Upload scripts/mk_p1_plots.py with huggingface_hub

scripts/mk_p1_plots.py

@@ -0,0 +1,145 @@
+#!/usr/bin/env python3
+"""Phase-1 (680M, 4k->16k) plots: training lineage + NIAH eval.
+
+Every number is read from a file. The 1.06B reference points come from the
+summary.csv files this project already published to HF (ref680/logs/...), not
+from anything transcribed by hand.
+
+  fig 1  p1_lineage.png   training loss across the 680M lineage
+  fig 2  p1_results.png   niah_single accuracy vs prompt length, Wilson 95% CI
+"""
+import csv, math, os, re, sys
+import matplotlib; matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+S = {"surface":"#fcfcfb","ink":"#0b0b0b","ink2":"#52514e","grid":"#e4e4e1"}
+C680, CBASE = "#12508f", "#9dc3ea"
+CPAR, CMIX, CPURE = "#eb6834", "#b3243b", "#8e8e88"
+klab = lambda v: (f"{round(v/1024)}k" if v >= 1024 else str(int(v)))
+
+def wilson(k, n, z=1.96):
+    if not n: return (0.0, 0.0, 0.0)
+    p = k/n; d = 1 + z*z/n
+    c = (p + z*z/(2*n))/d
+    h = z*math.sqrt(p*(1-p)/n + z*z/(4*n*n))/d
+    return (100*p, 100*max(0,c-h), 100*min(1,c+h))
+
+def style(ax):
+    ax.set_facecolor(S["surface"])
+    for s in ax.spines.values(): s.set_color(S["grid"])
+    ax.grid(True, color=S["grid"], lw=.7, alpha=.9)
+    ax.tick_params(colors=S["ink2"], labelsize=9)
+
+# ---------------------------------------------------------------- fig 1: loss
+rows = []
+CSVP = "CP_BluScreipt_Training_logs/CP_Training_log12.csv"
+if os.path.exists(CSVP):
+    for r in csv.DictReader(open(CSVP)):
+        try: rows.append((int(r["step"]), float(r["train_loss"]), float(r["val_loss"] or 0)))
+        except Exception: pass
+if not rows:  # fall back to the text log
+    t = open("p1_logs/train.log", errors="ignore").read()
+    rows = [(int(a), float(b), 0.0) for a,b in
+            re.findall(r"^step +(\d+) \| loss: ([0-9.]+)", t, re.M)]
+if not rows:
+    sys.exit("no training rows found -- nothing to plot")
+
+xs  = [r[0] for r in rows]; ys = [r[1] for r in rows]
+d   = {s:l for s,l,_ in rows}
+rx, ry = [], []
+for s in xs:
+    w = [d[j] for j in range(s-49, s+1) if j in d]
+    if len(w) == 50: rx.append(s); ry.append(sum(w)/50)
+# val only where it CHANGES (the trainer carries the last value forward)
+vx, vy, prev = [], [], None
+for s,_,v in rows:
+    if v and v != prev: vx.append(s); vy.append(v); prev = v
+
+fig, ax = plt.subplots(figsize=(11.5,5.6), facecolor=S["surface"])
+ax.plot(xs, ys, color=C680, lw=.6, alpha=.30)
+ax.plot(rx, ry, color=C680, lw=2.0, label="phase 1  4k -> 16k  (50-step mean)")
+if vx: ax.plot(vx, vy, "o--", color=CPAR, ms=4.5, lw=1.1, label="validation loss")
+style(ax)
+ax.set_xlabel("step", color=S["ink2"], fontsize=10)
+ax.set_ylabel("loss (nats)", color=S["ink2"], fontsize=10)
+ax.set_title("680M lineage -- phase 1 context extension, 4k -> 16k",
+             color=S["ink"], fontsize=13, loc="left", pad=12)
+last = ry[-1] if ry else ys[-1]
+ax.text(.995,.965, f"parent varybatch (10.5 tok/param)\nYaRN 4 x 4096 = 16384\n"
+                   f"{len(xs)} steps  |  last 50-step mean {last:.4f}",
+        transform=ax.transAxes, ha="right", va="top", fontsize=8.4, color=S["ink2"],
+        bbox=dict(fc="white", ec=S["grid"]))
+ax.text(.005,.02, "varybatch's own pretrain curve is not plotted: its training log was "
+                  "not published with the checkpoint, so the segment is unrecoverable.",
+        transform=ax.transAxes, fontsize=7.6, color=S["ink2"], style="italic")
+ax.legend(frameon=False, fontsize=9, labelcolor=S["ink2"], loc="upper right",
+          bbox_to_anchor=(.995,.80))
+fig.tight_layout(); fig.savefig("plots_p1/p1_lineage.png", dpi=155)
+print("  wrote plots_p1/p1_lineage.png")
+
+# -------------------------------------------------------------- fig 2: evals
+def read_sum(p, col=1):
+    if not os.path.exists(p): return None
+    for r in csv.reader(open(p)):
+        if r and r[0] == "niah_single":
+            try: return float(r[col])
+            except Exception: return None
+    return None
+
+LENS = [8192, 16360, 32744, 65512]
+REF = {
+ "1.06B ext64k parent (695)":      (CPAR,  "-",  [read_sum(f"ref680/logs/ext64k_len_{L}/summary.csv") for L in LENS]),
+ "1.06B mix64k (1200), best":      (CMIX,  "-",  [read_sum(f"ref680/logs/mix64k_1200_len_{L}/summary.csv") for L in LENS]),
+ "1.06B mixcont (1400), collapsed":(CPURE, "--", [read_sum(f"ref680/logs/mixcont/cliff1400_len_{L}/summary.csv") for L in LENS]),
+}
+P680 = {}
+for name, key in (("680M phase 1 (1200)","p1final"), ("680M varybatch, zero-shot YaRN","base680")):
+    pts = []
+    for L in (8192, 16360):
+        v = None
+        for c in (1,2,3):
+            v = read_sum(f"p1_eval_len{L}/summary.csv", c)
+            if v is not None: break
+        # summary.csv columns are per-model; find this model's column by header
+        hp = f"p1_eval_len{L}/summary.csv"
+        if os.path.exists(hp):
+            hdr = next(csv.reader(open(hp)), [])
+            if key in hdr:
+                v = read_sum(hp, hdr.index(key))
+        pts.append(v)
+    P680[name] = pts
+
+fig, ax = plt.subplots(figsize=(11.0,6.0), facecolor=S["surface"])
+X = list(range(len(LENS)))
+for name,(col,ls,vals) in REF.items():
+    xs_ = [i for i,v in zip(X,vals) if v is not None]
+    ys_ = [v for v in vals if v is not None]
+    if not ys_: continue
+    err = [[v-wilson(round(v),100)[1] for v in ys_],[wilson(round(v),100)[2]-v for v in ys_]]
+    ax.errorbar(xs_, ys_, yerr=err, color=col, ls=ls, lw=1.6, marker="o", ms=5,
+                capsize=3, alpha=.85, label=name)
+for name,(col,mk) in zip(P680, ((C680,"s"),(CBASE,"^"))):
+    vals = P680[name]
+    xs_ = [i for i,v in zip(X,vals) if v is not None]
+    ys_ = [v for v in vals if v is not None]
+    if not ys_:
+        continue
+    err = [[v-wilson(round(v),100)[1] for v in ys_],[wilson(round(v),100)[2]-v for v in ys_]]
+    ax.errorbar(xs_, ys_, yerr=err, color=col, lw=2.4, marker=mk, ms=8, capsize=3, label=name)
+ax.axvspan(1.5, 3.4, color="#f0f0ee", zorder=0)
+ax.text(2.45,.5,"beyond phase 1's 16k window\n(680M not evaluated here)",
+        ha="center", fontsize=8.2, color=S["ink2"], style="italic")
+style(ax)
+ax.set_xticks(X); ax.set_xticklabels([klab(L) for L in LENS])
+ax.set_xlim(-.3, 3.4); ax.set_ylim(0, 100)
+ax.set_xlabel("prompt length (tokens)", color=S["ink2"], fontsize=10)
+ax.set_ylabel("niah_single accuracy (%)", color=S["ink2"], fontsize=10)
+ax.set_title("Needle-in-a-haystack retrieval -- 680M phase 1 vs the 1.06B lineage",
+             color=S["ink"], fontsize=13, loc="left", pad=12)
+ax.text(.005,-.155,"N=100 per point, greedy, identical haystack (b2 val_shard_00000) and prompts "
+                   "across every model. Bars are Wilson 95% CI (+/-9-13pp at N=100): "
+                   "differences inside a bar are not differences.",
+        transform=ax.transAxes, fontsize=7.8, color=S["ink2"])
+ax.legend(frameon=False, fontsize=9, labelcolor=S["ink2"], loc="upper right")
+fig.tight_layout(); fig.savefig("plots_p1/p1_results.png", dpi=155, bbox_inches="tight")
+print("  wrote plots_p1/p1_results.png")