Update app.py

app.py

@@ -1,61 +1,50 @@
 """
-HF Space · WFGY 1-click Variance Gate
+HF Space · WFGY 1-click Variance Gate  (貼上就能部署)
 """
 
-import gradio as gr, numpy as np
+import io, numpy as np, gradio as gr, matplotlib.pyplot as plt
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from wfgy_sdk import get_engine
-from wfgy_sdk.evaluator import compare_logits
+from wfgy_sdk.evaluator import compare_logits, plot_histogram
 
 MODEL  = "sshleifer/tiny-gpt2"
 tok    = AutoTokenizer.from_pretrained(MODEL)
 mdl    = AutoModelForCausalLM.from_pretrained(MODEL)
 ENG    = get_engine()
 
-def run(prompt, enable, boost):
+def run(prompt:str):
     if not prompt.strip():
-        return "-", "-", "Please enter a prompt."
+        return "-", "-", "-", None
 
-    inputs = tok(prompt, return_tensors="pt")
-    rawL   = mdl(**inputs).logits[0, -1].detach().cpu().numpy()
+    inp   = tok(prompt, return_tensors="pt")
+    rawL  = mdl(**inp).logits[0, -1].detach().cpu().numpy()
+    I, G  = np.random.randn(2, 256).astype(np.float32)
+    modL  = ENG.run(I, G, rawL)
 
-    # demo-only fake semantic vectors
-    I = np.random.randn(256).astype(np.float32)
-    G = np.random.randn(256).astype(np.float32)
+    mets  = compare_logits(rawL, modL)
+    head  = f"▼ Var {mets['var_drop']*100:.1f}% | KL {mets['kl']:.2f}"
 
-    modL = ENG.run(I, G, rawL) if enable else rawL
-    mets = compare_logits(rawL, modL)
+    # ── 圖表轉成 PNG buffer ──
+    fig   = plot_histogram(rawL, modL)
+    buf   = io.BytesIO(); fig.savefig(buf, format="png"); buf.seek(0)
 
-    headline = f"variance ▼ {int(mets['var_drop']*100)} % | " \
-               f"KL {mets['kl']:.2f} | " \
-               f"top-1 {'✔' if mets['top1'] else '✘'}"
+    raw_txt = prompt + tok.decode(int(rawL.argmax()))
+    mod_txt = prompt + tok.decode(int(modL.argmax()))
+    return raw_txt, mod_txt, head, buf
 
-    return (
-        prompt + tok.decode(int(rawL.argmax())),
-        prompt + tok.decode(int(modL.argmax())),
-        headline,
-    )
-
-
-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")
-    enable = gr.Checkbox(True, label="Enable WFGY")
-    boost  = gr.Slider(0.5, 3.0, 1.0, label="Demo Boost (visual only)")
-
-    run_btn = gr.Button("Run", variant="primary")
+with gr.Blocks(title="WFGY 1-Click Variance Gate") as demo:
+    gr.Markdown("# 🧠 WFGY 模擬實驗\n*輸入任意 Prompt，立刻觀看 Logit 直方圖*")
+    prompt = gr.Textbox(label="Prompt", value="Explain Schrödinger's cat")
+    run_b  = gr.Button("🚀 Run")
 
     with gr.Row():
-        out_raw = gr.Textbox(label="Raw GPT-2")
-        out_mod = gr.Textbox(label="After WFGY")
+        raw  = gr.Textbox(label="Raw GPT-2")
+        mod  = gr.Textbox(label="After WFGY")
 
-    headline = gr.Markdown()
+    head = gr.Markdown()
+    img  = gr.Image(label="Logit Histogram")
 
-    run_btn.click(run, [prompt, enable, boost],
-                  [out_raw, out_mod, headline])
+    run_b.click(run, prompt, [raw, mod, head, img])
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.queue(default_concurrency_limit=2).launch()