icml rebuttal

scripts/main_exp/eval/d2l-plot-chunking-abalation.py

@@ -0,0 +1,248 @@
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import sys
+from pathlib import Path
+
+import matplotlib
+import matplotlib as mpl
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+sys.path.append(str(Path(__file__).resolve().parents[2] / "niah"))
+
+from plot_ablation_utils import (  # noqa: E402
+    choose_latest_runs_by_chunk,
+    ensure_output_path,
+    get_performance_value,
+    load_results_json,
+    load_script_text,
+    parse_checkpoint_path,
+    parse_eval_results_root,
+)
+
+LATTE_STYLE = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
+DATASET_LABELS = {
+    "longbench/qasper_e": "Qasper",
+    "longbench/2wikimqa_e": "2WikiMQA",
+    "longbench/multifieldqa_en_e": "MultiFieldQA-EN",
+    "longbench/gov_report_e": "GovReport",
+}
+REFERENCE_PERFORMANCE = {
+    "longbench/2wikimqa_e": 0.3387,
+    "longbench/multifieldqa_en_e": 0.3938,
+    "longbench/qasper_e": 0.3839,
+}
+METRIC_CANDIDATES = (
+    "qa_f1_score",
+    "rougeL.f1",
+    "qa_f1",
+    "f1",
+    "accuracy",
+)
+
+
+def chunk_size_tick_label(chunk_size: int) -> str:
+    return rf"$2^{{{int(chunk_size).bit_length() - 1}}}$"
+
+
+def approx_chunk_label(value: float) -> str:
+    rounded = round(value)
+    if abs(value - rounded) < 0.05:
+        return str(int(rounded))
+    return f"{value:.1f}"
+
+
+def configure_plot_style() -> None:
+    plt.style.use(["ggplot", LATTE_STYLE])
+    plt.rcParams["axes.facecolor"] = "F7FBFC"
+    plt.rcParams["figure.facecolor"] = "white"
+    plt.rcParams["savefig.facecolor"] = "white"
+    plt.rcParams["grid.color"] = "cccccc"
+    plt.rcParams["grid.linewidth"] = 1
+    plt.rcParams["axes.edgecolor"] = "ccd0da"
+    plt.rcParams["legend.facecolor"] = "white"
+    plt.rcParams["legend.fontsize"] = 14
+    plt.rcParams["xtick.labelsize"] = 13
+    plt.rcParams["ytick.labelsize"] = 13
+    plt.rcParams["axes.labelweight"] = "bold"
+    plt.rcParams["axes.prop_cycle"] = mpl.cycler(color=plt.cm.tab10.colors)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Plot one D2L chunking-ablation figure per dataset.",
+    )
+    parser.add_argument(
+        "--eval-script",
+        type=Path,
+        default=Path(__file__).with_name("d2l-chunking-abalation.sh"),
+        help="Path to the eval sweep shell script.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        default=None,
+        help="Optional directory for output PNGs. Defaults to eval-results-*/plots/.",
+    )
+    return parser
+
+
+def parse_checkpoint_from_script(script_text: str) -> Path:
+    try:
+        return parse_checkpoint_path(script_text)
+    except ValueError:
+        match = re.search(r"--checkpoint_path\s+(\S+)", script_text)
+        if match is None:
+            raise
+        return Path(match.group(1))
+
+
+def infer_metric_name(run_dir: Path, dataset_name: str) -> str:
+    result_json = load_results_json(run_dir, dataset_name)
+    for metric_name in METRIC_CANDIDATES:
+        if f"test_{dataset_name}_{metric_name}" in result_json:
+            return metric_name
+    raise ValueError(f"Could not infer a metric for dataset {dataset_name}.")
+
+
+def dataset_label(dataset_name: str) -> str:
+    return DATASET_LABELS.get(dataset_name, dataset_name.split("/")[-1])
+
+
+def parse_datasets_from_script(script_text: str) -> list[str]:
+    match = re.search(r"--datasets\s+(.+?)\s+--split", script_text, flags=re.DOTALL)
+    if match is None:
+        raise ValueError("Could not parse datasets from the eval script.")
+    return match.group(1).split()
+
+
+def mean_context_length(run_dir: Path, dataset_name: str) -> float:
+    path = run_dir / f"test_{dataset_name}_generated_text.jsonl"
+    total = 0.0
+    count = 0
+    with path.open(encoding="utf-8") as handle:
+        for line in handle:
+            record = json.loads(line)
+            ctx_ids_len = record.get("ctx_ids_len")
+            if isinstance(ctx_ids_len, (int, float)):
+                total += float(ctx_ids_len)
+                count += 1
+    if count == 0:
+        raise ValueError(f"Could not find ctx_ids_len values in {path}.")
+    return total / count
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    configure_plot_style()
+
+    script_text = load_script_text(args.eval_script)
+    checkpoint_path = parse_checkpoint_from_script(script_text)
+    eval_results_root = parse_eval_results_root(checkpoint_path)
+    expected_datasets = parse_datasets_from_script(script_text)
+    selected_runs = choose_latest_runs_by_chunk(eval_results_root, expected_datasets)
+
+    if not selected_runs:
+        raise SystemExit(f"No matching eval runs found under {eval_results_root}.")
+
+    output_dir = (
+        eval_results_root / "plots" if args.output_dir is None else args.output_dir
+    )
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    for dataset_name in expected_datasets:
+        metric_name = infer_metric_name(selected_runs[0].run_dir, dataset_name)
+        reference_performance = REFERENCE_PERFORMANCE.get(dataset_name)
+        chunk_sizes: list[int] = []
+        approx_num_chunks: list[float] = []
+        normalized_performance_values: list[float] = []
+
+        for run_info in selected_runs:
+            chunk_size = run_info.chunk_size
+            if chunk_size == 65536:
+                continue
+            performance = get_performance_value(
+                run_info.run_dir,
+                dataset_name,
+                metric_name=metric_name,
+            )
+            if performance is None:
+                continue
+
+            avg_context_length = mean_context_length(run_info.run_dir, dataset_name)
+            chunk_sizes.append(chunk_size)
+            approx_num_chunks.append(avg_context_length / chunk_size)
+            plot_value = (
+                performance / reference_performance
+                if reference_performance not in (None, 0)
+                else performance
+            )
+            normalized_performance_values.append(plot_value)
+
+        if not chunk_sizes:
+            continue
+
+        fig, ax_perf = plt.subplots(
+            1,
+            1,
+            figsize=(7.5, 5.5),
+            constrained_layout=False,
+        )
+        fig.subplots_adjust(bottom=0.18, top=0.86)
+
+        ax_perf.plot(
+            chunk_sizes,
+            normalized_performance_values,
+            marker="o",
+            linewidth=2,
+            color="#1f77b4",
+            markeredgecolor="black",
+            markeredgewidth=1,
+        )
+        for chunk_size, performance, approx_chunks in zip(
+            chunk_sizes,
+            normalized_performance_values,
+            approx_num_chunks,
+        ):
+            ax_perf.annotate(
+                approx_chunk_label(approx_chunks),
+                xy=(chunk_size, performance),
+                xytext=(0, 7),
+                textcoords="offset points",
+                ha="center",
+                va="bottom",
+                fontsize=11,
+                color="#4c4f69",
+            )
+
+        ax_perf.grid(True, alpha=0.25)
+        ax_perf.set_xscale("log", base=2)
+        ax_perf.set_xticks(chunk_sizes)
+        ax_perf.set_xticklabels(
+            [chunk_size_tick_label(chunk_size) for chunk_size in chunk_sizes]
+        )
+        ax_perf.xaxis.set_minor_locator(plt.NullLocator())
+
+        pretty_name = dataset_label(dataset_name)
+        ax_perf.set_title(f"Retrieval Performance\n({pretty_name})", fontweight="bold")
+        ax_perf.set_xlabel("Chunk Size (tokens)")
+        ax_perf.set_ylabel(
+            "Normalized Performance"
+            if reference_performance not in (None, 0)
+            else "Performance"
+        )
+        ax_perf.set_ylim(bottom=0.0)
+
+        output_path = ensure_output_path(
+            output_dir / f"d2l-chunking-abalation-{dataset_name.split('/')[-1]}.png"
+        )
+        fig.savefig(output_path, dpi=200, bbox_inches="tight")
+        print(f"Saved plot to {output_path}")
+
+
+if __name__ == "__main__":
+    main()