app.py

import io, numpy as np, gradio as gr
import matplotlib.pyplot as plt
from PIL import Image
from transformers import (
    AutoModelForCausalLM, AutoTokenizer, set_seed
)

import wfgy_sdk as w
from wfgy_sdk.evaluator import compare_logits
from wfgy_sdk.visual    import plot_histogram

# ─────────────── tiny GPT-2 (fits free HF Space) ────────────────
MODEL = "sshleifer/tiny-gpt2"
tokenizer = AutoTokenizer.from_pretrained(MODEL)
model     = AutoModelForCausalLM.from_pretrained(MODEL)
set_seed(42)

ENGINE = w.get_engine()

# ───────────────────────── core helper ──────────────────────────
def one_pass(prompt: str, boost: float):
    """Return (raw_txt, mod_txt, metrics, raw_l, mod_l)."""
    ids        = tokenizer(prompt, return_tensors="pt").input_ids
    raw_logits = model(ids).logits[0, -1].detach().cpu().numpy()

    # demo vectors – boost multiplies semantic distance
    G = np.random.randn(256); G /= np.linalg.norm(G)
    I = G + np.random.normal(scale=boost, size=256)

    mod_logits = ENGINE.run(I, G, raw_logits, bbmc_scale=boost)

    metrics = compare_logits(raw_logits, mod_logits)
    return (
        prompt + tokenizer.decode(int(raw_logits.argmax())),
        prompt + tokenizer.decode(int(mod_logits.argmax())),
        metrics, raw_logits, mod_logits
    )

def wfgy_pipeline(prompt: str, enable: bool, boost: float):
    if not prompt.strip():
        return "", "", "<i>Please enter a prompt.</i>", None

    try:
        raw_txt, mod_txt, met, rl, ml = one_pass(prompt, boost if enable else 0.0)

        # safety: if enable & variance drop < 5 %, force BBCR collapse once
        if enable and (1 - met["std_ratio"]) < .05:
            _, mod_txt, met, rl, ml = one_pass(prompt, boost * 1.8)

        stats = (
            f"<b>variance ▼ {(1-met['std_ratio'])*100:.0f}%</b> &nbsp;|&nbsp; "
            f"<b>KL {met['kl_divergence']:.2f}</b> &nbsp;|&nbsp; "
            f"top-1 {'✔' if met['top1_shift'] else '✘'}"
        )

        # histogram → PIL
        fig = plot_histogram(rl, ml) or plt.gcf()
        buf = io.BytesIO()
        fig.savefig(buf, format="png", bbox_inches="tight")
        plt.close(fig)
        img = Image.open(buf)

        return raw_txt, mod_txt, stats, img

    except Exception as exc:
        return "", "", f"<b style='color:red'>Error:</b> {exc}", None

# ────────────────────────── UI layout ───────────────────────────
with gr.Blocks(title="WFGY Variance Gate", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
### 🧠 WFGY 1-click Variance Gate  
**Turn any model—even GPT-2—into a calmer thinker.**  
Move the slider and watch variance dive.

| Metric | Meaning |
| --- | --- |
| **variance ▼** | logits get less noisy |
| **KL** | distribution really changed |
| **top-1** | most-likely token swapped ✔ / ✘ |
        """
    )

    prompt = gr.Textbox(label="Prompt", lines=2, placeholder="Ask anything…")
    enable = gr.Checkbox(label="Enable WFGY", value=True)
    boost  = gr.Slider(0, 3, value=1.2, step=.1,
                       label="Demo Boost (higher → bigger effect)")
    runbtn = gr.Button("Run")

    raw_box = gr.Textbox(label="Raw GPT-2")
    mod_box = gr.Textbox(label="After WFGY")
    metrics = gr.HTML()
    hist    = gr.Image(label="Logit distribution", width=460)

    runbtn.click(wfgy_pipeline,
                 inputs=[prompt, enable, boost],
                 outputs=[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,  
prepend <code>Use WFGY:</code> and enjoy sharper answers.

⭐ <a href="https://github.com/onestardao/WFGY" target="_blank">
GitHub repo – star to unlock WFGY 2.0 (10 k ⭐ before 2025-08-01)
</a>

📂 Hidden folder <b>I_am_not_lizardman/</b> holds 8 + 1 “Challenge-Einstein” papers — tweet a screenshot if you find them!
        """
    )

demo.launch()