Release ODM Mini v1 DecisionHead

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,176 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS

README.md

@@ -0,0 +1,145 @@
+---
+license: apache-2.0
+base_model:
+  - Qwen/Qwen3-0.6B
+library_name: pytorch
+tags:
+  - agent
+  - agentic-ai
+  - decision-model
+  - tool-selection
+  - system-one
+  - open-weights
+  - qwen
+---
+
+# ODM Mini v1
+
+**Open-weight local Choice decision head for agent tool selection, built on frozen Qwen3-0.6B.**
+
+ODM Mini takes an agent state plus candidate action descriptions and directly ranks the candidates. It produces candidate scores, grouped-softmax probabilities, a selected action, confidence, decision margin, and measured latency. It does **not** generate text.
+
+## Model architecture
+
+```text
+Frozen Qwen3-0.6B
+→ candidate-description mean pooling
+→ Linear(1024,256)
+→ GELU
+→ Linear(256,1)
+→ grouped softmax
+```
+
+Only description-token hidden states are pooled. The candidate ID, prompt headers, state tokens, and padding are excluded from the pooling mask. The Qwen backbone remains frozen.
+
+## Size and required backbone
+
+The ODM-specific trained DecisionHead is approximately **1.1 MB**. The complete runnable model is substantially larger because it additionally requires the separately downloaded [`Qwen/Qwen3-0.6B`](https://huggingface.co/Qwen/Qwen3-0.6B) backbone. This repository does not redistribute Qwen weights.
+
+## Installation and usage
+
+```bash
+pip install torch transformers safetensors huggingface_hub
+```
+
+The code below downloads this five-file release, imports its deterministic loader without `trust_remote_code=True`, separately obtains Qwen3-0.6B, and performs a Choice:
+
+```python
+import sys
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+
+release_dir = Path(
+    snapshot_download(
+        repo_id="samatv256/mini-Jev",
+        allow_patterns=[
+            "README.md",
+            "model.safetensors",
+            "config.json",
+            "odm_mini.py",
+            "LICENSE",
+        ],
+    )
+)
+sys.path.insert(0, str(release_dir))
+
+from odm_mini import ODMMiniModel
+
+model = ODMMiniModel.from_pretrained("samatv256/mini-Jev")
+
+state = {
+    "user_request": "Find the current weather in Boston.",
+    "available_context": "The user has not provided weather data.",
+}
+candidates = [
+    {
+        "id": "weather.lookup",
+        "description": "Look up current weather for a specified city.",
+    },
+    {
+        "id": "calendar.list",
+        "description": "List upcoming calendar events for the user.",
+    },
+    {
+        "id": "control.finish",
+        "description": "Finish because the request has already been completed.",
+    },
+]
+
+choice = model.predict_choice(state=state, candidates=candidates)
+print("selected candidate:", choice.selected)
+print("probabilities:", choice.probabilities)
+print("confidence:", choice.confidence)
+print("decision margin:", choice.decision_margin)
+print("latency (ms):", choice.latency_ms)
+```
+
+`ODMMiniModel.from_pretrained()` accepts either the Hub model ID or a local snapshot directory. On CUDA it defaults to BF16 backbone inference; on CPU it defaults to FP32. The DecisionHead is always evaluated in FP32.
+
+## Training
+
+- 50,000 synthetic training decisions
+- frozen `Qwen/Qwen3-0.6B` backbone
+- DecisionHead-only training
+- grouped cross-entropy objective
+- selected seed: **41**
+
+The selected head is the seed-41 epoch-4 checkpoint. Temperature remains at `1.0`; its maximum grouped-softmax output is reported as confidence, not as a universally calibrated probability.
+
+## Evaluation
+
+Synthetic held-out evaluation:
+
+| Metric | Result |
+|---|---:|
+| Semantic Choice accuracy | **72.97%** |
+| Stress Choice accuracy | **67.64%** |
+| Counterfactual pair consistency | **67.12%** |
+
+Performance measurements used an NVIDIA GH200, BF16 backbone inference, and the shared-prefix KV-cache serving path. For short and medium contexts of approximately 256–1,024 state tokens, total latency was approximately **76–85 ms** for **3–16 candidates** (measured range: 76.11–82.36 ms). Offline representation extraction reached **831.3 candidates/second** at candidate batch size 512 while caching 430,072 candidates. These local measurements are hardware- and workload-specific and are not a direct speed comparison with Jev or any hosted service.
+
+## Limitations and intended use
+
+> **ODM Mini v1 is a research/hackathon prototype and is not ready for unmonitored production agent control.**
+
+In a real shadow-agent evaluation covering **75 multi-step trajectories and 243 decisions**, ODM Mini achieved:
+
+- action accuracy: **27.98%**
+- controller agreement: **26.34%**
+
+The major known failure is **high-confidence premature completion on intermediate multi-step trajectories**. After partial progress, the model can over-index on successful receipts and choose `control.finish` before remaining steps have been completed. Training used synthetic, static decision snapshots, so the published metrics should not be assumed to transfer to arbitrary tools, domains, or agent loops.
+
+Recommended uses are research, offline evaluation, single-step Choice experiments, and shadow-mode analysis with an independent controller. Do not use this checkpoint as the sole decision-maker for consequential actions or autonomous production control.
+
+## Released files
+
+- `model.safetensors` — ODM Mini DecisionHead tensors only
+- `config.json` — architecture and backbone reference
+- `odm_mini.py` — inference-only loader
+- `README.md` — model card and working example
+- `LICENSE` — Apache License 2.0
+
+## License
+
+ODM Mini-specific code and weights in this repository are released under Apache-2.0. The separately downloaded Qwen backbone is governed by its own repository terms.

config.json

@@ -0,0 +1,11 @@
+{
+  "architectures": ["ODMMiniModel"],
+  "model_type": "odm_mini",
+  "base_model": "Qwen/Qwen3-0.6B",
+  "hidden_size": 1024,
+  "head_inner_size": 256,
+  "pooling": "candidate_description_mean",
+  "temperature": 1.0,
+  "max_state_length": 4096,
+  "max_suffix_length": 256
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52b9cdd647edef215d64cc156be98aef5bace6c4cce1f33f10cadfad08d59b16
+size 1050940

odm_mini.py

@@ -0,0 +1,354 @@
+"""Inference-only loader for ODM Mini v1.
+
+The released repository contains only the trained DecisionHead. The frozen
+Qwen/Qwen3-0.6B backbone is downloaded separately by ``from_pretrained``.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Mapping, Sequence
+
+import torch
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+from torch import nn
+from transformers import AutoModel, AutoTokenizer
+from transformers.cache_utils import DynamicCache
+
+
+DEFAULT_REPO_ID = "samatv256/mini-Jev"
+DEFAULT_QUESTION = "What is the best next action for the agent?"
+
+
+def canonical_state(state: Any) -> str:
+    if isinstance(state, str):
+        return state.rstrip()
+    return json.dumps(state, sort_keys=True, ensure_ascii=False)
+
+
+def format_state_prefix(state: Any) -> str:
+    return f"STATE:\n{canonical_state(state)}\n\n"
+
+
+def format_candidate_suffix_prefix(question: str, candidate_id: str) -> str:
+    question = question.strip() if question and question.strip() else DEFAULT_QUESTION
+    return f"QUESTION:\n{question}\n\nCANDIDATE:\n{str(candidate_id).strip()}\n\nDESCRIPTION:\n"
+
+
+def mean_pool_description(
+    hidden_states: torch.Tensor, description_mask: torch.Tensor
+) -> torch.Tensor:
+    mask = description_mask.to(
+        dtype=hidden_states.dtype, device=hidden_states.device
+    ).unsqueeze(-1)
+    return (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1.0)
+
+
+class DecisionHead(nn.Module):
+    """Linear(1024, 256) -> GELU -> Linear(256, 1)."""
+
+    def __init__(self, hidden_size: int = 1024, inner_size: int = 256) -> None:
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(hidden_size, inner_size),
+            nn.GELU(),
+            nn.Linear(inner_size, 1),
+        )
+
+    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
+        return self.net(hidden.to(self.net[0].weight.dtype)).squeeze(-1)
+
+
+@dataclass(frozen=True)
+class DecisionCandidate:
+    id: str
+    description: str
+
+    def __post_init__(self) -> None:
+        if not self.id.strip():
+            raise ValueError("candidate id must not be empty")
+        if not self.description.strip():
+            raise ValueError("candidate description must not be empty")
+
+
+CandidateLike = str | Mapping[str, Any] | DecisionCandidate
+
+
+def normalize_candidates(candidates: Sequence[CandidateLike]) -> list[DecisionCandidate]:
+    normalized: list[DecisionCandidate] = []
+    for item in candidates:
+        if isinstance(item, DecisionCandidate):
+            candidate = item
+        elif isinstance(item, Mapping):
+            candidate = DecisionCandidate(
+                id=str(item["id"]), description=str(item["description"])
+            )
+        elif isinstance(item, str):
+            candidate = DecisionCandidate(id=item, description=item)
+        else:
+            raise TypeError(f"unsupported candidate type: {type(item)!r}")
+        normalized.append(candidate)
+    if not normalized:
+        raise ValueError("candidates list must not be empty")
+    ids = [candidate.id for candidate in normalized]
+    if len(ids) != len(set(ids)):
+        raise ValueError("candidate ids must be unique within a decision")
+    return normalized
+
+
+@dataclass(frozen=True)
+class ChoiceResponse:
+    selected: str
+    probabilities: dict[str, float]
+    confidence: float
+    decision_margin: float
+    latency_ms: float
+
+
+def _resolve_file(
+    model_id_or_path: str | Path,
+    filename: str,
+    *,
+    revision: str | None,
+    token: str | bool | None,
+) -> Path:
+    local_path = Path(model_id_or_path)
+    if local_path.is_dir():
+        artifact = local_path / filename
+        if not artifact.is_file():
+            raise FileNotFoundError(f"missing release artifact: {artifact}")
+        return artifact
+    return Path(
+        hf_hub_download(
+            repo_id=str(model_id_or_path),
+            filename=filename,
+            revision=revision,
+            token=token,
+        )
+    )
+
+
+def replicate_cache(source_cache: DynamicCache, repeats: int) -> DynamicCache:
+    new_cache = DynamicCache()
+    for layer in source_cache.layers:
+        keys = layer.keys.clone() if repeats == 1 else layer.keys.repeat(repeats, 1, 1, 1)
+        values = (
+            layer.values.clone()
+            if repeats == 1
+            else layer.values.repeat(repeats, 1, 1, 1)
+        )
+        new_cache.update(keys, values, len(new_cache.layers))
+    return new_cache
+
+
+class ODMMiniModel(nn.Module):
+    """Frozen Qwen3-0.6B plus the trained ODM Mini DecisionHead."""
+
+    def __init__(
+        self,
+        backbone: nn.Module,
+        tokenizer: Any,
+        head: DecisionHead,
+        *,
+        hidden_size: int = 1024,
+        head_inner_size: int = 256,
+        max_state_length: int = 4096,
+        max_suffix_length: int = 256,
+        candidate_batch_size: int = 32,
+        temperature: float = 1.0,
+    ) -> None:
+        super().__init__()
+        self.backbone = backbone
+        self.tokenizer = tokenizer
+        self.head = head
+        self.hidden_size = hidden_size
+        self.head_inner_size = head_inner_size
+        self.max_state_length = max_state_length
+        self.max_suffix_length = max_suffix_length
+        self.candidate_batch_size = candidate_batch_size
+        self.temperature = float(temperature)
+        self.escalation_threshold = 0.85
+        self.backbone.eval()
+        for parameter in self.backbone.parameters():
+            parameter.requires_grad = False
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_id_or_path: str | Path = DEFAULT_REPO_ID,
+        *,
+        revision: str | None = None,
+        device: str | torch.device | None = None,
+        dtype: torch.dtype | None = None,
+        token: str | bool | None = None,
+        candidate_batch_size: int = 32,
+    ) -> "ODMMiniModel":
+        device = device or ("cuda:0" if torch.cuda.is_available() else "cpu")
+        dtype = dtype or (torch.bfloat16 if torch.cuda.is_available() else torch.float32)
+        config_path = _resolve_file(
+            model_id_or_path, "config.json", revision=revision, token=token
+        )
+        weights_path = _resolve_file(
+            model_id_or_path, "model.safetensors", revision=revision, token=token
+        )
+        config = json.loads(config_path.read_text(encoding="utf-8"))
+
+        base_model = str(config["base_model"])
+        tokenizer = AutoTokenizer.from_pretrained(base_model)
+        backbone = AutoModel.from_pretrained(base_model, dtype=dtype)
+        backbone.to(device)
+        backbone.eval()
+
+        hidden_size = int(config["hidden_size"])
+        actual_hidden_size = int(getattr(backbone.config, "hidden_size", hidden_size))
+        if actual_hidden_size != hidden_size:
+            raise ValueError(
+                f"backbone hidden size {actual_hidden_size} does not match release config {hidden_size}"
+            )
+
+        head = DecisionHead(
+            hidden_size=hidden_size,
+            inner_size=int(config["head_inner_size"]),
+        )
+        head.load_state_dict(load_file(str(weights_path), device="cpu"), strict=True)
+        head.to(device=device, dtype=torch.float32)
+        head.eval()
+
+        return cls(
+            backbone=backbone,
+            tokenizer=tokenizer,
+            head=head,
+            hidden_size=hidden_size,
+            head_inner_size=int(config["head_inner_size"]),
+            max_state_length=int(config["max_state_length"]),
+            max_suffix_length=int(config["max_suffix_length"]),
+            candidate_batch_size=candidate_batch_size,
+            temperature=float(config["temperature"]),
+        )
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.backbone.parameters()).device
+
+    @torch.inference_mode()
+    def encode_state(self, state: Any) -> tuple[DynamicCache, int]:
+        prefix_ids = self.tokenizer.encode(
+            format_state_prefix(state),
+            add_special_tokens=False,
+            truncation=True,
+            max_length=self.max_state_length,
+        )
+        if not prefix_ids:
+            prefix_ids = [self.tokenizer.pad_token_id or 0]
+        input_ids = torch.tensor([prefix_ids], dtype=torch.long, device=self.device)
+        output = self.backbone(input_ids=input_ids, use_cache=True)
+        return output.past_key_values, len(prefix_ids)
+
+    @torch.inference_mode()
+    def score_candidates_cached(
+        self,
+        past_key_values: DynamicCache,
+        prefix_length: int,
+        candidates: Sequence[DecisionCandidate],
+        question: str = DEFAULT_QUESTION,
+        candidate_batch_size: int | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        batch_size_limit = candidate_batch_size or self.candidate_batch_size
+        all_representations: list[torch.Tensor] = []
+
+        for start in range(0, len(candidates), batch_size_limit):
+            batch = candidates[start : start + batch_size_limit]
+            suffixes: list[list[int]] = []
+            description_spans: list[tuple[int, int]] = []
+            for candidate in batch:
+                suffix_prefix = format_candidate_suffix_prefix(question, candidate.id)
+                prefix_ids = self.tokenizer.encode(
+                    suffix_prefix,
+                    add_special_tokens=False,
+                    truncation=True,
+                    max_length=self.max_suffix_length // 2,
+                )
+                description_ids = self.tokenizer.encode(
+                    candidate.description.strip(),
+                    add_special_tokens=False,
+                    truncation=True,
+                    max_length=max(1, self.max_suffix_length - len(prefix_ids)),
+                )
+                suffix = prefix_ids + description_ids
+                suffixes.append(suffix)
+                description_spans.append((len(prefix_ids), len(suffix)))
+
+            max_length = max(len(suffix) for suffix in suffixes)
+            pad_id = int(self.tokenizer.pad_token_id or 0)
+            input_ids = torch.full(
+                (len(batch), max_length), pad_id, dtype=torch.long, device=self.device
+            )
+            attention_mask = torch.ones(
+                (len(batch), prefix_length + max_length),
+                dtype=torch.long,
+                device=self.device,
+            )
+            for index, suffix in enumerate(suffixes):
+                input_ids[index, : len(suffix)] = torch.tensor(
+                    suffix, dtype=torch.long, device=self.device
+                )
+                if len(suffix) < max_length:
+                    attention_mask[index, prefix_length + len(suffix) :] = 0
+
+            output = self.backbone(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                past_key_values=replicate_cache(past_key_values, len(batch)),
+                use_cache=False,
+            )
+            description_mask = torch.zeros(
+                (len(batch), max_length), dtype=torch.bool, device=self.device
+            )
+            for index, (description_start, description_end) in enumerate(description_spans):
+                description_mask[index, description_start:description_end] = True
+            all_representations.append(
+                mean_pool_description(output.last_hidden_state, description_mask)
+            )
+
+        representations = torch.cat(all_representations, dim=0)
+        return self.head(representations), representations
+
+    @torch.inference_mode()
+    def predict_choice(
+        self,
+        state: Any,
+        candidates: Sequence[CandidateLike],
+        question: str = DEFAULT_QUESTION,
+        temperature: float | None = None,
+    ) -> ChoiceResponse:
+        started = time.perf_counter()
+        normalized = normalize_candidates(candidates)
+        cache, prefix_length = self.encode_state(state)
+        logits, _ = self.score_candidates_cached(
+            cache, prefix_length, normalized, question=question
+        )
+        float_logits = logits.float()
+        sorted_logits, _ = torch.sort(float_logits, descending=True)
+        margin = float(
+            (sorted_logits[0] - sorted_logits[1]).item()
+            if len(sorted_logits) > 1
+            else sorted_logits[0].item()
+        )
+        effective_temperature = self.temperature if temperature is None else float(temperature)
+        probabilities = torch.softmax(float_logits / max(effective_temperature, 1e-4), dim=-1)
+        selected_index = int(probabilities.argmax().item())
+        probability_map = {
+            candidate.id: float(probability.item())
+            for candidate, probability in zip(normalized, probabilities, strict=True)
+        }
+        return ChoiceResponse(
+            selected=normalized[selected_index].id,
+            probabilities=probability_map,
+            confidence=float(probabilities[selected_index].item()),
+            decision_margin=margin,
+            latency_ms=(time.perf_counter() - started) * 1000.0,
+        )