Update app.py

app.py

@@ -8,35 +8,28 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 from wfgy_sdk import get_engine
 from wfgy_sdk.evaluator import compare_logits, plot_histogram
 
-# ────────────────────────────
-# tiny model + engine
-# ────────────────────────────
+# tiny model (CPU-friendly demo)
 tok = AutoTokenizer.from_pretrained("sshleifer/tiny-gpt2")
 mdl = AutoModelForCausalLM.from_pretrained("sshleifer/tiny-gpt2")
 eng = get_engine()
 
-# ────────────────────────────
-# runtime history (dummy zero so line is never empty)
-# ────────────────────────────
+# history buffer
 history = {"step": [0], "var": [0.0], "kl": [0.0]}
 
-# ────────────────────────────
-# paper benchmark table
-# ────────────────────────────
+# paper table
 paper_df = pd.DataFrame({
     "Benchmark": ["MMLU","GSM8K","BBH","MathBench","TruthfulQA",
                   "XNLI","MLQA","LongBench","VQAv2","OK-VQA"],
     "Baseline":  [61.0,78.0,79.3,72.2,62.4,59.5,78.1,51.4,69.1,65.7],
     "WFGY":      [89.8,98.7,100.7,87.4,90.4,77.3,106.6,69.6,86.6,86.8]
 })
-paper_df["Abs_gain"]  = (paper_df["WFGY"] - paper_df["Baseline"]).round(1)
-paper_df["Rel_gain%"] = ((paper_df["Abs_gain"] / paper_df["Baseline"]) * 100).round(0)
+paper_df["Abs_gain"]  = (paper_df["WFGY"]-paper_df["Baseline"]).round(1)
+paper_df["Rel_gain%"] = ((paper_df["Abs_gain"]/paper_df["Baseline"])*100).round(0)
 
 styled_df = (
     paper_df.style
-    .background_gradient(subset=["Abs_gain"], cmap="Greens")
-    .background_gradient(subset=["Rel_gain%"], cmap="Greens")
-    .format({"Abs_gain": "{:.1f}", "Rel_gain%": "{:.0f}"})
+    .background_gradient(cmap="Greens", subset=["Abs_gain","Rel_gain%"])
+    .format({"Abs_gain":"{:.1f}","Rel_gain%":"{:.0f}"})
 )
 
 paper_bar = px.bar(
@@ -45,68 +38,47 @@ paper_bar = px.bar(
     color_continuous_scale="Greens", height=300
 )
 
-# ────────────────────────────
 # helpers
-# ────────────────────────────
-def top5_tokens(logits: np.ndarray):
-    """return list of (token, prob) sorted desc"""
-    probs = torch.softmax(torch.tensor(logits), dim=0).numpy()
-    idx = probs.argsort()[-5:][::-1]
-    items = []
-    for i in idx:
-        token = tok.decode(int(i)).replace("\n", "\\n")
-        prob  = probs[i]
-        items.append(f"{token!r}: {prob:.3f}")
-    return "\n".join(items)
-
-def plot_history():
+def top5(logits: np.ndarray):
+    p = torch.softmax(torch.tensor(logits), dim=0).numpy()
+    idx = p.argsort()[-5:][::-1]
+    return "\n".join([f"{tok.decode(int(i))!r}: {p[i]:.2e}" for i in idx])
+
+def hist_plot():
     df = pd.DataFrame(history)
-    return px.line(df, x="step", y=["var", "kl"],
+    return px.line(df, x="step", y=["var","kl"],
                    labels={"value":"metric","step":"call"},
                    title="history (var% ↓  &  KL)").update_layout(height=260)
 
-def clear_history():
+def clear_hist():
     history["step"][:] = [0]; history["var"][:]=[0.0]; history["kl"][:]=[0.0]
-    return plot_history()
+    return hist_plot()
 
-# ────────────────────────────
-# main run
-# ────────────────────────────
 def run(prompt: str):
     p = prompt.strip()
     if not p:
-        return "", "", "", "", None, plot_history()
+        return "", "", "", "", None, hist_plot()
 
-    ids = tok(p, return_tensors="pt").input_ids
-    rawL = mdl(ids).logits[0, -1].detach().cpu().numpy()
+    ids  = tok(p, return_tensors="pt").input_ids
+    rawL = mdl(ids).logits[0,-1].detach().cpu().numpy()
     G    = np.random.randn(256).astype(np.float32)
-    I    = G + np.random.normal(scale=0.05, size=256).astype(np.float32)
-    modL = eng.run(I, G, rawL)
-
-    m = compare_logits(rawL, modL)
-    step = len(history["step"])
-    history["step"].append(step)
-    history["var"].append(m["var_drop"] * 100)
-    history["kl"].append(m["kl"])
-
-    fig = plot_histogram(rawL, modL)
-    buf = io.BytesIO(); fig.savefig(buf, format="png"); buf.seek(0)
-    img = Image.open(buf)
+    I    = G + np.random.normal(scale=0.05,size=256).astype(np.float32)
+    modL = eng.run(I,G,rawL)
 
-    headline = f"▼ var {m['var_drop']*100:4.1f}% | KL {m['kl']:.3f} | top-1 {'kept' if m['top1'] else 'changed'}"
+    m = compare_logits(rawL,modL)
+    n  = len(history["step"])
+    history["step"].append(n); history["var"].append(m["var_drop"]*100); history["kl"].append(m["kl"])
 
-    raw_top5 = top5_tokens(rawL)
-    mod_top5 = top5_tokens(modL)
+    fig = plot_histogram(rawL,modL); buf=io.BytesIO(); fig.savefig(buf,format="png"); buf.seek(0)
 
-    return raw_top5, mod_top5, headline, img, plot_history()
+    head = f"▼ var {m['var_drop']*100:4.1f}% | KL {m['kl']:.3f} | top-1 {'kept' if m['top1'] else 'changed'}"
+    return top5(rawL), top5(modL), head, Image.open(buf), hist_plot()
 
-# ────────────────────────────
-# UI layout
-# ────────────────────────────
+# UI
 with gr.Blocks(title="WFGY variance gate demo") as demo:
     gr.Markdown("# 🧠 WFGY simulation demo")
     prompt = gr.Textbox(label="Prompt", value="Explain Schrödinger's cat")
-    run_btn = gr.Button("🚀 Run")
+    run_b  = gr.Button("🚀 Run")
 
     with gr.Row():
         raw_box = gr.Textbox(label="Raw top-5 tokens", lines=6)
@@ -114,18 +86,17 @@ with gr.Blocks(title="WFGY variance gate demo") as demo:
 
     headline = gr.Markdown()
     hist_img = gr.Image(type="pil", label="Logit histogram")
-    hist_plot = gr.Plot()
-    clr_btn = gr.Button("Clear history")
+    hist_p   = gr.Plot()
+    clr_b    = gr.Button("Clear history")
 
     with gr.Accordion("Paper benchmarks", open=False):
-        gr.DataFrame(value=styled_df, interactive=False, wrap=True)
+        gr.DataFrame(styled_df, interactive=False, wrap=True)
         gr.Plot(paper_bar)
 
-    gr.Markdown("---\n⭐ **10 000 GitHub stars before 2025-08-01 unlock WFGY 2.0**")
+    gr.Markdown("---\n⭐ **10 k GitHub stars before 2025-08-01 unlock WFGY 2.0**")
 
-    run_btn.click(run, prompt,
-                  [raw_box, mod_box, headline, hist_img, hist_plot])
-    clr_btn.click(clear_history, None, hist_plot)
+    run_b.click(run, prompt, [raw_box,mod_box,headline,hist_img,hist_p])
+    clr_b.click(clear_hist, None, hist_p)
 
 if __name__ == "__main__":
     demo.queue().launch()