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