Upload scripts/mk_mix64k_plots.py with huggingface_hub

scripts/mk_mix64k_plots.py

@@ -0,0 +1,209 @@
+#!/usr/bin/env python3
+"""
+mk_mix64k_plots.py -- plots for the mixed-corpus 64k run, built from whatever
+eval cells exist so far. Safe to run repeatedly: it regenerates as more land.
+
+  python3 mk_mix64k_plots.py [--tag 1200] [--out mix64k_results]
+
+THE COMPARISON THIS RUN EXISTS TO MAKE
+  parent  ext64k step 695  67/65/52/29   <- the starting point of BOTH runs
+  pure-b5 cont64k 1651     45/17/ 2/ 1   <- 100% b5: collapsed, cliff at step 1199
+  mix     this run         ?             <- 50% b5 / 30% b3 / 20% b1, same parent,
+                                            same schedule, same window
+Only the data differs. The loss panel overlays both runs on the same axis, so a
+cliff in the mix run would sit exactly where the pure-b5 one fell.
+
+Numbers come from summary.csv / depth_summary.csv, never from prose.
+"""
+import argparse, csv, glob, json, math, os, re, sys
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from matplotlib.ticker import FuncFormatter
+
+S = {"surface":"#fcfcfb","ink":"#0b0b0b","ink2":"#52514e","grid":"#e4e4e1"}
+C = {"parent":"#eb6834","pure":"#b3243b","mix":"#12508f","warn":"#eda100","ok":"#7a9a3a"}
+LENS = [8192, 16360, 32744, 65512]
+
+def wilson(p, n=100):
+    if not n: return float("nan")
+    z=1.96; ph=p/100.0
+    return 100.0*z*math.sqrt(ph*(1-ph)/n + z*z/(4*n*n))/(1+z*z/n)
+
+def read_summary(path):
+    try:
+        for r in csv.reader(open(path)):
+            if r and r[0].strip()=="niah_single": return float(r[1])
+    except Exception: pass
+    return None
+
+def read_depth(path):
+    out=[]
+    try:
+        for r in csv.DictReader(open(path)):
+            out.append({"depth":float(r["depth"]), "hits":int(r["hits"]), "n":int(r["n"])})
+    except Exception: return []
+    return sorted(out, key=lambda d:d["depth"])
+
+def loss_rows(path):
+    try: t=open(path,errors="ignore").read()
+    except Exception: return []
+    return [(int(a),float(b)) for a,b in re.findall(r"^step +(\d+) \| loss: ([0-9.]+)", t, re.M)]
+
+def roll(rows, k=50):
+    xs=[];ys=[]
+    d=dict(rows)
+    for s,_ in rows:
+        w=[d[j] for j in range(s-k+1, s+1) if j in d]
+        if len(w)==k: xs.append(s); ys.append(sum(w)/k)
+    return xs, ys
+
+def style(ax, xl, yl, title=None):
+    ax.set_facecolor(S["surface"]); ax.grid(True,color=S["grid"],lw=.8,zorder=0); ax.set_axisbelow(True)
+    for k,sp in ax.spines.items(): sp.set_visible(k in ("left","bottom")); sp.set_color(S["grid"])
+    ax.tick_params(colors=S["ink2"],labelsize=9)
+    ax.set_xlabel(xl,color=S["ink2"],fontsize=10); ax.set_ylabel(yl,color=S["ink2"],fontsize=10)
+    if title: ax.set_title(title,color=S["ink"],fontsize=12,loc="left",pad=12)
+
+klab = lambda v,_=None: (f"{round(v/1024)}k" if v>=1024 else str(int(v)))
+
+def main():
+    ap=argparse.ArgumentParser()
+    ap.add_argument("--tag", default="1200")
+    ap.add_argument("--out", default="mix64k_results")
+    a=ap.parse_args()
+    R="mix64k_ref"
+
+    parent={L:read_summary(f"{R}/parent_{L}_summary.csv") for L in LENS}
+    pure  ={L:read_summary(f"{R}/purec5_{L}_summary.csv") for L in LENS}
+    mix   ={L:read_summary(f"mix64k_{a.tag}_len_{L}/summary.csv") for L in LENS}
+    parent={k:v for k,v in parent.items() if v is not None}
+    pure  ={k:v for k,v in pure.items()   if v is not None}
+    mix   ={k:v for k,v in mix.items()    if v is not None}
+    have=sorted(mix)
+    print(f"  parent {parent}\n  pure-b5 {pure}\n  mix@{a.tag} {mix}")
+    if not mix:
+        print("  no mix cells yet -- nothing to plot"); return 1
+
+    fig = plt.figure(figsize=(11, 17), facecolor=S["surface"])
+    gs = fig.add_gridspec(4, 1, hspace=0.34, height_ratios=[1,1,1.25,1])
+
+    # ---- 1. the three-way comparison -------------------------------------
+    ax=fig.add_subplot(gs[0])
+    for lab,d,col,mk in (("parent: ext64k step 695 (both runs start here)",parent,C["parent"],"D"),
+                         ("100% b5, step 1651 -- COLLAPSED",pure,C["pure"],"o"),
+                         (f"MIX 50/30/20, step {a.tag}",mix,C["mix"],"s")):
+        if not d: continue
+        xs=sorted(d); ys=[d[x] for x in xs]
+        ax.errorbar(xs,ys,yerr=[wilson(y) for y in ys],color=col,lw=2.4,capsize=3,
+                    marker=mk,ms=7,zorder=4,label=lab)
+    ax.set_xscale("log",base=2); ax.set_xticks(LENS)
+    ax.xaxis.set_major_formatter(FuncFormatter(klab))
+    ax.set_ylim(-4,100); ax.axhline(0,color=S["grid"],lw=1)
+    style(ax,"evaluation context length (tokens)","niah_single accuracy (%)",
+          "Same parent, same schedule, same window -- only the DATA differs")
+    ax.legend(frameon=False,fontsize=8.5,labelcolor=S["ink2"],loc="upper right")
+
+    # ---- 2. loss overlay: does the cliff reappear? ------------------------
+    ax=fig.add_subplot(gs[1])
+    pr=loss_rows(f"{R}/cont64k_purec5.log"); mr=loss_rows("mix64k_logs/train.log")
+    for rows,col,lab in ((pr,C["pure"],"100% b5 (cliffed at step 1199)"),
+                         (mr,C["mix"],"MIX 50/30/20 (this run)")):
+        if not rows: continue
+        ax.plot([s for s,_ in rows],[l for _,l in rows],color=col,lw=.5,alpha=.28,zorder=2)
+        xs,ys=roll(rows)
+        if xs: ax.plot(xs,ys,color=col,lw=2.1,zorder=4,label=lab+"  (50-step mean)")
+    if pr:
+        base=[l for s,l in pr if 200<=s<=700]
+        if len(base)>=100:
+            b=sum(base)/len(base)
+            ax.axhline(b-0.20,color=C["warn"],lw=1.2,ls="--",zorder=3)
+            ax.text(20,b-0.17,f"cliff threshold (baseline {b:.2f} - 0.20)",fontsize=8,color="#8a6400")
+    ax.axvline(1199,color=C["pure"],lw=1,ls=":",zorder=1)
+    ax.text(1205,1.35,"pure-b5 cliff\nonset (1199)",fontsize=8,color=C["pure"])
+    ax.set_ylim(1.0,3.4)
+    style(ax,"step (from the ext64k-695 parent)","training loss (nats)",
+          "A FALLING loss here is the failure signature, not success")
+    ax.legend(frameon=False,fontsize=8.5,labelcolor=S["ink2"],loc="lower left")
+
+    # ---- 3. the whole lineage, forking at ext64k-695 -----------------------
+    ax=fig.add_subplot(gs[2])
+    import json as _j
+    try: stages=_j.load(open(f"{R}/lineage/stages.json"))
+    except Exception: stages={}
+    ORDER=[("base pretrain",4096,"#8e8e88"),("l2 16k bucket",16384,"#9dc3ea"),
+           ("ext32k",32768,"#4a8fd4"),("b4 full epoch",32768,"#12508f"),
+           ("ext64k",65536,C["parent"])]
+    off=0; marks=[]
+    for name,ctx,col in ORDER:
+        st=stages.get(name)
+        if not st: continue
+        rows=st["rows"]; base=rows[0][0]
+        xs=[off+(r[0]-base) for r in rows]; ys=[r[1] for r in rows]
+        ax.plot(xs,ys,color=col,lw=.9,alpha=.85,zorder=3,label=f"{name} ({klab(ctx)})")
+        marks.append((off,xs[-1],name,col)); off=xs[-1]
+    fork=off   # both 64k runs start here, from ext64k step 695
+    for rows,col,lab,lw in ((pr,C["pure"],"100% b5 -> COLLAPSED",1.3),
+                            (mr,C["mix"],"MIX 50/30/20 -> this run",1.6)):
+        if not rows: continue
+        xs=[fork+s_ for s_,_ in rows]; ys=[l for _,l in rows]
+        ax.plot(xs,ys,color=col,lw=.5,alpha=.25,zorder=2)
+        rx,ry=roll(rows)
+        if rx: ax.plot([fork+x for x in rx],ry,color=col,lw=lw,zorder=5,label=lab)
+    ax.axvline(fork,color=S["ink2"],lw=1.1,ls="--",alpha=.7,zorder=4)
+    ax.annotate("both 64k runs fork here\n(ext64k step 695)",xy=(fork,9),xytext=(fork*0.52,9.2),
+                fontsize=8.5,color=S["ink2"],arrowprops=dict(arrowstyle="->",color=S["ink2"],lw=.9))
+    ax.set_yscale("log"); ax.set_ylim(0.9,13)
+    ax.set_yticks([1,2,3,5,10]); ax.set_yticklabels(["1","2","3","5","10"])
+    style(ax,"cumulative training step across the lineage","loss (nats, log)",
+          "Every stage this model has been through -- and where the two 64k runs diverge")
+    ax.legend(frameon=False,fontsize=7.8,labelcolor=S["ink2"],loc="lower left",ncol=2)
+
+    # ---- 4. depth profiles ------------------------------------------------
+    ax=fig.add_subplot(gs[3])
+    shades={8192:"#9dc3ea",16360:"#4a8fd4",32744:"#12508f",65512:"#08306b"}
+    for L in have:
+        rows=read_depth(f"mix64k_{a.tag}_len_{L}/depth_summary.csv")
+        if not rows: continue
+        xs=[r["depth"] for r in rows]; ys=[100.0*r["hits"]/r["n"] for r in rows]
+        ax.errorbar(xs,ys,yerr=[wilson(y,r["n"]) for y,r in zip(ys,rows)],color=shades.get(L,"#12508f"),
+                    lw=2.1,capsize=3,marker="s",ms=6,zorder=4,label=f"MIX @ {L:,}")
+    pd_=read_depth(f"{R}/purec5_32744_depth_summary.csv")
+    if pd_:
+        ax.errorbar([r["depth"] for r in pd_],[100.0*r["hits"]/r["n"] for r in pd_],
+                    color=C["pure"],lw=2.0,ls="--",marker="o",ms=5,zorder=3,
+                    label="100% b5 @ 32,744 (collapsed: 0/0/0/0/10)")
+    ax.set_ylim(-4,104)
+    style(ax,"needle depth  (0 = window start, 1 = next to the question)","accuracy (%)",
+          "Depth profile: a collapsed model retrieves only from the very end")
+    ax.legend(frameon=False,fontsize=8.5,labelcolor=S["ink2"],loc="upper left")
+
+    png=f"{a.out}_{a.tag}.png"
+    fig.savefig(png,dpi=130,facecolor=S["surface"],bbox_inches="tight")
+    print(f"  wrote {png} ({os.path.getsize(png)//1024} KB)")
+
+    # ---- a self-contained HTML with the numbers ---------------------------
+    import base64
+    b64=base64.b64encode(open(png,"rb").read()).decode()
+    rows_html=""
+    for L in LENS:
+        p=parent.get(L); u=pure.get(L); m=mix.get(L)
+        f=lambda v: "-" if v is None else f"{v:.1f}"
+        d="" if (m is None or p is None) else f"{m-p:+.1f}"
+        rows_html+=f"<tr><td>{L:,}</td><td>{f(p)}</td><td>{f(u)}</td><td><b>{f(m)}</b></td><td>{d}</td></tr>"
+    html=f"""<!doctype html><meta charset=utf-8><title>Mixed-corpus 64k run, step {a.tag}</title>
+<style>body{{max-width:1000px;margin:2rem auto;padding:0 1rem;font:15px/1.6 -apple-system,Segoe UI,Roboto,sans-serif;color:#111;background:#fcfcfb}}
+table{{border-collapse:collapse;margin:1rem 0;font-size:14px}}th,td{{border:1px solid #e4e4e1;padding:.4rem .7rem;text-align:left}}th{{background:#f3f3f1}}
+img{{max-width:100%;border:1px solid #e4e4e1}}code{{background:#f3f3f1;padding:.1rem .3rem;border-radius:3px}}</style>
+<h1>Mixed-corpus 64k extension, step {a.tag}</h1>
+<p>Parent <code>ext64k step 695</code>; 50% b5 / 30% b3 / 20% b1; plateau LR 1.8e-4; cp=4 dp=2 on 8x H200.
+The 100% b5 run from the <b>same parent and schedule</b> cliffed at its step 1199 and ended at 45/17/2/1.</p>
+<table><tr><th>context</th><th>parent (695)</th><th>100% b5 (1651)</th><th>MIX ({a.tag})</th><th>mix - parent</th></tr>{rows_html}</table>
+<img src="data:image/png;base64,{b64}">
+<p style="color:#666;font-size:13px">Generated from <code>mix64k_{a.tag}_len_*/summary.csv</code> and
+<code>mix64k_logs/train.log</code>. Cells showing "-" had not finished when this was built.</p>"""
+    hp=f"{a.out}_{a.tag}.html"; open(hp,"w").write(html)
+    print(f"  wrote {hp} ({os.path.getsize(hp)//1024} KB)")
+    return 0
+
+if __name__=="__main__": sys.exit(main())