File size: 10,609 Bytes
fa82a9d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 | #!/usr/bin/env python3
"""Plot training curves (loss, token accuracy, dev exact match) without new deps.
Reads the TRL ``trainer_state.json`` log history plus ``dev_metrics.jsonl`` from
one or more run directories and writes, per run:
* ``history.csv`` -- step, loss, learning_rate, grad_norm, token accuracy, eval loss
* ``loss.svg`` -- training/eval loss over steps
* ``accuracy.svg`` -- token accuracy and generation-based dev exact match
* ``curves.html`` -- both charts side by side for quick viewing
Pure standard library: the project lock stays unchanged and the same command
works on the training server.
Usage::
uv run --no-sync python scripts/plot_training_curves.py --run runs/formal/8b/spatial
uv run --no-sync python scripts/plot_training_curves.py --glob 'runs/formal/8b/*' --output-dir docs/curves
"""
from __future__ import annotations
import argparse
import csv
import glob
import json
from pathlib import Path
import sys
from typing import Any, Iterable, Sequence
ROOT = Path(__file__).resolve().parents[1]
WIDTH, HEIGHT = 900, 420
PAD_LEFT, PAD_RIGHT, PAD_TOP, PAD_BOTTOM = 70, 20, 40, 50
COLORS = ["#1f77b4", "#d62728", "#2ca02c", "#9467bd", "#ff7f0e", "#17becf"]
def find_state(run_dir: Path) -> Path | None:
candidates = sorted(run_dir.glob("checkpoint-*/trainer_state.json"), key=lambda p: int(p.parent.name.split("-")[-1]))
if candidates:
return candidates[-1]
fallback = run_dir / "trainer_state.json"
return fallback if fallback.exists() else None
def load_history(run_dir: Path) -> list[dict[str, Any]]:
state_path = find_state(run_dir)
if state_path is None:
return []
state = json.loads(state_path.read_text(encoding="utf-8"))
return list(state.get("log_history", []))
def load_dev_metrics(run_dir: Path) -> list[dict[str, Any]]:
path = run_dir / "dev_metrics.jsonl"
if not path.exists():
return []
rows = []
for line in path.read_text(encoding="utf-8").splitlines():
if line.strip():
rows.append(json.loads(line))
return rows
def collect_series(history: Iterable[dict[str, Any]], dev_rows: Iterable[dict[str, Any]]) -> dict[str, list[tuple[float, float]]]:
series: dict[str, list[tuple[float, float]]] = {
"train_loss": [],
"eval_loss": [],
"token_accuracy": [],
"eval_token_accuracy": [],
"grad_norm": [],
"learning_rate": [],
"dev_exact_match": [],
"dev_parseable_rate": [],
}
for entry in history:
step = entry.get("step")
if step is None:
continue
if "loss" in entry:
series["train_loss"].append((float(step), float(entry["loss"])))
if "eval_loss" in entry:
series["eval_loss"].append((float(step), float(entry["eval_loss"])))
if "mean_token_accuracy" in entry:
series["token_accuracy"].append((float(step), float(entry["mean_token_accuracy"])))
if "eval_mean_token_accuracy" in entry:
series["eval_token_accuracy"].append((float(step), float(entry["eval_mean_token_accuracy"])))
if "grad_norm" in entry:
series["grad_norm"].append((float(step), float(entry["grad_norm"])))
if "learning_rate" in entry:
series["learning_rate"].append((float(step), float(entry["learning_rate"])))
for row in dev_rows:
step = row.get("step")
if step is None:
continue
if "dev_exact_match" in row:
series["dev_exact_match"].append((float(step), float(row["dev_exact_match"])))
if "dev_parseable_rate" in row:
series["dev_parseable_rate"].append((float(step), float(row["dev_parseable_rate"])))
return {key: value for key, value in series.items() if value}
def _scale(values: Sequence[float], lo: float, hi: float, start: float, end: float) -> list[float]:
span = hi - lo or 1.0
return [start + (value - lo) / span * (end - start) for value in values]
def render_chart(title: str, panels: list[tuple[str, list[tuple[float, float]]]], *, y_min: float | None = None, y_max: float | None = None) -> str:
points = [p for _, series in panels for p in series]
if not points:
return f'<svg xmlns="http://www.w3.org/2000/svg" width="{WIDTH}" height="{HEIGHT}"><text x="20" y="30">no data for {title}</text></svg>'
xs = [p[0] for p in points]
ys = [p[1] for p in points]
x_lo, x_hi = min(xs), max(xs)
y_lo = y_min if y_min is not None else min(ys)
y_hi = y_max if y_max is not None else max(ys)
if y_lo == y_hi:
y_lo, y_hi = y_lo - 0.5, y_hi + 0.5
pad = (y_hi - y_lo) * 0.08
y_lo -= pad
y_hi += pad
plot_w = WIDTH - PAD_LEFT - PAD_RIGHT
plot_h = HEIGHT - PAD_TOP - PAD_BOTTOM
def sx(x: float) -> float:
return PAD_LEFT + (x - x_lo) / (x_hi - x_lo or 1.0) * plot_w
def sy(y: float) -> float:
return PAD_TOP + plot_h - (y - y_lo) / (y_hi - y_lo) * plot_h
parts = [
f'<svg xmlns="http://www.w3.org/2000/svg" width="{WIDTH}" height="{HEIGHT}" font-family="Helvetica,Arial,sans-serif">',
f'<rect width="{WIDTH}" height="{HEIGHT}" fill="white"/>',
f'<text x="{PAD_LEFT}" y="24" font-size="16" font-weight="bold">{title}</text>',
]
for i in range(5):
value = y_lo + (y_hi - y_lo) * i / 4
y = sy(value)
parts.append(f'<line x1="{PAD_LEFT}" y1="{y:.1f}" x2="{WIDTH - PAD_RIGHT}" y2="{y:.1f}" stroke="#e5e5e5"/>')
parts.append(f'<text x="{PAD_LEFT - 8}" y="{y + 4:.1f}" font-size="11" text-anchor="end" fill="#555">{value:.3g}</text>')
for i in range(5):
value = x_lo + (x_hi - x_lo) * i / 4
x = sx(value)
parts.append(f'<line x1="{x:.1f}" y1="{PAD_TOP}" x2="{x:.1f}" y2="{PAD_TOP + plot_h}" stroke="#f0f0f0"/>')
parts.append(f'<text x="{x:.1f}" y="{PAD_TOP + plot_h + 18}" font-size="11" text-anchor="middle" fill="#555">{value:.0f}</text>')
parts.append(f'<text x="{PAD_LEFT + plot_w / 2:.0f}" y="{HEIGHT - 8}" font-size="12" text-anchor="middle" fill="#333">training step</text>')
for index, (label, series) in enumerate(panels):
color = COLORS[index % len(COLORS)]
coords = " ".join(f"{sx(x):.1f},{sy(y):.1f}" for x, y in series)
parts.append(f'<polyline fill="none" stroke="{color}" stroke-width="2" points="{coords}"/>')
for x, y in series:
parts.append(f'<circle cx="{sx(x):.1f}" cy="{sy(y):.1f}" r="2.2" fill="{color}"/>')
legend_x = PAD_LEFT + 8 + index * 220
parts.append(f'<rect x="{legend_x}" y="{HEIGHT - 34}" width="12" height="12" fill="{color}"/>')
parts.append(f'<text x="{legend_x + 18}" y="{HEIGHT - 24}" font-size="12" fill="#333">{label}</text>')
parts.append("</svg>")
return "\n".join(parts)
def process_run(run_dir: Path, output_dir: Path) -> dict[str, Any]:
history = load_history(run_dir)
dev_rows = load_dev_metrics(run_dir)
series = collect_series(history, dev_rows)
name = run_dir.name
output_dir.mkdir(parents=True, exist_ok=True)
csv_path = output_dir / f"{name}-history.csv"
steps = sorted({int(step) for values in series.values() for step, _ in values})
with csv_path.open("w", encoding="utf-8", newline="") as handle:
writer = csv.writer(handle)
keys = [key for key in series]
writer.writerow(["step", *keys])
for step in steps:
row: list[Any] = [step]
for key in keys:
value = next((value for s, value in series[key] if int(s) == step), "")
row.append(value)
writer.writerow(row)
loss_panels = [(label, series[key]) for label, key in (("train loss", "train_loss"), ("eval loss", "eval_loss")) if key in series]
acc_panels = [
(label, series[key])
for label, key in (
("token accuracy", "token_accuracy"),
("eval token accuracy", "eval_token_accuracy"),
("dev exact match", "dev_exact_match"),
("dev parseable", "dev_parseable_rate"),
)
if key in series
]
loss_svg = render_chart(f"{name}: loss", loss_panels)
acc_svg = render_chart(f"{name}: accuracy", acc_panels, y_min=0.0, y_max=1.0)
(output_dir / f"{name}-loss.svg").write_text(loss_svg, encoding="utf-8")
(output_dir / f"{name}-accuracy.svg").write_text(acc_svg, encoding="utf-8")
html = (
"<!doctype html><meta charset='utf-8'>"
f"<title>{name} training curves</title>"
"<body style='font-family:Helvetica,Arial,sans-serif;margin:24px'>"
f"<h2>{name}</h2>"
f"<div>{loss_svg}</div><div style='margin-top:24px'>{acc_svg}</div>"
"</body>"
)
(output_dir / f"{name}-curves.html").write_text(html, encoding="utf-8")
final = {key: values[-1][1] for key, values in series.items()}
return {
"run": name,
"steps": len(steps),
"final": final,
"csv": str(csv_path),
"loss_svg": str(output_dir / f"{name}-loss.svg"),
"accuracy_svg": str(output_dir / f"{name}-accuracy.svg"),
"html": str(output_dir / f"{name}-curves.html"),
}
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--run", action="append", default=[], help="run directory (repeatable)")
parser.add_argument("--glob", default=None, help="glob of run directories")
parser.add_argument("--output-dir", default=None, help="where to write curves (default: each run directory)")
parser.add_argument("--summary", default=None, help="optional JSON file collecting all runs")
return parser.parse_args(argv)
def main(argv: Sequence[str] | None = None) -> int:
args = parse_args(argv)
runs = [Path(item) for item in args.run]
if args.glob:
runs += [Path(item) for item in sorted(glob.glob(args.glob)) if Path(item).is_dir()]
runs = [run for run in runs if run.exists()]
if not runs:
print("no run directories matched", file=sys.stderr)
return 2
results = []
for run in runs:
target = Path(args.output_dir) if args.output_dir else run
results.append(process_run(run, target))
print(json.dumps(results[-1], ensure_ascii=False))
if args.summary:
Path(args.summary).write_text(json.dumps(results, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|