#!/usr/bin/env python """Draw the figures for Labs 2, 4 and 5 from the finished runs. python scripts/summarize.py results/runs # first: writes results/summary.json uv run --with matplotlib python scripts/plot_runs.py # any machine, a few seconds Reads results/summary.json, the answer lengths in results/runs/**/merged.jsonl and Lab 7's results/lab7/side_effects.json, and writes results/figures/lab2_*.png, lab4_*.png and lab5_*.png. Runs that have not finished yet are left out. """ import json import os import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt # noqa: E402 ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) RES = os.path.join(ROOT, "results") FIG = os.path.join(RES, "figures") os.makedirs(FIG, exist_ok=True) S = {r["name"]: r["summary"] for r in json.load(open(os.path.join(RES, "summary.json")))} SIDE = json.load(open(os.path.join(RES, "lab7", "side_effects.json")))["models"] # the blog's think35 numbers: accuracy %, cut at the cap %, mean tokens PAPER = {"math500": {"base": (7.9, 93.0, 30868), "p50c": (48.1, 33.05, 22037)}, "gsm8k": {"base": (13.84, 87.5, 29289), "p50c": (67.1, 7.24, 11993)}, "ifeval": {"base": (21.0, 74.9, 25296), "p50c": (62.3, 2.2, 4241)}} GRAY, GREEN, RED, PURPLE, BLUE = "#8c959f", "#1a7f37", "#cf222e", "#8250df", "#0969da" METRICS = [("acc", "accuracy %", 100), ("cut_share", "answers cut at the token cap %", 100), ("mean_tokens", "mean answer length (thousand tokens)", 1e-3)] def val(s, k, scale): return None if s is None or s.get(k) is None else s[k] * scale def err(s, k): return 100 * s["acc_se"] if k == "acc" and s and s.get("acc_se") is not None else 0 def lengths(name): p = os.path.join(RES, "runs", name, "merged.jsonl") if not os.path.exists(p): return None return [s["tokens"] for line in open(p) for s in json.loads(line)["samples"]] made = [] # Lab 2: the paper's benchmarks, ours next to the blog's bench = [b for b in PAPER if f"bench/base/{b}" in S or f"bench/p50c/{b}" in S] if bench: fig, ax = plt.subplots(1, 3, figsize=(15, 4.2)) w = 0.2 for i, (k, lab, scale) in enumerate(METRICS): for j, b in enumerate(bench): for o, (model, src, col, hatch) in enumerate([("base", "ours", GRAY, None), ("base", "paper", GRAY, "//"), ("p50c", "ours", GREEN, None), ("p50c", "paper", GREEN, "//")]): if src == "ours": s = S.get(f"bench/{model}/{b}") v, e = val(s, k, scale), err(s, k) else: v, e = PAPER[b][model][i] * (1e-3 if k == "mean_tokens" else 1), 0 if v is None: continue ax[i].bar(j + (o - 1.5) * w, v, w, yerr=e or None, color=col, hatch=hatch, edgecolor="white", label=f"{'unedited' if model == 'base' else 'edited'}, {src}" if j == 0 else None) ax[i].set_xticks(range(len(bench)), [b.upper().replace("MATH500", "MATH-500") for b in bench]) ax[i].set_title(lab) ax[0].legend(fontsize=8) fig.suptitle("The paper's three benchmarks: our runs (solid) next to the blog's numbers (hatched)") fig.tight_layout() fig.savefig(os.path.join(FIG, "lab2_benchmarks.png"), dpi=130) made.append("lab2_benchmarks.png") pairs = [(n, c) for n, c in (("dev/base", GRAY), ("dev/p50c", GREEN)) if lengths(n)] if pairs: fig, ax = plt.subplots(figsize=(10, 3.8)) for n, c in pairs: ax.hist(lengths(n), bins=64, range=(0, 32768), color=c, alpha=0.6, label=f"{'unedited' if 'base' in n else 'edited'} ({n})") ax.set_yscale("log") ax.set_xlabel("answer length in tokens (the cap is 32,768)") ax.set_ylabel("answers") ax.set_title("Where the answers end: the tall bar at the right edge is answers that never stopped") ax.legend() fig.tight_layout() fig.savefig(os.path.join(FIG, "lab2_lengths.png"), dpi=130) made.append("lab2_lengths.png") # Lab 4: turn the knob, then the controls DOSE = [(-4, "a-4"), (0, "base"), (1, "a1"), (2, "a2"), (4, "a4"), (6, "p50c"), (8, "a8"), (12, "a12")] dose = [(x, S[f"dev/{n}"], SIDE.get(n)) for x, n in DOSE if f"dev/{n}" in S] if len(dose) >= 3: fig, ax = plt.subplots(1, 4, figsize=(17, 4)) xs = [d[0] for d in dose] for i, (k, lab, scale) in enumerate(METRICS): ax[i].errorbar(xs, [val(d[1], k, scale) for d in dose], yerr=[err(d[1], k) for d in dose], fmt="o-", color=BLUE, capsize=3, label="released u and 64 neurons") for n, dx, mk, col, lab2 in (("rebuilt", 0.35, "s", RED, "your Lab 3 rebuild"), ("capA-a6", 0.7, "^", PURPLE, "capstone, recipe A"), ("capB-a6", 1.05, "v", PURPLE, "capstone, recipe B")): if f"dev/{n}" in S: ax[i].errorbar([6 + dx], [val(S[f"dev/{n}"], k, scale)], yerr=[err(S[f"dev/{n}"], k)], fmt=mk, color=col, mfc="none", capsize=3, label=f"{lab2} (alpha 6, drawn to the right)") ax[i].set_title(lab) sx = [d[0] for d in dose if d[2]] ax[3].plot(sx, [d[2]["ppl_change_pct"] for d in dose if d[2]], "o-", color=RED) if SIDE.get("rebuilt"): ax[3].plot([6.35], [SIDE["rebuilt"]["ppl_change_pct"]], "s", color=RED, mfc="none") ax[3].axhline(0, color="black", lw=0.5) ax[3].set_title("side effect: perplexity change % (Lab 7)") for a_ in ax: a_.set_xlabel("alpha (0 = no edit, 1 = erase, 2 = mirror, 6 = released)") a_.set_xticks([-4, 0, 1, 2, 4, 6, 8, 12]) h, lb = ax[0].get_legend_handles_labels() fig.legend(h, lb, loc="lower center", ncol=len(lb), fontsize=9, frameon=False) fig.suptitle("Turning the knob: same direction and neurons, different alpha (256 unseen practice problems)") fig.tight_layout(rect=(0, 0.07, 1, 1)) fig.savefig(os.path.join(FIG, "lab4_dose.png"), dpi=130) made.append("lab4_dose.png") CTRL = [("base", "no edit", GRAY), ("p50c", "released edit", GREEN), ("rebuilt", "your Lab 3 rebuild", GREEN), ("capA-a6", "capstone A: own loops, just before", PURPLE), ("capB-a6", "capstone B: own loops, well before", PURPLE), ("l15only", "layer 15 only", BLUE), ("l19only", "layer 19 only", BLUE), ("all1", "erase u from all 3,584 neurons", BLUE), ("all2", "mirror u in all 3,584 neurons", BLUE), ("randdir", "control: random direction", RED), ("randcols", "control: 64 random neurons", RED)] ctrl = [(lab, S[f"dev/{n}"], c) for n, lab, c in CTRL if f"dev/{n}" in S] if len(ctrl) >= 3: fig, ax = plt.subplots(1, 3, figsize=(17, 0.45 * len(ctrl) + 1.5), sharey=True) ys = list(range(len(ctrl)))[::-1] for i, (k, lab, scale) in enumerate(METRICS): ax[i].barh(ys, [val(s, k, scale) for _, s, _ in ctrl], xerr=[err(s, k) for _, s, _ in ctrl] if k == "acc" else None, color=[c for *_, c in ctrl], capsize=2) ax[i].set_title(lab) ax[0].set_yticks(ys, [lab for lab, *_ in ctrl]) fig.suptitle("Which parts of the edit matter? Same 256 unseen practice problems, 2 answers each") fig.tight_layout() fig.savefig(os.path.join(FIG, "lab4_controls.png"), dpi=130) made.append("lab4_controls.png") # Lab 5: decoding fixes on the unedited model vs the weight edit DEC = [("base", "unedited, paper settings", GRAY), ("base+temp1.0", "temperature 1.0 (from 0.6)", BLUE), ("base+pp0", "presence penalty 0 (from 1.5)", BLUE), ("base+pp2", "presence penalty 2.0", BLUE), ("base+rep1.1", "repetition penalty 1.1", BLUE), ("base+max4k", "cap at 4,096 tokens", BLUE), ("base+nothink", "thinking off (cap 8,192)", BLUE), ("p50c", "edited, paper settings", GREEN), ("p50c+nothink", "edited, thinking off", GREEN)] dec = [(lab, S[f"dev/{n}"], c) for n, lab, c in DEC if f"dev/{n}" in S] if len(dec) >= 3: fig, ax = plt.subplots(1, 3, figsize=(17, 0.45 * len(dec) + 1.5), sharey=True) ys = list(range(len(dec)))[::-1] for i, (k, lab, scale) in enumerate(METRICS): ax[i].barh(ys, [val(s, k, scale) for _, s, _ in dec], xerr=[err(s, k) for _, s, _ in dec] if k == "acc" else None, color=[c for *_, c in dec], capsize=2) ax[i].set_title(lab) ax[0].set_yticks(ys, [lab for lab, *_ in dec]) fig.suptitle("Can decoding settings do what the weight edit does? Same 256 unseen practice problems") fig.tight_layout() fig.savefig(os.path.join(FIG, "lab5_decoding.png"), dpi=130) made.append("lab5_decoding.png") print("wrote", ", ".join(made) if made else "nothing yet (no finished runs)")