Update app.py

app.py

@@ -3,19 +3,26 @@ matplotlib.use("Agg")
 
 from PIL import Image
 import pandas as pd, plotly.express as px, gradio as gr
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from wfgy_sdk import get_engine
 from wfgy_sdk.evaluator import compare_logits, plot_histogram
 
-# tiny model for demo
+# ────────────────────────────
+# tiny model + engine
+# ────────────────────────────
 tok = AutoTokenizer.from_pretrained("sshleifer/tiny-gpt2")
 mdl = AutoModelForCausalLM.from_pretrained("sshleifer/tiny-gpt2")
 eng = get_engine()
 
-# runtime history (start with a dummy zero so the plot is never empty)
+# ────────────────────────────
+# runtime history (dummy zero so line is never empty)
+# ────────────────────────────
 history = {"step": [0], "var": [0.0], "kl": [0.0]}
 
-# paper benchmark absolute numbers
+# ────────────────────────────
+# paper benchmark table
+# ────────────────────────────
 paper_df = pd.DataFrame({
     "Benchmark": ["MMLU","GSM8K","BBH","MathBench","TruthfulQA",
                   "XNLI","MLQA","LongBench","VQAv2","OK-VQA"],
@@ -23,14 +30,55 @@ paper_df = pd.DataFrame({
     "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["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}"})
+)
+
+paper_bar = px.bar(
+    paper_df, x="Benchmark", y="Rel_gain%",
+    title="Relative gain (%)", color="Rel_gain%",
+    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():
+    df = pd.DataFrame(history)
+    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():
+    history["step"][:] = [0]; history["var"][:]=[0.0]; history["kl"][:]=[0.0]
+    return plot_history()
+
+# ────────────────────────────
+# main run
+# ────────────────────────────
 def run(prompt: str):
-    prompt = prompt.strip()
-    if not prompt:
-        return "", "", "", None, plot_history()
-    ids  = tok(prompt, return_tensors="pt").input_ids
-    rawL = mdl(ids).logits[0,-1].detach().cpu().numpy()
+    p = prompt.strip()
+    if not p:
+        return "", "", "", "", None, plot_history()
+
+    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)
@@ -38,52 +86,46 @@ def run(prompt: str):
     m = compare_logits(rawL, modL)
     step = len(history["step"])
     history["step"].append(step)
-    history["var"].append(m["var_drop"]*100)
+    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)
 
-    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'}*"
+    headline = f"▼ var {m['var_drop']*100:4.1f}% | KL {m['kl']:.3f} | top-1 {'kept' if m['top1'] else 'changed'}"
 
-    raw_text = prompt + tok.decode(int(rawL.argmax()))
-    mod_text = prompt + tok.decode(int(modL.argmax()))
+    raw_top5 = top5_tokens(rawL)
+    mod_top5 = top5_tokens(modL)
 
-    return raw_text, mod_text, headline + "  " + note, img, plot_history()
+    return raw_top5, mod_top5, headline, img, plot_history()
 
-def plot_history():
-    df = pd.DataFrame(history)
-    return px.line(df, x="step", y=["var","kl"],
-                   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:
+# ────────────────────────────
+# UI layout
+# ────────────────────────────
+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")
+
     with gr.Row():
-        raw_box = gr.Textbox(label="Raw GPT-2")
-        mod_box = gr.Textbox(label="After WFGY")
+        raw_box = gr.Textbox(label="Raw top-5 tokens", lines=6)
+        mod_box = gr.Textbox(label="WFGY top-5 tokens", lines=6)
+
     headline = gr.Markdown()
     hist_img = gr.Image(type="pil", label="Logit histogram")
-    hist_plot = gr.Plot(label="History")
+    hist_plot = gr.Plot()
     clr_btn = gr.Button("Clear history")
 
     with gr.Accordion("Paper benchmarks", open=False):
-        gr.DataFrame(paper_df, interactive=False, wrap=True)
+        gr.DataFrame(value=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 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)
+    run_btn.click(run, prompt,
+                  [raw_box, mod_box, headline, hist_img, hist_plot])
+    clr_btn.click(clear_history, None, hist_plot)
 
 if __name__ == "__main__":
     demo.queue().launch()