The Mathematician V2 — Beta
This is the beta release of The Mathematician V2, a pretrained mathematical language model.
This is a raw pretrained base model with no SFT, GRPO, or RL applied.
The inference script below automatically downloads the checkpoint from:
AlgoDriveAI/AIMO_EP_25/checkpoints_optimized2/latest_checkpoint.pt
Run in Google Colab
Copy and paste the following install command into a Colab cell and run it:
!pip install -q torch transformers gradio huggingface_hub sentencepiece
Then copy and paste the full inference code below into the next Colab cell and run it.
The checkpoint is downloaded automatically from Hugging Face. You do not need to manually upload the .pt file or edit a /content/checkpoints/... path.
Full Inference Code
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
THE MATHEMATICIAN V2 — BASE (PRETRAINED) INFERENCE
===============================================================================
Gradio + CLI inference for the *pretrained base* checkpoint. No GRPO, no SFT,
no reward machinery — this is raw next-token completion from the pretrain run.
ARCHITECTURE (older V2 MLA variant — matches the training script exactly):
• q_proj : single full projection, split at runtime into [q_nope | q_rope]
(NOT the canonical q_a_proj / q_a_norm / q_b_proj compression)
• k_rope_proj : PER-HEAD (d_model → n_heads * qk_rope_head_dim)
(NOT the shared single-head broadcast used in later versions)
• k_up / v_up : separate up-projections off the latent
(NOT a fused kv_b_proj)
• QK-norm : applied AFTER concat + AFTER RoPE, over the FULL head_dim
(NOT pre-RoPE on nope dims only)
• SwiGLU : fused gate_up_proj → chunk(2)
• lm_head weight-tied to embed
• RoPE over qk_rope_head_dim only, rotate_half convention, base 10000
If you point this at a checkpoint from a *newer* architecture it will fail
loudly on load_state_dict rather than silently mis-wire — that's intentional.
FEATURES:
• Incremental KV-cache decoding (prefill once, then 1 token/step)
• Streaming Gradio output + Stop button
• Architecture auto-detection from the checkpoint state dict (see
AUTO_ADOPT_CHECKPOINT_ARCH) — no more guessing whether it's 2048d or 2304d
• maj@k tab: batched self-consistency sampling with majority vote over
extracted answers
• CLI mode via RUN_CLI
CONFIG: all hardcoded at the top of this file. No argparse.
"""
import os
import re
import gc
import time
import random
import warnings
from collections import Counter
from dataclasses import dataclass
from typing import Optional, Tuple, List, Dict, Any
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ["HF_HUB_DISABLE_XET"] = "1"
warnings.filterwarnings("ignore", message=r".*UnsupportedFieldAttributeWarning.*")
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoTokenizer
from huggingface_hub import hf_hub_download
# =============================================================================
# HARDCODED CONFIGURATION
# =============================================================================
# ── Checkpoint ───────────────────────────────────────────────────────────────
# Download the beta checkpoint directly from Hugging Face.
CHECKPOINT_REPO = "AlgoDriveAI/AIMO_EP_25"
CHECKPOINT_SUBFOLDER = "checkpoints_optimized2"
CHECKPOINT_FILE = "latest_checkpoint.pt"
print(f"📥 Downloading checkpoint from Hugging Face: {CHECKPOINT_REPO}/{CHECKPOINT_SUBFOLDER}/{CHECKPOINT_FILE}")
CHECKPOINT_PATH = hf_hub_download(
repo_id=CHECKPOINT_REPO,
filename=f"{CHECKPOINT_SUBFOLDER}/{CHECKPOINT_FILE}",
)
# The rest of the inference code expects a directory + filename.
CHECKPOINT_DIR = os.path.dirname(CHECKPOINT_PATH)
print(f" ✅ Checkpoint ready: {CHECKPOINT_PATH}")
# ── Run mode ─────────────────────────────────────────────────────────────────
RUN_CLI = False # True → terminal REPL instead of Gradio
GRADIO_SHARE = True
GRADIO_PORT = 7860
# ── Architecture (must match the pretrain run) ───────────────────────────────
D_MODEL = 2048
N_LAYERS = 36
N_HEADS = 16
KV_LORA_RANK = 512 # MLA latent compression rank
QK_ROPE_HEAD_DIM = 64 # decoupled RoPE dims per head
FF_MULT = 3.5 # SwiGLU multiplier
CONTEXT_LEN = 4096
QK_NORM = True
ROPE_BASE = 10000.0
# Derived (recomputed after auto-detect)
HEAD_DIM = D_MODEL // N_HEADS
QK_NOPE_HEAD_DIM = HEAD_DIM - QK_ROPE_HEAD_DIM
# If True, read the true architecture out of the checkpoint's state dict and
# override the constants above before building the model. Every dim in this
# architecture is recoverable from tensor shapes, so this is exact, not a guess.
AUTO_ADOPT_CHECKPOINT_ARCH = True
# ── Tokenizer ────────────────────────────────────────────────────────────────
VOCAB_NAME = "mistralai/Mistral-7B-v0.1"
DOC_EOS_TOKEN = "<|endoftext|>"
VOCAB_PAD_MULTIPLE = 64
# ── Generation defaults ──────────────────────────────────────────────────────
DEFAULT_MAX_NEW_TOKENS = 512
DEFAULT_TEMPERATURE = 0.8
DEFAULT_TOP_K = 50
DEFAULT_TOP_P = 0.95
DEFAULT_REP_PENALTY = 1.05
DEFAULT_MAJ_K = 8 # samples for the maj@k tab
# Prompt templates offered in the UI. The base model saw `[SOLUTION]` section
# headers throughout the synthetic pretrain corpus, so that template is usually
# the strongest completion trigger. "Raw" sends your text untouched.
PROMPT_TEMPLATES = {
"Raw (no template)": "{p}",
"[SOLUTION] header": "{p}\n[SOLUTION]\n",
"Solve step by step": "{p} Solve step by step:\n",
}
DEFAULT_TEMPLATE = "[SOLUTION] header"
# =============================================================================
# RMSNorm — fused if available, custom fallback otherwise
# =============================================================================
try:
from torch.nn import RMSNorm
_rmsnorm_source = "nn.RMSNorm (fused)"
except ImportError:
class RMSNorm(nn.Module):
__constants__ = ["eps"]
def __init__(self, normalized_shape, eps: float = 1e-6, **kwargs):
super().__init__()
if isinstance(normalized_shape, int):
normalized_shape = (normalized_shape,)
self.eps = eps
self.weight = nn.Parameter(torch.ones(normalized_shape))
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.weight * (
x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
).to(x.dtype)
_rmsnorm_source = "custom fallback"
# =============================================================================
# ROTARY EMBEDDINGS — precomputed, rotate_half convention
# =============================================================================
class RotaryEmbedding(nn.Module):
def __init__(self, dim: int, base: float = ROPE_BASE, max_seq_len: int = 8192):
super().__init__()
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
t = torch.arange(max_seq_len, dtype=torch.float32)
freqs = torch.outer(t, inv_freq)
self.register_buffer("cos_cached", freqs.cos().repeat(1, 2), persistent=False)
self.register_buffer("sin_cached", freqs.sin().repeat(1, 2), persistent=False)
def forward(self, seq_len: int, dtype: torch.dtype):
return self.cos_cached[:seq_len].to(dtype), self.sin_cached[:seq_len].to(dtype)
def _rotate_half(x: torch.Tensor) -> torch.Tensor:
half = x.shape[-1] // 2
return torch.cat([-x[..., half:], x[..., :half]], dim=-1)
def apply_rotary_emb(q, k, cos, sin):
cos = cos.unsqueeze(0).unsqueeze(0)
sin = sin.unsqueeze(0).unsqueeze(0)
return (q * cos) + (_rotate_half(q) * sin), (k * cos) + (_rotate_half(k) * sin)
# =============================================================================
# MLA — verbatim from the training script (per-head k_rope, split q_proj)
# =============================================================================
class MLA(nn.Module):
def __init__(self, d_model, n_heads, kv_lora_rank, qk_rope_head_dim, rope,
qk_norm=False, attn_dropout=0.0):
super().__init__()
assert d_model % n_heads == 0
self.n_heads = n_heads
self.head_dim = d_model // n_heads
self.qk_rope_head_dim = qk_rope_head_dim
self.qk_nope_head_dim = self.head_dim - qk_rope_head_dim
self.kv_lora_rank = kv_lora_rank
self.attn_drop = attn_dropout
self.rope = rope
self.q_proj = nn.Linear(d_model, n_heads * self.head_dim, bias=False)
self.kv_down = nn.Linear(d_model, kv_lora_rank, bias=False)
self.kv_norm = RMSNorm(kv_lora_rank)
self.k_up = nn.Linear(kv_lora_rank, n_heads * self.qk_nope_head_dim, bias=False)
self.v_up = nn.Linear(kv_lora_rank, n_heads * self.head_dim, bias=False)
self.k_rope_proj = nn.Linear(d_model, n_heads * qk_rope_head_dim, bias=False)
self.o_proj = nn.Linear(d_model, d_model, bias=False)
self.o_proj._is_residual = True
self.qk_norm = qk_norm
if qk_norm:
self.q_norm = RMSNorm(self.head_dim)
self.k_norm = RMSNorm(self.head_dim)
def forward(self, x, past_kv=None, use_cache=False):
B, T, _ = x.shape
q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim)
q_nope = q[..., :self.qk_nope_head_dim]
q_rope = q[..., self.qk_nope_head_dim:]
c = self.kv_norm(self.kv_down(x))
k_nope = self.k_up(c).view(B, T, self.n_heads, self.qk_nope_head_dim)
v = self.v_up(c).view(B, T, self.n_heads, self.head_dim)
k_rope = self.k_rope_proj(x).view(B, T, self.n_heads, self.qk_rope_head_dim)
pos_offset = 0
if past_kv is not None:
pos_offset = past_kv[0].shape[2]
cos, sin = self.rope(pos_offset + T, x.dtype)
cos = cos[pos_offset:pos_offset + T]
sin = sin[pos_offset:pos_offset + T]
q_rope = q_rope.transpose(1, 2)
k_rope = k_rope.transpose(1, 2)
q_rope, k_rope = apply_rotary_emb(q_rope, k_rope, cos, sin)
q_rope = q_rope.transpose(1, 2)
k_rope = k_rope.transpose(1, 2)
q = torch.cat([q_nope, q_rope], dim=-1).transpose(1, 2).contiguous()
k = torch.cat([k_nope, k_rope], dim=-1).transpose(1, 2).contiguous()
v = v.transpose(1, 2).contiguous()
# QK-norm AFTER assembly, over full head_dim — matches training.
if self.qk_norm:
q = self.q_norm(q)
k = self.k_norm(k)
present_kv = None
if use_cache:
if past_kv is not None:
k = torch.cat([past_kv[0], k], dim=2)
v = torch.cat([past_kv[1], v], dim=2)
present_kv = (k, v)
drop_p = self.attn_drop if self.training else 0.0
is_causal = (T > 1) and (past_kv is None)
out = F.scaled_dot_product_attention(q, k, v, dropout_p=drop_p, is_causal=is_causal)
return self.o_proj(out.transpose(1, 2).reshape(B, T, -1)), present_kv
class SwiGLU(nn.Module):
def __init__(self, d_model, hidden_mult=3.5):
super().__init__()
inner = int(hidden_mult * d_model)
self.gate_up_proj = nn.Linear(d_model, 2 * inner, bias=False)
self.down_proj = nn.Linear(inner, d_model, bias=False)
self.down_proj._is_residual = True
def forward(self, x):
gate, up = self.gate_up_proj(x).chunk(2, dim=-1)
return self.down_proj(F.silu(gate) * up)
class Block(nn.Module):
def __init__(self, d_model, n_heads, kv_lora_rank, qk_rope_head_dim, rope,
ff_hidden_mult=3.5, qk_norm=False, attn_dropout=0.0, resid_dropout=0.0):
super().__init__()
self.ln_attn = RMSNorm(d_model)
self.ln_ff = RMSNorm(d_model)
self.attn = MLA(d_model, n_heads, kv_lora_rank, qk_rope_head_dim,
rope, qk_norm=qk_norm, attn_dropout=attn_dropout)
self.ff = SwiGLU(d_model, hidden_mult=ff_hidden_mult)
self.resid_drop = nn.Dropout(resid_dropout) if resid_dropout > 0 else nn.Identity()
def forward(self, x, past_kv=None, use_cache=False):
attn_out, present_kv = self.attn(self.ln_attn(x), past_kv=past_kv, use_cache=use_cache)
x = x + self.resid_drop(attn_out)
x = x + self.resid_drop(self.ff(self.ln_ff(x)))
return x, present_kv
@dataclass
class ModelConfig:
vocab_size: int
d_model: int
n_layers: int
n_heads: int
kv_lora_rank: int
qk_rope_head_dim: int
ff_hidden_mult: float
qk_norm: bool
max_seq_len: int = 8192
attn_dropout: float = 0.0
resid_dropout: float = 0.0
emb_dropout: float = 0.0
@property
def head_dim(self): return self.d_model // self.n_heads
@property
def qk_nope_head_dim(self): return self.head_dim - self.qk_rope_head_dim
class DenseLLM(nn.Module):
def __init__(self, cfg: ModelConfig):
super().__init__()
self.cfg = cfg
self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
self.emb_drop = nn.Identity()
self.rope = RotaryEmbedding(cfg.qk_rope_head_dim, max_seq_len=cfg.max_seq_len)
self.blocks = nn.ModuleList([
Block(cfg.d_model, cfg.n_heads, cfg.kv_lora_rank, cfg.qk_rope_head_dim,
self.rope, ff_hidden_mult=cfg.ff_hidden_mult, qk_norm=cfg.qk_norm,
attn_dropout=cfg.attn_dropout, resid_dropout=cfg.resid_dropout)
for _ in range(cfg.n_layers)
])
self.ln_f = RMSNorm(cfg.d_model)
self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
self.lm_head.weight = self.embed.weight
def forward(self, idx, past_kvs=None, use_cache=False):
x = self.emb_drop(self.embed(idx))
present_kvs = [] if use_cache else None
for i, blk in enumerate(self.blocks):
layer_past = past_kvs[i] if past_kvs is not None else None
x, present_kv = blk(x, past_kv=layer_past, use_cache=use_cache)
if use_cache:
present_kvs.append(present_kv)
logits = self.lm_head(self.ln_f(x))
if use_cache:
return logits, present_kvs
return logits, None
# =============================================================================
# TOKENIZER
# =============================================================================
def get_tokenizer(vocab_name, doc_eos_token, pad_to_multiple, checkpoint_dir):
if os.path.exists(os.path.join(checkpoint_dir, "tokenizer_config.json")):
print(f"📚 Tokenizer from checkpoint dir: {checkpoint_dir}")
tok = AutoTokenizer.from_pretrained(checkpoint_dir, use_fast=True)
else:
print(f"📚 Tokenizer from HuggingFace: {vocab_name}")
tok = AutoTokenizer.from_pretrained(vocab_name, use_fast=True)
if doc_eos_token not in tok.get_vocab():
tok.add_special_tokens({"additional_special_tokens": [doc_eos_token]})
if tok.pad_token is None:
tok.pad_token = doc_eos_token
tok.pad_token_id = tok.convert_tokens_to_ids(doc_eos_token)
if pad_to_multiple and (len(tok) % pad_to_multiple != 0):
n = pad_to_multiple - (len(tok) % pad_to_multiple)
tok.add_tokens([f"<|dummy_{i}|>" for i in range(n)], special_tokens=False)
tok.doc_eos_token = doc_eos_token
tok.doc_eos_token_id = tok.convert_tokens_to_ids(doc_eos_token)
if tok.pad_token_id is None:
tok.pad_token = doc_eos_token
tok.pad_token_id = tok.doc_eos_token_id
tok.model_max_length = int(1e9)
print(f" ✅ Vocab: {len(tok):,} | doc EOS '{doc_eos_token}' = id {tok.doc_eos_token_id}")
return tok
# =============================================================================
# ARCHITECTURE AUTO-DETECTION FROM STATE DICT
# =============================================================================
def infer_arch_from_state_dict(sd: Dict[str, torch.Tensor]) -> Dict[str, Any]:
"""
Every architectural dim in this MLA variant is recoverable from shapes:
embed.weight → (vocab, d_model)
blocks.N.* → n_layers (max N + 1)
blocks.0.attn.q_norm.weight → (head_dim,) [only if qk_norm]
blocks.0.attn.kv_down.weight → (kv_lora_rank, d_model)
blocks.0.attn.k_rope_proj.weight→ (n_heads * qk_rope, d_model)
blocks.0.attn.k_up.weight → (n_heads * qk_nope, kv_lora_rank)
blocks.0.ff.gate_up_proj.weight → (2 * inner, d_model)
Without qk_norm, head_dim is not directly observable, so n_heads falls back
to the hardcoded constant and qk_rope is derived from it.
"""
out: Dict[str, Any] = {}
vocab, d_model = sd["embed.weight"].shape
out["vocab_size"] = int(vocab)
out["d_model"] = int(d_model)
layer_ids = set()
for k in sd:
m = re.match(r"blocks\.(\d+)\.", k)
if m:
layer_ids.add(int(m.group(1)))
out["n_layers"] = max(layer_ids) + 1 if layer_ids else N_LAYERS
out["qk_norm"] = "blocks.0.attn.q_norm.weight" in sd
if out["qk_norm"]:
head_dim = int(sd["blocks.0.attn.q_norm.weight"].shape[0])
out["n_heads"] = out["d_model"] // head_dim
else:
out["n_heads"] = N_HEADS
head_dim = out["d_model"] // out["n_heads"]
out["head_dim"] = head_dim
out["kv_lora_rank"] = int(sd["blocks.0.attn.kv_down.weight"].shape[0])
out["qk_rope_head_dim"] = int(sd["blocks.0.attn.k_rope_proj.weight"].shape[0]) // out["n_heads"]
out["qk_nope_head_dim"] = int(sd["blocks.0.attn.k_up.weight"].shape[0]) // out["n_heads"]
inner = int(sd["blocks.0.ff.gate_up_proj.weight"].shape[0]) // 2
out["ff_inner"] = inner
out["ff_mult"] = inner / out["d_model"]
return out
def _print_arch_table(a: Dict[str, Any], hardcoded: Dict[str, Any]):
rows = [
("d_model", a["d_model"], hardcoded["d_model"]),
("n_layers", a["n_layers"], hardcoded["n_layers"]),
("n_heads", a["n_heads"], hardcoded["n_heads"]),
("head_dim", a["head_dim"], hardcoded["head_dim"]),
("qk_rope_head_dim", a["qk_rope_head_dim"], hardcoded["qk_rope_head_dim"]),
("qk_nope_head_dim", a["qk_nope_head_dim"], hardcoded["qk_nope_head_dim"]),
("kv_lora_rank", a["kv_lora_rank"], hardcoded["kv_lora_rank"]),
("ff_mult", round(a["ff_mult"], 4), hardcoded["ff_mult"]),
("qk_norm", a["qk_norm"], hardcoded["qk_norm"]),
("vocab_size", a["vocab_size"], hardcoded["vocab_size"]),
]
print(f" {'field':<20} {'checkpoint':>14} {'hardcoded':>14}")
print(f" {'-'*20} {'-'*14} {'-'*14}")
n_mismatch = 0
for name, ck, hc in rows:
marker = "" if ck == hc else " ← MISMATCH"
if ck != hc:
n_mismatch += 1
print(f" {name:<20} {str(ck):>14} {str(hc):>14}{marker}")
if n_mismatch:
print(f" ⚠️ {n_mismatch} field(s) differ from the hardcoded constants.")
else:
print(f" ✅ Checkpoint matches hardcoded constants exactly.")
# =============================================================================
# MODEL LOADING
# =============================================================================
def load_model(checkpoint_dir, checkpoint_file, tokenizer, device, dtype):
global D_MODEL, N_LAYERS, N_HEADS, KV_LORA_RANK, QK_ROPE_HEAD_DIM
global FF_MULT, QK_NORM, HEAD_DIM, QK_NOPE_HEAD_DIM
ckpt_path = os.path.join(checkpoint_dir, checkpoint_file)
print(f"\n{'='*68}\n🔧 Loading base model\n{'='*68}")
if not os.path.exists(ckpt_path):
avail = []
if os.path.isdir(checkpoint_dir):
avail = sorted(f for f in os.listdir(checkpoint_dir) if f.endswith(".pt"))
raise FileNotFoundError(
f"Checkpoint not found: {ckpt_path}\n"
f" .pt files in {checkpoint_dir}: {avail if avail else '(none)'}"
)
print(f" File: {ckpt_path}")
ckpt = torch.load(ckpt_path, map_location="cpu")
state_dict = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt
state_dict = {k.replace("_orig_mod.", ""): v for k, v in state_dict.items()}
for key in ("step", "global_step", "epoch", "best_val_loss", "val_loss", "accuracy"):
if isinstance(ckpt, dict) and key in ckpt:
print(f" {key}: {ckpt[key]}")
# ── Verify / adopt architecture from the actual tensors
print(f"\n 🔍 Architecture read from checkpoint tensors:")
arch = infer_arch_from_state_dict(state_dict)
_print_arch_table(arch, {
"d_model": D_MODEL, "n_layers": N_LAYERS, "n_heads": N_HEADS,
"head_dim": HEAD_DIM, "qk_rope_head_dim": QK_ROPE_HEAD_DIM,
"qk_nope_head_dim": QK_NOPE_HEAD_DIM, "kv_lora_rank": KV_LORA_RANK,
"ff_mult": FF_MULT, "qk_norm": QK_NORM, "vocab_size": len(tokenizer),
})
if AUTO_ADOPT_CHECKPOINT_ARCH:
D_MODEL = arch["d_model"]
N_LAYERS = arch["n_layers"]
N_HEADS = arch["n_heads"]
KV_LORA_RANK = arch["kv_lora_rank"]
QK_ROPE_HEAD_DIM = arch["qk_rope_head_dim"]
FF_MULT = arch["ff_mult"]
QK_NORM = arch["qk_norm"]
HEAD_DIM = arch["head_dim"]
QK_NOPE_HEAD_DIM = arch["qk_nope_head_dim"]
print(f" ↪ AUTO_ADOPT_CHECKPOINT_ARCH=True — building from checkpoint dims.")
vocab_size = arch["vocab_size"]
if vocab_size != len(tokenizer):
print(f" ⚠️ Checkpoint vocab {vocab_size:,} != tokenizer {len(tokenizer):,}. "
f"Building embedding at checkpoint size; token ids above "
f"{vocab_size-1} will be out of range.")
mcfg = ModelConfig(
vocab_size=vocab_size, d_model=D_MODEL, n_layers=N_LAYERS, n_heads=N_HEADS,
kv_lora_rank=KV_LORA_RANK, qk_rope_head_dim=QK_ROPE_HEAD_DIM,
ff_hidden_mult=FF_MULT, qk_norm=QK_NORM, max_seq_len=CONTEXT_LEN,
)
model = DenseLLM(mcfg)
missing, unexpected = model.load_state_dict(state_dict, strict=False)
if missing:
raise RuntimeError(f"{len(missing)} missing keys — architecture mismatch. "
f"First few: {missing[:6]}")
if unexpected:
print(f" ⚠️ {len(unexpected)} unexpected keys ignored: {unexpected[:6]}")
model = model.to(device=device, dtype=dtype).eval()
for p in model.parameters():
p.requires_grad_(False)
del ckpt, state_dict
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
total = sum(p.numel() for p in model.parameters())
param_str = f"{total/1e9:.2f}B" if total >= 1e9 else f"{total/1e6:.0f}M"
cache_d = N_HEADS * HEAD_DIM * 2 # this impl caches full k and v
print(f"\n ✅ {total:,} params ({param_str})")
print(f" ✅ {D_MODEL}d × {N_LAYERS}L × {N_HEADS}H | "
f"head={HEAD_DIM} = nope({QK_NOPE_HEAD_DIM}) + rope({QK_ROPE_HEAD_DIM})")
print(f" ✅ kv_lora_rank={KV_LORA_RANK} | ff_inner={int(FF_MULT*D_MODEL)} | "
f"qk_norm={QK_NORM}")
print(f" ✅ Decode cache: {cache_d}d/token/layer | RMSNorm: {_rmsnorm_source}")
print(f" ✅ Device: {device} | dtype: {dtype}")
print(f"{'='*68}\n")
return model, param_str
# =============================================================================
# SAMPLING HELPERS
# =============================================================================
def _apply_rep_penalty(logits: torch.Tensor, seen_ids: torch.Tensor, penalty: float):
"""In-place repetition penalty over already-emitted ids. logits: (B, V)."""
if penalty == 1.0:
return logits
for b in range(logits.shape[0]):
uniq = torch.unique(seen_ids[b])
s = logits[b, uniq]
logits[b, uniq] = torch.where(s > 0, s / penalty, s * penalty)
return logits
def _filter_and_sample(logits, temperature, top_k, top_p):
"""logits: (B, V) → next tokens (B, 1)."""
if temperature <= 0:
return torch.argmax(logits, dim=-1, keepdim=True)
logits = logits / max(temperature, 1e-8)
if top_k and top_k > 0:
v, _ = torch.topk(logits, min(int(top_k), logits.size(-1)))
logits = logits.masked_fill(logits < v[:, [-1]], -float("inf"))
if top_p and top_p < 1.0:
sorted_logits, sorted_idx = torch.sort(logits, descending=True)
cum = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
remove = cum > top_p
remove[:, 1:] = remove[:, :-1].clone()
remove[:, 0] = False
mask = torch.zeros_like(logits, dtype=torch.bool).scatter_(1, sorted_idx, remove)
logits = logits.masked_fill(mask, -float("inf"))
probs = F.softmax(logits, dim=-1)
return torch.multinomial(probs, num_samples=1)
# =============================================================================
# STREAMING GENERATION (batch of 1, KV cache)
# =============================================================================
@torch.inference_mode()
def generate_streaming(model, tokenizer, prompt, device, dtype,
max_new_tokens=512, temperature=0.8, top_k=50, top_p=0.95,
repetition_penalty=1.05, seed=-1, yield_every=4):
"""Yields (text_so_far, n_tokens, ended_with_eos, tok_per_sec, hit_limit)."""
if seed is not None and seed >= 0:
torch.manual_seed(int(seed))
random.seed(int(seed))
eos_id = tokenizer.doc_eos_token_id
ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to(device)
if ids.shape[1] >= CONTEXT_LEN:
ids = ids[:, -(CONTEXT_LEN - 8):]
prompt_len = ids.shape[1]
t0 = time.time()
logits, past_kvs = model(ids, use_cache=True)
next_logits = logits[:, -1, :].float()
gen_ids: List[int] = []
ended_with_eos = False
hit_limit = False
for i in range(int(max_new_tokens)):
if prompt_len + len(gen_ids) >= CONTEXT_LEN:
hit_limit = True
break
seen = torch.cat([ids, torch.tensor([gen_ids], dtype=torch.long, device=device)], dim=1) \
if gen_ids else ids
next_logits = _apply_rep_penalty(next_logits, seen, float(repetition_penalty))
next_token = _filter_and_sample(next_logits, temperature, top_k, top_p)
tid = int(next_token.item())
if tid == eos_id:
ended_with_eos = True
break
gen_ids.append(tid)
if (i % yield_every == 0) or (i == int(max_new_tokens) - 1):
text = tokenizer.decode(gen_ids, skip_special_tokens=True)
yield text, len(gen_ids), False, len(gen_ids) / max(time.time() - t0, 1e-9), False
logits, past_kvs = model(next_token, past_kvs=past_kvs, use_cache=True)
next_logits = logits[:, -1, :].float()
else:
hit_limit = True
del past_kvs
text = tokenizer.decode(gen_ids, skip_special_tokens=True)
yield text, len(gen_ids), ended_with_eos, len(gen_ids) / max(time.time() - t0, 1e-9), hit_limit
# =============================================================================
# BATCHED GENERATION (maj@k self-consistency)
# =============================================================================
@torch.inference_mode()
def generate_batch(model, tokenizer, prompt, device, dtype, k=8,
max_new_tokens=512, temperature=0.8, top_k=50, top_p=0.95,
repetition_penalty=1.05, seed=-1):
"""Generate k independent samples in parallel. Returns list of dicts."""
if seed is not None and seed >= 0:
torch.manual_seed(int(seed))
eos_id = tokenizer.doc_eos_token_id
prompt_ids = tokenizer(prompt, return_tensors="pt")["input_ids"][0].tolist()
if len(prompt_ids) >= CONTEXT_LEN:
prompt_ids = prompt_ids[-(CONTEXT_LEN - 8):]
prompt_len = len(prompt_ids)
k = int(k)
ids = torch.tensor([prompt_ids] * k, dtype=torch.long, device=device)
finished = torch.zeros(k, dtype=torch.bool, device=device)
finished_at_eos = torch.zeros(k, dtype=torch.bool, device=device)
all_gen: List[List[int]] = [[] for _ in range(k)]
t0 = time.time()
logits, past_kvs = model(ids, use_cache=True)
next_logits = logits[:, -1, :].float()
for step in range(int(max_new_tokens)):
if bool(finished.all()) or (prompt_len + step) >= CONTEXT_LEN:
break
if repetition_penalty != 1.0 and step > 0:
seen = torch.tensor(
[prompt_ids + g + [eos_id] * (step - len(g)) for g in all_gen],
dtype=torch.long, device=device)
next_logits = _apply_rep_penalty(next_logits, seen, float(repetition_penalty))
next_tokens = _filter_and_sample(next_logits, temperature, top_k, top_p)
next_tokens[finished] = eos_id
for g in range(k):
if not finished[g]:
tid = int(next_tokens[g].item())
if tid != eos_id:
all_gen[g].append(tid)
just_done = (next_tokens.squeeze(-1) == eos_id) & ~finished
finished_at_eos |= just_done
finished |= just_done
logits, past_kvs = model(next_tokens, past_kvs=past_kvs, use_cache=True)
next_logits = logits[:, -1, :].float()
del past_kvs
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
elapsed = time.time() - t0
total_tok = sum(len(g) for g in all_gen)
results = []
for g in range(k):
text = tokenizer.decode(all_gen[g], skip_special_tokens=True)
results.append({
"idx": g,
"text": text,
"n_tokens": len(all_gen[g]),
"ended_with_eos": bool(finished_at_eos[g].item()),
"answer": extract_answer(text),
})
return results, elapsed, total_tok / max(elapsed, 1e-9)
# =============================================================================
# ANSWER EXTRACTION (same logic as the training/eval scripts)
# =============================================================================
def extract_answer(text: str) -> Optional[str]:
if not text:
return None
boxed = []
for m in re.finditer(r'\\boxed\{', text):
start = m.end()
depth, i = 1, start
while i < len(text) and depth > 0:
if text[i] == '{': depth += 1
elif text[i] == '}': depth -= 1
i += 1
if depth == 0:
boxed.append(text[start:i-1].strip())
if boxed:
return boxed[-1]
m = re.search(r'(?:the\s+)?answer\s+is\s*[:\s]*([+-]?\d+(?:\.\d+)?(?:/\d+)?)', text, re.IGNORECASE)
if m: return m.group(1).strip()
m = re.search(r'answer\s*[:=]\s*([+-]?\d+(?:\.\d+)?(?:/\d+)?)', text, re.IGNORECASE)
if m: return m.group(1).strip()
eqs = re.findall(r'=\s*([+-]?\d+(?:\.\d+)?(?:/\d+)?)\s*[.\s]*$', text, re.MULTILINE)
if eqs: return eqs[-1].strip()
nums = re.findall(r'(?<![.\w])([+-]?\d+(?:\.\d+)?(?:/\d+)?)(?![.\w])', text)
if nums: return nums[-1].strip()
return None
def normalize_answer(answer: Optional[str]) -> Optional[str]:
if answer is None: return None
a = answer.strip()
a = re.sub(r'\\text\{([^}]*)\}', r'\1', a)
a = re.sub(r'\\mathrm\{([^}]*)\}', r'\1', a)
a = re.sub(r'[\$\\,\\;\\!\\>\\:]', '', a)
a = a.replace(' ', '').replace(',', '')
a = re.sub(r'%$', '', a)
a = re.sub(r'\^\{?circ\}?$', '', a)
a = re.sub(r'°$', '', a).strip()
try:
if '/' in a:
parts = a.split('/')
if len(parts) == 2:
val = float(parts[0]) / float(parts[1])
return str(int(val)) if val == int(val) else f"{val:.10f}".rstrip('0').rstrip('.')
val = float(a)
return str(int(val)) if val == int(val) else f"{val:.10f}".rstrip('0').rstrip('.')
except (ValueError, ZeroDivisionError, OverflowError):
pass
return a.lower().strip()
# =============================================================================
# GRADIO INTERFACE
# =============================================================================
def create_gradio_interface(model, tokenizer, device, dtype, param_str):
import gradio as gr
custom_css = """
.gradio-container { max-width: 1080px !important; margin: auto !important; }
.mono-out textarea {
font-family: 'JetBrains Mono','SF Mono','Courier New',monospace !important;
font-size: 13.5px !important; line-height: 1.65 !important;
}
.metric-strip {
font-family: 'JetBrains Mono','SF Mono',monospace; font-size: 13px;
padding: 6px 2px; opacity: 0.9;
}
"""
ff_inner = int(FF_MULT * D_MODEL)
def _metrics_md(answer, n_tok, eos, tps, hit_limit):
ans = f"`{answer}`" if answer else "—"
stop = "EOS ✅" if eos else ("length cap ⛔" if hit_limit else "…running")
return (f"<div class='metric-strip'>"
f"<b>answer:</b> {ans} · "
f"<b>tokens:</b> {n_tok} · "
f"<b>stop:</b> {stop} · "
f"<b>speed:</b> {tps:.1f} tok/s</div>")
def build_prompt(raw, template_name):
tmpl = PROMPT_TEMPLATES.get(template_name, "{p}")
return tmpl.format(p=raw.strip())
# ── Tab 1: single completion, streaming
def run_generate(raw_prompt, template_name, max_new, temp, tk, tp, rep, seed, echo):
if not raw_prompt.strip():
yield "Enter a prompt first.", _metrics_md(None, 0, False, 0.0, False)
return
prompt = build_prompt(raw_prompt, template_name)
head = prompt if echo else ""
last = ("", 0, False, 0.0, False)
for text, n_tok, eos, tps, hit in generate_streaming(
model, tokenizer, prompt, device, dtype,
max_new_tokens=max_new, temperature=temp, top_k=tk, top_p=tp,
repetition_penalty=rep, seed=seed):
last = (text, n_tok, eos, tps, hit)
yield head + text, _metrics_md(None, n_tok, eos, tps, hit)
text, n_tok, eos, tps, hit = last
yield head + text, _metrics_md(extract_answer(text), n_tok, eos, tps, hit)
# ── Tab 2: maj@k self-consistency
def run_majk(raw_prompt, template_name, k, max_new, temp, tk, tp, rep, seed):
if not raw_prompt.strip():
return "Enter a prompt first.", ""
prompt = build_prompt(raw_prompt, template_name)
results, elapsed, tps = generate_batch(
model, tokenizer, prompt, device, dtype, k=k, max_new_tokens=max_new,
temperature=temp, top_k=tk, top_p=tp, repetition_penalty=rep, seed=seed)
norm_counts = Counter()
display_for_norm = {}
for r in results:
n = normalize_answer(r["answer"])
if n is None:
continue
norm_counts[n] += 1
display_for_norm.setdefault(n, r["answer"])
lines = [f"### maj@{int(k)} vote", ""]
if norm_counts:
top_norm, votes = norm_counts.most_common(1)[0]
lines.append(f"**Consensus: `{display_for_norm[top_norm]}`** "
f"({votes}/{int(k)} samples, {votes/int(k)*100:.0f}%)")
lines.append("")
lines.append("| answer | votes |")
lines.append("|---|---|")
for n, c in norm_counts.most_common():
lines.append(f"| `{display_for_norm[n]}` | {c} |")
else:
lines.append("No answer could be extracted from any sample.")
n_eos = sum(1 for r in results if r["ended_with_eos"])
avg_len = sum(r["n_tokens"] for r in results) / max(len(results), 1)
lines += ["", f"_{n_eos}/{int(k)} terminated on EOS · avg {avg_len:.0f} tokens · "
f"{elapsed:.1f}s · {tps:.0f} tok/s aggregate_"]
blocks = []
for r in results:
flag = "EOS" if r["ended_with_eos"] else "cut"
blocks.append(f"{'='*64}\n# sample {r['idx']+1}/{int(k)} · {r['n_tokens']} tok · "
f"{flag} · answer: {r['answer']}\n{'='*64}\n{r['text']}\n")
return "\n".join(lines), "\n".join(blocks)
def set_preset(name):
return {
"Precise": (0.3, 40, 0.90, 1.05),
"Balanced": (0.8, 50, 0.95, 1.05),
"Creative": (1.0, 0, 0.98, 1.10),
"Greedy": (0.0, 0, 1.00, 1.00),
}.get(name, (0.8, 50, 0.95, 1.05))
with gr.Blocks(title="THE MATHEMATICIAN V2 — base",
css=custom_css,
theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="blue")) as demo:
gr.Markdown(
f"""
# 🧠 THE MATHEMATICIAN V2 — base ({param_str})
Pretrained base checkpoint, no instruction tuning and no RL. It completes
text; it does not follow instructions. Give it a problem statement and let
the template supply the continuation cue.
`{D_MODEL}d × {N_LAYERS}L × {N_HEADS}H` · MLA rank-`{KV_LORA_RANK}` ·
head `{HEAD_DIM}` = nope(`{QK_NOPE_HEAD_DIM}`) + rope(`{QK_ROPE_HEAD_DIM}`) ·
ff_inner `{ff_inner}` · ctx `{CONTEXT_LEN}` · `{device}`/`{str(dtype).replace('torch.','')}`
---
"""
)
with gr.Row():
with gr.Column(scale=3):
prompt_input = gr.Textbox(
label="Prompt", lines=4, max_lines=12,
placeholder="Find the remainder when 2^100 is divided by 7.",
)
template_dd = gr.Dropdown(
choices=list(PROMPT_TEMPLATES.keys()), value=DEFAULT_TEMPLATE,
label="Prompt template",
)
with gr.Column(scale=2):
preset_dd = gr.Dropdown(
choices=["Balanced", "Precise", "Creative", "Greedy"],
value="Balanced", label="Preset",
)
temp_sl = gr.Slider(0.0, 2.0, DEFAULT_TEMPERATURE, step=0.05, label="temperature")
topk_sl = gr.Slider(0, 200, DEFAULT_TOP_K, step=5, label="top_k (0 = off)")
topp_sl = gr.Slider(0.0, 1.0, DEFAULT_TOP_P, step=0.01, label="top_p")
rep_sl = gr.Slider(1.0, 1.5, DEFAULT_REP_PENALTY, step=0.01,
label="repetition penalty")
maxtok_sl = gr.Slider(16, 2048, DEFAULT_MAX_NEW_TOKENS, step=16,
label="max new tokens")
seed_num = gr.Number(value=-1, precision=0, label="seed (-1 = random)")
with gr.Tabs():
with gr.Tab("Completion"):
echo_cb = gr.Checkbox(value=True, label="Echo prompt in output")
out_box = gr.Textbox(label="Output", lines=20, max_lines=40,
interactive=False, elem_classes=["mono-out"])
metrics_md = gr.Markdown(_metrics_md(None, 0, False, 0.0, False))
with gr.Row():
gen_btn = gr.Button("Generate", variant="primary", scale=2)
stop_btn = gr.Button("Stop", variant="stop", scale=1)
clear_btn = gr.Button("Clear", scale=1)
gr.Examples(
examples=[
"Find the remainder when 2^100 is divided by 7.",
"In triangle ABC, AB = 13, BC = 14, and CA = 15. Find the area of triangle ABC.",
"If x + y = 7 and xy = 10, find the value of x^3 + y^3.",
"How many ways can the letters in the word BANANA be arranged?",
"Find the sum of all positive divisors of 360.",
"Find the number of trailing zeros in 50!.",
"Let f(x) = x^2 - 4x + 3. Find the sum of all integers n such that f(f(n)) = 3.",
],
inputs=prompt_input,
)
with gr.Tab("maj@k (self-consistency)"):
gr.Markdown(
"Samples k completions in parallel from the same prompt and takes a "
"majority vote over extracted answers. Temperature must be > 0 or "
"every sample will be identical."
)
k_sl = gr.Slider(2, 32, DEFAULT_MAJ_K, step=1, label="k samples")
majk_btn = gr.Button("Run maj@k", variant="primary")
vote_md = gr.Markdown()
samples_box = gr.Textbox(label="All samples", lines=24, max_lines=60,
interactive=False, elem_classes=["mono-out"])
with gr.Accordion("Model card", open=False):
gr.Markdown(
f"""
| field | value |
|---|---|
| checkpoint | `{os.path.join(CHECKPOINT_DIR, CHECKPOINT_FILE)}` |
| parameters | {param_str} |
| d_model | {D_MODEL} |
| layers | {N_LAYERS} |
| heads | {N_HEADS} |
| head_dim | {HEAD_DIM} = nope({QK_NOPE_HEAD_DIM}) + rope({QK_ROPE_HEAD_DIM}) |
| kv_lora_rank | {KV_LORA_RANK} |
| ff_mult | {FF_MULT} (inner {ff_inner}) |
| qk_norm | {QK_NORM} (post-assembly, full head_dim) |
| k_rope_proj | per-head ({N_HEADS} × {QK_ROPE_HEAD_DIM}) |
| rope base | {ROPE_BASE:,.0f} |
| context | {CONTEXT_LEN:,} |
| tokenizer | {VOCAB_NAME} ({len(tokenizer):,} tokens) |
| RMSNorm | {_rmsnorm_source} |
**Notes on this variant.** `q_proj` is a single full projection split at
runtime rather than the canonical q-compression path, `k_rope_proj` is
per-head rather than shared-and-broadcast, and `k_up`/`v_up` are separate
rather than a fused `kv_b_proj`. Because K is materialised in full for
SDPA, decode caches `n_heads × head_dim × 2` per token per layer — the
latent cache saving isn't realised in this implementation, which is fine
at this scale but worth remembering when comparing memory numbers.
**Base-model behaviour.** Expect textbook-style continuations, section
headers, and no reliable `\\boxed{{}}` + EOS discipline — that comes from
the SFT/GRPO stages. Repetition penalty above ~1.1 tends to hurt math
output because formulas legitimately repeat tokens.
"""
)
gen_inputs = [prompt_input, template_dd, maxtok_sl, temp_sl, topk_sl,
topp_sl, rep_sl, seed_num, echo_cb]
gen_event = gen_btn.click(fn=run_generate, inputs=gen_inputs,
outputs=[out_box, metrics_md])
submit_event = prompt_input.submit(fn=run_generate, inputs=gen_inputs,
outputs=[out_box, metrics_md])
stop_btn.click(fn=None, inputs=None, outputs=None,
cancels=[gen_event, submit_event])
clear_btn.click(fn=lambda: ("", "", _metrics_md(None, 0, False, 0.0, False)),
outputs=[prompt_input, out_box, metrics_md])
preset_dd.change(fn=set_preset, inputs=preset_dd,
outputs=[temp_sl, topk_sl, topp_sl, rep_sl])
majk_btn.click(
fn=run_majk,
inputs=[prompt_input, template_dd, k_sl, maxtok_sl, temp_sl, topk_sl,
topp_sl, rep_sl, seed_num],
outputs=[vote_md, samples_box],
)
return demo
# =============================================================================
# CLI MODE
# =============================================================================
def run_cli(model, tokenizer, device, dtype, param_str):
print(f"\n{'='*68}\n💬 CLI mode — THE MATHEMATICIAN V2 base ({param_str})")
print(f" Template: {DEFAULT_TEMPLATE}")
print(f" Commands: /quit /temp <f> /max <int> /raw (toggle template)")
print(f"{'='*68}\n")
temp = DEFAULT_TEMPERATURE
max_new = DEFAULT_MAX_NEW_TOKENS
use_template = True
while True:
try:
raw = input("\n> ").strip()
except (EOFError, KeyboardInterrupt):
print("\nbye")
return
if not raw:
continue
if raw in ("/quit", "/exit", "/q"):
print("bye")
return
if raw.startswith("/temp"):
try:
temp = float(raw.split()[1]); print(f" temperature = {temp}")
except Exception:
print(" usage: /temp 0.7")
continue
if raw.startswith("/max"):
try:
max_new = int(raw.split()[1]); print(f" max_new_tokens = {max_new}")
except Exception:
print(" usage: /max 512")
continue
if raw == "/raw":
use_template = not use_template
print(f" template = {'ON (' + DEFAULT_TEMPLATE + ')' if use_template else 'OFF'}")
continue
prompt = PROMPT_TEMPLATES[DEFAULT_TEMPLATE].format(p=raw) if use_template else raw
print("-" * 68)
printed = 0
final = ("", 0, False, 0.0, False)
for text, n_tok, eos, tps, hit in generate_streaming(
model, tokenizer, prompt, device, dtype,
max_new_tokens=max_new, temperature=temp, top_k=DEFAULT_TOP_K,
top_p=DEFAULT_TOP_P, repetition_penalty=DEFAULT_REP_PENALTY):
print(text[printed:], end="", flush=True)
printed = len(text)
final = (text, n_tok, eos, tps, hit)
text, n_tok, eos, tps, hit = final
stop = "EOS" if eos else ("length cap" if hit else "?")
print(f"\n{'-'*68}")
print(f"answer={extract_answer(text)} | {n_tok} tok | stop={stop} | {tps:.1f} tok/s")
# =============================================================================
# MAIN
# =============================================================================
def main():
print(f"""
╔══════════════════════════════════════════════════════════════════════╗
║ 🧠 THE MATHEMATICIAN V2 — BASE (PRETRAINED) INFERENCE ║
║ older MLA variant · per-head k_rope · split q_proj · no RL ║
╚══════════════════════════════════════════════════════════════════════╝
""")
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cuda":
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
torch.set_float32_matmul_precision("high")
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
props = torch.cuda.get_device_properties(0)
print(f"🖥️ {props.name} ({props.total_memory/(1024**3):.1f}GB)")
else:
dtype = torch.float32
print("⚠️ CPU only — generation will be slow.")
tokenizer = get_tokenizer(VOCAB_NAME, DOC_EOS_TOKEN, VOCAB_PAD_MULTIPLE, CHECKPOINT_DIR)
model, param_str = load_model(CHECKPOINT_DIR, CHECKPOINT_FILE, tokenizer, device, dtype)
if RUN_CLI:
run_cli(model, tokenizer, device, dtype, param_str)
return
print(f"🚀 Launching Gradio on port {GRADIO_PORT} (share={GRADIO_SHARE})\n")
demo = create_gradio_interface(model, tokenizer, device, dtype, param_str)
demo.queue().launch(share=GRADIO_SHARE, server_port=GRADIO_PORT,
debug=False, show_error=True)
if __name__ == "__main__":
main()