app.py

"""
HF Space · WFGY 1-click Variance Gate  (貼上就能部署)
"""

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, plot_histogram

MODEL  = "sshleifer/tiny-gpt2"
tok    = AutoTokenizer.from_pretrained(MODEL)
mdl    = AutoModelForCausalLM.from_pretrained(MODEL)
ENG    = get_engine()

def run(prompt:str):
    if not prompt.strip():
        return "-", "-", "-", None

    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)

    mets  = compare_logits(rawL, modL)
    head  = f"▼ Var {mets['var_drop']*100:.1f}% | KL {mets['kl']:.2f}"

    # ── 圖表轉成 PNG buffer ──
    fig   = plot_histogram(rawL, modL)
    buf   = io.BytesIO(); fig.savefig(buf, format="png"); buf.seek(0)

    raw_txt = prompt + tok.decode(int(rawL.argmax()))
    mod_txt = prompt + tok.decode(int(modL.argmax()))
    return raw_txt, mod_txt, head, buf

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():
        raw  = gr.Textbox(label="Raw GPT-2")
        mod  = gr.Textbox(label="After WFGY")

    head = gr.Markdown()
    img  = gr.Image(label="Logit Histogram")

    run_b.click(run, prompt, [raw, mod, head, img])

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=2).launch()