Update app.py

app.py

@@ -1,115 +1,68 @@
-import io, json, traceback, numpy as np, matplotlib
+import io, numpy as np, matplotlib
 matplotlib.use("Agg")
 
 from PIL import Image
-import gradio as gr
+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
 
-# model + engine
-MODEL_ID = "sshleifer/tiny-gpt2"
-tok  = AutoTokenizer.from_pretrained(MODEL_ID)
-mdl  = AutoModelForCausalLM.from_pretrained(MODEL_ID)
-eng  = get_engine()
-
-# runtime history
-hist = {"step": [], "var": [], "kl": []}
-
-# paper benchmark numbers
-paper = {
-    "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]
-}
-
-def run(prompt: str):
-    prompt = prompt.strip()
-    if not prompt:
-        return "", "", "", "", None, update_history()
-    try:
-        ids  = tok(prompt, 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)
-
-        # update history
-        step = len(hist["step"]) + 1
-        hist["step"].append(step)
-        hist["var"].append(m["var_drop"] * 100)
-        hist["kl"].append(m["kl"])
-
-        headline = f"▼ var {m['var_drop']*100:4.1f}% | KL {m['kl']:.3f}"
-        metrics_md = (
-            "| metric | value |\n"
-            "|--------|-------|\n"
-            f"| std_ratio | {m['std_ratio']:.3f} |\n"
-            f"| var_drop  | {m['var_drop']*100:.1f}% |\n"
-            f"| KL        | {m['kl']:.3f} |\n"
-            f"| top-1 same| {'yes' if m['top1'] else 'no'} |"
-        )
-
-        fig = plot_histogram(raw, mod)
-        buf = io.BytesIO(); fig.savefig(buf, format="png"); buf.seek(0)
-        img = Image.open(buf)
-
-        raw_txt = prompt + tok.decode(int(raw.argmax()))
-        mod_txt = prompt + tok.decode(int(mod.argmax()))
-        return raw_txt, mod_txt, headline, metrics_md, img, update_history()
-    except Exception:
-        tb = traceback.format_exc()
-        return "runtime error", tb, "runtime error", "", None, update_history()
-
-def update_history():
-    return {
-        "step": hist["step"],
-        "var":  hist["var"],
-        "kl":   hist["kl"]
-    }
-
-def clear_history():
-    hist["step"].clear(); hist["var"].clear(); hist["kl"].clear()
-    return update_history()
+tok = AutoTokenizer.from_pretrained("sshleifer/tiny-gpt2")
+mdl = AutoModelForCausalLM.from_pretrained("sshleifer/tiny-gpt2")
+eng = get_engine()
+
+history = {"step": [], "var": [], "kl": []}
+
+paper = 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],
+    "WFGY":      [89.8,98.7,100.7,87.4,90.4,77.3,106.6,69.6,86.6,86.8]
+})
+
+def run(prompt):
+    p = prompt.strip()
+    if not p:
+        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
+    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)
+    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()
+
+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)
 
 with gr.Blocks(title="WFGY variance gate") as demo:
     gr.Markdown("# 🧠 WFGY simulation demo")
-    prompt = gr.Textbox(label="Prompt", value="Explain Schrödinger's cat")
-    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()
-    metrics  = gr.Markdown()
-    img      = gr.Image(label="Logit histogram", type="pil")
-
-    hist_plot = gr.LinePlot(
-        label="History (var% ↓ & KL)",
-        x="step", y=["var", "kl"],
-        overlay=True, height=250
-    )
-    clear_btn = gr.Button("Clear history")
-
-    with gr.Accordion("Paper benchmarks", open=False):
-        bench_df  = gr.DataFrame(paper, interactive=False)
-        bench_bar = gr.BarPlot(paper, x="benchmark", y=["baseline", "wfgy"],
-                               overlay=False, height=300)
-
-    gr.Markdown(
-        "---\n"
-        "### ⭐ 10 000 GitHub stars before **2025-08-01** unlock **WFGY 2.0**"
-    )
-
-    btn.click(run, prompt,
-              [raw_box, mod_box, headline, metrics, img, hist_plot])
-    clear_btn.click(fn=clear_history, inputs=None, outputs=hist_plot)
+    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")
 
 if __name__ == "__main__":
     demo.queue().launch()