"""Core model architecture, token sequence construction, and confidence estimation for laya.""" import json import math import os import threading import warnings from contextlib import nullcontext from contextvars import ContextVar from functools import wraps from typing import Any, Dict, List, Optional, Sequence, Union import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.checkpoint import checkpoint QTYPES = {"choice": 0, "score": 1, "noul": 2} QTYPE_NAMES = {v: k for k, v in QTYPES.items()} _DEFAULT_NOUL_LABELS = {"false": "false", "true": "true"} def _autocast_enabled(device_type: str) -> bool: """Whether an autocast scope covers `device_type`. `torch.is_autocast_enabled` only accepted a `device_type` from torch 2.4; before that it reported CUDA alone and CPU and MPS had no autocast to report. This package declares `torch>=2.0.0`, so call the device-aware form where it exists and fall back otherwise, using torch's separate CPU predicate. """ try: return torch.is_autocast_enabled(device_type) except TypeError: # torch < 2.4 took no `device_type`: the no-argument call reports CUDA alone, CPU has # its own predicate, and MPS had no autocast to report. if device_type == "cpu": return torch.is_autocast_cpu_enabled() return torch.is_autocast_enabled() # A fast tokenizer is not read-only: `truncation=True` / `padding=True` make it call # `enable_truncation` / `enable_padding`, which mutates the shared Rust object. One tokenizer is # parsed per checkpoint directory and shared by every Agent that wants it, so concurrent # `predict()` calls -- on one Agent or on two sharing the cache -- raced and raised # `RuntimeError: Already borrowed`. Serialise encoding instead: it is a small fraction of a call # next to the forward pass, and this keeps the cache's single parse. _TOKENIZE_LOCK = threading.RLock() _QUESTION_TOKEN_CACHE = ContextVar("laya_question_token_cache", default=None) def encode_text(tok, text, **kwargs): """Tokenize `text` while holding the lock a shared fast tokenizer needs.""" with _TOKENIZE_LOCK: return tok(text, **kwargs) def _reuse_question_tokens(fn): @wraps(fn) def wrapped(*args, **kwargs): # One cache per prediction call: nested calls get their own scope, and exceptions # restore the outer scope. Nothing is retained on an Agent or shared across threads. scope = {"thread": threading.get_ident(), "tokens": {}} token = _QUESTION_TOKEN_CACHE.set(scope) try: return fn(*args, **kwargs) finally: # A timed hook can copy this context to a worker that outlives the call. scope["tokens"] = None _QUESTION_TOKEN_CACHE.reset(token) return wrapped def _disable_question_token_reuse(): scope = _QUESTION_TOKEN_CACHE.get() if scope is not None: scope["tokens"] = None def _encode_question_text(tok, text, **kwargs): scope = _QUESTION_TOKEN_CACHE.get() cache = scope["tokens"] if scope is not None and scope["thread"] == threading.get_ident() else None if cache is None: return encode_text(tok, text, **kwargs)["input_ids"] # Key the rendered, mask-sanitized text and encoding settings, not a JSON question: # option order, structured criteria and custom noul labels must keep their meaning. key = (id(tok), text, tuple(kwargs.items())) if key not in cache: # Keep the tokenizer alive so its identity cannot be reused within this scope. cache[key] = (tok, tuple(encode_text(tok, text, **kwargs)["input_ids"])) # Sequence assembly must not mutate token lists retained for later states. return list(cache[key][1]) def serialize_state(state: Union[str, dict, list]) -> str: if isinstance(state, str): return state return json.dumps(state, ensure_ascii=False) def render_criterion(value) -> str: """Render one criterion value as text. Strings pass through; anything structured (dict, list, number) becomes a single-line JSON document with the default separators -- ``", "`` between members, ``": "`` before a value -- so a rubric reads as JSON rather than a Python repr. Without this a dict-valued criterion crashed `noul` outright and leaked `{'desc': ...}` into `choice` and `score` prompts. """ if isinstance(value, str): return value return json.dumps(value, ensure_ascii=False, separators=(", ", ": "), default=str) def _resolve_noul_labels(labels=None): if labels is None: labels = _DEFAULT_NOUL_LABELS if not isinstance(labels, dict) or set(labels) != {"false", "true"}: raise ValueError("noul labels must map exactly 'false' and 'true' to distinct non-empty strings") false_label, true_label = labels["false"], labels["true"] if not isinstance(false_label, str) or not isinstance(true_label, str): raise ValueError("noul labels must map exactly 'false' and 'true' to distinct non-empty strings") false_label, true_label = false_label.strip(), true_label.strip() if not false_label or not true_label or false_label == true_label: raise ValueError("noul labels must map exactly 'false' and 'true' to distinct non-empty strings") return false_label, true_label def render_options(q: Dict) -> List[str]: """Render option texts in label-index order. Noul semantic order is always [false, true].""" t, crit = q["t"], q.get("crit") if t != "noul" and "labels" in q: raise ValueError("labels is only supported for noul questions") if t == "choice": # only None/"" mean "no description"; 0 and False are legitimate criterion values. # `str(k)` unconditionally: a label with no description is rendered as itself, so an int # label used to come back as an int from a function annotated `-> List[str]` and then # reached `build_sequence`, which calls `.replace` on it and raised an AttributeError # naming neither the question nor the label. With a description the same label already # went through `"%s: %s" %` and was a str, which is why only the undescribed form broke. # `structured._enum_field` stringifies labels the same way; the returned answer still # carries the caller's original label, which is unchanged. return [str(k) if v is None or v == "" else "%s: %s" % (k, render_criterion(v)) for k, v in crit.items()] if t == "score": return ["level %d: %s" % (i, render_criterion(c)) for i, c in enumerate(crit)] crit = crit or {} false_label, true_label = _resolve_noul_labels(q.get("labels")) false_crit, true_crit = crit.get("false"), crit.get("true") return [ false_label + ": " + (render_criterion(false_crit) if false_crit not in (None, "") else "no, the statement does not hold"), true_label + ": " + (render_criterion(true_crit) if true_crit not in (None, "") else "yes, the statement holds"), ] def build_sequence( tok, state: Union[str, dict, list], q: Dict, max_len: int = 512, head_max_len: int = 192, option_order: Optional[List[int]] = None, truncate_left: bool = False, state_ids: Optional[List[int]] = None, return_stats: bool = False, return_truncation_stats: bool = False, return_layout: bool = False, ): """Format: [CLS] instructions [SEP] [MASK] opt0 [MASK] opt1 ... [SEP] state [SEP]. `state_ids` lets a caller tokenize the shared state once and reuse it across every question, instead of re-serializing and re-tokenizing the same document per question. `return_stats` adds a third return value describing what the head budget did to the options: `options` (how many the question defines), `options_distinct` (how many still have a token span of their own) and `tokens_per_option` (the cap applied to each, or None when none was). The state is clamped to whatever room is left after the head, so a long state loses tokens here silently. `return_truncation_stats=True` adds one more return value, after the option stats when both are asked for, reporting that clamp: {"state_tokens": int, "state_tokens_used": int, "state_tokens_dropped": int, "truncated": bool} Callers cannot reconstruct this from the outside. The budget is in tokens, not characters, and the room left for the state depends on `max_len`, `head_max_len`, the instruction and the rendered options - so it moves per checkpoint and per question. A caller guessing with a fixed character threshold is wrong in both directions: it reports truncation that did not happen, and stays silent while evidence is being dropped (issue #174). The question half is `build_head`; `state_room` reports how much of `max_len` is left for the state after it, which is what a caller must size a window against. `return_layout=True` adds one last return value, the parallel option layout from `parallel_layout`: `{"position_ids": [...], "option_ids": [...]}`, one entry per token. """ ids, markers, stats = build_head(tok, q, head_max_len, option_order=option_order) head_len = len(ids) room = max(0, max_len - len(ids) - 1) if state_ids is None: state_ids = encode_text(tok, serialize_state(state).replace(tok.mask_token, " "), add_special_tokens=False)["input_ids"] # not state_ids[-room:]: with no room left, state_ids[-0:] is the whole state rather than none of it st = state_ids[max(0, len(state_ids) - room):] if truncate_left else state_ids[:room] ids = ids + st + [tok.sep_token_id] # Laid out before the clamp: a dropped trailing option must not widen the last surviving span layout = parallel_layout(markers, head_len, len(ids)) if return_layout else None ids, markers = ids[:max_len], [m for m in markers if m < max_len] extra = () if return_truncation_stats: # `room` leaves space for the closing [SEP], so every token in `st` survives the [:max_len] clamp extra = ({ "state_tokens": len(state_ids), "state_tokens_used": len(st), "state_tokens_dropped": len(state_ids) - len(st), "truncated": len(st) < len(state_ids), },) if return_layout: extra += ({k: v[:max_len] for k, v in layout.items()},) if not return_stats: return (ids, markers) + extra return (ids, markers, stats) + extra def build_head(tok, q: Dict, head_max_len: int = 192, option_order: Optional[List[int]] = None): """The question half of a sequence: `[CLS] instructions [SEP] [MASK] opt0 [MASK] opt1 ... [SEP]`. `build_sequence` appends the state to this and `state_room` measures what is left over for it, so neither can disagree with the other about what the head costs. Returns `(ids, markers, stats)`: the head token ids, the `[MASK]` position of each option in `option_order`, and the stats dict `build_sequence` returns for `return_stats`. """ mask_tok = tok.mask_token opts = render_options(q) order = option_order if option_order is not None else list(range(len(opts))) ins = str(q["ins"]).replace(mask_tok, " ") head_ids = _encode_question_text(tok, "%s question: %s" % (q["t"], ins), add_special_tokens=False) opt_ids = [] for i in order: # Cap at the tokenizer, not after the fact: `[:48]` still makes the tokenizer process the # whole (possibly long) description. truncation=True, max_length=48 keeps the first 48 # tokens, which is exactly what the previous slice produced. opt_tokens = _encode_question_text( tok, " " + opts[i].replace(mask_tok, " "), add_special_tokens=False, truncation=True, max_length=48, ) opt_ids.append([tok.mask_token_id] + opt_tokens) opt_budget = head_max_len - sum(len(o) for o in opt_ids) per_option = None if opt_budget < 16: per = max(4, (head_max_len - 16) // max(1, len(opt_ids))) per_option = per opt_ids = [o[:per] for o in opt_ids] opt_budget = head_max_len - sum(len(o) for o in opt_ids) head_ids = head_ids[: max(8, opt_budget)] ids = [tok.cls_token_id] + head_ids + [tok.sep_token_id] markers = [] for o in opt_ids: markers.append(len(ids)) ids.extend(o) ids.append(tok.sep_token_id) # Two options that share a prefix can come out of the cut as the same token span: the marker # count still matches the option count, so the guard in `Agent._encode_state` passes and # nothing downstream can tell that the question lost the ability to name them apart. Counted # on the capped option ids, before assembly: re-slicing the finished sequence cannot close # the last option's span -- it runs on into the serialized state, which differs per request, # so the last option always looks distinguishable however it collided (#538). return ids, markers, { "options": len(opt_ids), "options_distinct": len({tuple(o) for o in opt_ids}), "tokens_per_option": per_option, } def parallel_layout(markers: List[int], head_len: int, length: int) -> Dict[str, List[int]]: """Position ids and option ids that make the encoder blind to option order. In the sequential layout option `s` sits at positions after option `s-1`, and every option attends to every other, so where an option is listed changes its embedding: the slot logits of five identical options spread by 3.59 (`tests/test_option_order.py`). Here every option starts at the same position, the first one after the instruction's [SEP], and the head's closing [SEP] and the state continue after the longest option. Together with `parallel_option_masks`, which stops an option from attending to another, reordering the options only reorders the marker embeddings. `option_ids` is 0 for shared tokens (instruction, closing [SEP], state, padding) and `s + 1` for the tokens of slot `s`, which runs from its [MASK] up to the next marker or, for the last slot, up to the [SEP] at `head_len - 1`. """ if not markers: return {"position_ids": list(range(length)), "option_ids": [0] * length} start = markers[0] spans = list(zip(markers, markers[1:] + [head_len - 1])) position_ids, option_ids = list(range(start)), [0] * start for s, (a, b) in enumerate(spans): position_ids += range(start, start + b - a) option_ids += [s + 1] * (b - a) after = start + max(b - a for a, b in spans) position_ids += range(after, after + length - len(position_ids)) option_ids += [0] * (length - len(option_ids)) return {"position_ids": position_ids, "option_ids": option_ids} OPTION_LAYOUTS = ("sequential", "parallel") def uses_parallel_layout(cfg: Dict) -> bool: """Whether a checkpoint config asks for the parallel option layout (`parallel_layout`). The layout is a property of the weights -- a checkpoint trained on one layout reads the other as a different input -- so it lives in the config, defaulting to the sequential layout every published checkpoint was trained on. The parallel masks go through ModernBERT's per-layer-type mask argument, which transformers 4.x does not have. """ layout = cfg.get("option_layout", "sequential") if layout not in OPTION_LAYOUTS: raise ValueError("option_layout must be one of %s, got %r" % (", ".join(OPTION_LAYOUTS), layout)) if layout == "parallel": import transformers if int(transformers.__version__.split(".")[0]) < 5: raise RuntimeError("option_layout='parallel' needs transformers>=5, found %s" % transformers.__version__) return layout == "parallel" def parallel_option_masks(attention_mask: torch.Tensor, position_ids: torch.Tensor, option_ids: torch.Tensor, sliding_window: Optional[int]) -> Dict[str, torch.Tensor]: """Per-layer-type boolean attention masks for the parallel layout (`parallel_layout`). A token attends to every real token except those of a *different* option. ModernBERT's local layers keep their window, measured in position ids rather than sequence index: options share positions, so a sequence-index window would let slot 0 see more of the state than slot 3. Returned in the `{"full_attention": ..., "sliding_attention": ...}` form ModernBERT takes in place of a padding mask; each mask is `[batch, 1, query, key]`, True where attention is allowed. """ keys = attention_mask.bool()[:, None, :] other_option = (option_ids[:, :, None] > 0) & (option_ids[:, None, :] > 0) \ & (option_ids[:, :, None] != option_ids[:, None, :]) full = keys & ~other_option masks = {"full_attention": full[:, None]} if sliding_window is not None: near = (position_ids[:, :, None] - position_ids[:, None, :]).abs() <= sliding_window masks["sliding_attention"] = (full & near)[:, None] return masks def state_room(tok, q: Dict, max_len: int = 512, head_max_len: int = 192) -> int: """How many state tokens `q` leaves inside `max_len`, which is what `build_sequence` keeps. The head is the question's own -- its instructions plus one `[MASK]`-prefixed span per option -- so a question with many options leaves less room for the state than one with two, and two questions in the same request do not have to leave the same amount. Anything past the return value is cut off (the start is kept, or the end for a conversation turn list). """ head, _, _ = build_head(tok, q, head_max_len) return max(0, max_len - len(head) - 1) # -1 for the [SEP] that closes the state #: How far the DEFAULT window may be cut before `window_budget` says so. A small clamp is ordinary #: and warning about it would be noise; a large one multiplies the number of forward passes by the #: same factor, which a caller needs to be told about because nothing in their code implies it. 2x #: is the point where the scan costs at least twice what the caller would estimate from `max_len`. _WINDOW_CLAMP_WARN_RATIO = 2 # How much capping the window may multiply the window count before `window_batch_cap` starts # chunking. Under this, chunking would cost the single-shared-pass property for no real protection. _WINDOW_BATCH_BLOWUP = 2 def window_budget(tok, questions, max_len: int = 512, head_max_len: int = 192, window: Optional[int] = None, stride: Optional[int] = None): """Window size and stride for scanning a state that is longer than one sequence. `predict_long` decodes each token window back to text and scores it as an ordinary state, so a window wider than the room the questions leave is re-truncated by `build_sequence` on the way in: the tail of every window reaches no model, while the reported span says it did. The window is therefore capped at `state_room`, and at the *smallest* room of `questions`, because the windows are one list of states scored for every question in shared forward passes -- a window sized for the roomiest question would be cut short for the tightest one, and the offsets reported on its answers would mean something different per question. With no questions there is nothing to fit, so the caller's window (or the checkpoint default) stands. That default is `max(64, max_len - head_max_len - 8)`: the state budget the config leaves, with a 64-token floor under it, so a widened `head_max_len` stops shrinking it there. `questions` are internal question dicts, as `Agent._to_internal` returns them. An explicit `window` wider than the room is clamped to it with a `RuntimeWarning`, since a scan at the requested size cannot read what it claims to. A `stride` past the effective window is refused: the tokens between two windows would be read by neither, which is the failure `predict_long` exists to prevent. The default stride keeps its 50% overlap of the *effective* window, so a span near a boundary still lands whole inside some window; an explicit stride equal to the window still reads every token, with no overlap to catch a span that straddles a boundary. Returns `(window, stride, room)`. """ rooms = [state_room(tok, q, max_len, head_max_len) for q in questions] requested = window if (window and window > 0) else max(64, max_len - head_max_len - 8) size = requested room = min(rooms) if rooms else size if room <= 0: raise ValueError( "predict_long: the questions' options fill the whole sequence (max_len=%d," " head_max_len=%d), leaving no room for the state; no window can carry any of it." " A label set this large is what laya.shortlist.predict_shortlist is for" % (max_len, head_max_len)) if size > room: if window and window > 0: warnings.warn( "laya: predict_long: window=%d is wider than the %d state tokens these questions" " leave inside max_len=%d, so every window would be truncated to %d on the way to" " the model; scanning with window=%d instead" % (size, room, max_len, room, room), RuntimeWarning, stacklevel=3) elif size >= room * _WINDOW_CLAMP_WARN_RATIO: # The DEFAULT window was cut, and cut hard. Capping it is what stops the tail of every # window reaching no model, but it is not free and it must not be silent: the scan now # needs about `size / room` times as many windows, each one a full forward pass, and # nothing in the caller's code says why. Measured on the English checkpoint with 100 # four-word options: room 102 of max_len 512, so the question head alone is 409 tokens, # the window falls 312 -> 102 and the scan goes 11 windows -> 36, a 3.1x wall-clock # increase (1902 ms -> 5858 ms) for the same document. # # The cost is not the capping, it is the shape of the request: 80% of every sequence is # the question, and because the encoder is bidirectional the head cannot be computed # once and reused -- its representations depend on the state it is paired with. So the # warning names the real remedy rather than only reporting the clamp. warnings.warn( "laya: predict_long: these questions leave only %d of max_len=%d for the state" " (their heads take the rest), so the scan window is capped %d -> %d and roughly" " %.1fx as many windows -- each a full forward pass -- are needed to read the" " document. Fewer or shorter options, a larger max_len, or" " laya.shortlist.predict_shortlist for a large label set will all cost less than" " scanning at this width" % (room, max_len, size, room, size / room), RuntimeWarning, stacklevel=3) size = room step = stride if (stride and stride > 0) else max(1, size // 2) if step > size: if size < requested and stride and stride <= requested: # The window the caller asked for was reduced above, and their stride was valid for the # window they asked for -- so this is the library's clamp, not their mistake. Reducing # the stride to match keeps the no-gap guarantee without making a self-consistent pair # of arguments an error. A stride that overshot the *requested* window is still refused # below, because that one really is the caller's. warnings.warn( "laya: predict_long: stride=%d was a 50%% step for the window=%d you asked for, but" " the window was reduced to %d to fit the room these questions leave; scanning with" " stride=%d instead" % (step, requested, size, max(1, size // 2)), RuntimeWarning, stacklevel=3) step = max(1, size // 2) else: raise ValueError( "predict_long: stride=%d steps past the %d-token window%s, so %d tokens between" " every pair of windows would be read by no window at all; pass stride <= %d" % (step, size, " these questions leave room for" if size < requested else "", step - size, size)) return size, step, room def window_batch_cap(n_windows: int, window: int, config_budget: int, batch_size: Optional[int] = None) -> Optional[int]: """Keep one forward pass no wider than the un-capped scan's would have been -- but only when capping the window has multiplied the window count enough to matter. Capping the window at the room the questions leave multiplies the number of windows on exactly the inputs it targets: measured, a 4 561-token document at 120 options goes from 29 windows of 312 tokens to roughly 413 of 23. `predict_batch` with `batch_size=None` puts every state in one forward pass, so a caller who passed no `batch_size` would go from a 29-row pass to a 413-row one at `max_len` width -- a plausible out-of-memory on an input that used to fit. Sending every window in one run is also a deliberate property (`ONNXAgent` asserts it), and a mild cap -- a 64-token budget reduced to 43, say -- multiplies the count by well under two. Chunking those would trade a real property for no real protection. So the cap only applies past `_WINDOW_BATCH_BLOWUP`x, and then bounds the pass at the count the un-capped budget would have produced: peak memory stays at parity with the behaviour before the cap, and the scan still reads the whole document, just in more passes. An explicit `batch_size` is always honoured. """ if batch_size and batch_size > 0: return batch_size if window >= config_budget: return None # not capped: one pass, exactly as before unclamped = max(1, -(-n_windows * max(1, window // 2) // max(1, config_budget // 2))) if n_windows <= unclamped * _WINDOW_BATCH_BLOWUP: return None # a mild cap: keep the single shared pass return unclamped def collapsed_options(qids, items) -> Dict[str, Dict[str, Optional[int]]]: """The questions whose options no longer have a token span each, from per-item stats. `total` is the number of options the question defines, not the number of markers that reached the sequence: a report counted from the markers would say "43/58" about a request where 28 options never made it into the input at all. """ out = {} for qid, item in zip(qids, items): stats = item.get("options") if stats and stats["options_distinct"] < stats["options"]: out[qid] = {"total": stats["options"], "distinct": stats["options_distinct"], "tokens_per_option": stats["tokens_per_option"]} return out class _DynamicMultiheadAttention(nn.MultiheadAttention): """`nn.MultiheadAttention` that keeps the shapes it traces. The stock module reshapes the packed projection with sizes captured while tracing, so under the legacy ONNX exporter the traced sequence length becomes a constant and a model exported from a short dummy input only runs at that length. The parameters and the maths are the same here; the reshape uses only constant shape arguments (`chunk` / `unflatten` / `flatten`) and `scaled_dot_product_attention`, the kernel the stock path already uses when weights are not requested. """ def forward(self, query, key, value, key_padding_mask=None, need_weights=True, attn_mask=None, average_attn_weights=True, is_causal=False): if (need_weights or self.in_proj_weight is None or self.bias_k is not None or self.bias_v is not None or (attn_mask is not None and attn_mask.dtype != torch.bool)): # Weight averaging, additive masks and the optional k/v bias are not on the traced # path; the stock implementation keeps them correct. return super().forward(query, key, value, key_padding_mask=key_padding_mask, need_weights=need_weights, attn_mask=attn_mask, average_attn_weights=average_attn_weights, is_causal=is_causal) if self.batch_first: query, key, value = query.transpose(0, 1), key.transpose(0, 1), value.transpose(0, 1) # (T, B, E) from here, matching the stock module's internals; attention runs on (B, H, T, D). if query is key is value: q, k, v = (part.unflatten(-1, (self.num_heads, self.head_dim)).permute(1, 2, 0, 3) for part in F.linear(query, self.in_proj_weight, self.in_proj_bias).chunk(3, dim=-1)) else: embed_dim = query.shape[-1] wq, wk, wv = self.in_proj_weight.split(embed_dim, dim=0) bq, bk, bv = ((None, None, None) if self.in_proj_bias is None else self.in_proj_bias.split(embed_dim, dim=0)) q, k, v = ( F.linear(t, w, b).unflatten(-1, (self.num_heads, self.head_dim)).permute(1, 2, 0, 3) for t, w, b in ((query, wq, bq), (key, wk, bk), (value, wv, bv)) ) mask = None if attn_mask is not None: mask = ~attn_mask if key_padding_mask is not None: # `== 0` keeps this correct for a bool mask and for the 0 / -inf float mask the encoder # layer hands over (`F._canonical_mask`), where `~` would not be defined. keep = key_padding_mask[:, None, None, :] == 0 mask = keep if mask is None else mask & keep attn = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, is_causal=is_causal and mask is None, dropout_p=self.dropout if self.training else 0.0) attn = attn.permute(2, 0, 1, 3).flatten(-2) if self.batch_first: attn = attn.transpose(0, 1) return self.out_proj(attn), None def unpermute_probs(p: np.ndarray, option_order: Optional[List[int]]) -> np.ndarray: """Put a slot-ordered probability row back into the caller's option order. `build_sequence` puts option `option_order[s]` in slot `s`, so a model row comes back indexed by slot. Everything downstream indexes by option -- `zip(keys, p)` for a choice, `arange(k) * p` for a score level, `p[1]` for noul-true -- so the row has to be inverted first or the probabilities end up attached to the wrong options, which is silent. A missing or mismatched order returns `p` untouched, so the canonical path is unaffected. """ if option_order is None or len(option_order) != len(p): return p canonical = np.empty_like(p) canonical[np.asarray(option_order, dtype=int)] = p return canonical class DecisionModel(nn.Module): """Bidirectional transformer encoder backbone + typed decision head.""" def __init__(self, encoder: nn.Module, head_layers: int = 2, n_act: int = 2, dropout: float = 0.1, no_init: bool = False): super().__init__() self.encoder = encoder d = encoder.config.hidden_size # `no_init` means the caller is about to load every parameter from a checkpoint, so the # head's initial values are pure overhead -- and not just time. transformers' # `no_init_weights()` patches the torch.nn.init functions, but something inside # nn.TransformerEncoderLayer draws from the RNG outside them, so building it still # advanced the global generator and made `load()` a visible side effect. On the meta # device no initialisation kernel runs at all; the layers are materialised empty and the # load fills them. with torch.device("meta") if no_init else nullcontext(): nhead = max(1, d // 64) layer = nn.TransformerEncoderLayer(d, nhead, 4 * d, dropout, batch_first=True, norm_first=True) # The stock attention bakes the traced length into an exported graph; see # _DynamicMultiheadAttention. Same parameters, same maths, traceable shapes. layer.self_attn = _DynamicMultiheadAttention(d, nhead, dropout=dropout, batch_first=True) self.head = nn.TransformerEncoder(layer, head_layers, enable_nested_tensor=False) if head_layers > 0 else None self.type_emb = nn.Embedding(3, d) self.scorer = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, d), nn.GELU(), nn.Linear(d, 1)) self.act_head = nn.Sequential(nn.Linear(d + 4, 256), nn.GELU(), nn.Linear(256, n_act)) self.register_buffer("temperature", torch.ones(3)) if no_init: # Only the modules created above are on the meta device; the encoder is already real # and may hold non-persistent buffers (RoPE frequencies) that to_empty would wipe. for module in (self.head, self.type_emb, self.scorer, self.act_head): if module is not None: module.to_empty(device="cpu") self.temperature = torch.empty_like(self.temperature, device="cpu") self.head_checkpointing = False def forward(self, input_ids, attention_mask, marker_pos, marker_mask, qtype, detach_encoder: bool = False, position_ids=None, option_ids=None): if option_ids is None: h = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state else: # Parallel option layout (`parallel_layout`). The head below has no positional # encoding, so with order-blind marker embeddings the logits permute with the options. masks = parallel_option_masks(attention_mask, position_ids, option_ids, getattr(self.encoder.config, "sliding_window", None)) h = self.encoder(input_ids=input_ids, attention_mask=masks, position_ids=position_ids).last_hidden_state if detach_encoder: h = h.detach() h = h + self.type_emb(qtype)[:, None, :] if self.head is not None: pad = ~attention_mask.bool() for layer in self.head.layers: if self.head_checkpointing and self.training and torch.is_grad_enabled(): # Non-reentrant checkpointing also trains the head when its input # is frozen. Default RNG preservation keeps dropout consistent. h = checkpoint(layer, h, src_key_padding_mask=pad, use_reentrant=False) else: h = layer(h, src_key_padding_mask=pad) idx = marker_pos.clamp(min=0)[:, :, None].expand(-1, -1, h.size(-1)) m = torch.gather(h, 1, idx) logits = self.scorer(m).squeeze(-1).float() logits = logits.masked_fill(~marker_mask, -1e4) p = torch.softmax(logits.detach(), -1) k = marker_mask.sum(-1).clamp(min=2).float() ent = -(p * torch.log(p.clamp_min(1e-9))).sum(-1) / torch.log(k) if p.size(-1) >= 2: top2 = p.topk(2, -1).values else: # A single-option question has exactly one marker, so p.topk(2, ...) # has nothing to select for the second slot and raises. The answer # is still well-defined: softmax over one logit is 1.0 regardless of # its value, so pad the missing second entry with 0.0 - that gives # the act head top1 - top2 == 1.0, the same "fully decided" signal # it would see for any other unambiguous top-1-vs-rest gap. top1 = p.topk(1, -1).values top2 = torch.cat([top1, torch.zeros_like(top1)], dim=-1) feats = torch.stack([top2[:, 0], top2[:, 0] - top2[:, 1], ent, k / 255.0], -1) pooled = h[:, 0].float() act_input = torch.cat([pooled, feats], -1) # Explicit-dtype exports run without autocast; keep the confidence maths in fp32, # then match the head's weights. Autocast already chooses the linear's input dtype. if not _autocast_enabled(h.device.type): act_input = act_input.to(self.act_head[0].weight.dtype) act_logits = self.act_head(act_input) return logits, act_logits def _apply_rope_config(ecfg) -> None: """Carry transformers>=5 per-layer RoPE settings over to the attributes 4.x reads. A checkpoint re-saved by transformers 5 stores RoPE as `rope_parameters = {"full_attention": {"rope_theta": ...}, "sliding_attention": {...}}`. transformers 4.x does not know that key, so it keeps its own defaults (global 160000, local 10000) and any checkpoint whose sliding-attention theta differs silently runs the wrong RoPE base -- mmBERT is exactly that case, both of its thetas are 160000. Map the values onto `global_rope_theta` / `local_rope_theta`, which 4.x does read. On transformers 5 this is a no-op beyond re-setting the same numbers. """ rope = getattr(ecfg, "rope_parameters", None) if not isinstance(rope, dict): return flat = rope.get("rope_theta") for layer_type, attr in (("full_attention", "global_rope_theta"), ("sliding_attention", "local_rope_theta")): params = rope.get(layer_type) theta = params.get("rope_theta") if isinstance(params, dict) else flat if theta is not None and hasattr(ecfg, attr): setattr(ecfg, attr, float(theta)) def _no_init_weights(): """`no_init_weights` lives in different modules across transformers versions.""" try: from transformers.initialization import no_init_weights except ImportError: # transformers 4.x from transformers.modeling_utils import no_init_weights return no_init_weights() def build_model(cfg: Dict, encoder_dir: Optional[str] = None, pretrained: bool = True, revision: Optional[str] = None) -> DecisionModel: """Build the decision model described by `cfg`. With `pretrained=False`, or when `encoder_dir` holds a saved encoder config, nothing is downloaded and **no parameter is initialised**: the caller is expected to load a checkpoint into the result with `load_state_dict(..., strict=True)` immediately. Skipping initialisation keeps `load()` from spending time on, or consuming RNG for, weights it is about to overwrite. """ from transformers import AutoConfig, AutoModel head_layers, n_act = cfg.get("head_layers", 2), len(cfg.get("act_costs", {})) + 1 if not pretrained or (encoder_dir and os.path.exists(encoder_dir)): ecfg = AutoConfig.from_pretrained(encoder_dir or cfg["encoder"]) _apply_rope_config(ecfg) with _no_init_weights(): enc = AutoModel.from_config(ecfg, attn_implementation="sdpa") return DecisionModel(enc, head_layers, n_act, no_init=True) # Training-time Hub load of the base encoder; allow pinning it like the checkpoints. kw = {"attn_implementation": "sdpa"} if revision: kw["revision"] = revision enc = AutoModel.from_pretrained(cfg["encoder"], **kw) return DecisionModel(enc, head_layers, n_act) def proper_reward( q: torch.Tensor, target: torch.Tensor, qtype: torch.Tensor, mask: torch.Tensor, w_sph: float = 0.5, w_rps: float = 1.0, log_floor: float = -9.21, ) -> torch.Tensor: """Strictly proper scoring rule reward: log score + spherical score + ranked probability score. q: [..., N, K] reported distributions target: [N, K] (one-hot or soft target distributions) """ q = q * mask logq = torch.log(q.clamp_min(1e-12)).clamp_min(log_floor) log_score = (target * logq).sum(-1) sph = (target * q).sum(-1) / q.norm(dim=-1).clamp_min(1e-9) r = log_score + w_sph * sph is_score = (qtype == QTYPES["score"]).float() if is_score.any(): k = mask.sum(-1).clamp(min=2).float() cdf_q = torch.cumsum(q, -1) cdf_t = torch.cumsum(target, -1) rps = (((cdf_q - cdf_t) ** 2) * mask).sum(-1) / (k - 1) r = r - w_rps * rps * is_score return r def td_lambda_targets(p_true: torch.Tensor, batch: Dict, lam: float = 1.0) -> torch.Tensor: """TD(lambda) targets for multi-turn conversation trajectories.""" target = batch["target"].clone() groups = batch.get("ep_group") if groups is None: return target for g in torch.unique(groups[groups >= 0]).tolist(): idx = (groups == g).nonzero(as_tuple=True)[0] idx = idx[torch.argsort(batch["ep_step"][idx])] y = batch["target"][idx[-1], 1] G = y for j in range(len(idx) - 1, -1, -1): if j < len(idx) - 1: G = (1 - lam) * p_true[idx[j + 1]] + lam * G target[idx[j], 0], target[idx[j], 1] = 1 - G, G return target def ece_score(conf: np.ndarray, correct: np.ndarray, bins: int = 15) -> float: """Expected Calibration Error across confidence bins.""" if len(conf) == 0: return float("nan") edges = np.linspace(0, 1, bins + 1) e = 0.0 for i, (lo, hi) in enumerate(zip(edges[:-1], edges[1:])): sel = (conf >= lo if i == 0 else conf > lo) & (conf <= hi) if sel.any(): e += sel.mean() * abs(conf[sel].mean() - correct[sel].mean()) return float(e) def answer_confidence(p: np.ndarray, k: int) -> float: """Probability mass on the answer being reported: max(p). This is the quantity temperature scaling fits, and the quantity every calibration figure in this repository is computed on -- both benchmark harnesses take `conf = max(probs)` before calling `ece_score`. The README's gating section relies on the property that goes with it: of the answers returned at confidence c, about c of them are right. That property is conditional, and the condition is not met by default -- it holds only after the temperatures have been fitted and validated on held-out data for this checkpoint and this option count. The shipped checkpoints are over-confident: `choice:11+` is fitted at 0.1006, a ~10x sharpener that would publish a genuine 0.24 top probability as 0.9999. `clamp_temperature` refuses it and applies `TEMP_MIN` (0.5) instead, which still sharpens 2x and publishes that same 0.24 as 0.4993 -- so a threshold applied to these numbers still selects below model accuracy (issue #394). The clamp bounds the damage; it does not calibrate anything. `confidence_from_probs` below reports a different quantity on a different scale and carries no such guarantee, so the two must not be compared against the same threshold. """ if k < 1: return 1.0 return float(np.clip(np.max(p[:k]), 0.0, 1.0)) def confidence_from_probs(p: np.ndarray, k: int) -> float: """Normalized Shannon entropy confidence: 1 - H(p) / log(k). How concentrated the whole distribution is. Useful, but not calibrated: it is not what temperature scaling fits and not what the reported ECE measures. See `answer_confidence`. """ if k < 2: return 1.0 p = p[:k] ent = -(p * np.log(np.clip(p, 1e-12, 1.0))).sum() return float(np.clip(1.0 - ent / math.log(k), 0.0, 1.0)) def temp_bucket(qtype: int, k: int) -> str: size = "2" if k <= 2 else "3-5" if k <= 5 else "6-10" if k <= 10 else "11+" return "%s:%s" % (QTYPE_NAMES[int(qtype)], size) # A fitted temperature below 1 sharpens the logits instead of softening them. The shipped # `choice:11+` bucket is 0.1006, which multiplies them ~10x: a 0.24 top probability is published as # 0.99, so a caller gating on confidence is told a coin flip is a certainty. No honest calibration # needs to sharpen this hard, so refuse to apply one that does. TEMP_MIN = 0.5 TEMP_MAX = 5.0 def clamp_temperature(t, lo: float = TEMP_MIN, hi: float = TEMP_MAX) -> float: """A usable temperature: `t` confined to [lo, hi], falling back to 1.0 if it is not a number. A bool is not a number either: `True`/`False` used to float to 1.0/0.0 here and read as fitted/sharpening temperatures, the same class of quiet acceptance `check_min_confidence` already refuses. A custom `lo`/`hi` still bounds either way. """ if isinstance(t, bool): return 1.0 try: t = float(t) except (TypeError, ValueError): return 1.0 if t != t or t in (float("inf"), float("-inf")): # NaN / inf return 1.0 return min(hi, max(lo, t)) def resolve_lang_temperatures(raw: Optional[Dict[str, Any]], base_temperature: Sequence[float]) -> Dict[str, Dict[str, Any]]: """Parse the `lang_temperatures` option into `{language: {temperature, temperature_by_options}}`. One implementation, because `Agent` and `ONNXAgent` both accept this option and both promise the same confidences for it. Reading it with `cfg.get(...)` and `len(...)` before checking the shape of either raised `AttributeError` and `TypeError` for exactly the inputs the `ValueError` below is written for, after the whole checkpoint had loaded: {"de": {"temperature": 2}} -> TypeError: object of type 'int' has no len() {"de": {"temperature": None}} -> TypeError: object of type 'NoneType' has no len() {"de": None} -> AttributeError: 'NoneType' object has no attribute 'get' A `null` entry or a `null` temperature both mean "inherit the checkpoint's own", which is how the `laya-ts` port reads the same option (`agent.ts:322-327`). """ resolved: Dict[str, Dict[str, Any]] = {} for lang, cfg in (raw or {}).items(): if not isinstance(lang, str): raise ValueError("Language override keys must be strings, got %r" % (lang,)) norm = lang.split("-")[0].lower() if cfg is None: cfg = {} if not isinstance(cfg, dict): raise ValueError("Language override %r must be a mapping, got %s" % (lang, type(cfg).__name__)) t_raw = cfg.get("temperature") if t_raw is None: t_raw = base_temperature if not isinstance(t_raw, (list, tuple)) or len(t_raw) != 3: raise ValueError("Language override %r temperature must be a list of 3 floats, got %r" % (lang, t_raw)) tbo_raw = cfg.get("temperature_by_options") or {} if not isinstance(tbo_raw, dict): raise ValueError("Language override %r temperature_by_options must be a mapping of " "bucket -> float, got %s" % (lang, type(tbo_raw).__name__)) resolved[norm] = { "temperature": [clamp_temperature(t) for t in t_raw], "temperature_by_options": {k: clamp_temperature(v) for k, v in tbo_raw.items()}, } return resolved def amp_dtype(name: Optional[str]) -> torch.dtype: return torch.bfloat16 if name == "bf16" else torch.float16 def collate_items(batch, pad_id: int): items = [it for group in batch for it in group] if not items: return None n, L = len(items), max(len(it["ids"]) for it in items) kmax = max(len(it["markers"]) for it in items) ids = torch.full((n, L), pad_id, dtype=torch.long) att = torch.zeros((n, L), dtype=torch.long) mpos = torch.zeros((n, kmax), dtype=torch.long) mmask = torch.zeros((n, kmax), dtype=torch.bool) has_target = any("target" in it for it in items) target = torch.zeros((n, kmax), dtype=torch.float32) if has_target else None for i, it in enumerate(items): ids[i, : len(it["ids"])] = torch.tensor(it["ids"]) att[i, : len(it["ids"])] = 1 k = len(it["markers"]) mpos[i, :k] = torch.tensor(it["markers"]) mmask[i, :k] = True if has_target and "target" in it: if len(it["target"]) > k: # Otherwise this lands as "The expanded size of the tensor (k) must match the # existing size (kmax)" from inside the assignment, which says nothing about the # actual mistake: a target with more entries than the item has options. The limit # is this item's own marker count, not the batch-wide kmax: in a mixed-width # batch a longer sibling row must not legitimise extra entries (#311). raise ValueError( "collate_items: item %d has %d target entries but only %d marker positions; " "a target needs one entry per option" % (i, len(it["target"]), k)) target[i, : len(it["target"])] = torch.tensor(it["target"], dtype=torch.float32) res = { "input_ids": ids, "attention_mask": att, "marker_pos": mpos, "marker_mask": mmask, "qtype": torch.tensor([it["qtype"] for it in items]), "label": torch.tensor([it.get("label", -1) for it in items]), "meta": [{k: it[k] for k in it if k not in ("ids", "markers", "target", "layout")} for it in items], } if target is not None: res["target"] = target if any("layout" in it for it in items): # All or none: a batch that mixes layouts would run half its rows through the wrong masks if not all("layout" in it for it in items): raise ValueError("collate_items: some items carry a parallel layout and some do not") pos = torch.zeros((n, L), dtype=torch.long) opt = torch.zeros((n, L), dtype=torch.long) for i, it in enumerate(items): pos[i, : len(it["ids"])] = torch.tensor(it["layout"]["position_ids"]) opt[i, : len(it["ids"])] = torch.tensor(it["layout"]["option_ids"]) res["position_ids"], res["option_ids"] = pos, opt return res