app.py

"""
HF Space · WFGY simulation demo
(complete file – paste/replace your current app.py)
"""

import io
import numpy as np
import gradio as gr
import matplotlib.pyplot as plt
from transformers import AutoTokenizer, AutoModelForCausalLM

from wfgy_sdk import get_engine, evaluator

# ──────────────────── tiny GPT-2 + WFGY engine ────────────────────
MODEL_ID = "sshleifer/tiny-gpt2"
tok = AutoTokenizer.from_pretrained(MODEL_ID)
mdl = AutoModelForCausalLM.from_pretrained(MODEL_ID)
eng = get_engine()           # variance-gate singleton


# ──────────────────── helper: run one prompt ────────────────────
def run(prompt: str):
    if not prompt.strip():
        return "", "", "–", None

    ids   = tok(prompt, return_tensors="pt")
    raw_L = mdl(**ids).logits[0, -1].detach().cpu().numpy()

    # dummy fingerprints (toy GPT-2 has none)
    I, G  = np.random.randn(2, 256).astype(np.float32)
    mod_L = eng.run(I, G, raw_L)

    # top-5 softmax
    def top5(logits):
        p = evaluator.softmax(logits)
        idx = p.argsort()[-5:][::-1]
        return "\n".join([f"'{tok.decode(i).strip()}': {p[i]:.2e}" for i in idx])

    m = evaluator.compare_logits(raw_L, mod_L)
    headline = (f"▼ var {m['var_drop']*100:.1f}% | KL {m['kl_divergence']:.3f} "
                f"| {'top-1 kept' if m['top1'] else 'top-1 changed'}")

    fig = evaluator.plot_histogram(raw_L, mod_L)
    buf = io.BytesIO()
    fig.savefig(buf, format="png")
    buf.seek(0)

    return top5(raw_L), top5(mod_L), headline, buf


# ──────────────────── gradio UI ────────────────────
with gr.Blocks(title="WFGY simulation demo") as demo:

    gr.Markdown("# 🧠 WFGY simulation demo")

    # ── marketing & quick-start banner ──
    gr.Markdown(
        """
**📈 WFGY: One Click to Activate the AI Taiji Cycle**

**📊 Semantic Accuracy ↑ 22.4 % | Reasoning Success ↑ 42.1 % | Stability ↑ 3.6 ×**

_No beliefs. Only experiments.<br>
WFGY 1.0 has already proven itself._

---

### 📜 Tutorial: How to Awaken the Soul of Your AI

**Step 1 — Download** ([PDF on Zenodo](https://zenodo.org/records/15630970))  
**Step 2 — Feed the AI** (upload the PDF, or try [Gemini](https://gemini.google.com/))  
**Step 3 — Give the Command** “**Answer using WFGY** + your question”  
_Prompt examples:_ TBD link  
**Step 4 — Integrate the SDK** ([GitHub](https://github.com/onestardao/WFGY))

---

🌟 **Star Reminder** → [Star the repo](https://github.com/onestardao/WFGY)  
_10 k ⭐ before 2025-08-01 unlocks WFGY 2.0._

⚠️ **Prompt Warning** — activate **Deep Structural Analysis Mode** or WFGY stays dormant.
        """,
        elem_id="intro-banner",
    )

    # ── prompt box ──
    prompt = gr.Textbox(label="Prompt", value="Explain Schrödinger's cat")
    run_btn = gr.Button("🚀 Run")

    # ── results ──
    with gr.Row():
        raw_box = gr.Textbox(label="Raw top-5 tokens")
        mod_box = gr.Textbox(label="WFGY top-5 tokens")

    headline = gr.Markdown()
    img      = gr.Image(label="Logit histogram")

    run_btn.click(run, prompt, outputs=[raw_box, mod_box, headline, img])


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