| |
| """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 |
|
|
| 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"] |
| |
| 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 = [] |
|
|
| |
| 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") |
|
|
| |
| 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") |
|
|
| |
| 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)") |
|
|
