mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-26 17:11:02 +02:00
icml rebuttal
This commit is contained in:
51616
parent
22267c7666
commit
696b45c51b
44 changed files with 5300 additions and 40 deletions
389
scripts/main_exp/eval/d2l-plot-context-length-8192.py
Normal file
389
scripts/main_exp/eval/d2l-plot-context-length-8192.py
Normal file
|
|
@ -0,0 +1,389 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import math
|
||||
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,
|
||||
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",
|
||||
}
|
||||
BASE_MODEL_EVAL_ROOT = Path("eval_results/google/gemma-2-2b-it")
|
||||
METRIC_CANDIDATES = (
|
||||
"qa_f1_score",
|
||||
"rougeL.f1",
|
||||
"qa_f1",
|
||||
"f1",
|
||||
"accuracy",
|
||||
)
|
||||
LEN_KEY_RE = re.compile(r"^test_(.+)_(.+)_len_([0-9]+)-([0-9]+|inf)$")
|
||||
MIN_CONTEXT_LENGTH = 4096
|
||||
MAX_CONTEXT_LENGTH = 16384
|
||||
|
||||
|
||||
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 LongBench performance vs context length for the 8192 chunk-size D2L eval run.",
|
||||
)
|
||||
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(
|
||||
"--chunk-size",
|
||||
type=int,
|
||||
default=8192,
|
||||
help="Chunk size to select from the sweep.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Optional output directory. 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 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 infer_metric_name(result_json: dict[str, object], dataset_name: str) -> str:
|
||||
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 choose_latest_complete_base_run(expected_datasets: list[str]) -> Path:
|
||||
candidates: list[Path] = []
|
||||
for run_dir in sorted(BASE_MODEL_EVAL_ROOT.iterdir()):
|
||||
if not run_dir.is_dir():
|
||||
continue
|
||||
if all(
|
||||
(run_dir / f"test_{dataset_name}_results.json").exists()
|
||||
for dataset_name in expected_datasets
|
||||
):
|
||||
candidates.append(run_dir)
|
||||
|
||||
if not candidates:
|
||||
raise ValueError(
|
||||
f"Could not find a complete base-model run under {BASE_MODEL_EVAL_ROOT}."
|
||||
)
|
||||
return candidates[-1]
|
||||
|
||||
|
||||
def extract_length_series(
|
||||
result_json: dict[str, object],
|
||||
dataset_name: str,
|
||||
metric_name: str,
|
||||
) -> tuple[list[int], list[float]]:
|
||||
xs: list[int] = []
|
||||
ys: list[float] = []
|
||||
metric_prefix = f"test_{dataset_name}_{metric_name}_len_"
|
||||
count_prefix = f"test_{dataset_name}_num_samples_{metric_name}_len_"
|
||||
|
||||
for key, value in result_json.items():
|
||||
if not key.startswith(metric_prefix):
|
||||
continue
|
||||
match = LEN_KEY_RE.match(key)
|
||||
if match is None:
|
||||
continue
|
||||
|
||||
low = int(match.group(3))
|
||||
high_str = match.group(4)
|
||||
if low == 0 or high_str == "inf":
|
||||
continue
|
||||
if value in (None, "None", "N/A"):
|
||||
continue
|
||||
|
||||
high = int(high_str)
|
||||
upper_bound = high + 1
|
||||
if upper_bound < MIN_CONTEXT_LENGTH or upper_bound > MAX_CONTEXT_LENGTH:
|
||||
continue
|
||||
count_key = f"{count_prefix}{low}-{high_str}"
|
||||
count = result_json.get(count_key)
|
||||
if not isinstance(count, int) or count <= 0:
|
||||
continue
|
||||
|
||||
xs.append(upper_bound)
|
||||
ys.append(float(value))
|
||||
|
||||
pairs = sorted(zip(xs, ys), key=lambda pair: pair[0])
|
||||
return [x for x, _ in pairs], [y for _, y in pairs]
|
||||
|
||||
|
||||
def tick_label(length_value: int) -> str:
|
||||
exponent = int(round(math.log2(length_value)))
|
||||
return rf"$2^{{{exponent}}}$"
|
||||
|
||||
|
||||
def to_series_map(xs: list[int], ys: list[float]) -> dict[int, float]:
|
||||
return {x: y for x, y in zip(xs, ys)}
|
||||
|
||||
|
||||
def chunk_count_label(context_length: int, chunk_size: int) -> str:
|
||||
num_chunks = max(1, math.ceil(context_length / chunk_size))
|
||||
suffix = "" if num_chunks == 1 else "s"
|
||||
return f"{num_chunks} chunk{suffix}"
|
||||
|
||||
|
||||
def chunk_label_offset(dataset_name: str) -> tuple[int, int]:
|
||||
if dataset_name == "longbench/2wikimqa_e":
|
||||
return (0, -15)
|
||||
return (0, 11)
|
||||
|
||||
|
||||
def chunk_label_position(
|
||||
dataset_name: str,
|
||||
index: int,
|
||||
total: int,
|
||||
) -> tuple[tuple[int, int], str]:
|
||||
dx, dy = chunk_label_offset(dataset_name)
|
||||
ha = "center"
|
||||
if total > 1 and index == 0:
|
||||
dx += 2
|
||||
elif total > 1 and index == total - 1:
|
||||
dx -= 2
|
||||
return (dx, dy), ha
|
||||
|
||||
|
||||
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)
|
||||
base_run_dir = choose_latest_complete_base_run(expected_datasets)
|
||||
|
||||
target_run = next(
|
||||
(
|
||||
run_info
|
||||
for run_info in selected_runs
|
||||
if run_info.chunk_size == args.chunk_size
|
||||
),
|
||||
None,
|
||||
)
|
||||
if target_run is None:
|
||||
raise SystemExit(
|
||||
f"Could not find a run with chunk size {args.chunk_size} under {eval_results_root}."
|
||||
)
|
||||
|
||||
output_dir = eval_results_root / "plots" if args.output is None else args.output
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
color_map = plt.get_cmap("tab10", 2)
|
||||
for dataset_name in expected_datasets:
|
||||
result_json = load_results_json(target_run.run_dir, dataset_name)
|
||||
metric_name = infer_metric_name(result_json, dataset_name)
|
||||
xs, ys = extract_length_series(result_json, dataset_name, metric_name)
|
||||
base_result_json = load_results_json(base_run_dir, dataset_name)
|
||||
base_xs, base_ys = extract_length_series(
|
||||
base_result_json, dataset_name, metric_name
|
||||
)
|
||||
|
||||
d2l_series = to_series_map(xs, ys)
|
||||
base_series = to_series_map(base_xs, base_ys)
|
||||
normalized_xs = [
|
||||
x for x in sorted(set(d2l_series) & set(base_series)) if base_series[x] != 0
|
||||
]
|
||||
all_xs = sorted(set(xs) | set(base_xs))
|
||||
if not all_xs:
|
||||
continue
|
||||
|
||||
fig, (ax_raw, ax_norm) = plt.subplots(
|
||||
1,
|
||||
2,
|
||||
figsize=(14.5, 5.8),
|
||||
constrained_layout=False,
|
||||
)
|
||||
fig.subplots_adjust(bottom=0.16, top=0.84, wspace=0.28)
|
||||
|
||||
if xs:
|
||||
ax_raw.plot(
|
||||
xs,
|
||||
ys,
|
||||
marker="o",
|
||||
linewidth=2,
|
||||
label="D2L",
|
||||
markeredgecolor="black",
|
||||
markeredgewidth=1,
|
||||
color=color_map(0),
|
||||
)
|
||||
for i, (x, y) in enumerate(zip(xs, ys)):
|
||||
raw_label_offset, raw_ha = chunk_label_position(
|
||||
dataset_name, i, len(xs)
|
||||
)
|
||||
ax_raw.annotate(
|
||||
chunk_count_label(x, args.chunk_size),
|
||||
(x, y),
|
||||
textcoords="offset points",
|
||||
xytext=raw_label_offset,
|
||||
ha=raw_ha,
|
||||
fontsize=8,
|
||||
color="#6c6f85",
|
||||
)
|
||||
if base_xs:
|
||||
ax_raw.plot(
|
||||
base_xs,
|
||||
base_ys,
|
||||
linestyle="--",
|
||||
marker="o",
|
||||
linewidth=2,
|
||||
label="Base",
|
||||
color=color_map(1),
|
||||
alpha=0.9,
|
||||
markeredgecolor="black",
|
||||
markeredgewidth=1,
|
||||
)
|
||||
|
||||
ax_raw.set_xscale("log", base=2)
|
||||
ax_raw.set_xticks(all_xs)
|
||||
ax_raw.set_xticklabels([tick_label(x) for x in all_xs])
|
||||
ax_raw.xaxis.set_minor_locator(plt.NullLocator())
|
||||
ax_raw.grid(True, alpha=0.25)
|
||||
ax_raw.set_title("Raw Performance", fontweight="bold")
|
||||
ax_raw.set_xlabel("Context Length (tokens)")
|
||||
ax_raw.set_ylabel("QA F1 Score")
|
||||
ax_raw.set_ylim(-0.05, 1.05)
|
||||
ax_raw.legend(loc="lower left", frameon=True, fancybox=True)
|
||||
|
||||
if normalized_xs:
|
||||
normalized_d2l_ys = [d2l_series[x] / base_series[x] for x in normalized_xs]
|
||||
normalized_base_ys = [1.0 for _ in normalized_xs]
|
||||
ax_norm.plot(
|
||||
normalized_xs,
|
||||
normalized_d2l_ys,
|
||||
marker="o",
|
||||
linewidth=2,
|
||||
label="D2L",
|
||||
markeredgecolor="black",
|
||||
markeredgewidth=1,
|
||||
color=color_map(0),
|
||||
)
|
||||
for i, (x, y) in enumerate(zip(normalized_xs, normalized_d2l_ys)):
|
||||
norm_label_offset, norm_ha = chunk_label_position(
|
||||
dataset_name,
|
||||
i,
|
||||
len(normalized_xs),
|
||||
)
|
||||
ax_norm.annotate(
|
||||
chunk_count_label(x, args.chunk_size),
|
||||
(x, y),
|
||||
textcoords="offset points",
|
||||
xytext=norm_label_offset,
|
||||
ha=norm_ha,
|
||||
fontsize=8,
|
||||
color="#6c6f85",
|
||||
)
|
||||
ax_norm.plot(
|
||||
normalized_xs,
|
||||
normalized_base_ys,
|
||||
linestyle="--",
|
||||
marker="o",
|
||||
linewidth=2,
|
||||
label="Base",
|
||||
color=color_map(1),
|
||||
alpha=0.9,
|
||||
markeredgecolor="black",
|
||||
markeredgewidth=1,
|
||||
)
|
||||
else:
|
||||
ax_norm.text(
|
||||
0.5,
|
||||
0.5,
|
||||
"No overlapping bins for normalization.",
|
||||
ha="center",
|
||||
va="center",
|
||||
transform=ax_norm.transAxes,
|
||||
)
|
||||
|
||||
norm_ticks = normalized_xs if normalized_xs else all_xs
|
||||
ax_norm.set_xscale("log", base=2)
|
||||
ax_norm.set_xticks(norm_ticks)
|
||||
ax_norm.set_xticklabels([tick_label(x) for x in norm_ticks])
|
||||
ax_norm.xaxis.set_minor_locator(plt.NullLocator())
|
||||
ax_norm.grid(True, alpha=0.25)
|
||||
ax_norm.set_title("Normalized Performance", fontweight="bold")
|
||||
ax_norm.set_xlabel("Context Length (tokens)")
|
||||
ax_norm.set_ylabel("Normalized Performance")
|
||||
ax_norm.set_ylim(bottom=0.0)
|
||||
if normalized_xs:
|
||||
ax_norm.legend(loc="lower left", frameon=True, fancybox=True)
|
||||
|
||||
fig.suptitle(
|
||||
f"Performance Grouped by Length\n({dataset_label(dataset_name)}, chunk size = {args.chunk_size})",
|
||||
fontweight="bold",
|
||||
fontsize=16,
|
||||
)
|
||||
|
||||
output_path = (
|
||||
output_dir
|
||||
/ f"d2l-context-length-{args.chunk_size}-{dataset_name.split('/')[-1]}.png"
|
||||
)
|
||||
fig.savefig(output_path, dpi=200, bbox_inches="tight")
|
||||
print(f"Saved plot to {output_path}")
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Add table
Add a link
Reference in a new issue