Upload scripts/mk_abl_plot.py with huggingface_hub

scripts/mk_abl_plot.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python3
+"""Phase-1 LR-schedule ablation. Every number read from p1ev_*/summary.csv."""
+import csv,os,math
+import matplotlib; matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+S={"surface":"#fcfcfb","ink":"#0b0b0b","ink2":"#52514e","grid":"#e4e4e1"}
+C={"abl_const":"#b3243b","abl_cosine":"#eb6834","p1final":"#12508f","base680":"#9dc3ea"}
+LAB={"abl_const":"const 8e-5 (Muon+AdamW)","abl_cosine":"cosine 8e-5->8e-6 (Muon+AdamW)",
+     "p1final":"p1final  1.8e-4 WSD (AdamW)","base680":"base680  NO training"}
+def sty(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)
+def wilson(k,n=100,z=1.96):
+    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*max(0,c-h),100*min(1,c+h)
+PL=[4096,8192,16360]; D={}
+for pl in PL:
+    r=list(csv.reader(open(f"p1ev_{pl}/summary.csv"))); h=r[0]
+    row=[x for x in r[1:] if x and x[0]=="niah_single"][0]
+    for k,v in zip(h[1:],row[1:]): D.setdefault(k,{})[pl]=float(v)
+
+fig,(a,b,c)=plt.subplots(1,3,figsize=(17.6,5.6),facecolor=S["surface"])
+X=range(len(PL))
+for m in ("base680","p1final","abl_const","abl_cosine"):
+    ys=[D[m][p] for p in PL]
+    lo=[y-wilson(round(y))[0] for y in ys]; hi=[wilson(round(y))[1]-y for y in ys]
+    a.errorbar(list(X),ys,yerr=[lo,hi],color=C[m],lw=2.4,ms=8,marker="o",capsize=3,label=LAB[m])
+sty(a); a.set_xticks(list(X)); a.set_xticklabels(["4k","8k","16k"]); a.set_ylim(60,105)
+a.set_xlabel("prompt length",color=S["ink2"],fontsize=10)
+a.set_ylabel("niah_single (%)",color=S["ink2"],fontsize=10)
+a.set_title("A. All four models, Wilson 95% CI",color=S["ink"],fontsize=11.5,loc="left",pad=9)
+a.legend(frameon=False,fontsize=8,labelcolor=S["ink2"],loc="lower left")
+
+d=[D["abl_const"][p]-D["abl_cosine"][p] for p in PL]
+b.bar(list(X),d,color=["#b3243b" if x>0 else "#eb6834" for x in d],width=.5)
+b.axhline(0,color=S["ink2"],lw=1)
+b.axhspan(-11,11,color="#f0f0ee",zorder=0)
+b.text(1,8.4,"Wilson 95% CI band at N=100",ha="center",fontsize=8.4,color=S["ink2"],style="italic")
+sty(b); b.set_xticks(list(X)); b.set_xticklabels(["4k","8k","16k"]); b.set_ylim(-14,14)
+b.set_ylabel("const - cosine (pp)",color=S["ink2"],fontsize=10)
+b.set_title("B. THE ABLATION: schedule shape -- null",color=S["ink"],fontsize=11.5,loc="left",pad=9)
+for i,v in enumerate(d): b.text(i,v+(0.7 if v>=0 else -1.6),f"{v:+.0f}",ha="center",fontsize=10,color=S["ink2"])
+
+names=["base680","p1final","abl_const","abl_cosine"]
+tot=[sum(D[m][p] for p in PL) for m in names]
+c.bar(range(4),tot,color=[C[m] for m in names])
+for i,v in enumerate(tot): c.text(i,v+3,f"{v:.0f}",ha="center",fontsize=10.5,color=S["ink2"],weight="bold")
+sty(c); c.set_xticks(range(4))
+c.set_xticklabels(["base680\n(none)","p1final\n1.8e-4 AdamW","const\n8e-5 Muon","cosine\n8e-5 Muon"],fontsize=8.4)
+c.set_ylim(0,310); c.set_ylabel("sum over 4k/8k/16k (pp)",color=S["ink2"],fontsize=10)
+c.set_title("C. Both new arms sit BELOW the untrained base",color=S["ink"],fontsize=11.5,loc="left",pad=9)
+fig.text(.005,.005,"680M, 4k->16k, 1,200 steps, identical data (ablmix16, token-identical to p1mix) and identical haystack "
+  "(199,754,694 B). N=100/point. CAVEAT: the two new arms differ from p1final in BOTH optimizer (Muon vs AdamW) and LR "
+  "(8e-5 vs 1.8e-4) -- this run cannot separate them.",fontsize=7.8,color=S["ink2"])
+fig.tight_layout(); fig.savefig("plots_abl/phase1_ablation.png",dpi=150,bbox_inches="tight")
+print("  wrote plots_abl/phase1_ablation.png")