icml rebuttal

scripts/niah/plot_ablation_utils.py

@@ -0,0 +1,319 @@
+from __future__ import annotations
+
+import json
+import re
+import shlex
+from dataclasses import dataclass
+from pathlib import Path
+
+DATASET_RE = re.compile(r"ctx_magic_number_(\d+)_(\d+)")
+CHECKPOINT_RE = re.compile(r"^CHECKPOINT_PATH=(.+)$", re.MULTILINE)
+CMD_RE = re.compile(r"CMD:\s+(.*)$", re.MULTILINE)
+
+
+@dataclass(frozen=True)
+class RunInfo:
+    run_dir: Path
+    run_name: str
+    chunk_size: int
+    ctx_chunk_overlap: int
+    datasets: tuple[str, ...]
+    has_dense_logging: bool
+
+
+def load_script_text(script_path: Path) -> str:
+    return script_path.read_text(encoding="utf-8")
+
+
+def parse_checkpoint_path(script_text: str) -> Path:
+    match = CHECKPOINT_RE.search(script_text)
+    if match is None:
+        raise ValueError("Could not find CHECKPOINT_PATH in the eval script.")
+    return Path(match.group(1).strip().strip("\"'"))
+
+
+def parse_expected_datasets(script_text: str) -> list[str]:
+    datasets = {match.group(0) for match in DATASET_RE.finditer(script_text)}
+    if not datasets:
+        raise ValueError(
+            "Could not find any ctx_magic_number datasets in the eval script."
+        )
+    return sorted(datasets, key=dataset_context_length)
+
+
+def parse_eval_results_root(checkpoint_path: Path) -> Path:
+    checkpoint_match = re.search(r"checkpoint-(\d+)", checkpoint_path.as_posix())
+    if checkpoint_match is None:
+        raise ValueError(f"Could not parse checkpoint step from {checkpoint_path}.")
+    checkpoint_step = checkpoint_match.group(1)
+    return checkpoint_path.parent.parent / f"eval-results-{checkpoint_step}"
+
+
+def parse_dataset_context_bounds(dataset_name: str) -> tuple[int, int]:
+    match = DATASET_RE.fullmatch(dataset_name)
+    if match is None:
+        raise ValueError(f"Unsupported dataset name: {dataset_name}")
+    return int(match.group(1)), int(match.group(2))
+
+
+def dataset_context_length(dataset_name: str) -> int:
+    return parse_dataset_context_bounds(dataset_name)[1]
+
+
+def _extract_flag_value(tokens: list[str], flag: str) -> str | None:
+    for index, token in enumerate(tokens):
+        if token == flag and index + 1 < len(tokens):
+            return tokens[index + 1]
+        if token.startswith(f"{flag}="):
+            return token.split("=", 1)[1]
+    return None
+
+
+def _extract_datasets(tokens: list[str]) -> tuple[str, ...]:
+    if "--datasets" not in tokens:
+        return ()
+
+    index = tokens.index("--datasets") + 1
+    datasets: list[str] = []
+    while index < len(tokens) and not tokens[index].startswith("--"):
+        datasets.append(tokens[index])
+        index += 1
+    return tuple(datasets)
+
+
+def parse_run_info(run_dir: Path) -> RunInfo | None:
+    log_path = run_dir / "debug.log"
+    if not log_path.exists():
+        return None
+
+    log_text = log_path.read_text(encoding="utf-8", errors="replace")
+    match = CMD_RE.search(log_text)
+    if match is None:
+        return None
+
+    command = match.group(1).strip()
+    try:
+        tokens = shlex.split(command)
+    except ValueError:
+        tokens = command.split()
+
+    chunk_size_value = _extract_flag_value(tokens, "--max_ctx_chunk_len")
+    if chunk_size_value is None:
+        return None
+
+    datasets = _extract_datasets(tokens)
+    if not datasets:
+        return None
+
+    return RunInfo(
+        run_dir=run_dir,
+        run_name=run_dir.name,
+        chunk_size=int(chunk_size_value),
+        ctx_chunk_overlap=int(_extract_flag_value(tokens, "--ctx_chunk_overlap") or 0),
+        datasets=datasets,
+        has_dense_logging="--log_dense_lora_magnitude" in tokens,
+    )
+
+
+def _run_has_expected_outputs(run_dir: Path, datasets: list[str]) -> bool:
+    return all(
+        (run_dir / f"test_{dataset_name}_results.json").exists()
+        for dataset_name in datasets
+    )
+
+
+def choose_latest_runs_by_chunk(
+    eval_results_root: Path,
+    expected_datasets: list[str],
+) -> list[RunInfo]:
+    selected: dict[int, RunInfo] = {}
+    expected_dataset_set = set(expected_datasets)
+
+    for run_dir in sorted(eval_results_root.iterdir()):
+        if not run_dir.is_dir():
+            continue
+
+        run_info = parse_run_info(run_dir)
+        if run_info is None:
+            continue
+        if set(run_info.datasets) != expected_dataset_set:
+            continue
+        if not _run_has_expected_outputs(run_dir, expected_datasets):
+            continue
+
+        current = selected.get(run_info.chunk_size)
+        if current is None:
+            selected[run_info.chunk_size] = run_info
+            continue
+
+        current_key = (int(current.has_dense_logging), current.run_name)
+        candidate_key = (int(run_info.has_dense_logging), run_info.run_name)
+        if candidate_key > current_key:
+            selected[run_info.chunk_size] = run_info
+
+    return [selected[chunk_size] for chunk_size in sorted(selected)]
+
+
+def choose_latest_runs_by_dataset_and_chunk(
+    eval_results_root: Path,
+    expected_datasets: list[str],
+) -> dict[str, list[RunInfo]]:
+    selected: dict[tuple[str, int], RunInfo] = {}
+    expected_dataset_set = set(expected_datasets)
+
+    for run_dir in sorted(eval_results_root.iterdir()):
+        if not run_dir.is_dir():
+            continue
+
+        run_info = parse_run_info(run_dir)
+        if run_info is None or len(run_info.datasets) != 1:
+            continue
+
+        dataset_name = run_info.datasets[0]
+        if dataset_name not in expected_dataset_set:
+            continue
+        if not _run_has_expected_outputs(run_dir, [dataset_name]):
+            continue
+
+        key = (dataset_name, run_info.chunk_size)
+        current = selected.get(key)
+        if current is None:
+            selected[key] = run_info
+            continue
+
+        current_key = (int(current.has_dense_logging), current.run_name)
+        candidate_key = (int(run_info.has_dense_logging), run_info.run_name)
+        if candidate_key > current_key:
+            selected[key] = run_info
+
+    grouped: dict[str, list[RunInfo]] = {
+        dataset_name: [] for dataset_name in expected_datasets
+    }
+    for (dataset_name, _), run_info in selected.items():
+        grouped[dataset_name].append(run_info)
+
+    for dataset_name, run_infos in grouped.items():
+        grouped[dataset_name] = sorted(
+            run_infos, key=lambda run_info: run_info.chunk_size
+        )
+
+    return grouped
+
+
+def choose_latest_runs_by_dataset_chunk_and_overlap(
+    eval_results_root: Path,
+    expected_datasets: list[str],
+    ctx_chunk_overlap: int,
+) -> dict[str, list[RunInfo]]:
+    selected: dict[tuple[str, int], RunInfo] = {}
+    expected_dataset_set = set(expected_datasets)
+
+    for run_dir in sorted(eval_results_root.iterdir()):
+        if not run_dir.is_dir():
+            continue
+
+        run_info = parse_run_info(run_dir)
+        if run_info is None or len(run_info.datasets) != 1:
+            continue
+        if run_info.ctx_chunk_overlap != ctx_chunk_overlap:
+            continue
+
+        dataset_name = run_info.datasets[0]
+        if dataset_name not in expected_dataset_set:
+            continue
+        if not _run_has_expected_outputs(run_dir, [dataset_name]):
+            continue
+
+        key = (dataset_name, run_info.chunk_size)
+        current = selected.get(key)
+        if current is None:
+            selected[key] = run_info
+            continue
+
+        current_key = (int(current.has_dense_logging), current.run_name)
+        candidate_key = (int(run_info.has_dense_logging), run_info.run_name)
+        if candidate_key > current_key:
+            selected[key] = run_info
+
+    grouped: dict[str, list[RunInfo]] = {
+        dataset_name: [] for dataset_name in expected_datasets
+    }
+    for (dataset_name, _), run_info in selected.items():
+        grouped[dataset_name].append(run_info)
+
+    for dataset_name, run_infos in grouped.items():
+        grouped[dataset_name] = sorted(
+            run_infos, key=lambda run_info: run_info.chunk_size
+        )
+
+    return grouped
+
+
+def load_results_json(run_dir: Path, dataset_name: str) -> dict[str, object]:
+    path = run_dir / f"test_{dataset_name}_results.json"
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def load_dense_stats_json(run_dir: Path, dataset_name: str) -> dict[str, object] | None:
+    path = run_dir / f"test_{dataset_name}_dense_lora_magnitude.json"
+    if not path.exists():
+        return None
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def _to_float(value: object) -> float | None:
+    if value in (None, "None", "N/A"):
+        return None
+    if isinstance(value, (int, float)):
+        return float(value)
+    return None
+
+
+def get_performance_value(
+    run_dir: Path,
+    dataset_name: str,
+    metric_name: str = "rougeL.f1",
+) -> float | None:
+    result_json = load_results_json(run_dir, dataset_name)
+    return _to_float(result_json.get(f"test_{dataset_name}_{metric_name}"))
+
+
+def get_dense_value(
+    run_dir: Path,
+    dataset_name: str,
+    metric_name: str = "mean",
+) -> float | None:
+    result_json = load_results_json(run_dir, dataset_name)
+    metric_key = f"test_{dataset_name}_dense_lora_magnitude_{metric_name}"
+    value = _to_float(result_json.get(metric_key))
+    if value is not None:
+        return value
+
+    dense_stats = load_dense_stats_json(run_dir, dataset_name)
+    if dense_stats is None:
+        return None
+
+    overall = dense_stats.get("overall", {})
+    if not isinstance(overall, dict):
+        return None
+    return _to_float(overall.get(metric_name))
+
+
+def get_dense_values(
+    run_dir: Path,
+    dataset_name: str,
+    metric_names: list[str] | tuple[str, ...],
+) -> dict[str, float | None]:
+    return {
+        metric_name: get_dense_value(
+            run_dir,
+            dataset_name,
+            metric_name=metric_name,
+        )
+        for metric_name in metric_names
+    }
+
+
+def ensure_output_path(output_path: Path) -> Path:
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    return output_path