app.py

"""
v12 / eMoE — bare vibe-check harness.

ONE path only:  descriptor --BGE--> z --hyper--> adapter (scaled by alpha) --> generate.
No ladder, no RAG, no OOD controller, no multi-expert loop, no escalation.
This exists to feel the adapted base directly. alpha=0.0 -> bare base (no adapter),
so the honest comparison is sliding alpha 0.0 <-> 0.8 (the trained backstop) <-> 1.0
on the SAME prompt and watching what the adapter does (help / nothing / shatter).
"""

import gradio as gr
import numpy as np
import torch
from dataclasses import fields

from huggingface_hub import hf_hub_download
from model_hybrid import GPT, GPTConfig
import hyper_lora as Hmod
import tok_v9

# ============================================================
# EDIT THIS BLOCK WHEN SWAPPING ARTIFACTS
# ============================================================
MODEL_REPO = "Daxamite/V12Emoe"
BASE_FILE  = "ckpt_v12_190m_best.pt"        # 190M FFT-hybrid, frozen
HYPER_FILE = "hyper_ckpt_v12.best.pt"       # 9.20M hypernetwork
ST_MODEL   = "BAAI/bge-base-en-v1.5"        # z encoder, d_z=768
# v12 base is ChatML-pretrained and the hyper was trained with --chat_format,
# so we wrap the generation prompt in ChatML (unlike v11).
CHAT_FORMAT = True
# ============================================================

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


# --- frozen base (same load path as HyperExpertRunner) ---
print(f"Downloading base {BASE_FILE} ...")
base_path = hf_hub_download(repo_id=MODEL_REPO, filename=BASE_FILE)
ckpt = torch.load(base_path, map_location="cpu", weights_only=False)
margs = ckpt.get("model_args") or ckpt.get("config") or ckpt.get("args") or {}
if hasattr(margs, "__dict__"):
    margs = vars(margs)
valid = {f.name for f in fields(GPTConfig)}
gcfg = GPTConfig(**{k: v for k, v in margs.items() if k in valid})
sd = ckpt.get("model") or ckpt.get("state_dict") or ckpt.get("model_state_dict")
if sd is None and all(torch.is_tensor(v) for v in ckpt.values()):
    sd = ckpt
sd = {k.replace("_orig_mod.", ""): v for k, v in sd.items()}
base = GPT(gcfg)
miss, unexp = base.load_state_dict(sd, strict=False)
if miss or unexp:
    print(f"NOTE base state_dict: {len(miss)} missing / {len(unexp)} unexpected")
BLOCK = gcfg.block_size

# --- rebuild + load the hypernetwork exactly as trained ---
print(f"Downloading hyper {HYPER_FILE} ...")
hyper_path = hf_hub_download(repo_id=MODEL_REPO, filename=HYPER_FILE)
hst = torch.load(hyper_path, map_location="cpu", weights_only=False)
if "hyper" not in hst:
    raise SystemExit(f"{HYPER_FILE} is not a train_hyper_sft checkpoint")
acfg = Hmod.AdaptConfig(**hst["adapt_cfg"])
D_Z = int(hst["d_z"]); d_trunk = int(hst.get("d_trunk", 512))
adapted, hyper, sites = Hmod.build(base, d_z=D_Z, cfg=acfg, d_trunk=d_trunk)
assert len(sites) == len(hst["sites"]), \
    f"site mismatch {len(sites)} vs {len(hst['sites'])} — wrong base for this hyper"
hyper.load_state_dict(hst["hyper"])
adapted = adapted.to(DEVICE).eval()
hyper = hyper.to(DEVICE).eval()
print(f"base {base.get_num_params()/1e6:.1f}M frozen | "
      f"hyper {sum(p.numel() for p in hyper.parameters())/1e6:.2f}M | "
      f"{len(sites)} sites | d_z {D_Z} | {DEVICE}")

# --- tokenizer + EOS stops ---
tok = tok_v9.build()
VOCAB = tok.vocab_size
STOP_IDS = set()
for s in ("<|im_end|>", "<|endoftext|>"):
    e = tok.encode(s)
    if len(e) == 1:
        STOP_IDS.add(int(e[0]))

# --- BGE z encoder (lazy) ---
print(f"Loading encoder {ST_MODEL} ...")
from sentence_transformers import SentenceTransformer
_enc = SentenceTransformer(ST_MODEL, device=DEVICE)


def encode_z(text: str) -> np.ndarray:
    v = _enc.encode([text], normalize_embeddings=True, show_progress_bar=False)[0]
    return np.asarray(v, dtype=np.float32)


def scale_deltas(deltas, alpha):
    if alpha == 1.0:
        return deltas
    out = {}
    for name, p in deltas.items():
        if p[0] == "lora":
            out[name] = ("lora", p[1], p[2] * alpha)
        else:  # film: additive
            out[name] = ("film", p[1] * alpha, p[2] * alpha)
    return out


def build_prompt(request: str) -> str:
    if CHAT_FORMAT:
        return f"<|im_start|>user\n{request}<|im_end|>\n<|im_start|>assistant\n"
    return f"{request}\n"


# --- streaming mint + generate ---
def run(request, descriptor, alpha, max_new, temperature, top_k, penalty):
    if not request or not request.strip():
        yield ""
        return
    desc = descriptor.strip() or request          # z defaults to the request itself
    alpha = float(alpha)

    with torch.no_grad():
        if alpha == 0.0:
            adapted.set_deltas(None)               # bare base, no adapter
            z_note = "alpha=0 -> BARE BASE (no adapter)"
        else:
            z = torch.as_tensor(encode_z(desc), device=DEVICE)
            adapted.set_deltas(scale_deltas(hyper(z), alpha))
            z_note = f"adapter minted from descriptor, alpha={alpha}"

        ids = tok.encode(build_prompt(request))[:BLOCK]
        idx = torch.tensor([ids], dtype=torch.long, device=DEVICE)
        new_ids = []
        for _ in range(int(max_new)):
            cond = idx if idx.size(1) <= BLOCK else idx[:, -BLOCK:]
            logits, _ = adapted(cond)
            logits = logits[:, -1, :].float()
            if VOCAB < logits.size(-1):            # never emit untrained pad logits
                logits[:, VOCAB:] = -float("inf")
            if penalty != 1.0:
                seen = torch.unique(idx[0])
                vals = logits[0, seen]
                logits[0, seen] = torch.where(vals > 0, vals / penalty, vals * penalty)
            if temperature and temperature > 0:
                lg = logits / temperature
                if top_k:
                    v, _ = torch.topk(lg, min(int(top_k), lg.size(-1)))
                    lg[lg < v[:, [-1]]] = -float("inf")
                nxt = int(torch.multinomial(torch.softmax(lg, -1), 1))
            else:
                nxt = int(logits.argmax(-1))
            if nxt in STOP_IDS:
                break
            new_ids.append(nxt)
            idx = torch.cat([idx, torch.tensor([[nxt]], device=DEVICE)], dim=1)
            yield f"[{z_note}]\n\n" + tok.decode(new_ids)
        yield f"[{z_note}]\n\n" + tok.decode(new_ids)


DESC = (
    "**v12 / eMoE — bare vibe check.** Frozen 190M FFT-hybrid base + 9.20M "
    "hypernetwork. The **descriptor** is BGE-encoded to `z`; the hypernetwork mints "
    "a per-request adapter from `z`; it is scaled by **alpha** and applied, then the "
    "**prompt** is generated. That is the whole path — no ladder, no RAG, no OOD "
    "controller, no multi-expert loop. **alpha=0 = bare base** (no adapter); 0.8 is "
    "the trained backstop; 1.0 reproduces training. Slide alpha on a fixed prompt to "
    "feel help vs. nothing vs. shatter. Tiny model on CPU/no-KV-cache: streams slowly, "
    "confabulates. Not a reasoning demo."
)

with gr.Blocks(title="v12 eMoE — vibe check") as demo:
    gr.Markdown(f"# v12 / eMoE — vibe check\n\n{DESC}")
    request = gr.Textbox(label="Prompt (what gets generated)", lines=4,
                         value="Write a Python function that reverses a string.")
    descriptor = gr.Textbox(
        label="Task descriptor (encoded to z; blank = use the prompt)", lines=2,
        value="python coding: implement a small string utility function")
    with gr.Row():
        alpha = gr.Slider(0.0, 1.0, value=0.8, step=0.05, label="alpha (0 = bare base)")
        max_new = gr.Slider(16, 512, value=200, step=8, label="Max new tokens")
        temp = gr.Slider(0.0, 1.5, value=0.0, step=0.05, label="Temperature (0 = greedy)")
        topk = gr.Slider(0, 200, value=0, step=1, label="Top-k (0 = off)")
        pen = gr.Slider(1.0, 1.5, value=1.0, step=0.01, label="Repetition penalty")
    run_btn = gr.Button("Mint + generate", variant="primary")
    out = gr.Textbox(label="Output", lines=16)
    inputs = [request, descriptor, alpha, max_new, temp, topk, pen]
    run_btn.click(run, inputs, out)
    request.submit(run, inputs, out)
    gr.Examples(
        examples=[
            ["Write a Python function that reverses a string.",
             "python coding: implement a small string utility function"],
            ["What is the capital of France?", "factual question answering: geography"],
            ["Summarize the causes of the French Revolution.",
             "history: explain causes of a historical event"],
            ["def fibonacci(n):", "python coding: recursive numeric function"],
        ],
        inputs=[request, descriptor],
    )

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