File size: 8,670 Bytes
a2b8b4a a903ec8 a2b8b4a a903ec8 a2b8b4a a903ec8 a2b8b4a | 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 | """
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"} # fixed categorical order
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()
# Greedy label placement in pixel space: try nearby offsets, keep the first that overlaps nothing.
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: # markers are obstacles too
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}/")
|