Update app.py

app.py

@@ -1,32 +1,31 @@
-import base64, io, numpy as np, gradio as gr, wfgy_sdk as w
+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
 
-import torch
+from PIL import Image
+import matplotlib.pyplot as plt
 from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
 
-# ----------------------------------------------------------------------
-# tiny GPT-2 so the Space stays within free CPU limits
-# ----------------------------------------------------------------------
-MODEL = "sshleifer/tiny-gpt2"
+# ─────────────────────────── config ────────────────────────────
+MODEL = "sshleifer/tiny-gpt2"          # 124 MB, fits free tier
 tokenizer = AutoTokenizer.from_pretrained(MODEL)
 model     = AutoModelForCausalLM.from_pretrained(MODEL)
 set_seed(42)
 
 ENGINE = w.get_engine()
 
-# ----------------------------------------------------------------------
-# helper
-# ----------------------------------------------------------------------
+# ─────────────────────────── core fn ───────────────────────────
 def wfgy_pipeline(prompt: str, enable_wfgy: bool):
+    """Return raw text, mod text, HTML metrics and PIL histogram."""
     if not prompt.strip():
         return "", "", "<i>Please enter a prompt.</i>", None
 
     try:
+        # ── logits from tiny GPT-2
         ids        = tokenizer(prompt, return_tensors="pt").input_ids
         raw_logits = model(ids).logits[0, -1].detach().cpu().numpy()
 
-        # dummy semantic vectors (demo only)
+        # demo semantic vectors
         G = np.random.randn(256); G /= np.linalg.norm(G)
         I = G + np.random.normal(scale=0.05, size=256)
 
@@ -35,32 +34,30 @@ def wfgy_pipeline(prompt: str, enable_wfgy: bool):
             if enable_wfgy else raw_logits.copy()
         )
 
-        # metrics
-        m      = compare_logits(raw_logits, mod_logits)
-        top1   = "✔" if m["top1_shift"] else "✘"
-        metric = (f"<b>variance ▼ {(1-m['std_ratio'])*100:.0f}%</b> | "
-                  f"<b>KL {m['kl_divergence']:.2f}</b> | top-1 {top1}")
+        # ── metrics
+        m     = compare_logits(raw_logits, mod_logits)
+        top1  = "✔" if m["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; top-1 {top1}"
+        )
 
-        # histogram (support both “return fig” or “draw directly” versions)
-        maybe_fig = plot_histogram(raw_logits, mod_logits)   # **no show kwarg**
-        import matplotlib.pyplot as plt
-        fig = maybe_fig if maybe_fig is not None else plt.gcf()
-        buf = io.BytesIO(); fig.savefig(buf, format="png"); fig.clf()
-        hist_uri = "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
+        # ── 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)
 
-        # one-token continuations
+        # ── one-token continuations
         raw_txt = prompt + tokenizer.decode(int(raw_logits.argmax()))
         mod_txt = prompt + tokenizer.decode(int(mod_logits.argmax()))
 
-        return raw_txt, mod_txt, metric, hist_uri
+        return raw_txt, mod_txt, stats, hist_img
 
-    except Exception as e:
-        err = f"<b style='color:red'>Error:</b> {str(e)}"
+    except Exception as exc:
+        err = f"<b style='color:red'>Error:</b> {exc}"
         return "", "", err, None
 
-# ----------------------------------------------------------------------
-# UI
-# ----------------------------------------------------------------------
+# ─────────────────────────── UI ────────────────────────────────
 css = "#prompt-row{margin-bottom:1rem}.gr-box{font-size:.85rem}"
 
 with gr.Blocks(title="WFGY Variance Gate", css=css, theme=gr.themes.Soft()) as demo:
@@ -68,7 +65,7 @@ with gr.Blocks(title="WFGY Variance Gate", css=css, theme=gr.themes.Soft()) as d
         """
 ### 🧠 WFGY 1-click Variance Gate  
 
-Turn GPT-2 into a calmer thinker in seconds. **Bigger LLMs → even stronger gains.**
+Turn GPT-2 into a calmer thinker in seconds. **Bigger LLMs → stronger gains.**
 
 | Metric | Meaning |
 |--------|---------|
@@ -79,7 +76,7 @@ Turn GPT-2 into a calmer thinker in seconds. **Bigger LLMs → even stronger gai
 **Benchmarks (WFGY 1.0 vs base)**  
 
 | Task | Base % | WFGY % | Δ |
-|------|-------:|-------:|---:|
+|------|------:|------:|--:|
 | MMLU | 61.0 | **89.8** | +47 % |
 | TruthfulQA | 62.4 | **90.4** | +45 % |
 | GSM8K | 78.0 | **98.7** | +27 % |
@@ -95,15 +92,15 @@ Turn GPT-2 into a calmer thinker in seconds. **Bigger LLMs → even stronger gai
         raw_box = gr.Textbox(label="Raw GPT-2")
         mod_box = gr.Textbox(label="After WFGY")
 
-    metric_html = gr.HTML()
-    hist_img    = gr.Image(label="Logit distribution", width=440)
+    metrics = gr.HTML()
+    hist    = gr.Image(label="Logit distribution", width=440, show_label=True)
 
     runbtn.click(wfgy_pipeline, [prompt, enable],
-                 [raw_box, mod_box, metric_html, hist_img])
+                 [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,  
+**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!  
 
 ⭐ **10 000 GitHub stars before 2025-08-01** unlocks **WFGY 2.0**