Upload folder using huggingface_hub

README.md

@@ -0,0 +1,113 @@
+---
+license: apache-2.0
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+  - text-generation
+  - custom-code
+  - ml-engineering
+  - specialist
+  - metallum
+---
+
+## Metallum-1B
+
+A 1.0B-parameter **ML/LLM-engineering specialist**, pretrained **from random
+initialization on a single RTX 5090** — 16.0B tokens in ~13.5 days including
+gates and recoveries — then one bounded knowledge/format SFT stage. Not a
+general chatbot. By **Home Brewed Labs** (project codename MetaLLM).
+
+The engineering core of the project is measurement discipline: every training
+stage was gated by frozen batteries with exact checkpoint identity, replay
+accounting (≈1.04B replayed tokens from two hardware incidents, fully
+receipted), independent adversarial review of evaluation machinery, and a
+sealed never-trained-on final holdout (H9) reserved for a single one-time
+qualification.
+
+### Architecture
+26-layer decoder, d_model 1792, GQA 28Q/14KV, SwiGLU (d_ff 4864), RMSNorm
+pre-norm + QK-norm + z-loss, RoPE θ=500k with **NoPE every 4th layer**
+(length-generalizing retrieval heads), document-masked attention, 40k
+byte-level BPE with FIM sentinels, 2048 native context. Muon (2D) + AdamW,
+WSD schedule, curriculum with short-sequence early phases.
+
+### Loading
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+model = AutoModelForCausalLM.from_pretrained("HomeBrewedlabs/metallum-1b",
+                                             trust_remote_code=True)
+tok = AutoTokenizer.from_pretrained("HomeBrewedlabs/metallum-1b")
+```
+
+`trust_remote_code=True` is required: this is a custom architecture
+(QK-norm, NoPE interleave, document masking). **Do not** map to a stock
+Llama config — that silently drops QK-norm/NoPE and produces a different
+model than the one measured below. Stock GGUF/vLLM converters reject this
+architecture by design; a faithful llama.cpp port is on the roadmap.
+
+### Evaluation (champion s2b075, frozen batteries, repeated where noted)
+| Capability | Metric | Metallum-1B | Qwen3-1.7B | Qwen3-0.6B | SmolLM2-1.7B |
+|---|---|---|---|---|---|
+| In-domain ML-eng problems | ineval_v1, frozen 150-task exec suite | **7/150** | 7/150 | 3/150 | 0/150 |
+| In-domain code (internal suite) | code_v2 40-task, 2 exact repeats | **32/40** (disclosure: known training overlap) | — | — | — |
+| Long-context retrieval | RULER (90) / needles D1024+D2032 (18) | **86/90** and **18/18** | — | — | — |
+| Base LM efficiency | held-out ML-arXiv BPC | **0.7276** | — | — | — |
+| Final unseen generalization | H9 sealed holdout BPC (one-time run) | **1.3953** | — | — | — |
+
+All numbers were produced by frozen evaluators with sealed checkpoint
+identity; the ineval suite was used for model selection during development
+(selection-aware caveat), which is why the untouched H9 result is the number
+that matters for release claims.
+
+### Structured output: by decoding, not by hope
+Base weights score **0/60** on a strict whole-output format suite (JSON
+schema / tool-call / MCQ-letter), and three bounded SFT recipes plus a
+ReST-EM screen all failed to move that (all receipted). This model therefore
+ships with **decode-time structure enforcement**: `serve_metallum.py`
+(included) constrains generation with a token-level JSON grammar +
+schema-forced keys, and scores **60/60** on the same frozen suite. If you
+consume the raw weights without the wrapper, expect prose, not JSON.
+
+```bash
+python serve_metallum.py --port 8399   # localhost endpoint: /generate, modes free|json|tool|mcq
+```
+
+### Intended use
+ML/LLM-engineering assistant: PyTorch/training scaffolding, concept
+explanation, long-context retrieval over technical docs, structured-output
+endpoints via the wrapper. **Out of scope:** general chat, non-ML factual
+QA, general-purpose coding, safety-critical use, autonomous code execution.
+
+### Honest limitations
+- **neural_ops floor exception:** this checkpoint scores 4/5 on the internal
+  neural-ops category vs the 5/5 preregistered floor. The exception was
+  adjudicated explicitly by the owner; it is disclosed, not hidden.
+- **MCF (250q multiple-choice) is permanently non-promotable:** the teacher
+  generation prompts embedded real MCF items, so no MCF number from this
+  lineage is valid evidence. Do not report one.
+- **ineval is selection-aware:** 7/150 was measured on a suite used during
+  development. The sealed H9 result (**1.3953 BPC** on 400k post-cutoff
+  Wikipedia characters, single run, contamination-verified at 0.0099% shingle
+  containment) is the generalization claim.
+- Free generation makes local factual slips; verify specifics.
+- No preference/safety alignment; no RLHF.
+- Format behavior requires the constrained decoder (above).
+
+### Training data
+Pretraining: 19.399B effective tokens (8.473B unique): 53.6% code, 19.6%
+knowledge, 15.2% reasoning, 11.6% math. 1.24% (240.6M tokens) is
+provider-teacher-derived; full lineage map is maintained by the authors.
+SFT (12,496 rows): 40% locally self-hosted open-teacher MCQs
+(qwen2.5-coder:14b, Apache-2.0), 12% self-generated synthetic retrieval rows,
+48% provider-teacher-derived replay (disposition: owner rights review).
+
+**Attribution — Stack Exchange (`stack_ml_edu`, 2.973B pack tokens, 15.3% of
+the pack): CC BY-SA 4.0.** Attribution is owed to Stack Exchange contributors,
+and the share-alike obligation attaches to that portion of the training data.
+
+### License
+Apache-2.0 (model weights + code). See the attribution note above regarding
+training-data share-alike obligations.
+
+### Citation
+Technical report in preparation.

config.json

@@ -0,0 +1,34 @@
+{
+  "architectures": [
+    "MetaLLMForCausalLM"
+  ],
+  "attn_impl": "auto",
+  "auto_map": {
+    "AutoConfig": "configuration_metallm.MetaLLMConfig",
+    "AutoModelForCausalLM": "modeling_metallm.MetaLLMForCausalLM"
+  },
+  "bos_id": 1,
+  "bos_token_id": 1,
+  "d_ff": 4864,
+  "d_model": 1792,
+  "doc_mask": true,
+  "hidden_size": 1792,
+  "max_position_embeddings": 2048,
+  "max_seq_len": 2048,
+  "model_type": "metallm",
+  "n_heads": 28,
+  "n_kv_heads": 14,
+  "n_layers": 26,
+  "nope_every": 4,
+  "norm_eps": 1e-05,
+  "num_attention_heads": 28,
+  "num_hidden_layers": 26,
+  "num_key_value_heads": 14,
+  "qk_norm": true,
+  "rope_fp32": true,
+  "rope_theta": 500000.0,
+  "tie_embeddings": true,
+  "transformers_version": "4.45.0",
+  "vocab_size": 40000,
+  "z_loss_weight": 0.0001
+}

configuration_metallm.py

@@ -0,0 +1,62 @@
+"""HF config for MetaLLM / Vishvakarma (custom arch: NoPE-every-N, QK-norm, GQA, SwiGLU).
+
+Field names intentionally mirror configs/model.py:ModelConfig so this object can be
+passed directly to metallm_core.MetaLLMv2 as its `cfg` (duck-typed).
+"""
+from transformers import PretrainedConfig
+
+
+class MetaLLMConfig(PretrainedConfig):
+    model_type = "metallm"
+
+    def __init__(
+        self,
+        vocab_size: int = 40000,
+        n_layers: int = 26,
+        d_model: int = 1792,
+        n_heads: int = 28,
+        n_kv_heads: int = 14,
+        d_ff: int = 4864,
+        max_seq_len: int = 2048,
+        rope_theta: float = 500_000.0,
+        norm_eps: float = 1e-5,
+        tie_embeddings: bool = True,
+        qk_norm: bool = True,
+        z_loss_weight: float = 0.0,   # inference shim: aux loss unused, kept for fidelity
+        rope_fp32: bool = True,
+        doc_mask: bool = True,        # equals plain causal for single-document prompts
+        attn_impl: str = "sdpa",      # "sdpa" is the portable inference default
+        nope_every: int = 4,
+        bos_id: int = 1,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.n_layers = n_layers
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.d_ff = d_ff
+        self.max_seq_len = max_seq_len
+        self.rope_theta = rope_theta
+        self.norm_eps = norm_eps
+        self.tie_embeddings = tie_embeddings
+        self.qk_norm = qk_norm
+        self.z_loss_weight = z_loss_weight
+        self.rope_fp32 = rope_fp32
+        self.doc_mask = doc_mask
+        self.attn_impl = attn_impl
+        self.nope_every = nope_every
+        self.bos_id = bos_id
+        # standard-name aliases: transformers>=5.13 core reads these directly
+        self.num_hidden_layers = n_layers
+        self.hidden_size = d_model
+        self.num_attention_heads = n_heads
+        self.num_key_value_heads = n_kv_heads
+        self.max_position_embeddings = max_seq_len
+        kwargs.setdefault("tie_word_embeddings", tie_embeddings)
+        kwargs.setdefault("bos_token_id", bos_id)
+        super().__init__(**kwargs)
+
+    @property
+    def head_dim(self) -> int:
+        return self.d_model // self.n_heads

generation_config.json

@@ -0,0 +1,6 @@
+{
+ "bos_token_id": 1,
+ "eos_token_id": 2,
+ "pad_token_id": 0,
+ "max_length": 2048
+}

metallm_core.py

@@ -0,0 +1,374 @@
+"""
+MetaLLM v2 — Llama3-style decoder with the 2026 stability/retrieval bundle.
+
+Deltas vs v1 (model/llama.py), each gated by a ModelConfig flag so v1 behavior
+is reproducible by leaving the flags off:
+
+  * doc_mask        — document-masked causal attention. Tokens only attend within
+                      their own document (segment id = cumsum(ids == bos_id)). This
+                      is the structural fix for the v1 passkey=0 failure: v1 packed
+                      multiple documents into one sequence and attended across the
+                      BOS boundaries (model/llama.py:77 is_causal=True, no mask),
+                      teaching the model to ignore long-range in-context signal.
+                      See build_log A17/A18.
+  * qk_norm         — per-head RMSNorm on Q and K before RoPE (Qwen3/OLMo-2/Gemma-3).
+                      Caps attention-logit magnitude => training stability at higher
+                      LR / more tokens.
+  * rope_fp32       — apply RoPE rotation in float32 then cast back. v1 ran the whole
+                      model (and thus cos/sin) in bf16; bf16 rounding corrupts RoPE
+                      relative positions at long context (AnchorAttention 2411.13476).
+  * z_loss_weight   — auxiliary log(Z)^2 penalty on the LM logits (OLMo-2/PaLM); keeps
+                      final logits from drifting large under bf16. Returned separately
+                      for logging; folded into the training loss.
+
+attn_impl selects the masking kernel:
+  * "sdpa"  — F.scaled_dot_product_attention with an explicit boolean mask. Works on
+              CPU and CUDA, fine to moderate seq_len. Default; used by CPU smoke tests.
+  * "flex"  — torch FlexAttention with a (causal AND same-document) mask_mod. Scales to
+              long context without materializing a (t,t) mask. CUDA path for the real run.
+
+The model reuses RMSNorm / SwiGLU / build_rope_cache from model/llama.py.
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# --- inlined from model/llama.py (self-contained for HF trust_remote_code) ---
+class RMSNorm(nn.Module):
+    def __init__(self, d: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(d))
+
+    def forward(self, x):
+        dtype = x.dtype
+        x = x.float()
+        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return (self.weight * x).to(dtype)
+
+
+def build_rope_cache(head_dim: int, max_seq: int, theta: float, device, dtype=torch.float32):
+    freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=dtype) / head_dim))
+    t = torch.arange(max_seq, device=device, dtype=dtype)
+    freqs = torch.outer(t, freqs)                # (seq, head_dim/2)
+    return torch.cos(freqs), torch.sin(freqs)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.w_gate = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
+        self.w_up = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
+        self.w_down = nn.Linear(cfg.d_ff, cfg.d_model, bias=False)
+
+    def forward(self, x):
+        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))
+
+BOS_ID = 1  # token id 1 is BOS in the MetaLLM tokenizer (see data/shard.py validate_shard)
+FLEX_MIN_SEQ = 2048  # at/above this seq_len, "auto" uses FlexAttention (compact block mask)
+                     # instead of the dense (b,1,t,t) SDPA mask, which would OOM at long ctx.
+
+_FLEX_FN = None
+
+
+def _flex_attention_fn():
+    """Lazily build a torch.compile'd flex_attention. Compiling is what gives flex its
+    memory win (otherwise it materializes the full score matrix — grok review A18)."""
+    global _FLEX_FN
+    if _FLEX_FN is None:
+        from torch.nn.attention.flex_attention import flex_attention
+        _FLEX_FN = torch.compile(flex_attention)
+    return _FLEX_FN
+
+
+def resolve_impl(attn_impl: str, t: int, device_type: str) -> str:
+    """Pick the attention implementation for this forward. 'auto' uses flex on CUDA for
+    long sequences, sdpa otherwise. flex on CPU is unreliable, so CPU always uses sdpa."""
+    if device_type != "cuda":
+        return "sdpa"
+    if attn_impl == "flex":
+        return "flex"
+    if attn_impl == "auto":
+        return "flex" if t >= FLEX_MIN_SEQ else "sdpa"
+    return "sdpa"
+
+
+# -------- rope (fp32-capable) --------
+
+def apply_rope_v2(x, cos, sin, fp32: bool = True):
+    """Rotate x (b, n_heads, seq, head_dim). When fp32=True the rotation math runs
+    in float32 regardless of x.dtype, then casts back — avoids bf16 position drift."""
+    in_dtype = x.dtype
+    if fp32:
+        x = x.float()
+        cos = cos.float()
+        sin = sin.float()
+    d = x.shape[-1]
+    x1, x2 = x[..., : d // 2], x[..., d // 2:]
+    t = x.shape[-2]
+    cos = cos[:t].unsqueeze(0).unsqueeze(0)
+    sin = sin[:t].unsqueeze(0).unsqueeze(0)
+    out = torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+    return out.to(in_dtype)
+
+
+# -------- document-causal masking --------
+
+def segment_ids_from_tokens(ids: torch.Tensor, bos_id: int = BOS_ID) -> torch.Tensor:
+    """Per-token document segment id. cumsum over BOS occurrences: every token gets
+    the index of the document it belongs to within the packed sequence. The leading
+    partial document (window started mid-doc, before any BOS) is segment 0 and attends
+    only to itself — correct."""
+    return (ids == bos_id).cumsum(dim=1)
+
+
+def build_doc_causal_mask(seg: torch.Tensor) -> torch.Tensor:
+    """(b, 1, t, t) boolean mask, True = key participates. Allowed iff causal (j<=i)
+    AND same document (seg_i == seg_j)."""
+    b, t = seg.shape
+    causal = torch.tril(torch.ones(t, t, dtype=torch.bool, device=seg.device))
+    same_doc = seg.unsqueeze(2) == seg.unsqueeze(1)          # (b, t, t)
+    return (same_doc & causal).unsqueeze(1)                  # (b, 1, t, t)
+
+
+# -------- attention --------
+
+class GroupedQueryAttentionV2(nn.Module):
+    def __init__(self, cfg, use_rope: bool = True):
+        super().__init__()
+        self.n_heads = cfg.n_heads
+        self.n_kv = cfg.n_kv_heads
+        self.head_dim = cfg.head_dim
+        self.scale = 1.0 / math.sqrt(self.head_dim)
+        self.rope_fp32 = getattr(cfg, "rope_fp32", True)
+        self.use_rope = use_rope        # False => NoPE layer (full attn, no positional encoding)
+        self.wq = nn.Linear(cfg.d_model, cfg.n_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.head_dim, bias=False)
+        self.wv = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.head_dim, bias=False)
+        self.wo = nn.Linear(cfg.n_heads * self.head_dim, cfg.d_model, bias=False)
+
+        self.qk_norm = getattr(cfg, "qk_norm", False)
+        if self.qk_norm:
+            self.q_norm = RMSNorm(self.head_dim, cfg.norm_eps)
+            self.k_norm = RMSNorm(self.head_dim, cfg.norm_eps)
+
+    def forward(self, x, cos, sin, attn_mask=None, block_mask=None):
+        b, t, _ = x.shape
+        q = self.wq(x).view(b, t, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.wk(x).view(b, t, self.n_kv, self.head_dim).transpose(1, 2)
+        v = self.wv(x).view(b, t, self.n_kv, self.head_dim).transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        if self.use_rope:
+            q = apply_rope_v2(q, cos, sin, fp32=self.rope_fp32)
+            k = apply_rope_v2(k, cos, sin, fp32=self.rope_fp32)
+
+        rep = self.n_heads // self.n_kv
+        if rep > 1:
+            k = k.repeat_interleave(rep, dim=1)
+            v = v.repeat_interleave(rep, dim=1)
+
+        # The impl is chosen by which mask the model handed down: a flex BlockMask
+        # (long ctx, compact) or a dense SDPA boolean mask (short ctx) or neither (causal).
+        if block_mask is not None:
+            out = _flex_attention_fn()(q, k, v, block_mask=block_mask, scale=self.scale)
+        elif attn_mask is not None:
+            out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, scale=self.scale)
+        else:
+            out = F.scaled_dot_product_attention(q, k, v, is_causal=True, scale=self.scale)
+
+        out = out.transpose(1, 2).contiguous().view(b, t, -1)
+        return self.wo(out)
+
+
+# -------- block --------
+
+class BlockV2(nn.Module):
+    def __init__(self, cfg, use_ckpt: bool = False, use_rope: bool = True):
+        super().__init__()
+        self.use_ckpt = use_ckpt
+        self.norm1 = RMSNorm(cfg.d_model, cfg.norm_eps)
+        self.attn = GroupedQueryAttentionV2(cfg, use_rope=use_rope)
+        self.norm2 = RMSNorm(cfg.d_model, cfg.norm_eps)
+        self.mlp = SwiGLU(cfg)
+
+    def _fwd(self, x, cos, sin, attn_mask, block_mask):
+        x = x + self.attn(self.norm1(x), cos, sin, attn_mask=attn_mask, block_mask=block_mask)
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+    def forward(self, x, cos, sin, attn_mask=None, block_mask=None):
+        if self.use_ckpt and self.training:
+            return torch.utils.checkpoint.checkpoint(
+                self._fwd, x, cos, sin, attn_mask, block_mask, use_reentrant=False
+            )
+        return self._fwd(x, cos, sin, attn_mask, block_mask)
+
+
+# -------- model --------
+
+class MetaLLMv2(nn.Module):
+    def __init__(self, cfg, grad_ckpt: bool = False):
+        super().__init__()
+        self.cfg = cfg
+        self.grad_ckpt = grad_ckpt
+        self.doc_mask = getattr(cfg, "doc_mask", False)
+        self.attn_impl = getattr(cfg, "attn_impl", "sdpa")
+        self.z_loss_weight = getattr(cfg, "z_loss_weight", 0.0)
+        self.bos_id = getattr(cfg, "bos_id", BOS_ID)
+
+        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        # NoPE: every `nope_every`-th layer drops RoPE (full-attn, no positional encoding) to grow
+        # dedicated retrieval heads (RNoPE / HyPE). nope_every=0 => all layers use RoPE (v2 behavior).
+        nope_every = getattr(cfg, "nope_every", 0)
+        def _use_rope(i: int) -> bool:
+            return not (nope_every and (i + 1) % nope_every == 0)
+        self.nope_layers = [i for i in range(cfg.n_layers) if not _use_rope(i)]
+        self.blocks = nn.ModuleList(
+            [BlockV2(cfg, use_ckpt=grad_ckpt, use_rope=_use_rope(i)) for i in range(cfg.n_layers)]
+        )
+        self.norm = RMSNorm(cfg.d_model, cfg.norm_eps)
+        self.lm_head = None if cfg.tie_embeddings else nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+
+        # RoPE cache is built LAZILY in _rope() — registering it as a buffer breaks
+        # under transformers>=5.x meta-device loading (non-persistent buffers are
+        # materialized as uninitialized memory and never recomputed -> garbage/NaN).
+        self._rope_cache = None
+
+        self.apply(self._init)
+
+    @staticmethod
+    def _init(m):
+        if isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, mean=0.0, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Embedding):
+            nn.init.normal_(m.weight, mean=0.0, std=0.02)
+
+    def _make_masks(self, ids):
+        """Return (attn_mask, block_mask), or (None, None) when doc_mask off.
+        Resolves the impl per-forward: flex (compact block mask) for long ctx on CUDA,
+        dense SDPA bool mask for short ctx, so the long-context curriculum phase doesn't
+        OOM on a (b,1,t,t) mask (grok review A18)."""
+        if not self.doc_mask:
+            return None, None
+        b, t = ids.shape
+        seg = segment_ids_from_tokens(ids, self.bos_id)
+        impl = resolve_impl(self.attn_impl, t, ids.device.type)
+        if impl == "flex":
+            from torch.nn.attention.flex_attention import create_block_mask
+
+            def mask_mod(bi, hi, qi, ki):
+                return (qi >= ki) & (seg[bi, qi] == seg[bi, ki])
+
+            block_mask = create_block_mask(mask_mod, B=b, H=None, Q_LEN=t, KV_LEN=t, device=ids.device)
+            return None, block_mask
+        return build_doc_causal_mask(seg), None
+
+    def _rope(self, device):
+        # Fidelity: the native model registers cos/sin as buffers, so .to(bfloat16)
+        # rounds them to the model dtype — the network trained against THOSE values.
+        # Build fp32 then round through the param dtype (2026-07-10: pristine-fp32
+        # cache caused compounding layer drift, 4% greedy argmax flips, ~2x pass@1
+        # loss vs native; see build_log.md).
+        dtype = self.tok_emb.weight.dtype
+        if (self._rope_cache is None or self._rope_cache[0].device != device
+                or self._rope_cache[0].dtype != dtype):
+            cos, sin = build_rope_cache(self.cfg.head_dim, self.cfg.max_seq_len,
+                                        self.cfg.rope_theta, device=device)
+            self._rope_cache = (cos.to(dtype), sin.to(dtype))
+        return self._rope_cache
+
+    def forward(self, ids, targets=None):
+        x = self.tok_emb(ids)
+        cos, sin = self._rope(x.device)
+        attn_mask, block_mask = self._make_masks(ids)
+
+        for blk in self.blocks:
+            x = blk(x, cos, sin, attn_mask=attn_mask, block_mask=block_mask)
+        x = self.norm(x)
+
+        if targets is None:
+            logits = F.linear(x, self.tok_emb.weight) if self.lm_head is None else self.lm_head(x)
+            return logits, None
+
+        # Training path: fused, chunked head+loss. The dense [B,T,vocab] logits (plus the fp32
+        # upcast the z-loss needs) are the dominant activation cost at large vocab/seq and force
+        # tiny batches. Instead project + score in row-chunks, checkpointing each chunk so its
+        # logits are recomputed in backward rather than retained. Identical math to the dense
+        # path (only fp summation order differs). Training callers read only the loss, so logits
+        # are returned as None.
+        ce, z_loss = self._chunked_loss(x, targets)
+        loss = ce + z_loss
+        self._last_ce = ce.detach()
+        self._last_zloss = z_loss.detach()
+        return None, loss
+
+    def _chunked_loss(self, x, targets, chunk: int = 2048):
+        """Memory-bounded CE + z-loss. Peak head memory ~chunk*vocab regardless of B*T, so
+        large batches fit. Exact: each chunk is checkpointed (logits recomputed in backward)."""
+        W = self.tok_emb.weight if self.lm_head is None else self.lm_head.weight
+        xf = x.reshape(-1, x.size(-1))
+        tf = targets.reshape(-1)
+        N = xf.size(0)
+        zw = float(self.z_loss_weight)
+
+        def chunk_fn(xc, tc):
+            lc = F.linear(xc, W)
+            ce_sum = F.cross_entropy(lc, tc, ignore_index=-100, reduction="sum")
+            if zw > 0.0:
+                m = tc != -100
+                z_sum = torch.logsumexp(lc[m].float(), dim=-1).pow(2).sum() if m.any() \
+                    else lc.new_zeros(())
+            else:
+                z_sum = lc.new_zeros(())
+            return ce_sum, z_sum
+
+        ce_total = xf.new_zeros(())
+        z_total = xf.new_zeros(())
+        for i in range(0, N, chunk):
+            xc, tc = xf[i:i + chunk], tf[i:i + chunk]
+            if self.training and xc.requires_grad:
+                ce_sum, z_sum = torch.utils.checkpoint.checkpoint(
+                    chunk_fn, xc, tc, use_reentrant=False)
+            else:
+                ce_sum, z_sum = chunk_fn(xc, tc)
+            ce_total = ce_total + ce_sum
+            z_total = z_total + z_sum
+        n_valid = (tf != -100).sum().clamp(min=1)
+        ce = ce_total / n_valid
+        z_loss = (zw * z_total / n_valid) if zw > 0.0 else xf.new_zeros(())
+        return ce, z_loss
+
+    def encode(self, ids, pad_id: int = 0, bidirectional: bool = True):
+        """Pooled sentence embedding (b, d) for the Track B embedding model.
+        bidirectional=True runs full (non-causal) attention over the real tokens
+        (LLM2Vec-style — better embedding geometry than causal); pooling is masked
+        mean over non-pad positions. Returns grad-enabled embeddings for contrastive FT."""
+        x = self.tok_emb(ids)
+        cos, sin = self._rope(x.device)
+        nonpad = (ids != pad_id)                                  # (b, t)
+        b, t = ids.shape
+        if bidirectional:
+            attn_mask = nonpad[:, None, None, :].expand(b, 1, t, t)   # attend to all real keys
+        else:
+            attn_mask, _ = self._make_masks(ids)
+        for blk in self.blocks:
+            x = blk(x, cos, sin, attn_mask=attn_mask, block_mask=None)
+        x = self.norm(x)
+        m = nonpad.unsqueeze(-1).to(x.dtype)
+        return (x * m).sum(1) / m.sum(1).clamp(min=1.0)
+
+    def num_params(self, non_embedding: bool = False) -> int:
+        n = sum(p.numel() for p in self.parameters())
+        if non_embedding:
+            n -= self.tok_emb.weight.numel()
+        return n

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c754911b6ba65a98e1ed137488228d1613563c8db419a46046d398c1cd57988f
+size 2004481000

modeling_metallm.py

@@ -0,0 +1,72 @@
+"""HF modeling shim for MetaLLM / Vishvakarma.
+
+Wraps the native MetaLLMv2 (metallm_core.py, vendored + self-contained) in a
+PreTrainedModel so the checkpoint is `AutoModelForCausalLM.from_pretrained(...,
+trust_remote_code=True)`-loadable — the exact NoPE/QK-norm architecture, not a
+lossy Llama-config approximation.
+
+v1 scope: no KV cache (generation recomputes the prefix each step — correct, just
+slower). Fine for demos/evals; a cached path is future work.
+"""
+import torch
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_metallm import MetaLLMConfig
+from .metallm_core import MetaLLMv2
+
+
+class MetaLLMForCausalLM(PreTrainedModel):
+    config_class = MetaLLMConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["BlockV2"]
+    # weights are shared inside MetaLLMv2 (lm_head is None when tied) — nothing for
+    # HF to re-tie. transformers>=5.13 requires all_tied_weights_keys to exist.
+    _tied_weights_keys = []
+
+    @property
+    def all_tied_weights_keys(self):
+        return {}
+
+    def __init__(self, config: MetaLLMConfig):
+        super().__init__(config)
+        self.model = MetaLLMv2(config)
+        # post_init would re-run HF default init over our scaled init; the weights
+        # come from the checkpoint anyway, so skip it deliberately.
+
+    # --- embeddings plumbing (needed by resize/tie utilities) ---
+    def get_input_embeddings(self):
+        return self.model.tok_emb
+
+    def set_input_embeddings(self, value):
+        self.model.tok_emb = value
+
+    def get_output_embeddings(self):
+        return self.model.lm_head  # None when tied — HF treats tied as expected
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: torch.Tensor = None,  # accepted for API compat; causal-only v1
+        labels: torch.LongTensor = None,
+        past_key_values=None,
+        use_cache: bool = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        logits, _ = self.model(input_ids)  # native path, targets=None -> (logits, None)
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=None)
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        # No KV cache: feed the full sequence every step.
+        return {"input_ids": input_ids}

serve_metallum.py

@@ -0,0 +1,615 @@
+#!/usr/bin/env python3
+"""Local serving wrapper for the exported MetaLLM s2b075 HF bundle.
+
+Product answer to the measured format-acquisition failure: strict structured
+output is enforced AT DECODE TIME (token-level grammar/schema constraints),
+not hoped for from the weights. SFT got 0/60 on the frozen strict-format
+suite; this wrapper guarantees syntactic validity and schema conformance on
+every request by construction.
+
+Modes (POST /generate):
+  free  — plain continuation (greedy or temperature-sampled, seeded)
+  json  — flat-object schema forced: exact keys, typed values
+  tool  — harness contract: {"thought","tool","args"} or {"thought","final"}
+  mcq   — single letter A-D then EOS
+
+Binds 127.0.0.1 only. No auth, no network exposure: localhost product shim.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import sys
+import threading
+import time
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+
+import torch
+
+BUNDLE = Path(__file__).resolve().parent  # model files ship beside this script
+TOOLS = ["python_exec", "write_file", "read_file", "ls"]
+MAX_CTX = 2048
+BOS, EOS = 1, 2
+
+
+# --- strict whole-output validators (self-contained copies) -----------------
+def check_json(text: str, schema: dict) -> bool:
+    s = text.strip()
+    if s.startswith("```"):
+        s = s.strip("`").removeprefix("json").strip()
+    try:
+        obj = json.loads(s)
+    except Exception:
+        return False
+    if not isinstance(obj, dict):
+        return False
+    for k, t in schema.items():
+        if k not in obj or not isinstance(obj[k], t):
+            return False
+    return True
+
+
+def check_tool(text: str) -> bool:
+    s = text.strip()
+    try:
+        obj = json.loads(s)
+    except Exception:
+        return False
+    if not isinstance(obj, dict) or "thought" not in obj:
+        return False
+    if "final" in obj:
+        return True
+    return obj.get("tool") in TOOLS and isinstance(obj.get("args"), dict)
+
+
+def check_mcq(text: str) -> bool:
+    t = text.strip().rstrip(".").strip()
+    return len(t) == 1 and t in "ABCD"
+
+# ---------------------------------------------------------------------------
+# generic JSON prefix state machine (for free-form values such as tool args)
+# ---------------------------------------------------------------------------
+
+_INT_PART = re.compile(r"^-?(0|[1-9][0-9]*)?$")
+_INT_FULL = re.compile(r"^-?(0|[1-9][0-9]*)$")
+_NUM_PART = re.compile(r"^-?(0|[1-9][0-9]*)?(\.[0-9]*)?([eE][+-]?[0-9]*)?$")
+_NUM_FULL = re.compile(r"^-?(0|[1-9][0-9]*)(\.[0-9]+)?([eE][+-]?[0-9]+)?$")
+_LITS = ("true", "false", "null")
+_WS = " \t\n\r"
+_NUM_CHARS = set("0123456789.eE+-")
+_LIT_CHARS = set("abcdefghijklmnopqrstuvwxyz")
+
+
+class JsonError(Exception):
+    pass
+
+
+class JsonSM:
+    """Incremental validator for a single top-level JSON value."""
+
+    def __init__(self, depth_cap: int = 8):
+        self.stack: list[str] = []          # 'obj' | 'arr'
+        self.mode = "value"                 # value|key|key_or_end|colon|comma|string|number|lit|end
+        self.buf = ""                       # partial number / literal
+        self.esc = False
+        self.hex_left = 0
+        self.after_string = "comma"         # where to go when a string closes
+        self.depth_cap = depth_cap
+
+    def copy(self) -> "JsonSM":
+        s = JsonSM(self.depth_cap)
+        s.stack = list(self.stack)
+        s.mode, s.buf, s.esc, s.hex_left, s.after_string = (
+            self.mode, self.buf, self.esc, self.hex_left, self.after_string)
+        return s
+
+    # -- internal helpers --------------------------------------------------
+    def _value_done(self):
+        if not self.stack:
+            self.mode = "end"
+        else:
+            self.mode = "comma"
+
+    def _feed(self, ch: str):
+        m = self.mode
+        if m == "end":
+            if ch not in _WS:
+                raise JsonError("trailing content")
+            return
+        if m == "string":
+            if self.hex_left:
+                if ch not in "0123456789abcdefABCDEF":
+                    raise JsonError("bad unicode escape")
+                self.hex_left -= 1
+                return
+            if self.esc:
+                if ch == "u":
+                    self.hex_left = 4
+                elif ch not in '"\\/bfnrt':
+                    raise JsonError("bad escape")
+                self.esc = False
+                return
+            if ch == "\\":
+                self.esc = True
+            elif ch == '"':
+                self.mode = self.after_string
+            elif ord(ch) < 0x20:
+                raise JsonError("control char in string")
+            return
+        if m == "number":
+            if ch in _NUM_CHARS:
+                if not _NUM_PART.match(self.buf + ch):
+                    raise JsonError("bad number")
+                self.buf += ch
+                return
+            if not _NUM_FULL.match(self.buf):
+                raise JsonError("incomplete number")
+            self.buf = ""
+            self._value_done()
+            self._feed(ch)  # re-feed delimiter
+            return
+        if m == "lit":
+            if ch in _LIT_CHARS:
+                self.buf += ch
+                if not any(w.startswith(self.buf) for w in _LITS):
+                    raise JsonError("bad literal")
+                if self.buf in _LITS:
+                    self.buf = ""
+                    self._value_done()
+                return
+            if self.buf:  # incomplete literal followed by delimiter
+                raise JsonError("incomplete literal")
+            self._feed(ch)
+            return
+        # structural modes
+        if ch in _WS:
+            return
+        if m == "value":
+            if ch == "{":
+                if len(self.stack) >= self.depth_cap:
+                    raise JsonError("depth cap")
+                self.stack.append("obj")
+                self.mode = "key_or_end"
+            elif ch == "[":
+                if len(self.stack) >= self.depth_cap:
+                    raise JsonError("depth cap")
+                self.stack.append("arr")
+                self.mode = "value_or_end"
+            elif ch == '"':
+                self.after_string = "comma" if self.stack else "end"
+                self.mode = "string"
+            elif ch in "-0123456789":
+                self.mode = "number"
+                self.buf = ch
+            elif ch in "tfn":
+                self.mode = "lit"
+                self.buf = ch
+            else:
+                raise JsonError("expected value")
+            return
+        if m == "value_or_end":  # just after '['
+            if ch == "]":
+                self.stack.pop()
+                self._value_done()
+                return
+            self.mode = "value"
+            self._feed(ch)
+            return
+        if m == "key_or_end":  # just after '{' — empty object may close
+            if ch == "}":
+                self.stack.pop()
+                self._value_done()
+                return
+            self.mode = "key"
+            self._feed(ch)
+            return
+        if m == "key":
+            if ch == '"':
+                self.after_string = "colon"
+                self.mode = "string"
+                return
+            raise JsonError("expected key")
+        if m == "colon":
+            if ch == ":":
+                self.mode = "value"
+                return
+            raise JsonError("expected colon")
+        if m == "comma":
+            top = self.stack[-1] if self.stack else None
+            if top == "obj":
+                if ch == ",":
+                    self.mode = "key"
+                    return
+                if ch == "}":
+                    self.stack.pop()
+                    self._value_done()
+                    return
+            elif top == "arr":
+                if ch == ",":
+                    self.mode = "value"
+                    return
+                if ch == "]":
+                    self.stack.pop()
+                    self._value_done()
+                    return
+            raise JsonError("expected comma/close")
+        raise JsonError(f"bad mode {m}")
+
+    def feed(self, text: str):
+        for ch in text:
+            self._feed(ch)
+
+    def allows(self, text: str) -> bool:
+        try:
+            s = self.copy()
+            s.feed(text)
+            return True
+        except JsonError:
+            return False
+
+    @property
+    def complete(self) -> bool:
+        return self.mode == "end"
+
+
+# ---------------------------------------------------------------------------
+# constrained generator
+# ---------------------------------------------------------------------------
+
+class ConstrainedDecoder:
+    def __init__(self, model, tok, device: str):
+        self.model = model
+        self.tok = tok
+        self.device = device
+        self.lock = threading.Lock()
+        # vocab text cache
+        n = 40000
+        self.vtext: list[str] = [tok.decode([i]) for i in range(n)]
+        # token classes (special ids 0=pad,1=bos,2=eos excluded everywhere)
+        self.quote_ids = [i for i, t in enumerate(self.vtext) if t == '"' and i > 2]
+        self.closer_ids = [i for i, t in enumerate(self.vtext)
+                           if i > 2 and t and any(c in '}]"' for c in t)
+                           and all(c in '}]" \t\n\r' for c in t)]
+        self.num_cand_ids = [i for i, t in enumerate(self.vtext)
+                             if i > 2 and t and all(c in _NUM_CHARS for c in t)]
+        self.str_cand_ids = [i for i, t in enumerate(self.vtext)
+                             if i > 2 and t and '"' not in t and "\\" not in t
+                             and all(ord(c) >= 0x20 for c in t)]
+        self.mcq_ids = [i for i, t in enumerate(self.vtext)
+                        if i > 2 and t in ("A", "B", "C", "D", " A", " B", " C", " D")]
+        self._closer_set = set(self.closer_ids)
+        assert self.quote_ids and self.num_cand_ids and self.str_cand_ids and self.mcq_ids
+
+    # -- primitives ---------------------------------------------------------
+    def encode_lit(self, text: str) -> list[int]:
+        ids = self.tok.encode(text, add_special_tokens=False)
+        if self.tok.decode(ids) == text:
+            return ids
+        out = []
+        for ch in text:  # fall back to per-char forcing (roundtrip-safe)
+            ci = self.tok.encode(ch, add_special_tokens=False)
+            if self.tok.decode(ci) != ch:
+                raise RuntimeError(f"cannot force {text!r}")
+            out.extend(ci)
+        return out
+
+    @torch.no_grad()
+    def _logits(self, ids: list[int]) -> torch.Tensor:
+        x = torch.tensor([ids[-MAX_CTX:]], dtype=torch.long, device=self.device)
+        return self.model(input_ids=x).logits[0, -1].float()
+
+    @staticmethod
+    def _pick(logits: torch.Tensor, cand: list[int], gen: torch.Generator | None,
+              temperature: float) -> int:
+        t = logits[cand]
+        if gen is None or temperature <= 0:
+            return cand[int(t.argmax().item())]
+        p = torch.softmax(t / max(temperature, 1e-5), dim=-1)
+        return cand[int(torch.multinomial(p, 1, generator=gen).item())]
+
+    def _pick_with_stop(self, logits, cand, stop_id, gen, temperature):
+        """Choose among cand tokens plus a virtual STOP action (stop_id's logit)."""
+        ids = cand + [stop_id]
+        choice = self._pick(logits, ids, gen, temperature)
+        return None if choice == stop_id else choice
+
+    # -- typed values --------------------------------------------------------
+    def gen_number(self, ids, kind: str, next_struct: int, gen, temperature,
+                   budget: list[int]) -> str:
+        part, full = (_INT_PART, _INT_FULL) if kind == "int" else (_NUM_PART, _NUM_FULL)
+        so_far = ""
+        while budget[0] > 0 and len(so_far) < 18:
+            cand = [i for i in self.num_cand_ids if part.match(so_far + self.vtext[i])]
+            if not cand:
+                break
+            logits = self._logits(ids)
+            if full.match(so_far):
+                nxt = self._pick_with_stop(logits, cand, next_struct, gen, temperature)
+                if nxt is None:
+                    break
+            else:
+                nxt = self._pick(logits, cand, gen, temperature)
+            ids.append(nxt)
+            budget[0] -= 1
+            so_far += self.vtext[nxt]
+        if not full.match(so_far):
+            raise RuntimeError(f"number underflow: {so_far!r}")
+        return so_far
+
+    def gen_string(self, ids, gen, temperature, budget: list[int],
+                   max_chars: int = 160) -> str:
+        # Graceful close at max_chars/budget: the model is never required to
+        # prefer the bare-quote token, so strings are capped, never fatal.
+        out = ""
+        while budget[0] > 0 and len(out) < max_chars:
+            logits = self._logits(ids)
+            cand = self.str_cand_ids + self.quote_ids
+            nxt = self._pick(logits, cand, gen, temperature)
+            ids.append(nxt)
+            budget[0] -= 1
+            if nxt in self.quote_ids:
+                return out
+            out += self.vtext[nxt]
+        return out  # caller appends the closing quote
+
+    def gen_bool(self, ids, gen, temperature, budget: list[int]) -> str:
+        t_ids = self.encode_lit("true")
+        f_ids = self.encode_lit("false")
+        logits = self._logits(ids)
+        first = self._pick(logits, [t_ids[0], f_ids[0]], gen, temperature)
+        seq = t_ids if first == t_ids[0] else f_ids
+        for tid in seq:
+            ids.append(tid)
+            budget[0] -= 1
+        return "true" if seq is t_ids else "false"
+
+    def gen_json_free(self, ids, gen, temperature, budget: list[int],
+                      sm: JsonSM | None = None, require_container: bool = False) -> str:
+        sm = sm or JsonSM()
+        out = ""
+        mask_cache: dict = {}
+        while budget[0] > 0 and not sm.complete:
+            key = (sm.mode, tuple(sm.stack), sm.buf, sm.esc, sm.hex_left, sm.after_string)
+            cand = mask_cache.get(key)
+            if cand is None:
+                cand = [i for i, t in enumerate(self.vtext) if t and sm.allows(t)]
+                mask_cache[key] = cand
+            if require_container and not out.strip():
+                cand = [i for i in cand
+                        if self.vtext[i].strip() and self.vtext[i].lstrip()[:1] in "{["]
+            if budget[0] <= 24:  # force-close: steer to completion within reserve
+                closing = [i for i in cand if i in self._closer_set]
+                if closing:
+                    cand = closing
+            if not cand:
+                raise RuntimeError("json mask exhausted")
+            logits = self._logits(ids)
+            nxt = self._pick(logits, cand, gen, temperature)
+            ids.append(nxt)
+            budget[0] -= 1
+            txt = self.vtext[nxt]
+            sm.feed(txt)
+            out += txt
+        if not sm.complete:
+            raise RuntimeError("json did not complete")
+        return out
+
+    # -- top-level modes -----------------------------------------------------
+    def generate(self, prompt: str, mode: str = "free", schema: dict | None = None,
+                 max_new: int = 160, temperature: float = 0.0,
+                 seed: int | None = None) -> dict:
+        with self.lock:
+            return self._generate(prompt, mode, schema, max_new, temperature, seed)
+
+    def _generate(self, prompt, mode, schema, max_new, temperature, seed):
+        gen = None
+        if seed is not None:
+            gen = torch.Generator(device=self.device).manual_seed(seed)
+        ids = [BOS] + self.tok.encode(prompt, add_special_tokens=False)
+        n_prompt = len(ids)
+        if n_prompt + max_new > MAX_CTX:
+            raise ValueError(f"prompt({n_prompt})+max_new({max_new}) exceeds {MAX_CTX}")
+        budget = [max_new]
+        text = ""
+        t0 = time.time()
+
+        if mode == "free":
+            while budget[0] > 0:
+                logits = self._logits(ids)
+                if gen is None or temperature <= 0:
+                    nxt = int(logits.argmax().item())
+                else:
+                    p = torch.softmax(logits / temperature, dim=-1)
+                    nxt = int(torch.multinomial(p, 1, generator=gen).item())
+                if nxt == EOS:
+                    break
+                ids.append(nxt)
+                budget[0] -= 1
+                text += self.vtext[nxt]
+
+        elif mode == "json":
+            schema = schema or {}
+            items = list(schema.items())
+            if not items:  # schema-less: free-form JSON object/array, grammar-masked
+                text += self.gen_json_free(ids, gen, temperature, budget, JsonSM(),
+                                           require_container=True)
+                return self._finish(ids, n_prompt, t0, text, mode)
+            for tid in self.encode_lit("{"):
+                ids.append(tid); budget[0] -= 1
+            text += "{"
+            for i, (key, vtype) in enumerate(items):
+                lead = '"' if i == 0 else ', "'
+                lit = f'{lead}{key}": '
+                for tid in self.encode_lit(lit):
+                    ids.append(tid); budget[0] -= 1
+                text += lit
+                # structural token used as the STOP comparator
+                nxt_lit = "}" if i == len(items) - 1 else f', "{items[i+1][0]}": '
+                stop_id = self.encode_lit(nxt_lit)[0]
+                if vtype == "string":
+                    for tid in self.encode_lit('"'):
+                        ids.append(tid); budget[0] -= 1
+                    text += '"'
+                    val_cap = max(0, min(160, (budget[0] - 8) * 4))
+                    text += self.gen_string(ids, gen, temperature, budget, val_cap) + '"'
+                elif vtype in ("int", "number"):
+                    text += self.gen_number(ids, vtype, stop_id, gen, temperature, budget)
+                elif vtype == "bool":
+                    text += self.gen_bool(ids, gen, temperature, budget)
+                else:
+                    raise ValueError(f"unsupported schema type {vtype!r}")
+            for tid in self.encode_lit("}"):
+                ids.append(tid); budget[0] -= 1
+            text += "}"
+
+        elif mode == "tool":
+            for tid in self.encode_lit('{"thought": "'):
+                ids.append(tid); budget[0] -= 1
+            text += '{"thought": "'
+            thought_cap = max(0, min(120, (budget[0] - 40) * 4))
+            text += self.gen_string(ids, gen, temperature, budget, thought_cap) + '"'
+            tool_lit, final_lit = ', "tool": "', ', "final": "'
+            tl, fl = self.encode_lit(tool_lit), self.encode_lit(final_lit)
+            logits = self._logits(ids)
+            branch = self._pick(logits, [tl[0], fl[0]], gen, temperature)
+            if branch == fl[0]:
+                for tid in fl:
+                    ids.append(tid); budget[0] -= 1
+                text += final_lit
+                final_cap = max(0, min(600, (budget[0] - 4) * 4))
+                text += self.gen_string(ids, gen, temperature, budget, final_cap) + '"'
+            else:
+                for tid in tl:
+                    ids.append(tid); budget[0] -= 1
+                text += tool_lit
+                name_lits = {name: self.encode_lit(name) for name in TOOLS}
+                logits = self._logits(ids)
+                first = self._pick(logits, [nl[0] for nl in name_lits.values()],
+                                   gen, temperature)
+                chosen = next(n for n, nl in name_lits.items() if nl[0] == first)
+                for tid in name_lits[chosen]:
+                    ids.append(tid); budget[0] -= 1
+                text += chosen
+                mid = '", "args": '
+                for tid in self.encode_lit(mid):
+                    ids.append(tid); budget[0] -= 1
+                text += mid
+                sm = JsonSM()
+                sm.feed("{")
+                ids.extend(self.encode_lit("{")); budget[0] -= 1
+                text += "{"
+                text += self.gen_json_free(ids, gen, temperature, budget, sm)
+            for tid in self.encode_lit("}"):
+                ids.append(tid); budget[0] -= 1
+            text += "}"
+
+        elif mode == "mcq":
+            logits = self._logits(ids)
+            nxt = self._pick(logits, self.mcq_ids, gen, temperature)
+            ids.append(nxt)
+            budget[0] -= 1
+            text += self.vtext[nxt]
+
+        else:
+            raise ValueError(f"unknown mode {mode!r}")
+
+        return self._finish(ids, n_prompt, t0, text, mode)
+
+    @staticmethod
+    def _finish(ids, n_prompt, t0, text, mode):
+        dt = time.time() - t0
+        return {
+            "text": text, "mode": mode,
+            "usage": {"prompt_tokens": n_prompt, "gen_tokens": len(ids) - n_prompt,
+                      "seconds": round(dt, 2)},
+        }
+
+
+# ---------------------------------------------------------------------------
+# HTTP layer
+# ---------------------------------------------------------------------------
+
+def validate(mode: str, text: str, schema: dict | None) -> bool:
+    if mode == "json" and schema:
+        pysch = {k: {"int": int, "number": (int, float), "string": str,
+                     "bool": bool}[t] for k, t in schema.items()}
+        return check_json(text, pysch)
+    if mode == "json":
+        try:
+            return isinstance(json.loads(text), (dict, list))
+        except Exception:
+            return False
+    if mode == "tool":
+        return check_tool(text)
+    if mode == "mcq":
+        return check_mcq(text)
+    return True
+
+
+def make_handler(dec: ConstrainedDecoder):
+    class H(BaseHTTPRequestHandler):
+        def _send(self, code: int, obj: dict):
+            body = json.dumps(obj).encode()
+            self.send_response(code)
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(body)))
+            self.end_headers()
+            self.wfile.write(body)
+
+        def log_message(self, fmt, *a):  # quieter logs
+            sys.stderr.write("[serve] " + fmt % a + "\n")
+
+        def do_GET(self):
+            if self.path == "/health":
+                self._send(200, {"status": "ok", "bundle": str(BUNDLE),
+                                 "modes": ["free", "json", "tool", "mcq"]})
+            else:
+                self._send(404, {"error": "unknown path"})
+
+        def do_POST(self):
+            if self.path != "/generate":
+                self._send(404, {"error": "unknown path"})
+                return
+            try:
+                req = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
+                out = dec.generate(
+                    prompt=req["prompt"],
+                    mode=req.get("mode", "free"),
+                    schema=req.get("schema"),
+                    max_new=int(req.get("max_new", 160)),
+                    temperature=float(req.get("temperature", 0.0)),
+                    seed=req.get("seed"),
+                )
+                out["valid"] = validate(out["mode"], out["text"], req.get("schema"))
+                self._send(200, out)
+            except Exception as e:  # fail loudly, keep serving
+                self._send(400, {"error": f"{type(e).__name__}: {e}"})
+
+    return H
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--port", type=int, default=8399)
+    ap.add_argument("--device", default="cuda:0")
+    args = ap.parse_args()
+
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    tok = AutoTokenizer.from_pretrained(str(BUNDLE))
+    model = AutoModelForCausalLM.from_pretrained(
+        str(BUNDLE), trust_remote_code=True, torch_dtype=torch.bfloat16)
+    model = model.to(args.device).eval()
+    n_params = sum(p.numel() for p in model.parameters())
+    print(f"[serve] loaded {type(model).__name__} {n_params/1e6:.1f}M on {args.device}",
+          flush=True)
+
+    dec = ConstrainedDecoder(model, tok, args.device)
+    srv = ThreadingHTTPServer(("127.0.0.1", args.port), make_handler(dec))
+    print(f"[serve] listening on http://127.0.0.1:{args.port} (localhost only)", flush=True)
+    srv.serve_forever()
+
+
+if __name__ == "__main__":
+    main()

tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+ "tokenizer_class": "PreTrainedTokenizerFast",
+ "bos_token": "<s>",
+ "eos_token": "</s>",
+ "pad_token": "<pad>",
+ "model_input_names": [
+  "input_ids",
+  "attention_mask"
+ ]
+}