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