app.py

"""
Gradio demo for the JEPA rl_prep/final checkpoint -- a ~50M-param SpikeWhale chat
model. Inference is byte-identical to the bake-off/battle harness that produced
the good generations: ChatML prompt via format_chat, temperature + top-p 0.9
nucleus sampling, halt on <|im_end|>/<eos>, fixed seed per turn.

Run:  python app.py        (then open the printed local URL)
"""
import os, sys, torch, torch.nn.functional as F
import gradio as gr

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
from spike_tokenizer import SpikeTokenizer
from chat_format import format_chat, IM_END
from model_v2 import SpikeWhaleLM

DEV = "cuda" if torch.cuda.is_available() else "cpu"
REPO_ID = "Quazim0t0/Escarda-86M"            # HuggingFace source of truth (private)
LOCAL_CKPT = os.path.join(HERE, "checkpoints", "rl_prep", "final")  # offline fallback
# Read the HF token from the environment / Space secret store -- never hardcode it.
HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")


def _load_tokenizer():
    """Tokenizer from the HF repo; fall back to the bundled tokenizer.json."""
    try:
        from huggingface_hub import hf_hub_download
        path = hf_hub_download(repo_id=REPO_ID, filename="tokenizer.json", token=HF_TOKEN)
        print(f"Tokenizer pulled from {REPO_ID}")
        return SpikeTokenizer(path)
    except Exception as e:
        print(f"(HF tokenizer fetch failed: {e}; using local tokenizer.json)")
        return SpikeTokenizer(os.path.join(HERE, "tokenizer.json"))


def _load_model():
    """Model from the HF repo (private -> needs token); fall back to local dir."""
    try:
        print(f"Loading {REPO_ID} on {DEV} ...")
        return SpikeWhaleLM.from_pretrained(REPO_ID, token=HF_TOKEN).to(DEV).float().eval()
    except Exception as e:
        print(f"(HF model fetch failed: {e}; loading local {LOCAL_CKPT})")
        return SpikeWhaleLM.from_pretrained(LOCAL_CKPT).to(DEV).float().eval()


tok = _load_tokenizer()
END_ID = tok.convert_tokens_to_ids(IM_END)
EOS = getattr(tok, "eos_token_id", None)
MODEL = _load_model()
N_PARAMS = sum(p.numel() for p in MODEL.parameters())
print(f"Loaded. {N_PARAMS/1e6:.1f}M params.")


@torch.no_grad()
def generate(message, history, system_prompt, temperature, max_new):
    """EXACT bake-off/battle sampling: ChatML prompt, temp + top-p 0.9 nucleus,
    stop on <|im_end|>/<eos>, seed reset for reproducibility."""
    # Null-safe: example clicks / API calls may omit the extra inputs.
    temperature = float(temperature) if temperature else 0.3
    max_new = int(max_new) if max_new else 120
    system_prompt = system_prompt or None
    torch.manual_seed(0)
    msgs = []
    for u, a in (history or []):
        if u:
            msgs.append({"role": "user", "content": u})
        if a:
            msgs.append({"role": "assistant", "content": a})
    msgs.append({"role": "user", "content": message})
    prompt = format_chat(msgs, system_prompt=(system_prompt or None),
                         add_generation_prompt=True)
    ids = torch.tensor(tok.encode(prompt), device=DEV).unsqueeze(0)
    out = MODEL(ids, use_cache=True); past = out.past_key_values; last = out.logits[0, -1]
    gen_ids = []
    for _ in range(int(max_new)):
        logits = last.float() / max(float(temperature), 1e-5)
        p = F.softmax(logits, -1)
        sp, si = p.sort(descending=True); cum = sp.cumsum(0)
        cut = cum > 0.9; cut[1:] = cut[:-1].clone(); cut[0] = False; sp[cut] = 0
        nxt = si[torch.multinomial(sp / sp.sum(), 1)].item()
        if nxt == END_ID or (EOS is not None and nxt == EOS):
            break
        gen_ids.append(nxt)
        out = MODEL(torch.tensor([[nxt]], device=DEV), past_key_values=past, use_cache=True)
        past = out.past_key_values; last = out.logits[0, -1]
    return tok.decode(gen_ids, skip_special_tokens=True).strip()


# ---- pre-captured sample generations (from the actual battle harness) ----
SAMPLES = [
    ("Give one practical tip for staying focused while working.",
     "Staying focused while working can be a great way to stay motivated and avoid "
     "burnout. Start with a short break. Lean back, take a deep breath, and focus on "
     "your breath. This will help you stay focused and avoid distractions."),
    ("In one sentence, what is a prime number?",
     "A prime number is a positive integer greater than 1 that has exactly two "
     "distinct factors: 1 and itself."),
    ("What gas do plants absorb from the air for photosynthesis?",
     "Plants absorb carbon dioxide from the atmosphere through photosynthesis."),
    ("Write a haiku about the ocean.",
     "The ocean is a majestic place / With a rhythm that is both beautiful and strange / "
     "The waves crash against the shore."),
    ("Write a congratulations message for a friend's graduation.",
     "Happy graduation! I'm so glad you're celebrating this milestone with your friends!"),
    ("List three steps to make a cup of tea.",
     "1. Start by boiling water for your cup of tea.\n2. Add a tea bag and let it steep.\n"
     "3. Drizzle a splash of milk or honey to taste."),
]

ABOUT_MD = f"""
# 🐋 Escarda-86M · General Chat Demo
**Weights:** [`Quazim0t0/Escarda-86M`](https://huggingface.co/Quazim0t0/Escarda-86M) on HuggingFace —
chosen as the best chat model after a seed-controlled bake-off across **28 checkpoints**
and a head-to-head battle test.

- **~{N_PARAMS/1e6:.0f}M parameters** · 4096 token context · HRM-refine + ChatML SFT, RL-prepped
- Runs on **{DEV.upper()}** — and comfortably on a single consumer GPU (or CPU)
- Inference here is **byte-identical** to the harness that produced its winning generations:
  ChatML framing, temperature + top-p 0.9 nucleus sampling, stop on `<|im_end|>`

### Why this one won
It almost never degenerates into repetition loops, follows instructions, and writes
coherently — staying on-task across chat, how-to, common-sense, and short-form writing,
where larger sibling checkpoints collapsed.
"""

MISSION_MD = """
## 🌍 Built for the community

I made this model **for everyone in the community to use freely for general-purpose
tasks**. I strongly believe we'll figure out — and soon — how to build small models
that genuinely contend with the much bigger ones; this is a step in that direction
and an open invitation for others to build on it.

This model was trained using **Modal's credits** as part of the **Small Models, Big
Adventures Hackathon**. 🙏

**Model weights are available** at
[huggingface.co/Quazim0t0/Escarda-86M](https://huggingface.co/Quazim0t0/Escarda-86M),
and **benchmarks will be posted soon.**
"""

COST_MD = """
## 💸 Why a small model that *works* matters

A ~86M-param model is **roughly 80× smaller** than a 7B model and **~2,000×**
smaller than a 175B-class frontier model. That gap is the whole pitch:

| | This model (~86M) | 7B model | 175B model |
|---|---|---|---|
| **VRAM (fp16)** | ~0.17 GB | ~14 GB | ~350 GB |
| **Runs on** | a laptop / free-tier GPU / CPU | one high-end GPU | a multi-GPU server |
| **$ / 1M tokens (self-host)** | fractions of a cent | cents–dimes | dollars |
| **Cold-start latency** | sub-second load | seconds–minutes | minutes |
| **Train cost** | a weekend on one box | many GPU-days | GPU-*months* |

**The thesis:** for *bounded* assistant tasks — short chat, how-to steps, drafting,
classification, on-device helpers — you don't need a frontier model. A focused 86M
model that stays coherent and follows instructions delivers a usable experience at
**near-zero marginal cost**, no API bills, full data privacy, and it runs where big
models simply can't. Small + correct beats huge + overkill for the long tail of
everyday tasks.
"""

with gr.Blocks(title="Escarda-86M Chat Demo", theme=gr.themes.Soft()) as demo:
    gr.Markdown(ABOUT_MD)
    gr.Markdown(MISSION_MD)
    with gr.Row():
        with gr.Column(scale=3):
            chat = gr.ChatInterface(
                fn=generate,
                additional_inputs=[
                    gr.Textbox(value="", label="System prompt (optional)",
                               placeholder="e.g. You are a concise, friendly assistant."),
                    gr.Slider(0.1, 1.2, value=0.3, step=0.05, label="Temperature"),
                    gr.Slider(16, 256, value=120, step=8, label="Max new tokens"),
                ],
                examples=[[s[0]] for s in SAMPLES],
                cache_examples=False,
                title="💬 Talk to Escarda-86M",
                description="Same inference settings that produced its best generations.",
            )
        with gr.Column(scale=2):
            gr.Markdown("## 📋 Sample generations (pre-captured)")
            for q, a in SAMPLES:
                gr.Markdown(f"**Q:** {q}\n\n**A:** {a}")
    gr.Markdown(COST_MD)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)),
                show_api=False, ssr_mode=False)