code/minijev/model/minijev.py

"""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"]