| """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 |
|
|
| |
| _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)) |
| |
| 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()) |
|
|
| |
| 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(): |
| 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.") |
| 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 |
|
|
| |
| 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) |
|
|
| |
| 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"] |
| 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: |
| |
| 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"] |
|
|
