Update app.py

app.py

@@ -1,12 +1,15 @@
-import io, numpy as np, gradio as gr, wfgy_sdk as w
-from wfgy_sdk.evaluator import compare_logits
-from wfgy_sdk.visual import plot_histogram
-
-from PIL import Image
+import io, numpy as np, gradio as gr
 import matplotlib.pyplot as plt
-from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
+from PIL import Image
+from transformers import (
+    AutoModelForCausalLM, AutoTokenizer, set_seed
+)
+
+import wfgy_sdk as w
+from wfgy_sdk.evaluator import compare_logits
+from wfgy_sdk.visual    import plot_histogram
 
-# ─────────────────────── tiny GPT-2 (124 MB) ───────────────────────
+# ─────────────── tiny GPT-2 (fits free HF Space) ────────────────
 MODEL = "sshleifer/tiny-gpt2"
 tokenizer = AutoTokenizer.from_pretrained(MODEL)
 model     = AutoModelForCausalLM.from_pretrained(MODEL)
@@ -14,104 +17,96 @@ set_seed(42)
 
 ENGINE = w.get_engine()
 
-# ─────────────────────────── core demo fn ──────────────────────────
-def wfgy_pipeline(prompt, enable_wfgy, boost):
-    """
-    boost ∈ [0,1] controls semantic deviation.
-    Larger boost → larger variance drop & KL.
-    """
-    if not prompt.strip():
-        return "", "", "<i>Please enter a prompt.</i>", None
-
-    # — raw logits from GPT-2 (single token) —
+# ───────────────────────── 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()
 
-    # — synthetic semantic vectors (demo only) —
+    # 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)
 
-    # — apply WFGY —
-    mod_logits = (
-        ENGINE.run(input_vec=I, ground_vec=G, logits=raw_logits)
-        if enable_wfgy else raw_logits.copy()
-    )
+    mod_logits = ENGINE.run(I, G, raw_logits, bbmc_scale=boost)
 
-    # — metrics —
-    m     = compare_logits(raw_logits, mod_logits)
-    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 '✘'}"
+    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
     )
 
-    # — histogram → PIL —
-    fig = plot_histogram(raw_logits, mod_logits) or plt.gcf()
-    buf = io.BytesIO()
-    fig.savefig(buf, format="png", bbox_inches="tight"); plt.close(fig)
-    hist_img = Image.open(buf)
+def wfgy_pipeline(prompt: str, enable: bool, boost: float):
+    if not prompt.strip():
+        return "", "", "<i>Please enter a prompt.</i>", None
+
+    try:
+        raw_txt, mod_txt, met, rl, ml = one_pass(prompt, boost if enable else 0.0)
+
+        # 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)
 
-    # — one-token continuations —
-    raw_txt = prompt + tokenizer.decode(int(raw_logits.argmax()))
-    mod_txt = prompt + tokenizer.decode(int(mod_logits.argmax()))
+        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 '✘'}"
+        )
 
-    return raw_txt, mod_txt, stats, hist_img
+        # 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)
 
-# ──────────────────────────── UI layout ────────────────────────────
-css = "#prompt-row{margin-bottom:1rem}.gr-box{font-size:.85rem}"
+        return raw_txt, mod_txt, stats, img
 
-with gr.Blocks(title="WFGY Variance Gate", css=css, theme=gr.themes.Soft()) as demo:
+    except Exception as exc:
+        return "", "", f"<b style='color:red'>Error:</b> {exc}", None
+
+# ────────────────────────── UI layout ───────────────────────────
+with gr.Blocks(title="WFGY Variance Gate", theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
 ### 🧠 WFGY 1-click Variance Gate  
-
-Turn GPT-2 into a calmer thinker in seconds. **Bigger LLMs → stronger gains.**
+**Turn any model—even GPT-2—into a calmer thinker.**  
+Move the slider and watch variance dive.
 
 | Metric | Meaning |
-|--------|---------|
-| **variance ▼** | logits become less noisy |
-| **KL** | distribution reshaped |
+| --- | --- |
+| **variance ▼** | logits get less noisy |
+| **KL** | distribution really changed |
 | **top-1** | most-likely token swapped ✔ / ✘ |
-
-Benchmarks (WFGY 1.0 vs base) show up to **+47 %** accuracy on MMLU.
         """
     )
 
-    with gr.Row(elem_id="prompt-row"):
-        prompt = gr.Textbox(label="Prompt", lines=2, placeholder="Ask anything…")
-        enable = gr.Checkbox(label="Enable WFGY", value=True)
-        runbtn = gr.Button("Run")
-
-    boost = gr.Slider(
-        minimum=0.0, maximum=1.0, step=0.05, value=0.30,
-        label="Demo Boost (higher → bigger effect)"
-    )
-
-    with gr.Row():
-        raw_box = gr.Textbox(label="Raw GPT-2")
-        mod_box = gr.Textbox(label="After WFGY")
+    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,
+                       label="Demo Boost (higher → bigger effect)")
+    runbtn = gr.Button("Run")
 
+    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=440)
+    hist    = gr.Image(label="Logit distribution", width=460)
 
-    runbtn.click(
-        wfgy_pipeline,
-        inputs=[prompt, enable, boost],
-        outputs=[raw_box, mod_box, metrics, hist]
-    )
+    runbtn.click(wfgy_pipeline,
+                 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 watch replies get sharper. Prompt revolution!  
+**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.
 
 ⭐ <a href="https://github.com/onestardao/WFGY" target="_blank">
-GitHub repo – star to unlock WFGY 2.0 (10 000 ⭐ before 2025-08-01)
-</a>  
+GitHub repo – star to unlock WFGY 2.0 (10 k ⭐ before 2025-08-01)
+</a>
 
 📂 Hidden folder <b>I_am_not_lizardman/</b> holds 8 + 1 “Challenge-Einstein” papers — tweet a screenshot if you find them!
-        """,
-        elem_id="footer"
+        """
     )
 
 demo.launch()