Update app.py

app.py

@@ -1,112 +1,97 @@
-import io, numpy as np, gradio as gr
+import io, inspect, numpy as np, gradio as gr
 import matplotlib.pyplot as plt
 from PIL import Image
-from transformers import (
-    AutoModelForCausalLM, AutoTokenizer, set_seed
-)
+from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed
 
 import wfgy_sdk as w
 from wfgy_sdk.evaluator import compare_logits
 from wfgy_sdk.visual    import plot_histogram
 
-# ─────────────── tiny GPT-2 (fits free HF Space) ────────────────
+# ────────── tiny GPT-2 (最小模型也能看出落差) ──────────
 MODEL = "sshleifer/tiny-gpt2"
-tokenizer = AutoTokenizer.from_pretrained(MODEL)
-model     = AutoModelForCausalLM.from_pretrained(MODEL)
+tok   = AutoTokenizer.from_pretrained(MODEL)
+mdl   = AutoModelForCausalLM.from_pretrained(MODEL)
 set_seed(42)
 
 ENGINE = w.get_engine()
+BOOST  = 1.2                       # 預設 Demo 放大倍率
 
-# ───────────────────────── core helper ──────────────────────────
-def one_pass(prompt: str, boost: float):
-    """Return (raw_txt, mod_txt, metrics, raw_l, mod_l)."""
-    ids        = tokenizer(prompt, return_tensors="pt").input_ids
-    raw_logits = model(ids).logits[0, -1].detach().cpu().numpy()
+# 檢查 run() 是否支援 bbmc_scale
+_RUN_HAS_SCALE = "bbmc_scale" in inspect.signature(ENGINE.run).parameters
 
-    # demo vectors – boost multiplies semantic distance
-    G = np.random.randn(256); G /= np.linalg.norm(G)
-    I = G + np.random.normal(scale=boost, size=256)
-
-    mod_logits = ENGINE.run(I, G, raw_logits, bbmc_scale=boost)
-
-    metrics = compare_logits(raw_logits, mod_logits)
-    return (
-        prompt + tokenizer.decode(int(raw_logits.argmax())),
-        prompt + tokenizer.decode(int(mod_logits.argmax())),
-        metrics, raw_logits, mod_logits
-    )
-
-def wfgy_pipeline(prompt: str, enable: bool, boost: float):
+# ────────────────── 核心推論 ──────────────────
+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()
+
+    # 生成語意向量
+    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:
-        raw_txt, mod_txt, met, rl, ml = one_pass(prompt, boost if enable else 0.0)
+        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)
 
-        # safety: if enable & variance drop < 5 %, force BBCR collapse once
-        if enable and (1 - met["std_ratio"]) < .05:
-            _, mod_txt, met, rl, ml = one_pass(prompt, boost * 1.8)
+    m = compare_logits(rawL, modL)
 
-        stats = (
-            f"<b>variance ▼ {(1-met['std_ratio'])*100:.0f}%</b> &nbsp;|&nbsp; "
-            f"<b>KL {met['kl_divergence']:.2f}</b> &nbsp;|&nbsp; "
-            f"top-1 {'✔' if met['top1_shift'] else '✘'}"
-        )
+    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 '✘'}")
 
-        # histogram → PIL
-        fig = plot_histogram(rl, ml) or plt.gcf()
-        buf = io.BytesIO()
-        fig.savefig(buf, format="png", bbox_inches="tight")
-        plt.close(fig)
-        img = Image.open(buf)
+    # 產生直方圖
+    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)
 
-        return raw_txt, mod_txt, stats, img
+    raw_txt = prompt + tok.decode(int(rawL.argmax()))
+    mod_txt = prompt + tok.decode(int(modL.argmax()))
+    return raw_txt, mod_txt, stats, img
 
-    except Exception as exc:
-        return "", "", f"<b style='color:red'>Error:</b> {exc}", None
+# ──────────────────── Gradio 介面 ────────────────────
+with gr.Blocks(theme=gr.themes.Soft(),
+               title="WFGY Variance Gate") as demo:
 
-# ────────────────────────── UI layout ───────────────────────────
-with gr.Blocks(title="WFGY Variance Gate", theme=gr.themes.Soft()) as demo:
     gr.Markdown(
-        """
+    """
 ### 🧠 WFGY 1-click Variance Gate  
-**Turn any model—even GPT-2—into a calmer thinker.**  
-Move the slider and watch variance dive.
-
-| Metric | Meaning |
-| --- | --- |
-| **variance ▼** | logits get less noisy |
-| **KL** | distribution really changed |
-| **top-1** | most-likely token swapped ✔ / ✘ |
-        """
-    )
+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=1.2, step=.1,
+    boost  = gr.Slider(0, 3, value=BOOST, step=.1,
                        label="Demo Boost (higher → bigger effect)")
-    runbtn = gr.Button("Run")
+    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(wfgy_pipeline,
+    runbtn.click(run_wfgy,
                  inputs=[prompt, enable, boost],
                  outputs=[raw_box, mod_box, metrics, hist])
 
     gr.Markdown(
-        """
-**PDF mode ** – feed <code>I_am_not_lizardman/WFGY_1.0.pdf</code> to any chat-LLM,  
-prepend <code>Use WFGY:</code> and enjoy sharper answers.
+    """
+**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">
-GitHub repo – star to unlock WFGY 2.0 (10 k ⭐ before 2025-08-01)
-</a>
+10 000 GitHub&nbsp;stars before&nbsp;2025-08-01 unlock WFGY 2.0
+</a> – adaptive-gamma &amp; multimodal edition.  
 
-📂 Hidden folder <b>I_am_not_lizardman/</b> holds 8 + 1 “Challenge-Einstein” papers — tweet a screenshot if you find them!
-        """
-    )
+📂 Folder <b>I_am_not_lizardman/</b> hides eight “Challenge-Einstein” papers – find them!
+    """)
 
 demo.launch()