Update app.py

app.py

@@ -2,60 +2,54 @@
 HF Space · WFGY 1-click Variance Gate
 """
 
-import gradio as gr, numpy as np, torch
+import gradio as gr
+import numpy as np
+import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from wfgy_sdk import get_engine
 from wfgy_sdk.evaluator import compare_logits
 
-MODEL  = "sshleifer/tiny-gpt2"
-tok    = AutoTokenizer.from_pretrained(MODEL)
-mdl    = AutoModelForCausalLM.from_pretrained(MODEL)
-ENGINE = get_engine()
+# ---------- tiny demo backbone ----------
+MODEL      = "sshleifer/tiny-gpt2"
+tokenizer  = AutoTokenizer.from_pretrained(MODEL)
+model      = AutoModelForCausalLM.from_pretrained(MODEL)
+ENGINE     = get_engine()
 
-
-def run(prompt, enable, boost):
+# ---------- core inference ----------
+def wfgy_run(prompt: str, enable: bool, boost: float):
     if not prompt.strip():
-        return gr.update(value="-", visible=True)
+        return "—", "—", "Please enter a prompt."
 
-    # raw logits
-    inputs   = tok(prompt, return_tensors="pt")
-    rawL     = mdl(**inputs).logits[0, -1].detach().cpu().float().numpy()
+    # 1) raw logits from tiny-GPT-2
+    toks    = tokenizer(prompt, return_tensors="pt")
+    rawL    = model(**toks).logits[0, -1].detach().cpu().numpy()
 
-    # demo-only fake semantic vectors
-    I = np.random.randn(256).astype(np.float32)
+    # 2) dummy semantic vectors (demo only)
     G = np.random.randn(256).astype(np.float32)
+    I = (G + np.random.normal(scale=0.05, size=256).astype(np.float32)) * boost
 
-    if enable:
-        modL = ENGINE.run(
-            logits      = rawL,
-            input_vec   = I,
-            ground_vec  = G,
-            boost       = boost,
-        )
-    else:
-        modL = rawL
+    # 3) WFGY gate
+    modL = ENGINE.run(I, G, rawL) if enable else rawL
 
-    raw_txt = prompt + tok.decode(int(rawL.argmax()))
-    mod_txt = prompt + tok.decode(int(modL.argmax()))
+    # 4) text + metrics
+    raw_txt = prompt + tokenizer.decode(int(rawL.argmax()))
+    mod_txt = prompt + tokenizer.decode(int(modL.argmax()))
     m       = compare_logits(rawL, modL)
-
     headline = f"variance ▼ {int(m['var_drop']*100)} %  |  KL {m['kl']:.2f}  |  top-1 {'✔' if m['top1'] else '✘'}"
 
-    return (
-        raw_txt,
-        mod_txt,
-        headline,
-    )
+    return raw_txt, mod_txt, headline
 
+# ---------- UI ----------
+with gr.Blocks(title="WFGY 1-click Variance Gate", theme="soft") as demo:
+    gr.Markdown("## 🧠 **WFGY 1-click Variance Gate**  
+Turn GPT-2 into a calmer thinker. Move the slider → watch variance dive.\n")
 
-with gr.Blocks(title="WFGY 1-click Variance Gate") as demo:
-    gr.Markdown("## 🧠 WFGY 1-click Variance Gate\nTurn GPT-2 into a calmer thinker. Move the slider → watch variance dive.")
+    prompt  = gr.Textbox(label="Prompt")
+    enable  = gr.Checkbox(value=True, label="Enable WFGY")
+    boost   = gr.Slider(0.5, 3.0, value=1.0, step=0.1,
+                        label="Demo Boost (higher → bigger effect)")
 
-    prompt   = gr.Textbox(label="Prompt")
-    enable   = gr.Checkbox(value=True, label="Enable WFGY")
-    boost    = gr.Slider(0.5, 3.0, value=1.0, label="Demo Boost (higher → bigger effect)")
-
-    run_btn  = gr.Button("Run", variant="primary")
+    run_btn = gr.Button("Run", variant="primary")
 
     with gr.Row():
         out_raw = gr.Textbox(label="Raw GPT-2")
@@ -64,11 +58,12 @@ with gr.Blocks(title="WFGY 1-click Variance Gate") as demo:
     headline = gr.Markdown("")
 
     run_btn.click(
-        run,
+        wfgy_run,
         inputs=[prompt, enable, boost],
         outputs=[out_raw, out_mod, headline],
     )
+
+# ---------- launch ----------
 if __name__ == "__main__":
-    demo = build_ui()        # 你的 build_ui() 保持不變
-    demo.queue()             # 不帶任何參數 → 用預設 concurrency=2
-    demo.launch()            # Space 入口
+    demo.queue()      # default concurrency = 2
+    demo.launch()