""" 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|>/, 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|>/, 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)