mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
248 lines
7.8 KiB
Python
248 lines
7.8 KiB
Python
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()
|