"""MiniJEV network = input representation + MiniJEV Core + ChoiceEncoder (+ calibration). state ──StateEncoder (grid | text | numeric | log)──► state tokens ┐ question ──QuestionEncoder──────────────────────────► question tok ├─► MiniJEVCore ─► logit per choice choice_i ──ChoiceEncoder────────────────────────────► c_i ┘ (+ value, score, embedding) All encoders project into the same latent space of width `d_model`; the core only sees vectors. The text backend is either `byte` (the original from-scratch tokenizer, used by the Snake checkpoints) or a `pretrained` sentence encoder shared by question, choices and text states (see minijev/encoders/text.py). `forward(state, question_ids, choice_ids)` keeps the original signature. With a frozen pretrained backbone, callers may instead pass precomputed backbone outputs (`question_bb`, `choice_bb`) so the backbone runs once per string. """ from __future__ import annotations import torch from torch import nn from minijev.choices import ChoiceEncoder from minijev.core import NEG, MiniJEVCore from minijev.encoders import STATE_INPUTS, PretrainedTextBackbone, QuestionEncoder, build_state_encoder from minijev.tokenizer import PAD, VOCAB_SIZE # pre-v2 checkpoints: flat module names -> new modular names _LEGACY_PREFIXES = [("question.", "question_encoder.pooler."), ("choice.", "choice_encoder.pooler."), ("blocks.", "core.blocks."), ("ln_ctx.", "core.ln_ctx."), ("decision.", "core.decision."), ("score.", "core.choice_head."), ("value.", "core.value_head."), ("cls", "core.cls"), ("type_emb", "core.type_emb")] def upgrade_state_dict(sd: dict) -> dict: if not any(k.startswith("blocks.") for k in sd): return sd out = {} for k, v in sd.items(): for old, new in _LEGACY_PREFIXES: if k.startswith(old): k = new + k[len(old):] break out[k] = v return out class MiniJEVNet(nn.Module): def __init__(self, cfg: dict, encoder_classes: dict | None = None): super().__init__() self.cfg = cfg d = cfg["d_model"] tcfg = dict(cfg.get("text_encoder") or {"type": "byte"}) self.text_type = tcfg.get("type", "byte") if self.text_type == "byte": self.embed = nn.Embedding(VOCAB_SIZE, d, padding_idx=PAD) self.backbone, d_text, embed = None, None, self.embed self.pad_id = PAD else: from minijev.tokenizers import make_tokenizer self.pad_id = make_tokenizer(tcfg).pad_id self.backbone = PretrainedTextBackbone(tcfg["name"], tcfg.get("finetune", "frozen"), self.pad_id) d_text, embed = self.backbone.dim, None self.state_encoders = nn.ModuleDict({ kind: build_state_encoder(kind, d, enc_cfg, embed=embed, d_text=d_text, cls=(encoder_classes or {}).get(kind)) for kind, enc_cfg in cfg["state_encoders"].items() }) self.question_encoder = QuestionEncoder(d, embed, cfg.get("max_question_len", 32), d_text, tokens=tcfg.get("question_tokens", True)) self.choice_encoder = ChoiceEncoder(d, embed, cfg.get("max_choice_len", 16), d_text) self.core = MiniJEVCore(d, cfg["n_layers"], cfg.get("n_heads", 4), cfg.get("mlp_ratio", 4.0), cfg.get("dropout", 0.0), cfg.get("n_decision_layers", 1), cfg.get("scorer", "cross"), cfg.get("max_choices", 16)) # post-hoc calibration: logits are divided by T = exp(log_temperature) at inference self.register_buffer("log_temperature", torch.zeros(())) for name, m in self.named_children(): if name != "backbone": m.apply(self._init) @staticmethod def _init(m): if isinstance(m, nn.Linear): nn.init.trunc_normal_(m.weight, std=0.02) if m.bias is not None: nn.init.zeros_(m.bias) @property def temperature(self) -> float: return float(self.log_temperature.exp()) # ----------------------------------------------------------- checkpoints def _frozen_backbone(self) -> bool: return self.backbone is not None and self.backbone.frozen def state_dict(self, *a, **kw): sd = super().state_dict(*a, **kw) if self._frozen_backbone(): # frozen weights are reloaded from the hub/cache by name for k in [k for k in sd if k.startswith("backbone.")]: del sd[k] return sd def load_state_dict(self, sd, strict: bool = True, assign: bool = False): sd = upgrade_state_dict(dict(sd)) res = super().load_state_dict(sd, strict=False, assign=assign) missing = [k for k in res.missing_keys if not k.startswith("core.scalar_head.") # added in v2, untrained in old checkpoints and not (self._frozen_backbone() and k.startswith("backbone."))] if strict and (missing or res.unexpected_keys): raise RuntimeError(f"state_dict mismatch: missing={missing[:10]} unexpected={res.unexpected_keys[:10]}") return res # --------------------------------------------------------------- encoders def run_backbone(self, ids): """Backbone on the unique rows of ids (..., L) -> (tok, mask, pooled) shaped like ids.""" shape = ids.shape[:-1] flat = ids.reshape(-1, ids.shape[-1]) uniq, inv = torch.unique(flat, dim=0, return_inverse=True) tok, mask, pooled = self.backbone(uniq) return (tok[inv].reshape(*shape, *tok.shape[1:]), mask[inv].reshape(*shape, mask.shape[-1]), pooled[inv].reshape(*shape, pooled.shape[-1])) def encode_state(self, state: dict): for kind, enc in self.state_encoders.items(): key = STATE_INPUTS[kind] if key in state: x = state[key] if getattr(enc, "uses_backbone", False): x = state.get("state_bb") or self.backbone(x) return enc(x) if not state or "state_none" in state: return None, None raise ValueError(f"no state input for encoders {list(self.state_encoders)}; got {list(state)}") def encode_question(self, question_ids=None, question_bb=None): if self.backbone is None: return self.question_encoder(question_ids) return self.question_encoder(backbone_out=question_bb or self.run_backbone(question_ids)) def encode_choices(self, choice_ids=None, choice_bb=None, choice_mask=None): """-> (emb (B, N, d), mask (B, N)). choice_bb = (pooled (B, N, dt), mask (B, N)) precomputed.""" if self.backbone is None: return self.choice_encoder(choice_ids), (choice_ids != PAD).any(-1) if choice_bb is not None: pooled, choice_mask = choice_bb else: choice_mask = (choice_ids != self.pad_id).any(-1) pooled = self.run_backbone(choice_ids)[2] return self.choice_encoder(pooled=pooled), choice_mask def forward(self, state: dict, question_ids=None, choice_ids=None, question_bb=None, choice_bb=None): """choice_ids: (B, N, Lc); all-pad rows are absent choices. Returns raw (uncalibrated) logits.""" s_tok, s_mask = self.encode_state(state) q_tok, q_mask = self.encode_question(question_ids, question_bb) c_emb, c_mask = self.encode_choices(choice_ids, choice_bb) return self.core(s_tok, s_mask, q_tok, q_mask, c_emb, c_mask) # ------------------------------------------------------------ accounting def num_parameters(self) -> int: return sum(p.numel() for p in self.parameters()) def param_breakdown(self) -> dict: """Unique parameter counts per stage (shared byte embeddings counted once, in `text`).""" seen, out = set(), {} groups = [("text", ["embed", "backbone"]), ("question_encoder", ["question_encoder"]), ("choice_encoder", ["choice_encoder"]), ("state_encoders", ["state_encoders"]), ("core", ["core"])] for g, prefixes in groups: n = 0 for name, p in self.named_parameters(remove_duplicate=False): if name.split(".")[0] in prefixes and id(p) not in seen: seen.add(id(p)) n += p.numel() out[g] = n out["trainable"] = sum(p.numel() for p in self.parameters() if p.requires_grad) out["total"] = sum(v for k, v in out.items() if k != "trainable") out["core_total"] = out["total"] - out["text"] # everything but the (shared) text backbone return out class ExportWrapper(nn.Module): """Flat-signature module for ONNX: calibrated probabilities + logits + embedding.""" def __init__(self, model: MiniJEVNet, state_key: str): super().__init__() self.model, self.state_key = model, state_key def forward(self, state, question_ids, choice_ids): out = self.model({self.state_key: state}, question_ids, choice_ids) logits = out["logits"] / self.model.log_temperature.exp() return torch.softmax(logits, -1), logits, out["embedding"] def build_model(model_cfg: dict, task_encoder_cfgs: dict[str, dict] | None = None, encoder_classes: dict | None = None) -> MiniJEVNet: cfg = dict(model_cfg) encs = {k: dict(v or {}) for k, v in cfg.get("state_encoders", {}).items()} if task_encoder_cfgs is not None: # model-side overrides (e.g. stem width) apply only to modalities the task actually uses encs = {kind: {**(extra or {}), **encs.get(kind, {})} for kind, extra in task_encoder_cfgs.items()} cfg["state_encoders"] = encs return MiniJEVNet(cfg, encoder_classes) def human_params(n: int) -> str: if n >= 1e6: return f"{n / 1e6:.2f}M" return f"{n / 1e3:.0f}K" if n >= 1e3 else str(n) __all__ = ["MiniJEVNet", "ExportWrapper", "build_model", "human_params", "STATE_INPUTS", "NEG", "upgrade_state_dict"]