Update app.py

app.py

@@ -5,21 +5,28 @@ from wfgy_sdk.visual import plot_histogram
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
 
+# ----------------------------------------------------------------------
+# tiny GPT-2 so the Space stays within free CPU limits
+# ----------------------------------------------------------------------
 MODEL = "sshleifer/tiny-gpt2"
 tokenizer = AutoTokenizer.from_pretrained(MODEL)
-model = AutoModelForCausalLM.from_pretrained(MODEL)
+model     = AutoModelForCausalLM.from_pretrained(MODEL)
 set_seed(42)
-ENGINE = w.get_engine()
 
+ENGINE = w.get_engine()
 
+# ----------------------------------------------------------------------
+# helper
+# ----------------------------------------------------------------------
 def wfgy_pipeline(prompt: str, enable_wfgy: bool):
     if not prompt.strip():
         return "", "", "<i>Please enter a prompt.</i>", None
 
     try:
-        ids = tokenizer(prompt, return_tensors="pt").input_ids
-        raw_logits = model(ids).logits[0, -1].detach().numpy()
+        ids        = tokenizer(prompt, return_tensors="pt").input_ids
+        raw_logits = model(ids).logits[0, -1].detach().cpu().numpy()
 
+        # dummy semantic vectors (demo only)
         G = np.random.randn(256); G /= np.linalg.norm(G)
         I = G + np.random.normal(scale=0.05, size=256)
 
@@ -28,49 +35,51 @@ def wfgy_pipeline(prompt: str, enable_wfgy: bool):
             if enable_wfgy else raw_logits.copy()
         )
 
-        m = compare_logits(raw_logits, mod_logits)
-        top1 = "✔" if m["top1_shift"] else "✘"
-        metrics_html = (
-            f"<b>variance ▼ {(1-m['std_ratio'])*100:.0f}%</b> "
-            f"| <b>KL {m['kl_divergence']:.2f}</b> "
-            f"| top-1 {top1}"
-        )
+        # 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}")
 
-        fig = plot_histogram(raw_logits, mod_logits)          # <<< fixed
+        # 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()
-        img_uri = "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
+        hist_uri = "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
 
-        raw_next = tokenizer.decode(int(raw_logits.argmax()))
-        mod_next = tokenizer.decode(int(mod_logits.argmax()))
-        return prompt + raw_next, prompt + mod_next, metrics_html, img_uri
+        # one-token continuations
+        raw_txt = prompt + tokenizer.decode(int(raw_logits.argmax()))
+        mod_txt = prompt + tokenizer.decode(int(mod_logits.argmax()))
 
-    except Exception as e:
-        return "", "", f"<b style='color:red'>Error:</b> {str(e)}", None
+        return raw_txt, mod_txt, metric, hist_uri
 
+    except Exception as e:
+        err = f"<b style='color:red'>Error:</b> {str(e)}"
+        return "", "", err, None
 
-css = """
-#prompt-row {margin-bottom: 1.0rem}
-.gr-box {font-size: 0.85rem}
-"""
+# ----------------------------------------------------------------------
+# UI
+# ----------------------------------------------------------------------
+css = "#prompt-row{margin-bottom:1rem}.gr-box{font-size:.85rem}"
 
-with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="WFGY Variance Gate", css=css, theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
 ### 🧠 WFGY 1-click Variance Gate  
 
-Turn GPT-2 into a calmer thinker in seconds.<br>
-**Bigger LLMs → even stronger gains.**
+Turn GPT-2 into a calmer thinker in seconds. **Bigger LLMs → even stronger gains.**
 
 | Metric | Meaning |
 |--------|---------|
 | **variance ▼** | logits become less noisy |
 | **KL** | distribution reshaped |
-| **top-1** | most-likely token swapped ✔/✘ |
+| **top-1** | most-likely token swapped ✔ / ✘ |
 
 **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 % |
@@ -80,17 +89,17 @@ Turn GPT-2 into a calmer thinker in seconds.<br>
     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)
-        run_btn = gr.Button("Run")
+        runbtn = gr.Button("Run")
 
     with gr.Row():
         raw_box = gr.Textbox(label="Raw GPT-2")
         mod_box = gr.Textbox(label="After WFGY")
 
-    metrics = gr.HTML()
-    hist_img = gr.Image(label="Logit distribution", width=440)
+    metric_html = gr.HTML()
+    hist_img    = gr.Image(label="Logit distribution", width=440)
 
-    run_btn.click(wfgy_pipeline, [prompt, enable],
-                  [raw_box, mod_box, metrics, hist_img])
+    runbtn.click(wfgy_pipeline, [prompt, enable],
+                 [raw_box, mod_box, metric_html, hist_img])
 
     gr.Markdown(
         """