"""The scorer as a HuggingFace model: config, weights, and registration. A checkpoint used to be ``torch.save({"schema", "state_dict", "args"})``, which cost three things. Loading one ran arbitrary pickle, because the payload carried the training ``args`` beside the tensors and every reader had to pass ``weights_only=False``. Every load read the pretrained encoder and then threw it away -- 4.32 s to fetch 473 tensors that ``load_state_dict`` immediately overwrote, against 0.30 s for the same weights through ``from_pretrained``. And the layout was private, so four call sites each re-derived it and the encoder's identity lived in a free-text ``args`` field rather than in a config, which is why every entry point wanted ``--hf-model`` passed alongside the checkpoint. So the model is a ``PreTrainedModel``. ``save_pretrained`` writes ``config.json`` plus safetensors, ``from_pretrained`` reads them once, the encoder config nests inside the model config, and the schema travels with the weights. Registration below makes the directory loadable as ``AutoModel.from_pretrained(path, trust_remote_code=True)`` by anyone with transformers and no copy of this package. **Dtypes are mixed and that is load-bearing.** The encoder is bf16; the head -- ``option_proj``, ``context_proj``, ``norm`` -- is fp32, which is how training produced it and what ``layer_norm`` needs. Both failure modes here were measured: building the model in fp32 before loading upcast the whole encoder, ``F32`` went to disk for all 477 tensors, and predictions drifted 4.7e-03; passing ``dtype=torch.bfloat16`` to ``from_pretrained`` casts the head too and ``layer_norm`` then raises. The per-module dtype has to come from ``__init__``, which is what ``_dtype_for`` below is for. """ from __future__ import annotations import torch from torch import nn from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel from .schema import Schema ENCODER_DTYPE = torch.bfloat16 HEAD_DTYPE = torch.float32 DEFAULT_TRAINING_CONTEXT_TOKENS = 16_384 class PackedScorerConfig(PretrainedConfig): """Everything needed to rebuild the scorer, including the encoder's own config. ``schema`` is the JSON ``Schema.to_json`` already produces, so the schema a checkpoint was trained with travels with its weights. ``context_tokens`` records the 16K state budget used in training. Inference callers can select a larger state window without changing the model weights. Candidate text has no separate token limit. """ model_type = "mojev-scorer" sub_configs = {"encoder_config": AutoConfig} def __init__(self, encoder_config=None, rank: int = 512, schema=None, context_tokens: int = DEFAULT_TRAINING_CONTEXT_TOKENS, encoder_name: str | None = None, **kwargs) -> None: # from_dict hands back a plain dict; for_model rebuilds the right class # using the model_type that to_dict preserved. if isinstance(encoder_config, dict): encoder_config = AutoConfig.for_model(**encoder_config) self.encoder_config = encoder_config self.rank = rank self.schema = schema if context_tokens < 1: raise ValueError("context_tokens must be positive") self.context_tokens = context_tokens kwargs.pop("option_tokens", None) # Kept for provenance only. Loading does not consult it -- the weights are # in the checkpoint -- but it records which encoder the run started from. self.encoder_name = encoder_name super().__init__(**kwargs) class PackedScorer(PreTrainedModel): """Encode context and every option in one sequence, then score each option. Layout per row: [context tokens] [opt 0 tokens] [opt 1 tokens] ... [opt N-1 tokens] The attention mask is block-diagonal over the option spans: an option sees the context and itself, never another option. Without that, scores would depend on which rivals happened to be present and permutation invariance would be gone -- the property that separates this from putting the choices in a prompt. """ config_class = PackedScorerConfig base_model_prefix = "encoder" _supports_sdpa = True def __init__(self, config: PackedScorerConfig) -> None: super().__init__(config) self.encoder = AutoModel.from_config(config.encoder_config, dtype=ENCODER_DTYPE) width = self.hidden_size(config.encoder_config) self.width = width self.rank = config.rank self.option_proj = nn.Linear(width, config.rank, bias=False, dtype=HEAD_DTYPE) self.context_proj = nn.Linear(width, config.rank, bias=False, dtype=HEAD_DTYPE) self.norm = nn.LayerNorm(width, dtype=HEAD_DTYPE) # Sets all_tied_weights_keys and the rest of the composite-model state; # without it from_pretrained raises AttributeError on the first of them. self.post_init() @staticmethod def hidden_size(encoder_config) -> int: """Multimodal configs nest the text stack; single-stack ones do not.""" return getattr(encoder_config, "text_config", encoder_config).hidden_size @property def schema(self) -> Schema: return Schema.from_json(self.config.schema) @classmethod def from_encoder(cls, model_name: str, schema: Schema, rank: int = 512, context_tokens: int = DEFAULT_TRAINING_CONTEXT_TOKENS): """Build a fresh scorer on a pretrained encoder -- the training entry point. This is the one path that should read pretrained weights, because starting from them is the point. Loading a trained checkpoint goes through ``from_pretrained`` instead and reads the encoder exactly once. """ config = PackedScorerConfig( encoder_config=AutoConfig.from_pretrained(model_name), rank=rank, schema=schema.to_json(), context_tokens=context_tokens, encoder_name=str(model_name), ) model = cls(config) model.encoder = AutoModel.from_pretrained(model_name, dtype=ENCODER_DTYPE) return model def build_mask(self, context_span: torch.Tensor, field_span: torch.Tensor, option_span: torch.Tensor) -> torch.Tensor: """(B, 1, L, L) additive mask enforcing a state -> question -> candidate tree. Three levels, not two: state attends within itself question attends to the state and to itself candidate attends to the state, to *its own* question, and to itself What the two-level version got wrong is the middle row. With candidates hanging straight off the context, a question's wording reached every candidate in the request, so the eight questions of a customer-service row could see each other's text. Worse, question spans sat unmasked in the sequence, which let candidates of question A read question B -- scores then depend on which other questions happen to be in the same request, and the parallel-sampling guarantee (independent decisions, exact permutation invariance) no longer holds. Sibling isolation is the point: two candidates of the same question cannot see each other, and two questions of the same state cannot see each other. """ batch, fields, width, total = option_span.shape device = option_span.device state = context_span > 0 # (B, L) question = field_span > 0 # (B, F, L) candidate = option_span > 0 # (B, F, N, L) allow = torch.zeros(batch, total, total, dtype=torch.bool, device=device) # state -> state allow |= state[:, :, None] & state[:, None, :] # question -> state, question -> itself q_any = question.any(1) allow |= q_any[:, :, None] & state[:, None, :] allow |= torch.einsum("bfi,bfj->bij", question.float(), question.float()).bool() # candidate -> state c_any = candidate.any(1).any(1) allow |= c_any[:, :, None] & state[:, None, :] # candidate -> its own question (broadcast over that question's candidates) own_question = torch.einsum( "bfni,bfj->bij", candidate.float(), question.float() ).bool() allow |= own_question # candidate -> itself only, never a sibling allow |= torch.einsum("bfni,bfnj->bij", candidate.float(), candidate.float()).bool() # Padding positions belong to no node in the tree, so every one of their # rows would be entirely masked and softmax would divide by zero -- the # NaN this produced showed up within 200 steps. Let each position attend # to itself; the result is discarded because nothing pools from padding. eye = torch.eye(total, dtype=torch.bool, device=device) allow |= eye[None, :, :] return torch.where(allow, 0.0, torch.finfo(torch.float32).min).unsqueeze(1) def forward(self, batch: dict) -> torch.Tensor: # The tree mask, not a plain padding mask: without it every span in the # sequence is mutually visible and the per-decision independence this # design rests on is lost. mask = self.build_mask( batch["context_span"], batch["field_span"], batch["option_span"] ) # Padding columns are unreachable, but the diagonal must survive or the # padding rows go fully masked again and softmax divides by zero. floor = torch.finfo(mask.dtype).min total = mask.shape[-1] keep = batch["packed_mask"][:, None, None, :] | torch.eye( total, dtype=torch.bool, device=mask.device )[None, None] mask = mask.masked_fill(~keep, floor) extra = {} if "pixel_values" in batch: extra["pixel_values"] = batch["pixel_values"] extra["image_grid_thw"] = batch["image_grid_thw"] # The processor sizes mm_token_type_ids to the state alone; packing # appends questions and candidates, so it is padded with zeros (text) # out to the full sequence. token_types = batch["mm_token_type_ids"] if token_types.shape[1] < total: token_types = torch.cat([ token_types, token_types.new_zeros(token_types.shape[0], total - token_types.shape[1]), ], dim=1) extra["mm_token_type_ids"] = token_types[:, :total] # M-RoPE derives 3D positions by indexing the attention mask as a 2D # (B, L) padding mask. Ours is the 4D additive tree mask, which that # code cannot read -- it raised IndexError on attention_mask[b].bool(). # Compute the positions here from the real padding mask and hand them # over, so the encoder skips its own derivation. Position ids are # per-token, so the tree mask is irrelevant to them. extra["position_ids"] = self.encoder.get_rope_index( batch["packed_ids"], image_grid_thw=batch["image_grid_thw"], attention_mask=batch["packed_mask"].long(), mm_token_type_ids=extra["mm_token_type_ids"], )[0] hidden = self.encoder( input_ids=batch["packed_ids"], attention_mask=mask, **extra, ).last_hidden_state.float() hidden = self.norm(hidden) # Mean-pool the context span and each option span. context = (hidden * batch["context_span"].unsqueeze(-1)).sum(1) context = context / batch["context_span"].sum(-1, keepdim=True).clamp_min(1) query = self.context_proj(context) # (B, r) spans = batch["option_span"] # (B, F, N, L) weights = spans.sum(-1, keepdim=True).clamp_min(1) pooled = torch.einsum("bfnl,blw->bfnw", spans, hidden) / weights keys = self.option_proj(pooled) # (B, F, N, r) # The field's words are pooled from the sequence like everything else, # and added to the context query. No per-field parameters exist. field = torch.einsum("bfl,blw->bfw", batch["field_span"], hidden) field = field / batch["field_span"].sum(-1, keepdim=True).clamp_min(1) query = query[:, None, :] + self.context_proj(field) # (B, F, r) logits = (query[:, :, None, :] * keys).sum(-1) / self.rank ** 0.5 return logits.masked_fill(~batch["option_mask"], torch.finfo(logits.dtype).min) # Registering against the Auto classes, and recording auto_map in config.json, is # what makes a saved directory the whole model: someone with transformers and no # copy of this package can load it with # AutoModel.from_pretrained(path, trust_remote_code=True). These have to run at # module import in a real module -- defining the classes in __main__ silently # writes auto_map: null and the directory stops being portable. AutoConfig.register(PackedScorerConfig.model_type, PackedScorerConfig, exist_ok=True) AutoModel.register(PackedScorerConfig, PackedScorer, exist_ok=True) PackedScorerConfig.register_for_auto_class() PackedScorer.register_for_auto_class("AutoModel")