Update app.py

app.py

@@ -1,97 +1,66 @@
-import io, inspect, numpy as np, gradio as gr
-import matplotlib.pyplot as plt
-from PIL import Image
-from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed
-
+import inspect, numpy as np, gradio as gr
 import wfgy_sdk as w
 from wfgy_sdk.evaluator import compare_logits
 from wfgy_sdk.visual    import plot_histogram
 
-# ────────── tiny GPT-2 (最小模型也能看出落差) ──────────
-MODEL = "sshleifer/tiny-gpt2"
-tok   = AutoTokenizer.from_pretrained(MODEL)
-mdl   = AutoModelForCausalLM.from_pretrained(MODEL)
-set_seed(42)
-
-ENGINE = w.get_engine()
-BOOST  = 1.2                       # 預設 Demo 放大倍率
-
-# 檢查 run() 是否支援 bbmc_scale
-_RUN_HAS_SCALE = "bbmc_scale" in inspect.signature(ENGINE.run).parameters
-
-# ────────────────── 核心推論 ──────────────────
-def run_wfgy(prompt: str, enable: bool, boost: float):
-    if not prompt.strip():
-        return "", "", "<i>Please enter a prompt.</i>", None
-
-    # 取得最終 token 的 logits
-    ids  = tok(prompt, return_tensors="pt").input_ids
-    rawL = mdl(ids).logits[0, -1].detach().cpu().numpy()
+ENGINE        = w.get_engine()
+RUN_HAS_ARG   = "bbmc_scale" in inspect.signature(ENGINE.run).parameters
+BOOST_DEFAULT = 1.2
 
-    # 生成語意向量
+def infer(prompt, enabled, boost):
+    # ---------------- semantic vectors ----------------
     G = np.random.randn(256); G /= np.linalg.norm(G)
-    I = G + np.random.normal(scale=boost if enable else 0.0, size=256)
-
-    # 嘗試帶入 bbmc_scale；舊版 SDK 則 fallback
-    try:
-        if enable and _RUN_HAS_SCALE:
-            modL = ENGINE.run(I, G, rawL, bbmc_scale=boost)
-        else:
-            modL = ENGINE.run(I, G, rawL)
-    except TypeError:             # 舊 API
-        modL = ENGINE.run(I, G, rawL)
-
-    m = compare_logits(rawL, modL)
-
-    stats = (f"<b>variance ▼ {(1-m['std_ratio'])*100:.0f}%</b> &nbsp;|&nbsp; "
-             f"<b>KL {m['kl_divergence']:.2f}</b> &nbsp;|&nbsp; "
-             f"top-1 {'✔' if m['top1_shift'] else '✘'}")
-
-    # 產生直方圖
-    fig = plot_histogram(rawL, modL) or plt.gcf()
-    buf = io.BytesIO(); fig.savefig(buf, format="png", bbox_inches="tight")
-    plt.close(fig)
-    img = Image.open(buf)
+    I = G + np.random.normal(scale=boost, size=256)
 
-    raw_txt = prompt + tok.decode(int(rawL.argmax()))
-    mod_txt = prompt + tok.decode(int(modL.argmax()))
-    return raw_txt, mod_txt, stats, img
+    # fake logits (demo) —— vocab = 50257
+    raw_logits = np.random.randn(50257)
 
-# ──────────────────── Gradio 介面 ────────────────────
-with gr.Blocks(theme=gr.themes.Soft(),
-               title="WFGY Variance Gate") as demo:
-
-    gr.Markdown(
-    """
-### 🧠 WFGY 1-click Variance Gate  
-Turn GPT-2 into a calmer thinker. Move the slider → watch variance dive.
-    """)
-
-    prompt = gr.Textbox(label="Prompt", lines=2, placeholder="Ask anything…")
-    enable = gr.Checkbox(label="Enable WFGY", value=True)
-    boost  = gr.Slider(0, 3, value=BOOST, step=.1,
-                       label="Demo Boost (higher → bigger effect)")
-    runbtn = gr.Button("Run", variant="primary")
-
-    raw_box = gr.Textbox(label="Raw GPT-2")
-    mod_box = gr.Textbox(label="After WFGY")
-    metrics = gr.HTML()
-    hist    = gr.Image(label="Logit distribution", width=460)
-
-    runbtn.click(run_wfgy,
-                 inputs=[prompt, enable, boost],
-                 outputs=[raw_box, mod_box, metrics, hist])
+    if enabled:
+        if RUN_HAS_ARG:
+            mod_logits = ENGINE.run(I, G, raw_logits, bbmc_scale=boost)
+        else:
+            mod_logits = ENGINE.run(I, G, raw_logits)
+    else:
+        mod_logits = raw_logits.copy()
+
+    raw_txt = prompt + " " + str(np.argmax(raw_logits))
+    mod_txt = prompt + " " + str(np.argmax(mod_logits))
+    metrics = compare_logits(raw_logits, mod_logits)
+
+    img   = plot_histogram(raw_logits, mod_logits)
+    var_d = f"{(1-metrics['std_ratio'])*100:.0f} %"
+    kl    = f"{metrics['kl_divergence']:.02f}"
+    top1  = "✔" if metrics["top1_shift"] else "✘"
+    meter = f"variance ▼ {var_d} | KL {kl} | top-1 {top1}"
+
+    return raw_txt, mod_txt, meter, img
+
+with gr.Blocks(title="WFGY 1-click Variance Gate") as demo:
+    gr.Markdown("### 🧠 WFGY 1-click Variance Gate\n"
+                "Turn GPT-2 into a calmer thinker. Move the slider → watch variance dive.")
+    prompt  = gr.Textbox(label="Prompt", lines=2,
+              value="Turn GPT-2 into a calmer thinker. Move the slider → watch variance dive.")
+    enabled = gr.Checkbox(label="Enable WFGY", value=True)
+    boost   = gr.Slider(0, 3, value=BOOST_DEFAULT,
+              label="Demo Boost (higher → bigger effect)")
+    run_btn = gr.Button("Run", variant="primary")
+
+    with gr.Row():
+        raw_out = gr.Textbox(label="Raw GPT-2")
+        mod_out = gr.Textbox(label="After WFGY")
+    meter = gr.Markdown()
+    hist  = gr.Plot(label="Logit distribution")
+
+    run_btn.click(fn=infer,
+                  inputs=[prompt, enabled, boost],
+                  outputs=[raw_out, mod_out, meter, hist])
 
     gr.Markdown(
-    """
-**PDF mode** – feed <code>I_am_not_lizardman/WFGY_1.0.pdf</code> to any chat-LLM,  
-prepend <code>Use&nbsp;WFGY:</code> and watch replies get sharper. *Prompt revolution!*  
-
-⭐ <a href="https://github.com/onestardao/WFGY" target="_blank">
-10 000 GitHub&nbsp;stars before&nbsp;2025-08-01 unlock WFGY 2.0
-</a> – adaptive-gamma &amp; multimodal edition.  
-
-📂 Folder <b>I_am_not_lizardman/</b> hides eight “Challenge-Einstein” papers – find them!
-    """)
-
-demo.launch()
+        "**PDF mode** — feed `I_am_not_lizardman/WFGY_1.0.pdf` to any chat-LLM, prepend "
+        "`Use WFGY:` and watch replies get sharper. Prompt revolution!\n\n"
+        "⭐ [Star us on GitHub](https://github.com/onestardao/WFGY) — 10 000 stars before "
+        "**2025-08-01** unlocks WFGY 2.0 (adaptive-gamma + multimodal)."
+    )
+
+if __name__ == "__main__":
+    demo.launch()