Update app.py

app.py

@@ -3,66 +3,87 @@ matplotlib.use("Agg")
 
 from PIL import Image
 import pandas as pd, plotly.express as px, gradio as gr
-
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from wfgy_sdk import get_engine
 from wfgy_sdk.evaluator import compare_logits, plot_histogram
 
+# tiny model for demo
 tok = AutoTokenizer.from_pretrained("sshleifer/tiny-gpt2")
 mdl = AutoModelForCausalLM.from_pretrained("sshleifer/tiny-gpt2")
 eng = get_engine()
 
-history = {"step": [], "var": [], "kl": []}
+# runtime history (start with a dummy zero so the plot is never empty)
+history = {"step": [0], "var": [0.0], "kl": [0.0]}
 
-paper = pd.DataFrame({
+# paper benchmark absolute numbers
+paper_df = pd.DataFrame({
     "Benchmark": ["MMLU","GSM8K","BBH","MathBench","TruthfulQA",
                   "XNLI","MLQA","LongBench","VQAv2","OK-VQA"],
-    "Baseline":  [61,78,79.3,72.2,62.4,59.5,78.1,51.4,69.1,65.7],
+    "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)
 
-def run(prompt):
-    p = prompt.strip()
-    if not p:
+def run(prompt: str):
+    prompt = prompt.strip()
+    if not prompt:
         return "", "", "", None, plot_history()
-    ids = tok(p, return_tensors="pt").input_ids
-    raw = 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)
-    mod = eng.run(I,G,raw)
-    m = compare_logits(raw,mod)
-    step = len(history["step"])+1
+    ids  = tok(prompt, 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(raw,mod)
-    buf = io.BytesIO(); fig.savefig(buf,format="png"); buf.seek(0)
+
+    fig = plot_histogram(rawL, modL)
+    buf = io.BytesIO(); fig.savefig(buf, format="png"); buf.seek(0)
     img = Image.open(buf)
-    head = f"▼ var {m['var_drop']*100:4.1f}% | KL {m['kl']:.3f}"
-    raw_t = p + tok.decode(int(raw.argmax()))
-    mod_t = p + tok.decode(int(mod.argmax()))
-    return raw_t, mod_t, head, img, plot_history()
+
+    headline = f"▼ var {m['var_drop']*100:4.1f}% | KL {m['kl']:.3f}"
+    note = f"*top-1 token {'changed' if not m['top1'] else 'kept'}*"
+
+    raw_text = prompt + tok.decode(int(rawL.argmax()))
+    mod_text = prompt + tok.decode(int(modL.argmax()))
+
+    return raw_text, mod_text, headline + "  " + note, img, plot_history()
 
 def plot_history():
-    if not history["step"]:
-        return px.line(title="history").update_layout(height=250)
     df = pd.DataFrame(history)
     return px.line(df, x="step", y=["var","kl"],
-                   labels={"value":"metric","step":"call"}).update_layout(height=250)
+                   labels={"value":"metric","step":"call"},
+                   title="history (var% ↓  &  KL)").update_layout(height=260)
+
+def clear_hist():
+    history["step"][:] = [0]
+    history["var"][:]  = [0.0]
+    history["kl"][:]   = [0.0]
+    return plot_history()
 
 with gr.Blocks(title="WFGY variance gate") as demo:
     gr.Markdown("# 🧠 WFGY simulation demo")
-    inp = gr.Textbox(label="Prompt", value="Explain Schrödinger's cat")
-    btn = gr.Button("🚀 Run")
-    raw = gr.Textbox(label="Raw GPT-2")
-    mod = gr.Textbox(label="After WFGY")
-    head= gr.Markdown()
-    img = gr.Image(type="pil")
-    line= gr.Plot()
-    btn.click(run, inp, [raw, mod, head, img, line])
-    with gr.Accordion("Paper benchmark", open=False):
-        gr.DataFrame(paper, interactive=False)
-    gr.Markdown("---\n⭐ 10 000 stars → unlock **WFGY 2.0** by 2025-08-01")
+    prompt = gr.Textbox(label="Prompt", value="Explain Schrödinger's cat")
+    run_btn = gr.Button("🚀 Run")
+    with gr.Row():
+        raw_box = gr.Textbox(label="Raw GPT-2")
+        mod_box = gr.Textbox(label="After WFGY")
+    headline = gr.Markdown()
+    hist_img = gr.Image(type="pil", label="Logit histogram")
+    hist_plot = gr.Plot(label="History")
+    clr_btn = gr.Button("Clear history")
+
+    with gr.Accordion("Paper benchmarks", open=False):
+        gr.DataFrame(paper_df, interactive=False, wrap=True)
+
+    gr.Markdown("---\n⭐ **10 000 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_hist, None, hist_plot)
 
 if __name__ == "__main__":
     demo.queue().launch()