| """ |
| Make plots from a run's results (JSON) — used in the README. |
| |
| python plot_results.py runs/colab-l4-2026-10-08 |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import os |
| import sys |
|
|
| import matplotlib |
|
|
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import numpy as np |
|
|
| SURFACE = "#fcfcfb" |
| INK = "#0b0b0b" |
| INK_2 = "#52514e" |
| GRID = "#e4e3df" |
| SERIES = {"base": "#2a78d6", "sft": "#eb6834", "dpo": "#1baf7a"} |
|
|
|
|
| def _style(ax, title: str): |
| ax.set_facecolor(SURFACE) |
| ax.set_title(title, loc="left", color=INK, fontsize=12, fontweight="bold", pad=12) |
| for s in ("top", "right"): |
| ax.spines[s].set_visible(False) |
| for s in ("left", "bottom"): |
| ax.spines[s].set_color(GRID) |
| ax.tick_params(colors=INK_2, labelsize=9) |
| ax.grid(color=GRID, linewidth=0.8, which="major") |
| ax.set_axisbelow(True) |
|
|
|
|
| def plot_pretrain(run: str): |
| d = json.load(open(os.path.join(run, "pretrain.json"))) |
| tok = np.array([h["tokens"] for h in d["history"]]) / 1e6 |
| ppl = np.array([h["val_ppl"] for h in d["history"]]) |
| fig, ax = plt.subplots(figsize=(7, 4), facecolor=SURFACE) |
| _style(ax, "Pretraining — validation perplexity") |
| ax.plot(tok, ppl, color=SERIES["base"], linewidth=2) |
| ax.scatter(tok, ppl, s=28, color=SERIES["base"], edgecolor=SURFACE, linewidth=2, zorder=3) |
| ax.annotate(f"{ppl[-1]:.2f}", (tok[-1], ppl[-1]), xytext=(6, 6), textcoords="offset points", color=INK, fontsize=9) |
| ax.annotate(f"{ppl[0]:.2f}", (tok[0], ppl[0]), xytext=(6, 0), textcoords="offset points", color=INK, fontsize=9) |
| ax.set_xlabel("Tokens seen (millions)", color=INK_2) |
| ax.set_ylabel("Perplexity", color=INK_2) |
| fig.tight_layout() |
| fig.savefig(os.path.join(run, "pretrain.png"), dpi=160, facecolor=SURFACE) |
|
|
|
|
| def plot_scaling(run: str): |
| d = json.load(open(os.path.join(run, "scaling.json"))) |
| N = np.array([r["N"] for r in d["runs"]], dtype=float) |
| L = np.array([r["loss"] for r in d["runs"]]) |
| E, A, a = d["fit"]["E"], d["fit"]["A"], d["fit"]["alpha"] |
| ex = d["extrapolation"] |
| fig, ax = plt.subplots(figsize=(7, 4), facecolor=SURFACE) |
| _style(ax, "Scaling law — loss vs model size (D = 20·N)") |
| xs = np.logspace(np.log10(N.min() * 0.8), np.log10(ex["N_main"] * 2), 200) |
| ax.plot(xs, E + A * xs ** -a, color=INK_2, linewidth=1.5, linestyle="--", |
| label=f"fit: L = {E:.2f} + {A:.0f}·N^-{a:.2f}") |
| ax.scatter(N, L, s=48, color=SERIES["base"], edgecolor=SURFACE, linewidth=2, zorder=3, label="scaling runs") |
| ax.scatter([ex["N_main"]], [ex["actual_loss"]], s=64, marker="D", color=SERIES["dpo"], edgecolor=SURFACE, |
| linewidth=2, zorder=4, label="main model (actual)") |
| ax.annotate(f"predicted {ex['predicted_loss']:.3f}\nactual {ex['actual_loss']:.3f}", |
| (ex["N_main"], ex["actual_loss"]), xytext=(-20, 70), textcoords="offset points", ha="center", |
| color=INK, fontsize=9, arrowprops=dict(arrowstyle="-", color=INK_2, lw=0.8)) |
| ax.set_xscale("log") |
| ax.set_xlabel("Non-embedding params (N)", color=INK_2) |
| ax.set_ylabel("Validation loss", color=INK_2) |
| ax.legend(frameon=False, fontsize=9, labelcolor=INK_2) |
| fig.tight_layout() |
| fig.savefig(os.path.join(run, "scaling.png"), dpi=160, facecolor=SURFACE) |
|
|
|
|
| def plot_eval(run: str): |
| ev = json.load(open(os.path.join(run, "eval.json"))) |
| m = ev["metrics"] |
| models = ["base", "sft", "dpo"] |
| metrics = [("keyword_hit_rate", "Keyword hit"), ("all_keywords_rate", "All 3 words"), |
| ("finish_rate", "Finished (EOS)")] |
| fig, (ax, ax2) = plt.subplots(1, 2, figsize=(10, 4), facecolor=SURFACE, gridspec_kw={"width_ratios": [3, 1.3]}) |
| _style(ax, f"Instruction following (held-out prompts, n={ev.get('n_prompts', 200)})") |
| x = np.arange(len(metrics)) |
| w = 0.26 |
| for i, name in enumerate(models): |
| vals = [m[name][k] * 100 for k, _ in metrics] |
| bars = ax.bar(x + (i - 1) * w, vals, w * 0.92, color=SERIES[name], label=name) |
| for b, v in zip(bars, vals): |
| ax.text(b.get_x() + b.get_width() / 2, v + 1.5, (f"{v:.0f}" if abs(v - round(v)) < 0.05 else f"{v:.1f}"), ha="center", fontsize=8, color=INK_2) |
| ax.set_xticks(x, [lbl for _, lbl in metrics]) |
| ax.set_ylim(0, 110) |
| ax.set_ylabel("%", color=INK_2) |
| ax.legend(frameon=False, fontsize=9, labelcolor=INK_2, ncol=3, loc="upper left") |
|
|
| _style(ax2, "Val perplexity ↓") |
| ppl = [m[n]["val_ppl"] for n in models] |
| bars = ax2.bar(models, ppl, 0.6, color=[SERIES[n] for n in models]) |
| for b, v in zip(bars, ppl): |
| ax2.text(b.get_x() + b.get_width() / 2, v + 0.08, f"{v:.2f}", ha="center", fontsize=8, color=INK_2) |
| ax2.set_ylim(0, max(ppl) * 1.2) |
| fig.tight_layout() |
| fig.savefig(os.path.join(run, "eval.png"), dpi=160, facecolor=SURFACE) |
|
|
|
|
| def plot_sweep(run: str): |
| path = os.path.join(run, "dpo_sweep.json") |
| if not os.path.exists(path): |
| return |
| d = json.load(open(path)) |
| fig, ax = plt.subplots(figsize=(7.5, 4.4), facecolor=SURFACE) |
| _style(ax, "DPO sweep — instruction following vs fluency") |
| colors = {"v1": SERIES["base"], "v2": SERIES["sft"]} |
| sft = d["variants"]["sft (reference)"]["metrics"] |
| tol = d["selection_rule"]["fluency_ppl_max_ratio_vs_sft"] |
| ax.axvline(sft["fluency_ppl"] * tol, color=INK_2, linewidth=1, linestyle=":") |
| ax.annotate(f"guardrail ({tol:g}× SFT)", (sft["fluency_ppl"] * tol, 0), xytext=(4, 4), |
| textcoords="offset points", fontsize=8, color=INK_2) |
| ax.scatter([sft["fluency_ppl"]], [sft["all_keywords_rate"] * 100], s=60, color=INK_2, marker="s", |
| edgecolor=SURFACE, linewidth=2, zorder=3) |
| ax.annotate("SFT", (sft["fluency_ppl"], sft["all_keywords_rate"] * 100), xytext=(6, -12), |
| textcoords="offset points", fontsize=8, color=INK_2) |
| pts = [(n, v) for n, v in d["variants"].items() if not n.startswith("sft")] |
| for name, v in pts: |
| m = v["metrics"] |
| star = name == d["best"] |
| ax.scatter([m["fluency_ppl"]], [m["all_keywords_rate"] * 100], s=260 if star else 50, |
| marker="*" if star else "o", color=colors[v["reward"]], edgecolor=SURFACE, |
| linewidth=1.5 if star else 2, zorder=4) |
| ax.margins(x=0.12, y=0.15) |
| fig.canvas.draw() |
| |
| cands = [(8, 6), (8, -14), (-8, 6), (-8, -14), (30, 18), (-30, 18), (30, -30), (-30, -30), (55, 0), (-55, 0)] |
| boxes = [] |
| for _, v in pts: |
| mx, my = ax.transData.transform((v["metrics"]["fluency_ppl"], v["metrics"]["all_keywords_rate"] * 100)) |
| r = 7 * fig.dpi / 72 |
| boxes.append((mx - r, my - r, mx + r, my + r)) |
| for name, v in sorted(pts, key=lambda t: -t[1]["metrics"]["all_keywords_rate"]): |
| m = v["metrics"] |
| label = name + (" (best)" if name == d["best"] else "") |
| px, py = ax.transData.transform((m["fluency_ppl"], m["all_keywords_rate"] * 100)) |
| w, h = 5.2 * len(label) * fig.dpi / 100, 11 * fig.dpi / 100 |
| for dx, dy in cands: |
| dxp, dyp = dx * fig.dpi / 72, dy * fig.dpi / 72 |
| x0 = px + dxp if dx > 0 else px + dxp - w |
| box = (x0, py + dyp, x0 + w, py + dyp + h) |
| if not any(box[0] < b[2] and b[0] < box[2] and box[1] < b[3] and b[1] < box[3] for b in boxes): |
| break |
| boxes.append(box) |
| far = (dx, dy) not in cands[:4] |
| ax.annotate(label, (m["fluency_ppl"], m["all_keywords_rate"] * 100), xytext=(dx, dy), |
| textcoords="offset points", fontsize=8, color=INK, ha="left" if dx > 0 else "right", |
| arrowprops=dict(arrowstyle="-", color=INK_2, lw=0.6) if far else None) |
| for r, c in colors.items(): |
| ax.scatter([], [], s=40, color=c, label=f"reward {r}") |
| ax.set_xlabel("Fluency ppl under base model (lower = more natural)", color=INK_2) |
| ax.set_ylabel("All 3 words used (%)", color=INK_2) |
| ax.legend(frameon=False, fontsize=9, labelcolor=INK_2, loc="lower right") |
| fig.tight_layout() |
| fig.savefig(os.path.join(run, "dpo_sweep.png"), dpi=160, facecolor=SURFACE) |
|
|
| if __name__ == "__main__": |
| run_dir = sys.argv[1] if len(sys.argv) > 1 else "results" |
| for fn, needs in ((plot_pretrain, "pretrain.json"), (plot_scaling, "scaling.json"), |
| (plot_eval, "eval.json"), (plot_sweep, "dpo_sweep.json")): |
| if os.path.exists(os.path.join(run_dir, needs)): |
| fn(run_dir) |
| else: |
| print(f"skip {fn.__name__}: {needs} not found") |
| print(f"Plots saved to {run_dir}/") |
|
|
