apunkt/doc-to-lora
1
0
Fork 0
mirror of https://github.com/SakanaAI/doc-to-lora.git synced 2026-07-23 17:01:04 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

icml rebuttal

Browse source
This commit is contained in:
51616 2026-06-15 04:31:47 +00:00
parent 22267c7666
commit 696b45c51b
44 changed files with 5300 additions and 40 deletions
Download patch file Download diff file Expand all files Collapse all files

6
scripts/main_exp/eval/base_model.sh
View file

@ -1,8 +1,8 @@
# no truncation
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen 1
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e longbench/gov_report_e --split test --eval_batch_size_gen 1
# w/ truncation
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen 1 --truncate_if_too_long_inp
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e longbench/gov_report_e --split test --eval_batch_size_gen 1 --truncate_if_too_long_inp
# no context
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen 1 --remove_context
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e longbench/gov_report_e --split test --eval_batch_size_gen 1 --remove_context

8
scripts/main_exp/eval/base_model_long_ctx.sh Executable file
View file

@ -0,0 +1,8 @@
# no truncation
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e longbench/gov_report_e --split test --eval_batch_size_gen 1
# w/ truncation
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e longbench/gov_report_e --split test --eval_batch_size_gen 1 --truncate_if_too_long_inp
# no context
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e longbench/gov_report_e --split test --eval_batch_size_gen 1 --remove_context

2
scripts/main_exp/eval/cd.sh
View file

@ -3,4 +3,4 @@ WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it
# longbench
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/qasper_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --cd_use_gen_q --q_gen_rounds=1
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/qasper_e longbench/gov_report_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --cd_use_gen_q --q_gen_rounds=1

2
scripts/main_exp/eval/cd_oracle.sh
View file

@ -3,4 +3,4 @@ WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it
# longbench
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/qasper_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/qasper_e longbench/gov_report_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp

8
scripts/main_exp/eval/clipper.sh Normal file
View file

@ -0,0 +1,8 @@
# plain base-model baseline truncate
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets oolong-synth --split test --eval_batch_size_gen 1 --truncate_if_too_long_inp --max_test_samples_per_ds 500
# Context-free sanity check.
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets oolong-synth --split test --eval_batch_size_gen 1 --remove_context --max_test_samples_per_ds 500
# Iterative D2L variant.
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/checkpoint-$step/pytorch_model.bin --datasets oolong-synth --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 1 --max_test_samples_per_ds 500 --use_iterative_mode

10
scripts/main_exp/eval/d2l-chunking-abalation.sh Executable file
View file

@ -0,0 +1,10 @@
# main results
# batched
for max_ctx_chunk_len in 256 512 1024 2048 4096 8192; do
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path trained_d2l/gemma_2b_d2l/checkpoint-20000/pytorch_model.bin --datasets longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e longbench/gov_report_e --split test --max_ctx_chunk_len $max_ctx_chunk_len --eval_batch_size_gen 1
done
# # iterative
# WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/checkpoint-$step/pytorch_model.bin --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e longbench/gov_report_e --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 1 --use_iterative_mode

248
scripts/main_exp/eval/d2l-plot-chunking-abalation.py Normal file
View file

@ -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()

389
scripts/main_exp/eval/d2l-plot-context-length-8192.py Normal file
View 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()

740
scripts/main_exp/eval/d2l-plot-rag-vs-hybrid.py Normal file
View file

@ -0,0 +1,740 @@
from __future__ import annotations
import argparse
import csv
import json
import re
import shlex
from collections import defaultdict, deque
from dataclasses import dataclass
from pathlib import Path
import matplotlib
import matplotlib as mpl
matplotlib.use("Agg")
import matplotlib.pyplot as plt
LATTE_STYLE = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
BASE_MODEL_EVAL_ROOT = Path("eval_results/google/gemma-2-2b-it")
DATASET_LABELS = {
"oolong-synth": "Oolong-Synth",
"longbench/qasper_e": "Qasper",
"longbench/2wikimqa_e": "2WikiMQA",
"longbench/multifieldqa_en_e": "MultiFieldQA-EN",
"longbench/gov_report_e": "GovReport",
}
METRIC_CANDIDATES = (
"oolong_score",
"qa_f1_score",
"rougeL.f1",
"qa_f1",
"f1",
"accuracy",
)
METRIC_LABELS = {
"oolong_score": "Oolong Score",
"qa_f1_score": "QA F1 Score",
"rougeL.f1": "ROUGE-L F1",
"qa_f1": "QA F1",
"f1": "F1",
"accuracy": "Accuracy",
}
COMPARISON_TOLERANCE = 0.2
CMD_RE = re.compile(r"CMD:\s+(.*)$", re.MULTILINE)
BASELINE_EXCLUDE_FLAGS = {
"--use_rag",
"--use_hybrid_rag",
"--use_cd",
"--use_llmlingua",
"--use_t2l",
}
NO_CONTEXT_COLUMN = "No context"
CSV_SUBCOLUMNS = (NO_CONTEXT_COLUMN, "top-1", "top-4")
@dataclass(frozen=True)
class RagRunInfo:
run_dir: Path
run_name: str
datasets: tuple[str, ...]
is_hybrid: bool
top_k: int
chunk_size: int
chunk_overlap: int
max_retrieved_tokens: int
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"] = 12
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 raw-RAG vs hybrid-RAG comparison figure per dataset.",
)
parser.add_argument(
"--rag-root",
type=Path,
default=Path("eval_results/google/gemma-2-2b-it"),
help="Directory containing raw RAG eval runs.",
)
parser.add_argument(
"--hybrid-root",
type=Path,
default=Path("trained_d2l/gemma_2b_d2l/eval-results-20000"),
help="Directory containing hybrid RAG eval runs.",
)
parser.add_argument(
"--rag-chunk-size",
type=int,
default=256,
help="Filter runs by rag chunk size.",
)
parser.add_argument(
"--rag-chunk-overlap",
type=int,
default=64,
help="Filter runs by rag chunk overlap.",
)
parser.add_argument(
"--rag-max-retrieved-tokens",
type=int,
default=1536,
help="Filter runs by rag max retrieved tokens.",
)
parser.add_argument(
"--datasets",
nargs="*",
default=None,
help="Optional dataset allowlist. Defaults to all datasets with matched raw/hybrid runs.",
)
parser.add_argument(
"--top-ks",
nargs="*",
type=int,
default=None,
help="Optional top-k allowlist. Defaults to all matched top-k values.",
)
parser.add_argument(
"--output-dir",
type=Path,
default=None,
help="Optional output directory. Defaults to HYBRID_ROOT/plots.",
)
return parser
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 _has_max_test_samples_flag(tokens: list[str]) -> bool:
return _extract_flag_value(tokens, "--max_test_samples_per_ds") is not None
def parse_rag_run_info(run_dir: Path) -> RagRunInfo | 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()
use_rag = "--use_rag" in tokens
use_hybrid_rag = "--use_hybrid_rag" in tokens
if use_rag == use_hybrid_rag:
return None
if _has_max_test_samples_flag(tokens):
return None
datasets = _extract_datasets(tokens)
if not datasets:
return None
top_k_value = _extract_flag_value(tokens, "--rag_top_k")
chunk_size_value = _extract_flag_value(tokens, "--rag_chunk_size")
chunk_overlap_value = _extract_flag_value(tokens, "--rag_chunk_overlap")
max_tokens_value = _extract_flag_value(tokens, "--rag_max_retrieved_tokens")
if None in (
top_k_value,
chunk_size_value,
chunk_overlap_value,
max_tokens_value,
):
return None
return RagRunInfo(
run_dir=run_dir,
run_name=run_dir.name,
datasets=datasets,
is_hybrid=use_hybrid_rag,
top_k=int(top_k_value),
chunk_size=int(chunk_size_value),
chunk_overlap=int(chunk_overlap_value),
max_retrieved_tokens=int(max_tokens_value),
)
def parse_plain_base_run_datasets(run_dir: Path) -> tuple[str, ...] | 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()
if "--model_name_or_path" not in tokens:
return None
if any(flag in tokens for flag in BASELINE_EXCLUDE_FLAGS):
return None
if _has_max_test_samples_flag(tokens):
return None
datasets = _extract_datasets(tokens)
if not datasets:
return None
return datasets
def parse_plain_d2l_run_datasets(run_dir: Path) -> tuple[str, ...] | 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()
if "--checkpoint_path" not in tokens:
return None
if any(flag in tokens for flag in BASELINE_EXCLUDE_FLAGS):
return None
if "--add_ctx_to_input" in tokens:
return None
if _has_max_test_samples_flag(tokens):
return None
max_ctx_chunk_len = _extract_flag_value(tokens, "--max_ctx_chunk_len")
if max_ctx_chunk_len != "8192":
return None
datasets = _extract_datasets(tokens)
if not datasets:
return None
return datasets
def result_json_path(run_dir: Path, dataset_name: str) -> Path:
return run_dir / f"test_{dataset_name}_results.json"
def load_results_json(run_dir: Path, dataset_name: str) -> dict[str, object]:
return json.loads(
result_json_path(run_dir, dataset_name).read_text(encoding="utf-8")
)
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 metric_label(metric_name: str) -> str:
return METRIC_LABELS.get(metric_name, metric_name.replace("_", " ").title())
def dataset_label(dataset_name: str) -> str:
return DATASET_LABELS.get(dataset_name, dataset_name.split("/")[-1])
def select_latest_runs(
root: Path,
*,
want_hybrid: bool,
allowed_datasets: set[str] | None,
allowed_top_ks: set[int] | None,
rag_chunk_size: int,
rag_chunk_overlap: int,
rag_max_retrieved_tokens: int,
) -> dict[tuple[str, int], RagRunInfo]:
selected: dict[tuple[str, int], RagRunInfo] = {}
if not root.exists():
return selected
for run_dir in sorted(root.iterdir()):
if not run_dir.is_dir():
continue
run_info = parse_rag_run_info(run_dir)
if run_info is None or run_info.is_hybrid != want_hybrid:
continue
if run_info.chunk_size != rag_chunk_size:
continue
if run_info.chunk_overlap != rag_chunk_overlap:
continue
if run_info.max_retrieved_tokens != rag_max_retrieved_tokens:
continue
if allowed_top_ks is not None and run_info.top_k not in allowed_top_ks:
continue
for dataset_name in run_info.datasets:
if allowed_datasets is not None and dataset_name not in allowed_datasets:
continue
if not result_json_path(run_info.run_dir, dataset_name).exists():
continue
key = (dataset_name, run_info.top_k)
current = selected.get(key)
if current is None or run_info.run_name > current.run_name:
selected[key] = run_info
return selected
def get_metric_value(
run_dir: Path, dataset_name: str, metric_name: str
) -> float | None:
result_json = load_results_json(run_dir, dataset_name)
value = result_json.get(f"test_{dataset_name}_{metric_name}")
if value in (None, "None", "N/A"):
return None
if isinstance(value, (int, float)):
return float(value)
return None
def select_latest_plain_base_runs(
root: Path,
datasets: set[str],
) -> dict[str, Path]:
selected: dict[str, Path] = {}
if not root.exists():
return selected
for run_dir in sorted(root.iterdir()):
if not run_dir.is_dir():
continue
run_datasets = parse_plain_base_run_datasets(run_dir)
if run_datasets is None:
continue
for dataset_name in run_datasets:
if dataset_name not in datasets:
continue
if not result_json_path(run_dir, dataset_name).exists():
continue
selected[dataset_name] = run_dir
return selected
def select_latest_plain_d2l_runs(
root: Path,
datasets: set[str],
) -> dict[str, Path]:
selected: dict[str, Path] = {}
if not root.exists():
return selected
for run_dir in sorted(root.iterdir()):
if not run_dir.is_dir():
continue
run_datasets = parse_plain_d2l_run_datasets(run_dir)
if run_datasets is None:
continue
for dataset_name in run_datasets:
if dataset_name not in datasets:
continue
if not result_json_path(run_dir, dataset_name).exists():
continue
selected[dataset_name] = run_dir
return selected
CSV_TOP_KS = (1, 4)
def csv_output_path(output_dir: Path) -> Path:
return output_dir / "d2l-rag-vs-hybrid.csv"
def confusion_output_path_for_dataset(output_dir: Path, dataset_name: str) -> Path:
return output_dir / f"d2l-rag-vs-hybrid-confusion-{dataset_name.split('/')[-1]}.png"
def generated_text_path(run_dir: Path, dataset_name: str) -> Path:
return run_dir / f"test_{dataset_name}_generated_text.jsonl"
def sample_key(sample: dict[str, object], dataset_name: str) -> str:
if dataset_name == "oolong-synth":
return "oolong:" + str(sample.get("oolong_id"))
input_text = str(sample.get("input", ""))
label_text = str(sample.get("label", ""))
ctx_ids_len = str(sample.get("ctx_ids_len", ""))
return f"{dataset_name}\n{input_text}\n{label_text}\n{ctx_ids_len}"
def sample_metric_value(sample: dict[str, object], metric_name: str) -> float:
value = sample.get(metric_name)
if not isinstance(value, (int, float)):
return 0.0
return float(value)
def load_metric_by_sample(
run_dir: Path,
dataset_name: str,
metric_name: str,
) -> dict[str, deque[float]]:
path = generated_text_path(run_dir, dataset_name)
grouped: dict[str, deque[float]] = defaultdict(deque)
with path.open(encoding="utf-8") as handle:
for line in handle:
sample = json.loads(line)
grouped[sample_key(sample, dataset_name)].append(
sample_metric_value(sample, metric_name)
)
return grouped
def relative_bucket(score: float, other_score: float) -> int:
delta = score - other_score
if delta > COMPARISON_TOLERANCE:
return 2
if delta < -COMPARISON_TOLERANCE:
return 0
return 1
def build_confusion_counts(
raw_run_dir: Path,
hybrid_run_dir: Path,
dataset_name: str,
metric_name: str,
) -> list[list[int]]:
raw_grouped = load_metric_by_sample(raw_run_dir, dataset_name, metric_name)
hybrid_grouped = load_metric_by_sample(hybrid_run_dir, dataset_name, metric_name)
raw_counts = {key: len(values) for key, values in raw_grouped.items()}
hybrid_counts = {key: len(values) for key, values in hybrid_grouped.items()}
if raw_counts != hybrid_counts:
raise ValueError(f"Could not align raw and hybrid samples for {dataset_name}.")
matrix = [[0, 0, 0], [0, 0, 0], [0, 0, 0]]
for key in sorted(raw_grouped):
raw_values = raw_grouped[key]
hybrid_values = hybrid_grouped[key]
while raw_values:
raw_score = raw_values.popleft()
hybrid_score = hybrid_values.popleft()
raw_bucket = relative_bucket(raw_score, hybrid_score)
hybrid_bucket = relative_bucket(hybrid_score, raw_score)
matrix[raw_bucket][hybrid_bucket] += 1
return matrix
def plot_confusion_matrices(
output_dir: Path,
dataset_name: str,
metric_name: str,
paired_points: list[tuple[int, float, float]],
raw_runs: dict[tuple[str, int], RagRunInfo],
hybrid_runs: dict[tuple[str, int], RagRunInfo],
) -> None:
ncols = len(paired_points)
fig, axes = plt.subplots(
1,
ncols,
figsize=(4.8 * ncols, 4.6),
constrained_layout=False,
)
if ncols == 1:
axes = [axes]
fig.subplots_adjust(bottom=0.16, top=0.80, wspace=0.35)
for ax, (top_k, _, _) in zip(axes, paired_points):
matrix = build_confusion_counts(
raw_runs[(dataset_name, top_k)].run_dir,
hybrid_runs[(dataset_name, top_k)].run_dir,
dataset_name,
metric_name,
)
image = ax.imshow(matrix, cmap="Blues")
total = sum(sum(row) for row in matrix)
max_cell = max(max(row) for row in matrix) if total else 0
for raw_index in range(3):
for hybrid_index in range(3):
count = matrix[raw_index][hybrid_index]
pct = 100.0 * count / total if total else 0.0
color = "white" if count > max_cell / 2 else "#4c4f69"
ax.text(
hybrid_index,
raw_index,
f"{count}\n({pct:.1f}%)",
ha="center",
va="center",
color=color,
fontsize=11,
fontweight="bold" if count == max_cell and count > 0 else None,
)
ax.set_xticks([0, 1, 2])
ax.set_yticks([0, 1, 2])
ax.set_xticklabels(["More\nIncorrect", "Same", "More\nCorrect"])
ax.set_yticklabels(["More\nIncorrect", "Same", "More\nCorrect"])
ax.set_xlabel("Hybrid RAG")
if ax is axes[0]:
ax.set_ylabel("Raw RAG")
ax.set_title(f"top-k = {top_k}", fontweight="bold")
ax.grid(False)
for spine in ax.spines.values():
spine.set_visible(True)
spine.set_color("#ccd0da")
fig.colorbar(image, ax=ax, fraction=0.046, pad=0.04)
fig.suptitle(
f"Relative Score Agreement\n({dataset_label(dataset_name)}, delta threshold = {COMPARISON_TOLERANCE:.1f})",
fontweight="bold",
fontsize=16,
)
output_path = confusion_output_path_for_dataset(output_dir, dataset_name)
fig.savefig(output_path, dpi=200, bbox_inches="tight")
print(f"Saved plot to {output_path}")
plt.close(fig)
def write_summary_csv(
output_dir: Path,
datasets: list[str],
summary: dict[str, dict[str, dict[str, float | str | None]]],
) -> None:
output_path = csv_output_path(output_dir)
with output_path.open("w", encoding="utf-8", newline="") as handle:
writer = csv.writer(handle)
top_header = [""]
sub_header = ["Method"]
for dataset_name in datasets:
pretty_name = dataset_label(dataset_name)
top_header.extend([pretty_name] * len(CSV_SUBCOLUMNS))
sub_header.extend(CSV_SUBCOLUMNS)
writer.writerow(top_header)
writer.writerow(sub_header)
for method_name in ("RAG", "RAG + D2L", "D2L"):
row = [method_name]
for dataset_name in datasets:
method_values = summary.get(method_name, {}).get(dataset_name, {})
for subcolumn in CSV_SUBCOLUMNS:
value = method_values.get(subcolumn)
if isinstance(value, (int, float)):
row.append(f"{float(value):.4f}")
elif value is None:
row.append("")
else:
row.append(str(value))
writer.writerow(row)
print(f"Saved CSV to {output_path}")
def main() -> None:
args = build_parser().parse_args()
configure_plot_style()
allowed_datasets = set(args.datasets) if args.datasets else None
allowed_top_ks = set(args.top_ks) if args.top_ks else None
output_dir = (
args.hybrid_root / "plots" if args.output_dir is None else args.output_dir
)
output_dir.mkdir(parents=True, exist_ok=True)
raw_runs = select_latest_runs(
args.rag_root,
want_hybrid=False,
allowed_datasets=allowed_datasets,
allowed_top_ks=allowed_top_ks,
rag_chunk_size=args.rag_chunk_size,
rag_chunk_overlap=args.rag_chunk_overlap,
rag_max_retrieved_tokens=args.rag_max_retrieved_tokens,
)
hybrid_runs = select_latest_runs(
args.hybrid_root,
want_hybrid=True,
allowed_datasets=allowed_datasets,
allowed_top_ks=allowed_top_ks,
rag_chunk_size=args.rag_chunk_size,
rag_chunk_overlap=args.rag_chunk_overlap,
rag_max_retrieved_tokens=args.rag_max_retrieved_tokens,
)
shared_keys = set(raw_runs) & set(hybrid_runs)
if not shared_keys:
raise SystemExit("Could not find any matched raw-RAG and hybrid-RAG runs.")
datasets = sorted({dataset_name for dataset_name, _ in shared_keys})
base_runs = select_latest_plain_base_runs(BASE_MODEL_EVAL_ROOT, set(datasets))
d2l_runs = select_latest_plain_d2l_runs(args.hybrid_root, set(datasets))
summary: dict[str, dict[str, dict[str, float | str | None]]] = {
"RAG": {},
"RAG + D2L": {},
"D2L": {},
}
for dataset_name in datasets:
base_run_dir = base_runs.get(dataset_name)
if base_run_dir is None:
continue
shared_top_ks = sorted(
top_k
for candidate_dataset, top_k in shared_keys
if candidate_dataset == dataset_name
)
if not shared_top_ks:
continue
raw_result_json = load_results_json(
raw_runs[(dataset_name, shared_top_ks[0])].run_dir, dataset_name
)
hybrid_result_json = load_results_json(
hybrid_runs[(dataset_name, shared_top_ks[0])].run_dir,
dataset_name,
)
metric_name = infer_metric_name(raw_result_json, dataset_name)
hybrid_metric_name = infer_metric_name(hybrid_result_json, dataset_name)
if metric_name != hybrid_metric_name:
raise ValueError(
f"Metric mismatch for {dataset_name}: raw uses {metric_name}, hybrid uses {hybrid_metric_name}."
)
base_result_json = load_results_json(base_run_dir, dataset_name)
base_metric_name = infer_metric_name(base_result_json, dataset_name)
if metric_name != base_metric_name:
raise ValueError(
f"Metric mismatch for {dataset_name}: comparison uses {metric_name}, baseline uses {base_metric_name}."
)
base_value = get_metric_value(base_run_dir, dataset_name, metric_name)
if base_value in (None, 0):
continue
d2l_run_dir = d2l_runs.get(dataset_name)
d2l_value = None
if d2l_run_dir is not None:
d2l_result_json = load_results_json(d2l_run_dir, dataset_name)
d2l_metric_name = infer_metric_name(d2l_result_json, dataset_name)
if d2l_metric_name == metric_name:
candidate_value = get_metric_value(
d2l_run_dir, dataset_name, metric_name
)
if candidate_value is not None:
d2l_value = candidate_value / base_value
paired_points: list[tuple[int, float, float]] = []
for top_k in shared_top_ks:
raw_value = get_metric_value(
raw_runs[(dataset_name, top_k)].run_dir, dataset_name, metric_name
)
hybrid_value = get_metric_value(
hybrid_runs[(dataset_name, top_k)].run_dir,
dataset_name,
metric_name,
)
if raw_value is None or hybrid_value is None:
continue
paired_points.append(
(top_k, raw_value / base_value, hybrid_value / base_value)
)
if not paired_points:
continue
summary["RAG"][dataset_name] = {
NO_CONTEXT_COLUMN: "N/A",
**{f"top-{top_k}": raw_value for top_k, raw_value, _ in paired_points},
}
summary["RAG + D2L"][dataset_name] = {
NO_CONTEXT_COLUMN: "N/A",
**{
f"top-{top_k}": hybrid_value for top_k, _, hybrid_value in paired_points
},
}
summary["D2L"][dataset_name] = {
NO_CONTEXT_COLUMN: d2l_value,
"top-1": "N/A",
"top-4": "N/A",
}
plot_confusion_matrices(
output_dir,
dataset_name,
metric_name,
paired_points,
raw_runs,
hybrid_runs,
)
write_summary_csv(output_dir, datasets, summary)
if __name__ == "__main__":
main()

4
scripts/main_exp/eval/d2l.sh
View file

@ -1,9 +1,9 @@
# main results
# batched
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/checkpoint-$step/pytorch_model.bin --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 1
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/checkpoint-$step/pytorch_model.bin --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e longbench/gov_report_e --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 1
# iterative
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/checkpoint-$step/pytorch_model.bin --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 1 --use_iterative_mode
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/checkpoint-$step/pytorch_model.bin --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e longbench/gov_report_e --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 1 --use_iterative_mode
# query internalization
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/checkpoint-$step/pytorch_model.bin --datasets squad --split test --eval_batch_size_gen=1 --flip_ctx_inp

13
scripts/main_exp/eval/d2l_rag.sh Normal file
View file

@ -0,0 +1,13 @@
# Hybrid mode: Doc-to-LoRA internalizes the full document and BM25-style RAG adds
# retrieved evidence to the prompt-side input.
WANDB_MODE=disabled uv run run_eval.py \
--checkpoint_path train_outputs/runs/$RUN_NAME/checkpoint-$step/pytorch_model.bin \
--datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e longbench/gov_report_e \
--split test \
--eval_batch_size_gen=1 \
--use_hybrid_rag \
--rag_chunk_size 256 \
--rag_chunk_overlap 64 \
--rag_top_k 4 \
--rag_max_retrieved_tokens 1536

2
scripts/main_exp/eval/llmlingua.sh
View file

@ -1,4 +1,4 @@
for dataset in squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e; do
for dataset in squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e longbench/gov_report_e; do
for rate in 0.9 0.8 0.6 0.4 0.2 0.1; do
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets "$dataset" --split test --eval_batch_size_gen=1 --use_llmlingua --llmlingua_compression_rate "$rate" --truncate_if_too_long_ctx
done

3
scripts/main_exp/eval/llmlingua_oolong.sh Executable file
View file

@ -0,0 +1,3 @@
for rate in 0.9 0.8 0.6 0.4 0.2 0.1; do
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets oolong-synth --split test --eval_batch_size_gen=1 --use_llmlingua --llmlingua_compression_rate "$rate" --truncate_if_too_long_ctx
done

26
scripts/main_exp/eval/rag.sh Normal file
View file

@ -0,0 +1,26 @@
# BM25-style per-example RAG over the dataset `context` field.
# Keep `--max_ctx_chunk_len` disabled here; the RAG wrapper does its own chunking.
for top_k in 1 4; do
WANDB_MODE=disabled uv run run_eval.py \
--model_name_or_path google/gemma-2-2b-it \
--datasets oolong-synth longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e longbench/gov_report_e \
--split test \
--eval_batch_size_gen=1 \
--use_rag \
--rag_chunk_size 256 \
--rag_chunk_overlap 64 \
--rag_top_k 4 \
--rag_max_retrieved_tokens 1536
WANDB_MODE=disabled run uv run run_eval.py \
--checkpoint_path trained_d2l/gemma_2b_d2l/checkpoint-20000/pytorch_model.bin \
--datasets longbench/qasper_e \
--split test \
--eval_batch_size_gen 1 \
--use_hybrid_rag \
--rag_chunk_size 256 \
--rag_chunk_overlap 64 \
--rag_top_k 4 \
--rag_max_retrieved_tokens 1536 --use_iterative_mode
done

2
scripts/main_exp/eval/t2l.sh
View file

@ -1,4 +1,4 @@
# download t2l checkpoint
uv run huggingface-cli download SakanaAI/text-to-lora --local-dir . --include "trained_t2l/gemma_2b_t2l"
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen=1 --use_t2l
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e longbench/gov_report_e --split test --eval_batch_size_gen=1 --use_t2l

6
scripts/niah/3-eval-vary-chunk-size.sh Normal file
View file

@ -0,0 +1,6 @@
CHECKPOINT_PATH=train_outputs/runs/Oct03_10-58-26_slurm0-a3nodeset-8_91076_c4dac553/checkpoint-1482/pytorch_model.bin
CTX_CHUNK_OVERLAP=${CTX_CHUNK_OVERLAP:-0}
for ctx_chunk_len in 32 64 128 256 512 1024 2048 4096; do
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ${CHECKPOINT_PATH} --datasets ctx_magic_number_32_1024 ctx_magic_number_1024_2048 ctx_magic_number_3072_4096 ctx_magic_number_7168_8192 ctx_magic_number_15360_16384 ctx_magic_number_28672_32768 ctx_magic_number_57344_65536 --max_ctx_chunk_len=${ctx_chunk_len} --ctx_chunk_overlap=${CTX_CHUNK_OVERLAP} --split test --eval_batch_size_gen=4 --log_dense_lora_magnitude
done

258
scripts/niah/3-plot-vary-chunk-size.py Normal file
View file

@ -0,0 +1,258 @@
from __future__ import annotations
import argparse
from pathlib import Path
import matplotlib
import matplotlib as mpl
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from plot_ablation_utils import (
choose_latest_runs_by_chunk,
dataset_context_length,
ensure_output_path,
get_dense_values,
get_performance_value,
load_script_text,
parse_checkpoint_path,
parse_eval_results_root,
parse_expected_datasets,
)
LATTE_STYLE = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
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 the context-length ablation from 3-eval-vary-chunk-size.sh.",
)
parser.add_argument(
"--eval-script",
type=Path,
default=Path(__file__).with_name("3-eval-vary-chunk-size.sh"),
help="Path to the eval sweep shell script.",
)
parser.add_argument(
"--output",
type=Path,
default=None,
help="Optional output PNG path. Defaults to eval-results-*/plots/3-vary-chunk-size.png.",
)
parser.add_argument(
"--metric",
default="rougeL.f1",
help="Performance metric to read from each *_results.json file.",
)
parser.add_argument(
"--dense-metric",
default="mean",
help="Dense LoRA metric to read when available.",
)
parser.add_argument(
"--dense-std-metric",
default="std",
help="Dense LoRA std metric to use for the shaded area when available.",
)
parser.add_argument(
"--dense-max-metric",
default="max",
help="Dense LoRA max metric to plot as a separate line when available.",
)
return parser
def main() -> None:
args = build_parser().parse_args()
configure_plot_style()
script_text = load_script_text(args.eval_script)
checkpoint_path = parse_checkpoint_path(script_text)
eval_results_root = parse_eval_results_root(checkpoint_path)
expected_datasets = parse_expected_datasets(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_path = (
eval_results_root / "plots" / "3-vary-chunk-size.png"
if args.output is None
else args.output
)
output_path = ensure_output_path(output_path)
context_lengths = [
dataset_context_length(dataset_name) for dataset_name in expected_datasets
]
color_map = plt.get_cmap("tab10", len(selected_runs))
fig, (ax_perf, ax_dense, ax_dense_max) = plt.subplots(
1,
3,
figsize=(20, 5.5),
constrained_layout=False,
)
fig.subplots_adjust(bottom=0.30, top=0.88, wspace=0.28)
has_any_dense_mean_values = False
has_any_dense_max_values = False
for index, run_info in enumerate(selected_runs):
color = color_map(index)
performance_xs: list[int] = []
performance_ys: list[float] = []
dense_mean_xs: list[int] = []
dense_mean_ys: list[float] = []
dense_std_lows: list[float] = []
dense_std_highs: list[float] = []
dense_max_xs: list[int] = []
dense_max_ys: list[float] = []
for dataset_name in expected_datasets:
context_length = dataset_context_length(dataset_name)
performance = get_performance_value(
run_info.run_dir,
dataset_name,
metric_name=args.metric,
)
if performance is not None:
performance_xs.append(context_length)
performance_ys.append(performance)
dense_stats = get_dense_values(
run_info.run_dir,
dataset_name,
metric_names=(
args.dense_metric,
args.dense_std_metric,
args.dense_max_metric,
),
)
dense_mean = dense_stats[args.dense_metric]
dense_std = dense_stats[args.dense_std_metric]
dense_max = dense_stats[args.dense_max_metric]
if dense_mean is not None:
dense_mean_xs.append(context_length)
dense_mean_ys.append(dense_mean)
if dense_max is not None:
dense_max_xs.append(context_length)
dense_max_ys.append(dense_max)
label = str(run_info.chunk_size)
ax_perf.plot(
performance_xs,
performance_ys,
marker="o",
linewidth=2,
color=color,
label=label,
markeredgecolor="black",
markeredgewidth=1,
)
if dense_mean_xs:
has_any_dense_mean_values = True
ax_dense.plot(
dense_mean_xs,
dense_mean_ys,
marker="o",
linewidth=2,
color=color,
markeredgecolor="black",
markeredgewidth=1,
)
if dense_max_xs:
has_any_dense_max_values = True
ax_dense_max.plot(
dense_max_xs,
dense_max_ys,
marker="o",
linewidth=2,
color=color,
markeredgecolor="black",
markeredgewidth=1,
)
for axis in (ax_perf, ax_dense, ax_dense_max):
axis.set_xscale("log", base=2)
axis.set_xticks(context_lengths)
axis.set_xticklabels(
[rf"$2^{{{int(value).bit_length() - 1}}}$" for value in context_lengths],
rotation=0,
ha="center",
)
axis.grid(True, alpha=0.25)
axis.xaxis.set_minor_locator(plt.NullLocator())
ax_perf.set_title("Retrieval Performance", fontweight="bold")
ax_perf.set_xlabel("Haystack Length (tokens)")
ax_perf.set_ylabel("ROUGE-L F1 Score")
ax_perf.set_ylim(-0.05, 1.05)
ax_dense.set_title("LoRA Mean Magnitude", fontweight="bold")
ax_dense.set_xlabel("Haystack Length (tokens)")
ax_dense.set_ylabel("Mean Per-Element Magnitude")
if has_any_dense_mean_values:
ax_dense.set_ylim(bottom=0.0)
else:
ax_dense.text(
0.5,
0.5,
"Dense LoRA stats are missing for these runs.",
ha="center",
va="center",
transform=ax_dense.transAxes,
)
ax_dense_max.set_title("LoRA Max Magnitude", fontweight="bold")
ax_dense_max.set_xlabel("Haystack Length (tokens)")
ax_dense_max.set_ylabel("Maximum Per-Element Magnitude")
if has_any_dense_max_values:
ax_dense_max.set_ylim(bottom=0.0)
else:
ax_dense_max.text(
0.5,
0.5,
"Dense LoRA max stats are missing for these runs.",
ha="center",
va="center",
transform=ax_dense_max.transAxes,
)
handles, labels = ax_perf.get_legend_handles_labels()
chunk_legend = fig.legend(
handles,
labels,
loc="lower center",
bbox_to_anchor=(0.5, -0.005),
ncol=min(6, len(labels)),
frameon=True,
fancybox=True,
title="Chunk size",
title_fontsize=16,
)
chunk_legend.get_frame().set_edgecolor("#bdbdbd")
chunk_legend.get_frame().set_linewidth(1.0)
chunk_legend.get_frame().set_alpha(1.0)
fig.savefig(output_path, dpi=200, bbox_inches="tight")
print(f"Saved plot to {output_path}")
if __name__ == "__main__":
main()

6
scripts/niah/4-eval-vary-num-chunks-overlap.sh Executable file
View file

@ -0,0 +1,6 @@
CHECKPOINT_PATH=train_outputs/runs/Oct03_10-58-26_slurm0-a3nodeset-8_91076_c4dac553/checkpoint-1482/pytorch_model.bin
CTX_CHUNK_OVERLAP=8
for ctx_chunk_len in 16 32 64 128 256 512 1024; do
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ${CHECKPOINT_PATH} --datasets ctx_magic_number_32_1024 --max_ctx_chunk_len=${ctx_chunk_len} --ctx_chunk_overlap=${CTX_CHUNK_OVERLAP} --split test --eval_batch_size_gen=4 --log_dense_lora_magnitude
done

26
scripts/niah/4-eval-vary-num-chunks.sh Executable file
View file

@ -0,0 +1,26 @@
CHECKPOINT_PATH=train_outputs/runs/Oct03_10-58-26_slurm0-a3nodeset-8_91076_c4dac553/checkpoint-1482/pytorch_model.bin
CTX_CHUNK_OVERLAP=0
for ctx_chunk_len in 128 256 512 1024 2048 4096 8192 16384 32768; do
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ${CHECKPOINT_PATH} --datasets ctx_magic_number_28672_32768 --max_ctx_chunk_len=${ctx_chunk_len} --ctx_chunk_overlap=${CTX_CHUNK_OVERLAP} --split test --eval_batch_size_gen=4 --log_dense_lora_magnitude
done
for ctx_chunk_len in 64 128 256 512 1024 2048 4096 8192 16384; do
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ${CHECKPOINT_PATH} --datasets ctx_magic_number_15360_16384 --max_ctx_chunk_len=${ctx_chunk_len} --ctx_chunk_overlap=${CTX_CHUNK_OVERLAP} --split test --eval_batch_size_gen=4 --log_dense_lora_magnitude
done
for ctx_chunk_len in 32 64 128 256 512 1024 2048 4096 8192; do
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ${CHECKPOINT_PATH} --datasets ctx_magic_number_7168_8192 --max_ctx_chunk_len=${ctx_chunk_len} --ctx_chunk_overlap=${CTX_CHUNK_OVERLAP} --split test --eval_batch_size_gen=4 --log_dense_lora_magnitude
done
for ctx_chunk_len in 16 32 64 128 256 512 1024 2048 4096; do
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ${CHECKPOINT_PATH} --datasets ctx_magic_number_3072_4096 --max_ctx_chunk_len=${ctx_chunk_len} --ctx_chunk_overlap=${CTX_CHUNK_OVERLAP} --split test --eval_batch_size_gen=4 --log_dense_lora_magnitude
done
for ctx_chunk_len in 8 16 32 64 128 256 512 1024 2048; do
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ${CHECKPOINT_PATH} --datasets ctx_magic_number_1024_2048 --max_ctx_chunk_len=${ctx_chunk_len} --ctx_chunk_overlap=${CTX_CHUNK_OVERLAP} --split test --eval_batch_size_gen=4 --log_dense_lora_magnitude
done
for ctx_chunk_len in 4 8 16 32 64 128 256 512 1024; do
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ${CHECKPOINT_PATH} --datasets ctx_magic_number_32_1024 --max_ctx_chunk_len=${ctx_chunk_len} --ctx_chunk_overlap=${CTX_CHUNK_OVERLAP} --split test --eval_batch_size_gen=4 --log_dense_lora_magnitude
done

247
scripts/niah/4-plot-vary-num-chunks-overlap.py Normal file
View file

@ -0,0 +1,247 @@
from __future__ import annotations
import argparse
import math
from pathlib import Path
import matplotlib
import matplotlib as mpl
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from plot_ablation_utils import (
choose_latest_runs_by_dataset_chunk_and_overlap,
dataset_context_length,
ensure_output_path,
get_performance_value,
load_script_text,
parse_checkpoint_path,
parse_eval_results_root,
parse_expected_datasets,
)
LATTE_STYLE = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
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="Compare overlapping vs non-overlapping chunking for the 4-* NIAH sweeps.",
)
parser.add_argument(
"--non-overlap-script",
type=Path,
default=Path(__file__).with_name("4-eval-vary-num-chunks.sh"),
help="Path to the non-overlapping eval sweep shell script.",
)
parser.add_argument(
"--overlap-script",
type=Path,
default=Path(__file__).with_name("4-eval-vary-num-chunks-overlap.sh"),
help="Path to the overlapping eval sweep shell script.",
)
parser.add_argument(
"--output",
type=Path,
default=None,
help="Optional output PNG path. Defaults to eval-results-*/plots/4-vary-num-chunks-overlap-comparison.png.",
)
parser.add_argument(
"--metric",
default="rougeL.f1",
help="Performance metric to read from each *_results.json file.",
)
parser.add_argument(
"--non-overlap-value",
type=int,
default=0,
help="Overlap value used for the non-overlapping sweep.",
)
parser.add_argument(
"--overlap-value",
type=int,
default=8,
help="Overlap value used for the overlapping sweep.",
)
return parser
def tick_label(chunk_size: int) -> str:
exponent = math.log2(chunk_size)
if exponent.is_integer():
return rf"$2^{{{int(exponent)}}}$"
return rf"$2^{{{exponent:.1f}}}$"
def dataset_title(dataset_name: str) -> str:
context_length = dataset_context_length(dataset_name)
exponent = int(math.log2(context_length))
return rf"Haystack Length $2^{{{exponent}}}$"
def line_label(overlap_value: int) -> str:
if overlap_value == 0:
return "Non-overlapping"
return f"Overlap = {overlap_value}"
def main() -> None:
args = build_parser().parse_args()
configure_plot_style()
non_overlap_text = load_script_text(args.non_overlap_script)
overlap_text = load_script_text(args.overlap_script)
checkpoint_path = parse_checkpoint_path(non_overlap_text)
eval_results_root = parse_eval_results_root(checkpoint_path)
overlap_eval_results_root = parse_eval_results_root(
parse_checkpoint_path(overlap_text)
)
if overlap_eval_results_root != eval_results_root:
raise SystemExit(
"The overlap and non-overlap scripts point to different eval-results roots."
)
non_overlap_datasets = set(parse_expected_datasets(non_overlap_text))
overlap_datasets = set(parse_expected_datasets(overlap_text))
expected_datasets = sorted(
non_overlap_datasets & overlap_datasets,
key=dataset_context_length,
)
if not expected_datasets:
raise SystemExit("No shared datasets were found between the two sweep scripts.")
non_overlap_runs = choose_latest_runs_by_dataset_chunk_and_overlap(
eval_results_root,
expected_datasets,
ctx_chunk_overlap=args.non_overlap_value,
)
overlap_runs = choose_latest_runs_by_dataset_chunk_and_overlap(
eval_results_root,
expected_datasets,
ctx_chunk_overlap=args.overlap_value,
)
output_path = (
eval_results_root / "plots" / "4-vary-num-chunks-overlap-comparison.png"
if args.output is None
else args.output
)
output_path = ensure_output_path(output_path)
num_datasets = len(expected_datasets)
fig, axes = plt.subplots(
1,
num_datasets,
figsize=(6.4 * num_datasets, 5.5),
squeeze=False,
constrained_layout=False,
)
fig.subplots_adjust(bottom=0.26, top=0.84, wspace=0.28)
color_map = plt.get_cmap("tab10", 2)
legend_handles = []
legend_labels = []
for axis, dataset_name in zip(axes[0], expected_datasets):
plotted_anything = False
dataset_chunk_sizes: set[int] = set()
for overlap_value, run_infos, color in (
(
args.non_overlap_value,
non_overlap_runs.get(dataset_name, []),
color_map(0),
),
(args.overlap_value, overlap_runs.get(dataset_name, []), color_map(1)),
):
chunk_sizes: list[int] = []
performance_values: list[float] = []
for run_info in run_infos:
performance = get_performance_value(
run_info.run_dir,
dataset_name,
metric_name=args.metric,
)
if performance is None:
continue
chunk_sizes.append(run_info.chunk_size)
performance_values.append(performance)
dataset_chunk_sizes.add(run_info.chunk_size)
if not performance_values:
continue
plotted_anything = True
(line,) = axis.plot(
chunk_sizes,
performance_values,
marker="o",
linewidth=2,
color=color,
label=line_label(overlap_value),
markeredgecolor="black",
markeredgewidth=1,
)
if line.get_label() not in legend_labels:
legend_handles.append(line)
legend_labels.append(line.get_label())
if not plotted_anything:
axis.text(
0.5,
0.5,
"No matching runs found.",
ha="center",
va="center",
transform=axis.transAxes,
)
else:
sorted_chunk_sizes = sorted(dataset_chunk_sizes)
axis.set_xscale("log", base=2)
axis.set_xticks(sorted_chunk_sizes)
axis.set_xticklabels([tick_label(value) for value in sorted_chunk_sizes])
axis.xaxis.set_minor_locator(plt.NullLocator())
axis.grid(True, alpha=0.25)
axis.set_title(dataset_title(dataset_name), fontweight="bold")
axis.set_xlabel("Chunk Size")
axis.set_ylabel("ROUGE-L F1 Score")
axis.set_ylim(-0.05, 1.05)
if legend_handles:
legend = fig.legend(
legend_handles,
legend_labels,
loc="lower center",
bbox_to_anchor=(0.5, 0.01),
ncol=len(legend_labels),
frameon=True,
fancybox=True,
)
legend.get_frame().set_edgecolor("#bdbdbd")
legend.get_frame().set_linewidth(1.0)
legend.get_frame().set_alpha(1.0)
fig.suptitle("Overlapping vs Non-overlapping Chunking", fontweight="bold")
fig.savefig(output_path, dpi=200, bbox_inches="tight")
print(f"Saved plot to {output_path}")
if __name__ == "__main__":
main()

277
scripts/niah/4-plot-vary-num-chunks.py Normal file
View file

@ -0,0 +1,277 @@
from __future__ import annotations
import argparse
import math
from pathlib import Path
import matplotlib
import matplotlib as mpl
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from plot_ablation_utils import (
choose_latest_runs_by_dataset_and_chunk,
dataset_context_length,
ensure_output_path,
get_dense_values,
get_performance_value,
load_script_text,
parse_checkpoint_path,
parse_eval_results_root,
parse_expected_datasets,
)
LATTE_STYLE = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
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 the num-chunks ablation from 4-eval-vary-num-chunks.sh.",
)
parser.add_argument(
"--eval-script",
type=Path,
default=Path(__file__).with_name("4-eval-vary-num-chunks.sh"),
help="Path to the eval sweep shell script.",
)
parser.add_argument(
"--output",
type=Path,
default=None,
help="Optional output PNG path. Defaults to eval-results-*/plots/4-vary-num-chunks.png.",
)
parser.add_argument(
"--metric",
default="rougeL.f1",
help="Performance metric to read from each *_results.json file.",
)
parser.add_argument(
"--dense-metric",
default="mean",
help="Dense LoRA metric to read when available.",
)
parser.add_argument(
"--dense-std-metric",
default="std",
help="Dense LoRA std metric to use for the shaded area when available.",
)
parser.add_argument(
"--dense-max-metric",
default="max",
help="Dense LoRA max metric to plot as a separate line when available.",
)
return parser
def main() -> None:
args = build_parser().parse_args()
configure_plot_style()
script_text = load_script_text(args.eval_script)
checkpoint_path = parse_checkpoint_path(script_text)
eval_results_root = parse_eval_results_root(checkpoint_path)
expected_datasets = parse_expected_datasets(script_text)
selected_runs_by_dataset = choose_latest_runs_by_dataset_and_chunk(
eval_results_root,
expected_datasets,
)
if not any(selected_runs_by_dataset.values()):
raise SystemExit(f"No matching eval runs found under {eval_results_root}.")
output_path = (
eval_results_root / "plots" / "4-vary-num-chunks.png"
if args.output is None
else args.output
)
output_path = ensure_output_path(output_path)
fig, (ax_perf, ax_dense, ax_dense_max) = plt.subplots(
1,
3,
figsize=(20, 5.5),
constrained_layout=False,
)
fig.subplots_adjust(bottom=0.27, top=0.88, wspace=0.28)
color_map = plt.get_cmap("tab10", len(expected_datasets))
all_num_chunks: set[int] = set()
has_any_dense_mean_values = False
has_any_dense_max_values = False
for index, dataset_name in enumerate(expected_datasets):
color = color_map(index)
context_length = dataset_context_length(dataset_name)
label = rf"$2^{{{int(context_length).bit_length() - 1}}}$"
num_chunks: list[int] = []
performance_values: list[float] = []
dense_mean_xs: list[int] = []
dense_mean_ys: list[float] = []
dense_max_xs: list[int] = []
dense_max_ys: list[float] = []
for run_info in selected_runs_by_dataset.get(dataset_name, []):
if run_info.chunk_size == 24:
continue
approx_num_chunks = max(1, context_length // run_info.chunk_size)
performance = get_performance_value(
run_info.run_dir,
dataset_name,
metric_name=args.metric,
)
if performance is None:
continue
num_chunks.append(approx_num_chunks)
performance_values.append(performance)
all_num_chunks.add(approx_num_chunks)
dense_stats = get_dense_values(
run_info.run_dir,
dataset_name,
metric_names=(
args.dense_metric,
args.dense_std_metric,
args.dense_max_metric,
),
)
dense_mean = dense_stats[args.dense_metric]
dense_max = dense_stats[args.dense_max_metric]
if dense_mean is not None:
dense_mean_xs.append(approx_num_chunks)
dense_mean_ys.append(dense_mean)
if dense_max is not None:
dense_max_xs.append(approx_num_chunks)
dense_max_ys.append(dense_max)
if performance_values:
ax_perf.plot(
num_chunks,
performance_values,
marker="o",
linewidth=2,
color=color,
label=label,
markeredgecolor="black",
markeredgewidth=1,
)
if dense_mean_ys:
has_any_dense_mean_values = True
ax_dense.plot(
dense_mean_xs,
dense_mean_ys,
marker="o",
linewidth=2,
color=color,
label=label,
markeredgecolor="black",
markeredgewidth=1,
)
if dense_max_ys:
has_any_dense_max_values = True
ax_dense_max.plot(
dense_max_xs,
dense_max_ys,
marker="o",
linewidth=2,
color=color,
label=label,
markeredgecolor="black",
markeredgewidth=1,
)
if not all_num_chunks:
raise SystemExit("No performance values were found for the selected runs.")
sorted_num_chunks = sorted(all_num_chunks)
for axis in (ax_perf, ax_dense, ax_dense_max):
axis.set_xscale("log", base=2)
axis.set_xticks(sorted_num_chunks)
axis.set_xticklabels(
[
rf"$2^{{{math.log2(value):.1f}}}$"
if not math.log2(value).is_integer()
else rf"$2^{{{int(math.log2(value))}}}$"
for value in sorted_num_chunks
],
rotation=0,
ha="center",
)
axis.grid(True, alpha=0.25)
axis.xaxis.set_minor_locator(plt.NullLocator())
ax_perf.set_title("Retrieval Performance", fontweight="bold")
ax_perf.set_xlabel("Number of Chunks")
ax_perf.set_ylabel("ROUGE-L F1 Score")
ax_perf.set_ylim(-0.05, 1.05)
ax_dense.set_title("LoRA Mean Magnitude", fontweight="bold")
ax_dense.set_xlabel("Number of Chunks")
ax_dense.set_ylabel("Mean Per-Element Magnitude")
if has_any_dense_mean_values:
ax_dense.set_ylim(bottom=0.0)
else:
ax_dense.text(
0.5,
0.5,
"Dense LoRA stats are missing for these runs.",
ha="center",
va="center",
transform=ax_dense.transAxes,
)
ax_dense_max.set_title("LoRA Max Magnitude", fontweight="bold")
ax_dense_max.set_xlabel("Number of Chunks")
ax_dense_max.set_ylabel("Maximum Per-Element Magnitude")
if has_any_dense_max_values:
ax_dense_max.set_ylim(bottom=0.0)
else:
ax_dense_max.text(
0.5,
0.5,
"Dense LoRA max stats are missing for these runs.",
ha="center",
va="center",
transform=ax_dense_max.transAxes,
)
handles, labels = ax_perf.get_legend_handles_labels()
haystack_legend = fig.legend(
handles,
labels,
loc="lower center",
bbox_to_anchor=(0.5, -0.005),
ncol=min(6, len(labels)),
frameon=True,
fancybox=True,
title="Haystack length",
title_fontsize=16,
)
haystack_legend.get_frame().set_edgecolor("#bdbdbd")
haystack_legend.get_frame().set_linewidth(1.0)
haystack_legend.get_frame().set_alpha(1.0)
fig.savefig(output_path, dpi=200, bbox_inches="tight")
print(f"Saved plot to {output_path}")
if __name__ == "__main__":
main()

319
scripts/niah/plot_ablation_utils.py Normal file
View file

@ -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

10
scripts/sbatch_1_gpu.sh Normal file
View file

@ -0,0 +1,10 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#SBATCH --gpus=1
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
"$@"