"""Live demo: the weight-space-programmed Qwen3.5-0.8B vs its unedited base model.

From "Sparks of Weight Space Programming" (https://rikiyat.github.io/blog/weight-space-programming/):
an LLM hand-edited 131,072 weights (MLP down_proj, layers 15 & 19) to stop the model's
repetition-loop pathology. No training. This Space runs both checkpoints on your problem.
"""
import json

import gradio as gr

BASE = "Qwen/Qwen3.5-0.8B"
EDITED = "akhaliq/Qwen3.5-0.8B-weight-space-programming"
INSTRUCTION = "Please reason step by step, and put your final answer within \\boxed{}."

try:
    import spaces
    GPU = spaces.GPU
except Exception:  # local / CPU fallback
    def GPU(fn=None, **_kw):
        def deco(g):
            return g
        return deco(fn) if fn else deco

_MODELS = {}


def load_models():
    if _MODELS:
        return
    import torch
    from transformers import AutoModelForMultimodalLM, AutoTokenizer
    for name, repo in (("base", BASE), ("edited", EDITED)):
        tok = AutoTokenizer.from_pretrained(repo)
        mdl = AutoModelForMultimodalLM.from_pretrained(repo, dtype=torch.bfloat16, device_map="cuda")
        mdl.eval()
        _MODELS[name] = (tok, mdl)
    if torch.cuda.is_available():
        # Pre-warm torch.linalg on the main thread: its CUDA lazy loader is not thread-safe on
        # first use (pytorch/pytorch#90613) and the two generation threads below both hit
        # solve_triangular in Qwen3.5's gated delta rule.
        eye = torch.eye(2, device="cuda", dtype=torch.float32).unsqueeze(0)
        torch.linalg.solve_triangular(eye, eye.clone(), upper=False, unitriangular=True)


@GPU(duration=300)
def compare(problem, max_new_tokens):
    import queue
    import threading

    from transformers import TextIteratorStreamer
    load_models()
    prompt = problem.strip() + "\n" + INSTRUCTION
    updates = queue.Queue()

    def gen(name):
        tok, mdl = _MODELS[name]
        enc = tok.apply_chat_template([{"role": "user", "content": prompt}],
                                      add_generation_prompt=True, enable_thinking=True,
                                      return_tensors="pt", return_dict=True)
        kwargs = {"input_ids": enc["input_ids"].to("cuda"), "max_new_tokens": int(max_new_tokens),
                  "do_sample": True, "temperature": 0.6, "top_p": 0.95, "top_k": 20}
        if enc.get("attention_mask") is not None:
            kwargs["attention_mask"] = enc["attention_mask"].to("cuda")
        streamer = TextIteratorStreamer(tok, skip_prompt=True, skip_special_tokens=False)
        kwargs["streamer"] = streamer
        threading.Thread(target=mdl.generate, kwargs=kwargs, daemon=True).start()
        for chunk in streamer:
            updates.put((name, chunk))
        updates.put((name, None))  # done sentinel

    for n in ("base", "edited"):
        threading.Thread(target=gen, args=(n,), daemon=True).start()

    texts = {"base": "", "edited": ""}
    done = 0
    yield "", ""
    while done < 2:
        try:
            name, chunk = updates.get(timeout=0.5)
        except queue.Empty:
            continue
        if chunk is None:
            done += 1
        else:
            texts[name] += chunk
        yield texts["base"], texts["edited"]


def compare_safe(problem, max_new_tokens):
    try:
        yield from compare(problem, max_new_tokens)
    except Exception as e:
        msg = (f"Live generation isn't available on this Space's current hardware ({e}). "
               "See the 'Recorded verification runs' tab — those generations came from the "
               "actual 32,768-token verification job.")
        yield msg, msg


def fetch_recorded():
    """Recorded generations from the actual verification run (full 32,768-token protocol)."""
    from huggingface_hub import hf_hub_download
    try:
        base_f = hf_hub_download(EDITED, "eval/samples_base_math500.jsonl")
        edit_f = hf_hub_download(EDITED, "eval/samples_edited_math500.jsonl")
        base = {}
        for line in open(base_f):
            r = json.loads(line)
            base[r["id"]] = r
        rows = []
        for line in open(edit_f):
            r = json.loads(line)
            if r["id"] in base:
                rows.append((r["id"], base[r["id"]].get("answer", ""), base[r["id"]]["text"], r["text"]))
        return rows
    except Exception:
        return []


RECORDED = fetch_recorded()


def show_recorded(pid):
    for rid, answer, btext, etext in RECORDED:
        if rid == pid:
            return f"[gold answer: {answer}]\n\n{btext}", f"[gold answer: {answer}]\n\n{etext}"
    return "", ""


HEADER = """# 🧠 The model an LLM programmed by hand — try it

From **["Sparks of Weight Space Programming"](https://rikiyat.github.io/blog/weight-space-programming/)**:
a frontier LLM read Qwen3.5-0.8B's weights, found a direction that makes it loop until the token
limit, and **rewrote 131,072 weights to remove it — no training, no gradients**. The blog post
never released the model; this is a bit-exact, hash-verified rebuild
([akhaliq/Qwen3.5-0.8B-weight-space-programming](https://huggingface.co/akhaliq/Qwen3.5-0.8B-weight-space-programming)).

**Left: the original model** — watch it circle and repeat. **Right: the edited model** — same
weights everywhere else, same prompt, same sampling. Verified on the full MATH-500 (500 problems
× 4 samples): **51.0% accuracy** (post: 48.1) with **31.6%** of answers truncated at the token cap
(post: 33.1); the base model manages 5% with 95% truncated.

*Live runs cap the answer length for demo speed (default 4096 tokens; the verification used 32,768).
The base model's loop is usually visible within the first ~2,000 tokens. First call may take a
minute while both models load.*
"""


with gr.Blocks(title="Weight-space-programmed Qwen3.5-0.8B") as demo:
    gr.Markdown(HEADER)
    with gr.Tab("Live: base vs edited"):
        inp = gr.Textbox(label="Math problem", lines=2,
                         value="Let $x$ and $y$ be positive integers such that $x + y = 20$ and $xy = 91$. What is $x^2 + y^2$?")
        maxn = gr.Slider(1024, 8192, value=4096, step=1024,
                         label="Max new tokens (demo cap — verification used 32,768)")
        btn = gr.Button("Run both models", variant="primary")
        with gr.Row():
            out_b = gr.Textbox(label="Base Qwen3.5-0.8B", lines=26)
            out_e = gr.Textbox(label="Weight-edited (by an LLM, no training)", lines=26)
        btn.click(compare_safe, [inp, maxn], [out_b, out_e])
        gr.Examples(examples=[
            "What is $2^{10} + 24$?",
            "Compute $\\gcd(1071, 462)$.",
            "If $f(x) = 3x^2 - 2x + 1$, what is $f(5)$?",
            "A fair coin is flipped 8 times. What is the probability of getting exactly 3 heads? Express your answer as a common fraction.",
        ], inputs=inp)
    with gr.Tab("Recorded verification runs (full protocol)"):
        gr.Markdown("Generations from the verification job itself — 32,768-token cap, temperature 0.6, "
                    "thinking on, graded with math-verify. Pick a MATH-500 problem:")
        choices = [r[0] for r in RECORDED]
        drop = gr.Dropdown(choices=choices, value=choices[0] if choices else None, label="Problem")
        with gr.Row():
            rec_b = gr.Textbox(label="Base (original weights)", lines=30)
            rec_e = gr.Textbox(label="Edited (weight-space programmed)", lines=30)
        drop.change(show_recorded, drop, [rec_b, rec_e])
        demo.load(show_recorded, drop, [rec_b, rec_e])

if __name__ == "__main__":
    demo.launch()