| """MIMIR-1 as a transformers model: load with `AutoModel`, forward on encoded batches. |
| |
| Module and parameter names are the training checkpoint's state-dict keys. Renaming |
| one breaks `from_pretrained` for everybody who downloaded the release, so the head |
| below keeps the training code's attribute names exactly (`projection`, `roles`, |
| `typed_values`, `state_memory`, `candidate_processor`, `reasoner` with `.cell`, |
| `readout` with `utility`/`abstain`/`ordinal_score`/`threshold`/`histogram`), and the |
| encoder is a stock `ModernBertModel` under `self.encoder` for zero remapping. |
| |
| The computation is the same one the shipped `model.onnx` already carries, so this |
| file ports the head's forward only: no loss, no training machinery. Text handling, |
| calibration, the server and MCP stay in the `mimir-decisions` package; this path is load |
| plus forward on the 21 flat named tensors the graph reads, returning every depth's |
| readouts, evidence and workspace. |
| |
| The typed-value quantile edges live in the checkpoint buffers, so the constructor |
| takes them as optional: without them the buffers are zeros of the right shape and |
| `load_state_dict` fills the real values. A forward before loading is meaningless. |
| |
| The encoder forward replicates the training encoder call exactly (embeddings, per |
| layer-kind rotary, direct layer calls with prepared dense masks, typed-value |
| injection at `injection_layer`, final norm). SDPA only; tested transformers==5.17. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| from dataclasses import dataclass |
| from typing import Any |
|
|
| import numpy as np |
| import numpy.typing as npt |
| import torch |
| from torch import nn |
| from transformers import ModernBertModel, PreTrainedModel |
| from transformers.utils.generic import ModelOutput |
|
|
| try: |
| from .configuration_mimir import DECISION_TYPES, MimirConfig |
| except ImportError: |
| |
| from configuration_mimir import ( |
| DECISION_TYPES, |
| MimirConfig, |
| ) |
|
|
| STATE_ROLE = 0 |
| QUESTION_ROLE = 1 |
| CANDIDATE_ROLE = 2 |
|
|
| FULL_ATTENTION = "full_attention" |
| SLIDING_ATTENTION = "sliding_attention" |
|
|
| NUMBER = 0 |
| DATE = 1 |
| CYCLES = 4 |
| SECONDS_PER_DAY = 86_400 |
| FEBRUARY = 2 |
| EPOCH_WEEKDAY = 3 |
|
|
| RELATIONS = ("supports", "contradicts", "irrelevant") |
|
|
| INPUTS = ( |
| "chunk_ids", |
| "chunk_mask", |
| "chunk_record", |
| "question_length", |
| "key_chunk", |
| "key_position", |
| "key_mask", |
| "candidate_ids", |
| "candidate_mask", |
| "candidate_text_mask", |
| "candidate_index", |
| "candidate_present", |
| "decision_type", |
| "typed_kind", |
| "typed_number", |
| "typed_year", |
| "typed_month", |
| "typed_day", |
| "typed_second", |
| "typed_chunk", |
| "typed_position", |
| ) |
| """The forward's named inputs, in order: the same 21 the ONNX graph reads.""" |
|
|
|
|
| class Norm(nn.RMSNorm): |
| """RMSNorm in float32, returned in the input's dtype.""" |
|
|
| def forward(self, states: torch.Tensor) -> torch.Tensor: |
| normed: torch.Tensor = super().forward(states.float()) |
| return normed.to(states.dtype) |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class Memory: |
| """Keys (normalised) and values over [batch, heads, length, size], with the key mask.""" |
|
|
| key: torch.Tensor |
| value: torch.Tensor |
| mask: torch.Tensor |
|
|
| def extend(self, other: Memory) -> Memory: |
| return Memory( |
| torch.cat([self.key, other.key], dim=2), |
| torch.cat([self.value, other.value], dim=2), |
| torch.cat([self.mask, other.mask], dim=1), |
| ) |
|
|
|
|
| def masked_attention(query: torch.Tensor, memory: Memory) -> torch.Tensor: |
| """Fused attention; a row with no valid key reads zeros instead of NaN.""" |
| bias = torch.zeros(memory.mask.shape, dtype=query.dtype, device=query.device) |
| bias = bias.masked_fill(~memory.mask, torch.finfo(query.dtype).min)[:, None, None, :] |
| read: torch.Tensor = nn.functional.scaled_dot_product_attention( |
| query, memory.key, memory.value.to(query.dtype), attn_mask=bias |
| ) |
| return read * memory.mask.any(dim=1)[:, None, None, None] |
|
|
|
|
| def attention_weights(query: torch.Tensor, memory: Memory) -> torch.Tensor: |
| """The same attention's weights in float32, zero on masked keys and on keyless rows.""" |
| logits = (query @ memory.key.transpose(-1, -2)).float() / math.sqrt(query.shape[-1]) |
| mask = memory.mask[:, None, None, :] |
| logits = logits.masked_fill(~mask, torch.finfo(logits.dtype).min) |
| return logits.softmax(dim=-1) * mask.any(dim=-1, keepdim=True) |
|
|
|
|
| class Attention(nn.Module): |
| """Multi-head attention from queries onto a memory, with RMSNorm on queries and keys.""" |
|
|
| def __init__(self, config: MimirConfig, *, zero_output: bool) -> None: |
| super().__init__() |
| self.heads = config.heads |
| self.head_size = config.head_size |
| self.query = nn.Linear(config.width, config.width, bias=False) |
| self.key_value = nn.Linear(config.width, 2 * config.width, bias=False) |
| self.output = nn.Linear(config.width, config.width, bias=False) |
| self.query_norm = Norm(config.head_size) |
| self.key_norm = Norm(config.head_size) |
| if zero_output: |
| nn.init.zeros_(self.output.weight) |
|
|
| def _split(self, states: torch.Tensor) -> torch.Tensor: |
| batch, length, _ = states.shape |
| return states.view(batch, length, self.heads, self.head_size).transpose(1, 2) |
|
|
| def _merge(self, read: torch.Tensor) -> torch.Tensor: |
| batch, _, length, _ = read.shape |
| width = self.heads * self.head_size |
| merged: torch.Tensor = self.output(read.transpose(1, 2).reshape(batch, length, width)) |
| return merged |
|
|
| def memory(self, states: torch.Tensor, mask: torch.Tensor) -> Memory: |
| key, value = self.key_value(states).chunk(2, dim=-1) |
| return Memory(self.key_norm(self._split(key)), self._split(value), mask) |
|
|
| def queries(self, states: torch.Tensor) -> torch.Tensor: |
| split: torch.Tensor = self.query_norm(self._split(self.query(states))) |
| return split |
|
|
| def read(self, states: torch.Tensor, memory: Memory) -> torch.Tensor: |
| return self._merge(masked_attention(self.queries(states), memory)) |
|
|
| def read_with_weights( |
| self, states: torch.Tensor, memory: Memory |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| weights = attention_weights(self.queries(states), memory) |
| return self._merge(weights.to(memory.value.dtype) @ memory.value), weights |
|
|
|
|
| class AttentionBlock(nn.Module): |
| """states + Attention(norm(states), norm(memory)); self-attention when memory is states.""" |
|
|
| def __init__(self, config: MimirConfig) -> None: |
| super().__init__() |
| self.query_norm = Norm(config.width) |
| self.memory_norm = Norm(config.width) |
| self.attention = Attention(config, zero_output=True) |
|
|
| def memory(self, states: torch.Tensor, mask: torch.Tensor) -> Memory: |
| return self.attention.memory(self.memory_norm(states), mask) |
|
|
| def read(self, states: torch.Tensor, memory: Memory) -> torch.Tensor: |
| return states + self.attention.read(self.query_norm(states), memory) |
|
|
| def forward( |
| self, states: torch.Tensor, memory: torch.Tensor, memory_mask: torch.Tensor |
| ) -> torch.Tensor: |
| return self.read(states, self.memory(memory, memory_mask)) |
|
|
|
|
| class FeedForwardBlock(nn.Module): |
| """states + SwiGLU(norm(states)).""" |
|
|
| def __init__(self, config: MimirConfig) -> None: |
| super().__init__() |
| self.norm = Norm(config.width) |
| self.gate_up = nn.Linear(config.width, 2 * config.feedforward_hidden, bias=False) |
| self.down = nn.Linear(config.feedforward_hidden, config.width, bias=False) |
| nn.init.zeros_(self.down.weight) |
|
|
| def forward(self, states: torch.Tensor) -> torch.Tensor: |
| gate, up = self.gate_up(self.norm(states)).chunk(2, dim=-1) |
| update: torch.Tensor = self.down(nn.functional.silu(gate) * up) |
| return states + update |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class TypedValues: |
| kind: torch.Tensor |
| number: torch.Tensor |
| year: torch.Tensor |
| month: torch.Tensor |
| day: torch.Tensor |
| second: torch.Tensor |
|
|
|
|
| def signed_log(values: torch.Tensor) -> torch.Tensor: |
| return values.sign() * values.abs().log1p() |
|
|
|
|
| def days_from_civil(year: torch.Tensor, month: torch.Tensor, day: torch.Tensor) -> torch.Tensor: |
| """Days since 1970-01-01 in the proleptic Gregorian calendar (Hinnant's algorithm).""" |
| shifted = year - (month <= FEBRUARY).long() |
| era = torch.div(shifted, 400, rounding_mode="floor") |
| year_of_era = shifted - era * 400 |
| month_index = (month + 9) % 12 |
| day_of_year = torch.div(153 * month_index + 2, 5, rounding_mode="floor") + day - 1 |
| day_of_era = ( |
| year_of_era * 365 |
| + torch.div(year_of_era, 4, rounding_mode="floor") |
| - torch.div(year_of_era, 100, rounding_mode="floor") |
| + day_of_year |
| ) |
| return era * 146_097 + day_of_era - 719_468 |
|
|
|
|
| def piecewise_linear(values: torch.Tensor, edges: torch.Tensor) -> torch.Tensor: |
| """[n] values over T+1 edges -> [n, T]; the first bin is open below, the last above.""" |
| low, high = edges[:-1], edges[1:] |
| width = high - low |
| safe = torch.where(width > 0, width, torch.ones_like(width)) |
| fraction = (values[:, None] - low) / safe |
| fraction = torch.where(width > 0, fraction, (values[:, None] >= high).to(values.dtype)) |
| bins = edges.shape[0] - 1 |
| first = torch.arange(bins, device=values.device) == 0 |
| last = torch.arange(bins, device=values.device) == bins - 1 |
| fraction = torch.where(first, fraction, fraction.clamp(min=0)) |
| return torch.where(last, fraction, fraction.clamp(max=1)) |
|
|
|
|
| class TypedValueEncoder(nn.Module): |
| """Number bins, day bins and 4 date cycles -> MLP -> encoder width; zero at initialisation.""" |
|
|
| def __init__( |
| self, |
| config: MimirConfig, |
| number_edges: npt.NDArray[np.float64] | None = None, |
| day_edges: npt.NDArray[np.float64] | None = None, |
| ) -> None: |
| super().__init__() |
| for name, edges in (("number_edges", number_edges), ("day_edges", day_edges)): |
| if edges is None: |
| continue |
| if edges.shape != (config.value_bins + 1,) or np.any(np.diff(edges) < 0): |
| message = f"{name} of shape {edges.shape} for {config.value_bins} sorted bins" |
| raise ValueError(message) |
| missing = np.zeros(config.value_bins + 1, dtype=np.float64) |
| self.number_edges: torch.Tensor |
| self.day_edges: torch.Tensor |
| self.register_buffer( |
| "number_edges", |
| torch.from_numpy(number_edges if number_edges is not None else missing).float(), |
| ) |
| self.register_buffer( |
| "day_edges", |
| torch.from_numpy(day_edges if day_edges is not None else missing).float(), |
| ) |
| features = 2 * config.value_bins + 2 * CYCLES |
| self.hidden = nn.Linear(features, config.width, bias=False) |
| self.output = nn.Linear(config.width, config.encoder_width, bias=False) |
| nn.init.zeros_(self.output.weight) |
|
|
| def features(self, values: TypedValues) -> torch.Tensor: |
| is_number = values.kind == NUMBER |
| is_date = values.kind == DATE |
| number = torch.nan_to_num(values.number.float(), nan=0.0, posinf=0.0, neginf=0.0) |
| number_part = piecewise_linear(signed_log(number), self.number_edges) |
| days = days_from_civil(values.year, values.month.clamp(min=1), values.day.clamp(min=1)) |
| day_part = piecewise_linear(days.float(), self.day_edges) |
| weekday = (days + EPOCH_WEEKDAY) % 7 |
| angles = ( |
| 2 |
| * math.pi |
| * torch.stack( |
| [ |
| (values.month.float() - 1) / 12, |
| (values.day.float() - 1) / 31, |
| weekday.float() / 7, |
| values.second.float() / SECONDS_PER_DAY, |
| ], |
| dim=-1, |
| ) |
| ) |
| cycles = torch.cat([angles.sin(), angles.cos()], dim=-1) |
| return torch.cat( |
| [ |
| number_part * is_number[:, None], |
| day_part * is_date[:, None], |
| cycles * is_date[:, None], |
| ], |
| dim=-1, |
| ) |
|
|
| def forward(self, values: TypedValues) -> torch.Tensor: |
| encoded: torch.Tensor = self.output(nn.functional.silu(self.hidden(self.features(values)))) |
| return encoded |
|
|
|
|
| class StateMemory(nn.Module): |
| """Compress a record's chunk tokens into question-conditioned latents Z.""" |
|
|
| def __init__(self, config: MimirConfig) -> None: |
| super().__init__() |
| self.latents = nn.Parameter(torch.empty(config.latents, config.width)) |
| nn.init.orthogonal_(self.latents) |
| self.question = nn.Linear(config.width, config.width, bias=False) |
| self.read = AttentionBlock(config) |
| self.read_feedforward = FeedForwardBlock(config) |
| self.mix = AttentionBlock(config) |
| self.mix_feedforward = FeedForwardBlock(config) |
|
|
| def forward( |
| self, question: torch.Tensor, tokens: torch.Tensor, token_mask: torch.Tensor |
| ) -> torch.Tensor: |
| latents: torch.Tensor = self.latents + self.question(question)[:, None] |
| latents = self.read_feedforward(self.read(latents, tokens, token_mask)) |
| every = torch.ones(latents.shape[:2], dtype=torch.bool, device=latents.device) |
| memory: torch.Tensor = self.mix_feedforward(self.mix(latents, latents, every)) |
| return memory |
|
|
|
|
| class CandidateProcessor(nn.Module): |
| """Candidates as a set: each reads the question, then all attend to all, with no positions.""" |
|
|
| def __init__(self, config: MimirConfig) -> None: |
| super().__init__() |
| self.read_question = AttentionBlock(config) |
| self.interact = AttentionBlock(config) |
| self.feedforward = FeedForwardBlock(config) |
|
|
| def forward( |
| self, |
| candidates: torch.Tensor, |
| candidate_mask: torch.Tensor, |
| question: torch.Tensor, |
| question_mask: torch.Tensor, |
| ) -> torch.Tensor: |
| candidates = self.read_question(candidates, question, question_mask) |
| candidates = self.interact(candidates, candidates, candidate_mask) |
| refined: torch.Tensor = self.feedforward(candidates) |
| return refined |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class Context: |
| """The fixed memories every iteration reads, already projected.""" |
|
|
| recall: Memory |
| tokens: Memory |
| candidate_state: Memory |
| candidate_mask: torch.Tensor |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class Iteration: |
| workspace: torch.Tensor |
| candidates: torch.Tensor |
| evidence: torch.Tensor |
| gap_weights: torch.Tensor |
|
|
|
|
| def _every(states: torch.Tensor) -> torch.Tensor: |
| return torch.ones(states.shape[:2], dtype=torch.bool, device=states.device) |
|
|
|
|
| class ReasonerCell(nn.Module): |
| """The shared reasoning cell over workspace slots, applied once per iteration.""" |
|
|
| def __init__(self, config: MimirConfig) -> None: |
| super().__init__() |
| self.recall = AttentionBlock(config) |
| self.slot_norm = Norm(config.width) |
| self.token_norm = Norm(config.width) |
| self.evidence_norm = Norm(config.width) |
| self.gap = Attention(config, zero_output=False) |
| self.relation = nn.Linear(3 * config.width, len(RELATIONS), bias=False) |
| nn.init.zeros_(self.relation.weight) |
| self.support = nn.Linear(config.width, config.width, bias=False) |
| self.contradict = nn.Linear(config.width, config.width, bias=False) |
| nn.init.zeros_(self.support.weight) |
| nn.init.zeros_(self.contradict.weight) |
| self.workspace_feedforward = FeedForwardBlock(config) |
| self.read_candidates = AttentionBlock(config) |
| self.candidates_read = AttentionBlock(config) |
| self.interact = AttentionBlock(config) |
| self.candidate_feedforward = FeedForwardBlock(config) |
|
|
| def context( |
| self, |
| memory: torch.Tensor, |
| question: torch.Tensor, |
| question_mask: torch.Tensor, |
| tokens: torch.Tensor, |
| token_mask: torch.Tensor, |
| candidate_mask: torch.Tensor, |
| ) -> Context: |
| recall = torch.cat([memory, question], dim=1) |
| state = torch.cat([memory, tokens], dim=1) |
| return Context( |
| recall=self.recall.memory(recall, torch.cat([_every(memory), question_mask], dim=1)), |
| tokens=self.gap.memory(self.token_norm(tokens), token_mask), |
| candidate_state=self.candidates_read.memory( |
| state, torch.cat([_every(memory), token_mask], dim=1) |
| ), |
| candidate_mask=candidate_mask, |
| ) |
|
|
| def forward( |
| self, |
| workspace: torch.Tensor, |
| candidates: torch.Tensor, |
| evidence: torch.Tensor, |
| context: Context, |
| ) -> Iteration: |
| own = torch.cat([workspace, evidence], dim=1) |
| workspace = self.recall.read( |
| workspace, self.recall.memory(own, _every(own)).extend(context.recall) |
| ) |
| query = self.slot_norm(workspace) |
| found, gap_weights = self.gap.read_with_weights(query, context.tokens) |
| found_norm = self.evidence_norm(found) |
| relation = self.relation(torch.cat([query, found_norm, query * found_norm], dim=-1)) |
| weights = relation.float().softmax(dim=-1).to(workspace.dtype) |
| workspace = ( |
| workspace |
| + weights[..., 0:1] * self.support(found_norm) |
| + weights[..., 1:2] * self.contradict(found_norm) |
| ) |
| evidence = torch.cat([evidence, found], dim=1) |
| workspace = self.workspace_feedforward(workspace) |
| workspace = self.read_candidates(workspace, candidates, context.candidate_mask) |
| own = torch.cat([workspace, evidence], dim=1) |
| candidates = self.candidates_read.read( |
| candidates, |
| self.candidates_read.memory(own, _every(own)).extend(context.candidate_state), |
| ) |
| candidates = self.interact(candidates, candidates, context.candidate_mask) |
| candidates = self.candidate_feedforward(candidates) |
| return Iteration(workspace, candidates, evidence, gap_weights) |
|
|
|
|
| class Reasoner(nn.Module): |
| """Workspace slots seeded by the schema and the question, then the cell once per iteration.""" |
|
|
| def __init__(self, config: MimirConfig) -> None: |
| super().__init__() |
| self.iterations = config.iterations |
| self.slots = nn.Parameter(torch.empty(config.slots, config.width)) |
| nn.init.orthogonal_(self.slots) |
| self.schemas = nn.Embedding(len(DECISION_TYPES), config.width) |
| self.question = nn.Linear(config.width, config.width, bias=False) |
| self.cell = ReasonerCell(config) |
|
|
| def forward( |
| self, |
| question_mean: torch.Tensor, |
| decision_type: torch.Tensor, |
| candidates: torch.Tensor, |
| context: Context, |
| ) -> list[Iteration]: |
| seed = self.schemas(decision_type) + self.question(question_mean) |
| workspace = self.slots + seed[:, None] |
| evidence = workspace.new_zeros((workspace.shape[0], 0, workspace.shape[2])) |
| iterations: list[Iteration] = [] |
| for _ in range(self.iterations): |
| step = self.cell(workspace, candidates, evidence, context) |
| iterations.append(step) |
| workspace, candidates, evidence = step.workspace, step.candidates, step.evidence |
| return iterations |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class Readout: |
| utilities: torch.Tensor |
| abstain: torch.Tensor |
| ordinal_score: torch.Tensor |
| thresholds: torch.Tensor |
| histogram_logits: torch.Tensor |
|
|
|
|
| def readout_mlp(inputs: int, width: int, outputs: int) -> nn.Sequential: |
| return nn.Sequential( |
| Norm(inputs), |
| nn.Linear(inputs, width, bias=False), |
| nn.SiLU(), |
| nn.Linear(width, outputs), |
| ) |
|
|
|
|
| class DecisionReadout(nn.Module): |
| def __init__(self, config: MimirConfig) -> None: |
| super().__init__() |
| width = config.width |
| self.utility = readout_mlp(width, width, 1) |
| self.abstain = readout_mlp(width, width, 1) |
| self.ordinal_score = readout_mlp(width, width, 1) |
| self.threshold = readout_mlp(2 * width, width, 1) |
| self.histogram = readout_mlp(width, width, config.histogram_bins) |
|
|
| def forward(self, pooled: torch.Tensor, candidates: torch.Tensor) -> Readout: |
| """`pooled` is the workspace averaged over its slots.""" |
| pairs = torch.cat([candidates[:, :-1], candidates[:, 1:]], dim=-1) |
| raw = self.threshold(pairs).squeeze(-1) |
| steps = nn.functional.softplus(raw[:, 1:]) |
| thresholds = torch.cat([raw[:, :1], raw[:, :1] + steps.cumsum(dim=1)], dim=1) |
| return Readout( |
| utilities=self.utility(candidates).squeeze(-1), |
| abstain=self.abstain(pooled).squeeze(-1), |
| ordinal_score=self.ordinal_score(pooled).squeeze(-1), |
| thresholds=thresholds, |
| histogram_logits=self.histogram(pooled), |
| ) |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class HeadInputs: |
| """Encoder states per record: state keys, question tokens and pooled candidates.""" |
|
|
| tokens: torch.Tensor |
| token_mask: torch.Tensor |
| question: torch.Tensor |
| question_mask: torch.Tensor |
| candidates: torch.Tensor |
| candidate_mask: torch.Tensor |
| decision_type: torch.Tensor |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class DepthOutput: |
| """One depth's decision; `workspace` is the slot-averaged workspace the readout reads.""" |
|
|
| readout: Readout |
| evidence_scores: torch.Tensor |
| workspace: torch.Tensor |
|
|
|
|
| class DecisionHead(nn.Module): |
| def __init__( |
| self, |
| config: MimirConfig, |
| number_edges: npt.NDArray[np.float64] | None = None, |
| day_edges: npt.NDArray[np.float64] | None = None, |
| ) -> None: |
| super().__init__() |
| self.config = config |
| self.projection = nn.Linear(config.encoder_width, config.width, bias=False) |
| self.roles = nn.Embedding(3, config.width) |
| self.typed_values = TypedValueEncoder(config, number_edges, day_edges) |
| self.state_memory = StateMemory(config) |
| self.candidate_processor = CandidateProcessor(config) |
| self.reasoner = Reasoner(config) |
| self.readout = DecisionReadout(config) |
|
|
| def _project(self, states: torch.Tensor, role: int) -> torch.Tensor: |
| projected: torch.Tensor = self.projection(states) |
| return projected + self.roles.weight[role] |
|
|
| def forward(self, inputs: HeadInputs) -> list[DepthOutput]: |
| tokens = self._project(inputs.tokens, STATE_ROLE) |
| question = self._project(inputs.question, QUESTION_ROLE) |
| candidates = self._project(inputs.candidates, CANDIDATE_ROLE) |
| weight = inputs.question_mask[..., None].to(question.dtype) |
| question_mean = (question * weight).sum(dim=1) / weight.sum(dim=1).clamp(min=1) |
| memory = self.state_memory(question_mean, tokens, inputs.token_mask) |
| candidates = self.candidate_processor( |
| candidates, inputs.candidate_mask, question, inputs.question_mask |
| ) |
| context = self.reasoner.cell.context( |
| memory, |
| question, |
| inputs.question_mask, |
| tokens, |
| inputs.token_mask, |
| inputs.candidate_mask, |
| ) |
| iterations = self.reasoner(question_mean, inputs.decision_type, candidates, context) |
| outputs: list[DepthOutput] = [] |
| for step in iterations: |
| pooled = step.workspace.mean(dim=1) |
| outputs.append( |
| DepthOutput( |
| readout=self.readout(pooled, step.candidates), |
| evidence_scores=step.gap_weights.mean(dim=(1, 2)), |
| workspace=pooled, |
| ) |
| ) |
| return outputs |
|
|
|
|
| def padded_dense_masks( |
| attention_mask: torch.Tensor, sliding_window: int |
| ) -> dict[str, torch.Tensor]: |
| positions = torch.arange(attention_mask.shape[1], device=attention_mask.device) |
| keys = attention_mask.bool()[:, None, None, :] |
| near = (positions[:, None] - positions[None, :]).abs() <= sliding_window |
| return {FULL_ATTENTION: keys, SLIDING_ATTENTION: keys & near} |
|
|
|
|
| def _pool_candidates(hidden: torch.Tensor, text_mask: torch.Tensor) -> torch.Tensor: |
| """Each candidate row's mean state over its own text tokens.""" |
| weight = text_mask[..., None].to(hidden.dtype) |
| return (hidden * weight).sum(dim=1) / weight.sum(dim=1).clamp(min=1) |
|
|
|
|
| def _mean_question( |
| question: torch.Tensor, chunk_record: torch.Tensor, question_length: torch.Tensor |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Question positions 1.. of every chunk, averaged over the chunks of each record.""" |
| positions = torch.arange(question.shape[1], device=question.device) |
| per_chunk = (positions < question_length[:, None]).to(question.dtype) |
| records = torch.arange(chunk_record.max().item() + 1, device=question.device) |
| owner = (chunk_record[None, :] == records[:, None]).to(question.dtype) |
| summed = torch.einsum("rc,cqd->rqd", owner, question * per_chunk[..., None]) |
| counts = owner @ per_chunk |
| return summed / counts.clamp(min=1)[..., None], counts > 0 |
|
|
|
|
| @dataclass |
| class MimirModelOutput(ModelOutput): |
| """Every depth's readouts stacked first, plus the final pooled workspace.""" |
|
|
| utilities: torch.Tensor | None = None |
| thresholds: torch.Tensor | None = None |
| abstain: torch.Tensor | None = None |
| ordinal_score: torch.Tensor | None = None |
| histogram_logits: torch.Tensor | None = None |
| evidence: torch.Tensor | None = None |
| workspace: torch.Tensor | None = None |
|
|
|
|
| class MimirModel(PreTrainedModel): |
| """The decision model behind `AutoModel`: the stock encoder plus the head. |
| |
| The encoder lives under `self.encoder`, the head under `self.head`. Forward takes |
| the 21 flat named tensors the release graph reads and returns every depth's |
| readouts. The batch must hold at least one candidate row, and the question states |
| are `question_cap` wide whatever the chunk width, exactly like the exported path |
| the graph was cut from. |
| """ |
|
|
| config_class = MimirConfig |
| config: MimirConfig |
| base_model_prefix = "mimir" |
| _supports_sdpa = True |
|
|
| def __init__(self, config: MimirConfig) -> None: |
| super().__init__(config) |
| self.encoder = ModernBertModel(config) |
| self.head = DecisionHead(config) |
| self.post_init() |
| self.encoder.set_attn_implementation("sdpa") |
|
|
| def _init_weights(self, module: nn.Module) -> None: |
| |
| |
| |
| if isinstance(module, Norm) and isinstance(module.weight, torch.nn.Parameter): |
| nn.init.ones_(module.weight) |
|
|
| def get_input_embeddings(self) -> nn.Module: |
| return self.encoder.get_input_embeddings() |
|
|
| def set_input_embeddings(self, value: nn.Module) -> None: |
| self.encoder.set_input_embeddings(value) |
|
|
| def _encode( |
| self, |
| input_ids: torch.Tensor, |
| mask: torch.Tensor, |
| typed: TypedValues, |
| typed_chunk: torch.Tensor, |
| typed_position: torch.Tensor, |
| *, |
| inject: bool, |
| ) -> torch.Tensor: |
| positions = torch.arange(input_ids.shape[1], device=input_ids.device) |
| positions = positions.expand(input_ids.shape[0], -1) |
| masks = padded_dense_masks(mask.long(), int(self.config.sliding_window)) |
| injection: torch.Tensor | None = None |
| injection_layer = -1 |
| if inject and typed.kind.numel(): |
| values = self.head.typed_values(typed) |
| hidden = values.new_zeros((*input_ids.shape, values.shape[-1])) |
| hidden = hidden.index_put((typed_chunk, typed_position), values, accumulate=True) |
| injection = hidden |
| injection_layer = self.config.injection_layer |
| layer_types = self.encoder.config.layer_types |
| if layer_types is None: |
| message = "encoder config names no layer types" |
| raise ValueError(message) |
| hidden = self.encoder.embeddings(input_ids=input_ids) |
| rotary = { |
| kind: self.encoder.rotary_emb(hidden, positions, kind) for kind in set(layer_types) |
| } |
| for index, (layer, kind) in enumerate(zip(self.encoder.layers, layer_types, strict=True)): |
| if injection is not None and index == injection_layer: |
| hidden = hidden + injection.to(hidden.dtype) |
| hidden = layer(hidden, attention_mask=masks[kind], position_embeddings=rotary[kind]) |
| output: torch.Tensor = self.encoder.final_norm(hidden) |
| return output |
|
|
| def forward( |
| self, |
| chunk_ids: torch.Tensor, |
| chunk_mask: torch.Tensor, |
| chunk_record: torch.Tensor, |
| question_length: torch.Tensor, |
| key_chunk: torch.Tensor, |
| key_position: torch.Tensor, |
| key_mask: torch.Tensor, |
| candidate_ids: torch.Tensor, |
| candidate_mask: torch.Tensor, |
| candidate_text_mask: torch.Tensor, |
| candidate_index: torch.Tensor, |
| candidate_present: torch.Tensor, |
| decision_type: torch.Tensor, |
| typed_kind: torch.Tensor, |
| typed_number: torch.Tensor, |
| typed_year: torch.Tensor, |
| typed_month: torch.Tensor, |
| typed_day: torch.Tensor, |
| typed_second: torch.Tensor, |
| typed_chunk: torch.Tensor, |
| typed_position: torch.Tensor, |
| return_dict: bool | None = None, |
| **kwargs: Any, |
| ) -> MimirModelOutput | tuple[torch.Tensor, ...]: |
| if candidate_ids.shape[0] == 0: |
| message = "the batch holds no candidate row; the exported path needs at least one" |
| raise ValueError(message) |
| question_cap = self.config.question_cap |
| typed = TypedValues( |
| kind=typed_kind, |
| number=typed_number, |
| year=typed_year, |
| month=typed_month, |
| day=typed_day, |
| second=typed_second, |
| ) |
| hidden = self._encode( |
| chunk_ids, chunk_mask.bool(), typed, typed_chunk, typed_position, inject=True |
| ) |
| question = nn.functional.pad(hidden, (0, 0, 0, question_cap + 1))[:, 1 : 1 + question_cap] |
| candidates = self._encode( |
| candidate_ids, |
| candidate_mask.bool(), |
| typed, |
| typed_chunk, |
| typed_position, |
| inject=False, |
| ) |
| pooled = _pool_candidates(candidates, candidate_text_mask.bool()) |
| question, question_mask = _mean_question(question, chunk_record, question_length) |
| tokens = hidden[key_chunk, key_position] |
| if pooled.shape[0]: |
| selected = pooled[candidate_index] |
| else: |
| selected = pooled.new_zeros((*candidate_index.shape, pooled.shape[-1])) |
| inputs = HeadInputs( |
| tokens=tokens, |
| token_mask=key_mask.bool(), |
| question=question, |
| question_mask=question_mask, |
| candidates=selected, |
| candidate_mask=candidate_present.bool(), |
| decision_type=decision_type, |
| ) |
| outputs = self.head(inputs) |
| per_depth: dict[str, list[torch.Tensor]] = { |
| "utilities": [], |
| "thresholds": [], |
| "abstain": [], |
| "ordinal_score": [], |
| "histogram_logits": [], |
| "evidence": [], |
| } |
| for output in outputs: |
| readout = output.readout |
| per_depth["utilities"].append(readout.utilities.float()) |
| per_depth["thresholds"].append(readout.thresholds.float()) |
| per_depth["abstain"].append(readout.abstain.float()) |
| per_depth["ordinal_score"].append(readout.ordinal_score.float()) |
| per_depth["histogram_logits"].append(readout.histogram_logits.float()) |
| per_depth["evidence"].append(output.evidence_scores.float()) |
| stacked = {name: torch.stack(values) for name, values in per_depth.items()} |
| utilities = stacked["utilities"] |
| thresholds = stacked["thresholds"] |
| abstain = stacked["abstain"] |
| ordinal_score = stacked["ordinal_score"] |
| histogram_logits = stacked["histogram_logits"] |
| evidence = stacked["evidence"] |
| workspace = outputs[-1].workspace.float() |
| result = MimirModelOutput( |
| utilities=utilities, |
| thresholds=thresholds, |
| abstain=abstain, |
| ordinal_score=ordinal_score, |
| histogram_logits=histogram_logits, |
| evidence=evidence, |
| workspace=workspace, |
| ) |
| if return_dict is False: |
| return ( |
| utilities, |
| thresholds, |
| abstain, |
| ordinal_score, |
| histogram_logits, |
| evidence, |
| workspace, |
| ) |
| return result |
|
|
|
|
| __all__ = ["INPUTS", "MimirModel", "MimirModelOutput"] |
|
|
