| """Typed-decision core from STEAV's published System One inference implementation. |
| |
| Apache-2.0. The encoder layout and decision head follow Convai Laya's |
| Apache-2.0 implementation. No remote model code is executed. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import math |
| import os |
| import re |
| import sys |
| import warnings |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Any, Iterable, Sequence |
|
|
| import numpy as np |
| import torch |
| from torch import nn |
|
|
| __all__ = ["SystemOneDecisionModel", "LayaDecisionHead", "render_state", "offset_logit", "apply_prior_shift", |
| "encode_decision", "load_laya_encoder", "ModelError"] |
|
|
| |
| PROMPT_TEMPLATE = "cid/system-one-encoder-prompt/laya-v1" |
| HEAD_SCHEMA = "cid/system-one-head/v1" |
| CALIBRATION_SCHEMA = "cid/system-one-calibration/v1" |
| RELEASE_SCHEMA = "system-one-release/v1" |
| LAYA_QTYPES = {"choice": 0, "score": 1, "noul": 2} |
| LAYA_HEAD_MAX_LEN = 192 |
| LAYA_OPTION_TOKENS = 48 |
| MAX_SLOTS = 128 |
| OFFSET_CLIP = 1e-6 |
| PRIOR_CLIP = 1e-9 |
| DEFAULT_DESCRIPTIONS = {"choice": "Which answer is correct?", "score": "Which level applies?", "noul": "Is the proposition true?"} |
| NOUL_SLOTS = ({"id": "false", "description": "The proposition is false."}, |
| {"id": "true", "description": "The proposition is true."}) |
| NOUL_OPTION_TEXTS = ("false: no, the statement does not hold", "true: yes, the statement holds") |
| _NUMBER = re.compile(r"[+-]?(?:[0-9]+(?:\.[0-9]*)?|\.[0-9]+)(?:[eE][+-]?[0-9]+)?") |
|
|
| |
| |
| |
| LAYA_BASE_REPO = "convaiinnovations/laya" |
| LAYA_BASE_REVISION = "7b928d828b7b0e022f929d9bd2e44165aa270148" |
| LAYA_BASE_SHA256 = "891102d372688fc2a094dac56a384bc537b87c63f21f9f3dac0be2b7cbc8d86c" |
| LAYA_BASE_FILES = ("model.safetensors", "encoder/config.json", "tokenizer/tokenizer.json", "tokenizer/tokenizer_config.json") |
|
|
|
|
| class ModelError(ValueError): |
| """Raised for invalid model folders, inputs that break the input contract, or failed integrity checks.""" |
|
|
|
|
| |
|
|
| def question_slots(question: dict) -> list[dict]: |
| """Answer slots in scoring order: choice options, score levels, or [false, true] for noul.""" |
| qtype = str(question.get("type", "")).strip() |
| if qtype == "choice": |
| slots = question.get("options") or [] |
| elif qtype == "score": |
| slots = question.get("levels") or [] |
| elif qtype == "noul": |
| if question.get("options") or question.get("levels"): |
| raise ModelError("noul takes no options or levels") |
| slots = NOUL_SLOTS |
| else: |
| raise ModelError("question type must be choice, score or noul") |
| if not 2 <= len(slots) <= MAX_SLOTS: |
| raise ModelError("questions must carry between 2 and %d answer slots" % MAX_SLOTS) |
| ids = [str(slot.get("id", "")).strip() for slot in slots] |
| if any(not identity for identity in ids) or len(set(ids)) != len(ids): |
| raise ModelError("answer ids must be nonempty and distinct") |
| return [{"id": identity, "description": str(slot.get("description") or "")} for identity, slot in zip(ids, slots)] |
|
|
|
|
| def option_texts(question: dict, slots: Sequence[dict]) -> list[str]: |
| """Answer texts as rendered after each [MASK] marker.""" |
| qtype = str(question.get("type", "")).strip() |
| if qtype == "score": |
| return ["level %d: %s" % (i, " ".join(slot["description"].split()) or slot["id"]) for i, slot in enumerate(slots)] |
| if qtype == "noul": |
| return list(NOUL_OPTION_TEXTS) |
| return [(slot["id"] + ": " + " ".join(slot["description"].split())) if slot["description"].strip() else slot["id"] |
| for slot in slots] |
|
|
|
|
| def render_state(state: Any, offset_field: str | None = None, offset_hidden: bool = False) -> str: |
| """Evidence text exactly as trained. |
| |
| A dict is rendered as compact JSON with sorted keys and raw Unicode. Key order therefore decides which fields |
| survive end truncation: digits < uppercase (Context_*, Tree_*) < "_" < lowercase. A str state is used as given |
| (stripped). With offset_hidden, only the literal offset key is removed from the text. |
| """ |
| if offset_hidden and offset_field and isinstance(state, dict): |
| state = {key: value for key, value in state.items() if key != offset_field} |
| if isinstance(state, str): |
| return state.strip() |
| return json.dumps(state, sort_keys=True, separators=(",", ":"), ensure_ascii=False) |
|
|
|
|
| def offset_logit(state: Any, question: dict, field: str | None) -> float: |
| """Fixed upstream log-odds added to the "true" slot of noul questions. |
| |
| The value must be a JSON number or a numeric string in [0, 1]; booleans are rejected. It is clipped to |
| [1e-6, 1 - 1e-6], so the offset lies within about +/-13.8. |
| """ |
| if not field or str(question.get("type", "")).strip() != "noul": |
| return 0.0 |
| if not isinstance(state, dict) or field not in state: |
| raise ModelError("offset field %r missing from state" % field) |
| value = state[field] |
| if isinstance(value, bool) or not isinstance(value, (int, float, str)): |
| raise ModelError("offset field %r must be a probability in [0, 1]" % field) |
| if isinstance(value, str) and not _NUMBER.fullmatch(value.strip()): |
| raise ModelError("offset field %r must be a probability in [0, 1]" % field) |
| try: |
| p = float(value) |
| except (TypeError, ValueError, OverflowError): |
| raise ModelError("offset field %r must be a probability in [0, 1]" % field) from None |
| if not math.isfinite(p) or p < 0 or p > 1: |
| raise ModelError("offset field %r must be a probability in [0, 1]" % field) |
| p = min(max(p, OFFSET_CLIP), 1 - OFFSET_CLIP) |
| return math.log(p / (1 - p)) |
|
|
|
|
| def bucket_key(question: dict, count: int) -> str: |
| """Calibration bucket: '<type>:2', ':3-4', ':5-8' or ':9+'.""" |
| width = "2" if count <= 2 else "3-4" if count <= 4 else "5-8" if count <= 8 else "9+" |
| return str(question.get("type", "")).strip() + ":" + width |
|
|
|
|
| @dataclass |
| class Encoded: |
| ids: list[int] |
| positions: list[int] |
| qtype: int |
| offset: float = 0.0 |
| state_tokens: int = 0 |
| dropped_state_tokens: int = 0 |
|
|
|
|
| def encode_decision(tokenizer, state: Any, question: dict, limit: int, offset_field: str | None = None, |
| offset_hidden: bool = False) -> Encoded: |
| """Laya sequence (template cid/system-one-encoder-prompt/laya-v1): |
| |
| [CLS] "<type> question: <description>" [SEP] [MASK] " <option 0>" [MASK] " <option 1>" ... [SEP] <evidence> [SEP] |
| |
| Answers are never dropped; the question text and, if needed, the answers are shortened to the answer budget; |
| the evidence keeps its first tokens. Raw text is tokenized as is: special-token strings such as "[MASK]" |
| inside the evidence become special tokens, as they did in training. |
| """ |
| mask_id, cls_id, sep_id = tokenizer.mask_token_id, tokenizer.cls_token_id, tokenizer.sep_token_id |
| if mask_id is None or cls_id is None or sep_id is None: |
| raise ModelError("the tokenizer needs mask, cls and sep tokens") |
| slots = question_slots(question) |
| qtype = str(question.get("type", "")).strip() |
| description = " ".join(str(question.get("description") or "").split()) or DEFAULT_DESCRIPTIONS[qtype] |
| head_ids = tokenizer.encode(qtype + " question: " + description, add_special_tokens=False) |
| option_ids = [[mask_id] + tokenizer.encode(" " + text, add_special_tokens=False)[:LAYA_OPTION_TOKENS] |
| for text in option_texts(question, slots)] |
| head_max_len = max(LAYA_HEAD_MAX_LEN, min(8 * len(option_ids), limit // 2)) |
| budget = head_max_len - sum(len(o) for o in option_ids) |
| if budget < 16: |
| per = max(4, (head_max_len - 16) // max(1, len(option_ids))) |
| option_ids = [o[:per] for o in option_ids] |
| budget = head_max_len - sum(len(o) for o in option_ids) |
| head_ids = head_ids[:max(8, budget)] |
| ids = [cls_id] + head_ids + [sep_id] |
| positions = [] |
| for option in option_ids: |
| positions.append(len(ids)) |
| ids.extend(option) |
| ids.append(sep_id) |
| room = limit - len(ids) - 1 |
| if room < 8: |
| raise ModelError("decision prompt leaves no room for evidence; raise max_seq_length or shorten the answers") |
| evidence = tokenizer.encode(render_state(state, offset_field, offset_hidden), add_special_tokens=False) |
| ids = ids + evidence[:room] + [sep_id] |
| if len(ids) > limit: |
| raise ModelError("decision prompt exceeds max_seq_length") |
| return Encoded(ids, positions, LAYA_QTYPES[qtype], offset_logit(state, question, offset_field), |
| len(evidence), max(0, len(evidence) - room)) |
|
|
|
|
| def apply_prior_shift(p_true: float, training_prior: float, target_prior: float) -> float: |
| """Bayes prior correction of a calibrated P(true), as the benchmark scorer applied it.""" |
| for prior in (training_prior, target_prior): |
| if not 0 < prior < 1: |
| raise ModelError("priors must lie strictly between 0 and 1") |
| shift = math.log(target_prior / (1 - target_prior)) - math.log(training_prior / (1 - training_prior)) |
| log_odds = math.log(max(p_true, PRIOR_CLIP) / max(1 - p_true, PRIOR_CLIP)) + shift |
| return 1 / (1 + math.exp(-log_odds)) |
|
|
|
|
| |
|
|
| class LayaDecisionHead(nn.Module): |
| """Laya's typed decision head (Apache-2.0 `DecisionModel` minus the encoder), run in float32. |
| |
| Parameter names match head.safetensors: head.layers.{0,1}.*, type_emb.weight, scorer.*, |
| act_head.* (frozen act/escalate head, unused for scoring) and the inert `temperature` buffer (Laya's own |
| per-type temperatures; candidates are calibrated by calibration.json instead). |
| """ |
|
|
| def __init__(self, hidden: int, layers: int = 2, actions: int = 2, dropout: float = 0.1): |
| super().__init__() |
| layer = nn.TransformerEncoderLayer(hidden, max(1, hidden // 64), 4 * hidden, dropout, batch_first=True, norm_first=True) |
| self.head = nn.TransformerEncoder(layer, layers, enable_nested_tensor=False) if layers > 0 else None |
| self.type_emb = nn.Embedding(3, hidden) |
| self.scorer = nn.Sequential(nn.LayerNorm(hidden), nn.Linear(hidden, hidden), nn.GELU(), nn.Linear(hidden, 1)) |
| self.act_head = nn.Sequential(nn.Linear(hidden + 4, 256), nn.GELU(), nn.Linear(256, actions)) |
| self.register_buffer("temperature", torch.ones(3)) |
|
|
| def score(self, hidden: torch.Tensor, attention_mask: torch.Tensor, positions: Sequence[int], qtype: int) -> torch.Tensor: |
| """hidden [L, d] float32 of one padded sequence -> one logit per answer marker. |
| |
| Rows are scored one at a time, as in training, so batching never changes the head's arithmetic. |
| """ |
| h = hidden[None, :, :] + self.type_emb(torch.tensor([qtype], device=hidden.device))[:, None, :] |
| if self.head is not None: |
| pad = ~attention_mask[None, :].bool() |
| for layer in self.head.layers: |
| h = layer(h, src_key_padding_mask=pad) |
| markers = h[0, list(positions)] |
| return self.scorer(markers).squeeze(-1).float() |
|
|
|
|
| def _sha256(path: Path) -> str: |
| h = hashlib.sha256() |
| with open(path, "rb") as stream: |
| for chunk in iter(lambda: stream.read(1 << 22), b""): |
| h.update(chunk) |
| return h.hexdigest() |
|
|
|
|
| def _local_dir(source: str | os.PathLike, revision: str | None, allow_patterns: Iterable[str] | None = None) -> Path: |
| """A local directory as is; otherwise a Hugging Face repo id downloaded at `revision`.""" |
| path = Path(source).expanduser() |
| if path.is_dir(): |
| return path |
| if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]{0,95}/[A-Za-z0-9][A-Za-z0-9._-]{0,95}", str(source)): |
| raise ModelError("%s is neither a local directory nor a Hugging Face repo id" % source) |
| from huggingface_hub import snapshot_download |
| return Path(snapshot_download(repo_id=str(source), revision=revision, allow_patterns=list(allow_patterns) if allow_patterns else None)) |
|
|
|
|
| def base_layout(base_dir: Path) -> tuple[Path, Path, Path]: |
| """(weights, encoder config dir, tokenizer dir) for the HF repo layout or the flat bundle layout used in training. |
| |
| HF convaiinnovations/laya: model.safetensors + encoder/config.json + tokenizer/ (its root config.json is a |
| `model_type: laya` stub that transformers cannot load). Flat bundle: model.safetensors + config.json |
| (model_type modernbert) + tokenizer files in the same directory. |
| """ |
| weights = base_dir / "model.safetensors" |
| if not weights.is_file(): |
| raise ModelError("base directory has no model.safetensors: %s" % base_dir) |
| if (base_dir / "encoder" / "config.json").is_file(): |
| config_dir = base_dir / "encoder" |
| tokenizer_dir = base_dir / "tokenizer" if (base_dir / "tokenizer" / "tokenizer.json").is_file() else base_dir |
| elif (base_dir / "config.json").is_file(): |
| config_dir = tokenizer_dir = base_dir |
| else: |
| raise ModelError("base directory has neither encoder/config.json nor config.json: %s" % base_dir) |
| return weights, config_dir, tokenizer_dir |
|
|
|
|
| def load_laya_encoder(base_dir: Path, dtype: torch.dtype, device: torch.device, expected_sha256: str | None = None, |
| attn_implementation: str = "sdpa", rope_dtype: torch.dtype | None = torch.bfloat16) -> nn.Module: |
| """Laya's ModernBERT encoder, loaded the way the reference scorer loads it. |
| |
| from_config + strict key accounting over the `encoder.`-prefixed tensors, then ONE `.to(dtype)` for the whole |
| module. The checkpoint stores float16 tensors; they are copied into float32 parameters and, on CUDA, cast to |
| bfloat16. `.to(dtype)` also casts the rotary inv_freq buffers (from_pretrained(dtype=...) keeps them float32), |
| so do not replace this with from_pretrained when reproducing stored scores. |
| |
| rope_dtype: The original candidates trained and evaluated on CUDA in bfloat16, i.e. with bfloat16-rounded rotary |
| frequencies (up to 0.4% relative error). By default the frequencies are rounded the same way when computing in |
| another dtype (float32 on CPU), so positions are encoded as in training. None keeps exact float32 frequencies, |
| which is what the reference CPU path does. |
| """ |
| from safetensors.torch import load_file |
| from transformers import AutoConfig, AutoModel |
| weights, config_dir, _ = base_layout(base_dir) |
| if expected_sha256 and _sha256(weights) != expected_sha256: |
| raise ModelError("base weights sha256 differs from the release manifest (expected %s)" % expected_sha256) |
| config = AutoConfig.from_pretrained(str(config_dir), trust_remote_code=False, local_files_only=True) |
| if config.model_type != "modernbert" or getattr(config, "is_encoder_decoder", False): |
| raise ModelError("the base encoder must be a ModernBERT model") |
| model = AutoModel.from_config(config, dtype=torch.float32, attn_implementation=attn_implementation) |
| state = load_file(str(weights)) |
| encoder_state = {key[len("encoder."):]: value for key, value in state.items() if key.startswith("encoder.")} or state |
| missing, unexpected = model.load_state_dict(encoder_state, strict=False) |
| unexpected = [key for key in unexpected if not key.startswith(("head.", "decoder.", "cls.", "classifier."))] |
| if missing or unexpected: |
| raise ModelError("base encoder weights do not match the configuration (missing %d, unexpected %d)" % (len(missing), len(unexpected))) |
| del state, encoder_state |
| model = model.to(dtype=dtype) |
| rotary = getattr(model, "rotary_emb", None) |
| if rope_dtype is not None and rotary is not None: |
| with torch.no_grad(): |
| for name, buffer in rotary.named_buffers(): |
| if name.endswith("inv_freq") and buffer.is_floating_point(): |
| buffer.copy_(buffer.to(rope_dtype).to(buffer.dtype)) |
| return model.to(device) |
|
|
|
|
| def _resolve_dtype(dtype: str | torch.dtype, device: torch.device) -> torch.dtype: |
| if isinstance(dtype, torch.dtype): |
| return dtype |
| if dtype == "auto": |
| return torch.bfloat16 if device.type == "cuda" else torch.float32 |
| try: |
| return {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[dtype] |
| except KeyError: |
| raise ModelError("dtype must be 'auto', 'float32', 'bfloat16' or 'float16'") from None |
|
|
|
|
| def _tf32_enabled() -> bool: |
| """True when float32 matmuls may run in TF32 ("high"/"medium"); the stored scores used the default "highest".""" |
| try: |
| return torch.get_float32_matmul_precision() != "highest" |
| except Exception: |
| return False |
|
|
|
|
| def _read_json(path: Path) -> dict: |
| with open(path, "rb") as stream: |
| return json.loads(stream.read()) |
|
|
|
|
| |
|
|
| class SystemOneDecisionModel: |
| """A typed-decision encoder, adapter and calibrated marker-scoring head.""" |
|
|
| def __init__(self, encoder, tokenizer, head: LayaDecisionHead, calibration: dict, head_config: dict, release: dict, |
| device: torch.device, dtype: torch.dtype, max_seq_length: int, question: dict | None, |
| offset_field: str | None, offset_hidden: bool): |
| self.encoder, self.tokenizer, self.head = encoder, tokenizer, head |
| self.calibration, self.head_config, self.release = calibration, head_config, release |
| self.device, self.dtype, self.max_seq_length = device, dtype, max_seq_length |
| self.question, self.offset_field, self.offset_hidden = question, offset_field, offset_hidden |
|
|
| def __repr__(self) -> str: |
| return "SystemOneDecisionModel(%s, device=%s, dtype=%s, max_seq_length=%d, offset_field=%r)" % ( |
| self.release.get("model", "unnamed"), self.device, str(self.dtype).replace("torch.", ""), self.max_seq_length, self.offset_field) |
|
|
| @classmethod |
| def from_pretrained(cls, path: str | os.PathLike, *, base: str | os.PathLike | None = None, base_revision: str | None = None, |
| revision: str | None = None, question: dict | None = None, device: str | torch.device | None = None, |
| dtype: str | torch.dtype = "auto", verify_base: bool = True, max_seq_length: int | None = None, |
| offset_field: str | None | bool = True, rope_dtype: str | torch.dtype | None = "bfloat16") -> "SystemOneDecisionModel": |
| """Load a published model folder (or a training-run `adapter/` directory) plus the Laya base. |
| |
| path: model folder or HF repo id with adapter_config.json, adapter_model.safetensors and |
| system_one/{config.json, head.safetensors, calibration.json[, release.json]}. |
| base: Laya base (local dir in HF or flat layout, or a repo id). Default: release.json's base, else |
| convaiinnovations/laya at the pinned revision. |
| question: the exact question the model was trained with. Default: release.json's question; required for |
| a raw training-run adapter directory. |
| dtype: "auto" reproduces the reference scorer (bfloat16 on CUDA, float32 on CPU). |
| verify_base: check the base weights' sha256 against the release manifest (or the pinned Laya hash). |
| offset_field: True uses the trained setting; None/False disables the offset (ablation only); a string |
| overrides it (testing only). Trained models with an offset need state[offset_field] in [0, 1]. |
| rope_dtype: "bfloat16" (default) rounds the rotary frequencies as in training and benchmarking, also |
| on float32 CPU runs; None keeps float32 frequencies (the reference CPU behaviour). |
| """ |
| from peft import PeftModel |
| from safetensors.torch import load_file |
| from transformers import AutoTokenizer |
| model_dir = _local_dir(path, revision) |
| s1 = model_dir / "system_one" |
| head_config = _read_json(s1 / "config.json") |
| calibration = _read_json(s1 / "calibration.json") |
| release = _read_json(s1 / "release.json") if (s1 / "release.json").is_file() else {} |
| if head_config.get("schema") != HEAD_SCHEMA or calibration.get("schema") != CALIBRATION_SCHEMA: |
| raise ModelError("unsupported decision head or calibration schema") |
| if head_config.get("head") != "laya" or head_config.get("backbone") != "encoder" or head_config.get("prompt_template") != PROMPT_TEMPLATE: |
| raise ModelError("only encoder candidates with the Laya head and the laya-v1 prompt are supported") |
| if release and release.get("schema") != RELEASE_SCHEMA: |
| raise ModelError("unsupported release manifest schema") |
| device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu")) |
| dtype = _resolve_dtype(dtype, device) |
| if device.type == "cuda" and _tf32_enabled(): |
| warnings.warn("TF32 matmul is enabled; the float32 head then differs from the stored benchmark scores") |
|
|
| base_info = release.get("base", {}) |
| base_source = base if base is not None else base_info.get("repo_id", LAYA_BASE_REPO) |
| base_dir = _local_dir(base_source, base_revision or base_info.get("revision", LAYA_BASE_REVISION), LAYA_BASE_FILES) |
| expected = base_info.get("sha256", LAYA_BASE_SHA256) if verify_base else None |
| encoder = load_laya_encoder(base_dir, dtype, device, expected, |
| rope_dtype=None if rope_dtype is None else _resolve_dtype(rope_dtype, device)) |
| _, _, tokenizer_dir = base_layout(base_dir) |
| if not (tokenizer_dir / "tokenizer.json").is_file() and (model_dir / "tokenizer.json").is_file(): |
| tokenizer_dir = model_dir |
| tokenizer = AutoTokenizer.from_pretrained(str(tokenizer_dir), trust_remote_code=False, local_files_only=True) |
| if tokenizer.pad_token_id is None: |
| tokenizer.pad_token = tokenizer.eos_token |
| tokenizer.padding_side = "right" |
|
|
| encoder = PeftModel.from_pretrained(encoder, str(model_dir), is_trainable=False, local_files_only=True, autocast_adapter_dtype=True) |
| encoder.eval() |
| hidden = int(encoder.get_input_embeddings().weight.shape[1]) |
| if head_config.get("hidden_size") != hidden: |
| raise ModelError("decision head width differs from the base encoder") |
| weights = load_file(str(s1 / "head.safetensors")) |
| actions = int(weights["act_head.2.weight"].shape[0]) if "act_head.2.weight" in weights else 2 |
| head = LayaDecisionHead(hidden, int(head_config.get("laya_head_layers") or 2), actions).float() |
| expected_keys = set(head.state_dict()) |
| if set(weights) != expected_keys: |
| raise ModelError("head weights do not match the head configuration (missing %d, unexpected %d)" |
| % (len(expected_keys - set(weights)), len(set(weights) - expected_keys))) |
| head.load_state_dict({key: value.to(torch.float32) for key, value in weights.items()}, strict=True) |
| head.to(device).eval() |
|
|
| if offset_field is True: |
| field = head_config.get("offset_field") or None |
| elif offset_field in (None, False): |
| field = None |
| else: |
| field = str(offset_field) |
| limit = int(max_seq_length or head_config.get("max_seq_length") or 2048) |
| return cls(encoder, tokenizer, head, calibration, head_config, release, device, dtype, limit, |
| question or release.get("question"), field, bool(head_config.get("offset_hidden", False))) |
|
|
| |
|
|
| def _question(self, question: dict | None) -> dict: |
| question = question or self.question |
| if not isinstance(question, dict): |
| raise ModelError("no question: pass question={'type': 'noul', 'description': ...} or ship system_one/release.json") |
| return question |
|
|
| def temperature(self, question: dict | None = None) -> float: |
| question = self._question(question) |
| count = len(question_slots(question)) |
| return float(self.calibration.get("buckets", {}).get(bucket_key(question, count), self.calibration.get("global_temperature", 1.0))) |
|
|
| def encode(self, state: Any, question: dict | None = None) -> Encoded: |
| """Token ids, marker positions and offset for one state (for golden tests and truncation checks).""" |
| return encode_decision(self.tokenizer, state, self._question(question), self.max_seq_length, self.offset_field, self.offset_hidden) |
|
|
| @torch.inference_mode() |
| def _logits(self, rows: Sequence[Encoded]) -> list[torch.Tensor]: |
| """Raw logits (offset included, temperature not applied).""" |
| length = max(len(row.ids) for row in rows) |
| pad = self.tokenizer.pad_token_id |
| ids = torch.tensor([row.ids + [pad] * (length - len(row.ids)) for row in rows], device=self.device) |
| mask = torch.tensor([[1] * len(row.ids) + [0] * (length - len(row.ids)) for row in rows], device=self.device) |
| hidden = self.encoder(input_ids=ids, attention_mask=mask).last_hidden_state.float() |
| out = [] |
| for index, row in enumerate(rows): |
| z = self.head.score(hidden[index], mask[index], row.positions, row.qtype) |
| if row.offset: |
| |
| z = z + torch.tensor([0.0, row.offset], device=z.device, dtype=z.dtype) |
| out.append(z) |
| return out |
|
|
|
|
