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