ruling 0.2.0__py3-none-any.whl

ruling/distill.py

@@ -0,0 +1,131 @@
+"""Distillation data: a large open model writes and answers typed questions about many states.
+
+The served model is the fast decision encoder, which generalizes poorly when
+trained only on a few labeled sources. Here a large open-weight decoder writes
+varied Choice, Score, and Noul questions about each state, then answers them
+through the decoder engine with option-order averaging. Its probabilities
+become soft targets, so the encoder learns the larger model's judgment,
+uncertainty included, over far more kinds of question than any labeled
+dataset holds.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from pathlib import Path
+
+import mlx.core as mx
+from mlx_lm import batch_generate
+from mlx_lm.sample_utils import make_sampler
+from pydantic import TypeAdapter, ValidationError
+
+from ruling.calibration import softmax
+from ruling.engine import TEMPLATE_KWARGS, Engine, InputTooLong
+from ruling.evals import Record, load_dataset
+from ruling.questions import InvalidOptions, Question, State, check_askable
+
+WRITER_SYSTEM = ("You design evaluation questions for software that makes fast judgments about documents. "
+                 "Reply with one JSON object and nothing else.")
+QUESTION = TypeAdapter(Question)
+
+
+def writer_messages(state: State, count: int) -> list[dict[str, str]]:
+    body = state if isinstance(state, str) else json.dumps(state, indent=2, ensure_ascii=False)
+    instructions = f"""Document:
+{body}
+
+Write {count} different questions a program might need answered about this document. Mix three types:
+- "noul": a statement that is true or false about the document.
+- "choice": pick one of 2 to 8 named options; give each option a short description.
+- "score": an ordered scale of 3 to 5 levels from lowest to highest.
+Cover different aspects: intent, tone, facts stated, facts not stated, risk, next action, category. Include some questions
+whose answer is genuinely uncertain from the document, and some whose answer is clearly no. Use snake_case ids.
+
+Reply with exactly this JSON shape:
+{{"questions": {{
+  "some_id": {{"type": "noul", "instructions": "..."}},
+  "other_id": {{"type": "choice", "instructions": "...", "criteria": {{"option_a": "description", "option_b": "description"}}}},
+  "third_id": {{"type": "score", "instructions": "...", "criteria": ["lowest level", "middle level", "highest level"]}}
+}}}}"""
+    return [{"role": "system", "content": WRITER_SYSTEM}, {"role": "user", "content": instructions}]
+
+
+def parse_questions(text: str, max_options: int) -> dict:
+    """The valid questions in a writer reply; anything malformed is dropped rather than repaired."""
+    text = re.sub(r"<think>.*?</think>", "", text, flags=re.S)
+    start, end = text.find("{"), text.rfind("}")
+    if start < 0 or end <= start:
+        return {}
+    try:
+        payload = json.loads(text[start:end + 1])
+    except json.JSONDecodeError:
+        return {}
+    candidates = payload.get("questions") if isinstance(payload, dict) else None
+    if not isinstance(candidates, dict):
+        return {}
+    questions = {}
+    for qid, spec in candidates.items():
+        try:
+            question = QUESTION.validate_python(spec)
+            check_askable({str(qid): question}, max_options)
+        except (ValidationError, InvalidOptions):
+            continue
+        questions[str(qid)] = question
+    return questions
+
+
+def write_questions(states: list[tuple[str, State]], generate, questions_per_state: int, batch_size: int,
+                    max_options: int, log=print):
+    """Yield, batch by batch, records carrying the questions a writer model produced for each state.
+
+    `generate` takes a list of chat message lists and returns one reply text per prompt.
+    """
+    kept = 0
+    for start in range(0, len(states), batch_size):
+        batch = states[start:start + batch_size]
+        texts = generate([writer_messages(state, questions_per_state) for _, state in batch])
+        records = [Record(family=f"distill_{source}", state=state, questions=questions)
+                   for (source, state), text in zip(batch, texts)
+                   if (questions := parse_questions(text, max_options))]
+        kept += len(records)
+        log(f"{min(start + batch_size, len(states))}/{len(states)} states, {kept} with questions")
+        yield records
+
+
+def local_generator(engine: Engine, seed: int, temperature: float):
+    """A `generate` function backed by a local MLX model with batched sampling."""
+    mx.random.seed(seed)
+    sampler = make_sampler(temp=temperature, top_p=0.95)
+
+    def generate(prompts: list[list[dict]]) -> list[str]:
+        tokens = [engine.tokenizer.apply_chat_template(messages, add_generation_prompt=True, **TEMPLATE_KWARGS) for messages in prompts]
+        return batch_generate(engine.model, engine.tokenizer, tokens, max_tokens=1200, sampler=sampler).texts
+
+    return generate
+
+
+def unanswered(records: list[Record], written: Path) -> list[Record]:
+    """The records whose answers are not yet in `written`, matched by state and questions.
+
+    A record too long to answer leaves no line, so counting lines would misplace
+    the resume point and answer some records twice.
+    """
+    if not written.exists():
+        return records
+    key = lambda r: json.dumps([r.state, {qid: q.model_dump() for qid, q in r.questions.items()}], sort_keys=True)
+    done = {key(r) for r in load_dataset(written)}
+    return [r for r in records if key(r) not in done]
+
+
+def answer_locally(records: list[Record], engine, log=print):
+    """Yield each record with the engine's answer distributions as `target` references; records too long are skipped."""
+    for index, record in enumerate(records):
+        try:
+            scored = engine.score(record.state, record.questions)
+        except InputTooLong:
+            continue
+        targets = {qid: dict(zip(raw.keys, (float(p) for p in softmax(raw.logits)))) for qid, raw in scored.raw.items()}
+        if (index + 1) % 100 == 0:
+            log(f"{index + 1}/{len(records)} answered")
+        yield record.model_copy(update={"references": {qid: {"target": t} for qid, t in targets.items()}})

ruling/engine.py

@@ -0,0 +1,518 @@
+"""Score typed questions against a state with a local MLX model.
+
+No text is generated. The state prefix is encoded once into a KV cache. Every
+question suffix (one per rotation of its answer order) is then scored in
+batched forward passes against replicas of that cache, and the logits at each
+suffix's answer position are restricted to the allowed answer codes.
+
+Positional bias is removed two ways: a content-free letter prior measured once
+per answer count is subtracted from every log-probability, and the remaining
+orderings are averaged in log space.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import threading
+from dataclasses import dataclass
+from datetime import date
+from pathlib import Path
+
+import mlx.core as mx
+import numpy as np
+from huggingface_hub import HfApi, snapshot_download
+from huggingface_hub.errors import HfHubHTTPError, HFValidationError, LocalEntryNotFoundError
+from mlx.utils import tree_flatten
+from mlx_lm import load
+from mlx_lm.models.cache import make_prompt_cache
+from mlx_lm.tuner.utils import linear_to_lora_layers
+
+from ruling.calibration import Calibration, Provenance, choice_confidence, log_softmax, score_confidence, softmax
+from ruling.codes import AnswerCodes
+from ruling.prompt import (
+    PRIOR_STATE,
+    SPLIT_MARKER,
+    SYSTEM_PROMPT,
+    prior_question,
+    render_question,
+    render_state,
+)
+from ruling.questions import (
+    MAX_CHOICE_OPTIONS,
+    Choice,
+    ChoiceAnswer,
+    Noul,
+    NoulAnswer,
+    Score,
+    ScoreAnswer,
+    State,
+    SystemOneRequest,
+    SystemOneResponse,
+    Usage,
+    option_keys,
+    render_text,
+)
+
+# Prefix tokens are pushed through the model in chunks so the full-vocabulary
+# logits for a long state never have to exist all at once.
+PREFILL_STEP = 512
+
+# Suffix passes are batched until batch_size * (prefix + padded suffix) tokens
+# would exceed this budget, which bounds the memory of the replicated cache.
+BATCH_TOKEN_BUDGET = 65_536
+
+# For a Choice with more options than single-token codes, every option is first
+# scored on its own as a Noul and the strongest few go to a final Choice.
+STAGE_TWO_CANDIDATES = 8
+
+# Qwen3-family templates emit an empty thinking block when thinking is off.
+# Templates without the variable ignore it.
+TEMPLATE_KWARGS = {"enable_thinking": False}
+
+
+class InputTooLong(ValueError):
+    """Raised with a message that names the maximum context length, the phrase clients
+    such as the Decision Index harness look for to classify a rejection as a capacity
+    limit rather than a failure. `question_id` names the question at fault when one is."""
+
+    def __init__(self, message: str, question_id: str | None = None):
+        super().__init__(message)
+        self.question_id = question_id
+
+
+@dataclass(frozen=True)
+class RawScore:
+    """Debiased log-probabilities over `keys`, in canonical key order, before calibration."""
+
+    keys: list[str]
+    logits: np.ndarray
+
+
+@dataclass(frozen=True)
+class Scored:
+    raw: dict[str, RawScore]
+    input_tokens: int
+
+
+@dataclass(frozen=True)
+class _Pass:
+    question_id: str
+    keys: list[str]
+    tokens: list[int]
+
+
+@dataclass(frozen=True)
+class _Prefix:
+    """The tokens through the state, the template tail every question ends with, and what a
+    vision model needs beside them: its processor inputs for the images in the state, and how
+    far its position ids sit below the token count once the images are read."""
+    tokens: list[int]
+    tail: str
+    vision_inputs: dict | None = None
+    position_delta: int = 0
+
+
+class TwoStageChoice:
+    """Template method for engines that read at most `pass_options` options in one pass.
+
+    A Choice wider than that is first scored as a Noul per option; the strongest
+    few go to a final Choice, and every other option gets zero probability.
+    Subclasses supply `_score_plain` and `pass_options`.
+    """
+
+    pass_options: int
+
+    def _score_plain(self, state: State, questions: dict[str, Choice | Score | Noul]) -> Scored:
+        raise NotImplementedError
+
+    def score(self, state: State, questions: dict[str, Choice | Score | Noul]) -> Scored:
+        large = {qid: q for qid, q in questions.items()
+                 if isinstance(q, Choice) and len(q.criteria) > self.pass_options}
+        tokens = 0
+        reduced = dict(questions)
+        finalists: dict[str, list[str]] = {}
+        if large:
+            stage_one = self._score_plain(state, {
+                f"{qid}\x00{key}": Noul(instructions={"question": q.instructions, "candidate": {key: q.criteria[key]}})
+                for qid, q in large.items() for key in q.criteria
+            })
+            tokens += stage_one.input_tokens
+            for qid, q in large.items():
+                ranked = sorted(q.criteria, key=lambda key: -stage_one.raw[f"{qid}\x00{key}"].logits[0])
+                finalists[qid] = ranked[:STAGE_TWO_CANDIDATES]
+                reduced[qid] = Choice(instructions=q.instructions, criteria={k: q.criteria[k] for k in finalists[qid]})
+        scored = self._score_plain(state, reduced)
+        tokens += scored.input_tokens
+        raw = dict(scored.raw)
+        for qid, keys in finalists.items():
+            full = np.full(len(large[qid].criteria), -np.inf)
+            for key, value in zip(keys, scored.raw[qid].logits):
+                full[list(large[qid].criteria).index(key)] = value
+            raw[qid] = RawScore(keys=option_keys(large[qid]), logits=full)
+        return Scored(raw=raw, input_tokens=tokens)
+
+
+def answer(question: Choice | Score | Noul, raw: RawScore, calibration: Calibration) -> ChoiceAnswer | ScoreAnswer | NoulAnswer:
+    """A typed answer from calibrated probabilities over the question's canonical keys."""
+    probs = softmax(raw.logits, calibration.temperature(question.type))
+    distribution = {key: float(p) for key, p in zip(raw.keys, probs)}
+    if isinstance(question, Choice):
+        return ChoiceAnswer(choice=raw.keys[int(probs.argmax())], probabilities=distribution,
+                            confidence=choice_confidence(probs))
+    if isinstance(question, Score):
+        return ScoreAnswer(
+            score=float((probs * np.arange(len(probs))).sum()),
+            legend={str(i): render_text(level) for i, level in enumerate(question.criteria)},
+            probabilities=distribution,
+            confidence=score_confidence(probs),
+        )
+    return NoulAnswer(noul=distribution["yes"])
+
+
+def option_logits(model, tokens: mx.array, cache, last: mx.array, option_ids: mx.array,
+                  position_ids: mx.array | None = None, inputs_embeds: mx.array | None = None) -> mx.array:
+    """Logits for each row's answer codes at its last real position, shape [batch, options].
+
+    Only the hidden state at the answer position goes through the output head,
+    so a batch never materializes full-vocabulary logits for every token. This
+    follows the mlx-lm causal layout: an optional `language_model` wrapper, a
+    `.model` that returns final hidden states, and either a separate `lm_head`
+    or tied input embeddings.
+    """
+    text = getattr(model, "language_model", model)
+    # Embeddings with image features and the positions images moved are passed only when there are images;
+    # text-only models take neither.
+    extra = {key: value for key, value in (("position_ids", position_ids), ("inputs_embeds", inputs_embeds))
+             if value is not None}
+    hidden = text.model(tokens, cache=cache, **extra)
+    picked = hidden[mx.arange(tokens.shape[0]), last]
+    tied = getattr(text, "tie_word_embeddings", None)
+    if tied is None:
+        tied = text.args.tie_word_embeddings
+    logits = text.model.embed_tokens.as_linear(picked) if tied else text.lm_head(picked)
+    return mx.take_along_axis(logits, option_ids, axis=-1).astype(mx.float32)
+
+
+def adapter_digest(adapter: Path | None) -> str | None:
+    """A short content hash of LoRA weights, so calibration and reports name the exact adapter."""
+    if adapter is None:
+        return None
+    return hashlib.sha256((adapter / ADAPTER_WEIGHTS).read_bytes()).hexdigest()[:16]
+
+
+ADAPTER_WEIGHTS = "adapters.safetensors"
+ADAPTER_CONFIG = "adapter_config.json"
+
+
+def lora_parameters(model) -> dict[str, mx.array]:
+    """Every LoRA weight in the model, keyed as mlx-lm saves and loads them."""
+    return {key: value for key, value in tree_flatten(model.parameters()) if key.rsplit(".", 1)[-1] in ("lora_a", "lora_b")}
+
+
+class _SharedModel:
+    """What engines built on one loaded model have in common: the lock that serializes the GPU,
+    the LoRA shape its layers were converted to, and which engine's weights are in them now."""
+
+    def __init__(self):
+        self.lock = threading.Lock()
+        self.lora: dict | None = None
+        self.active: "Engine | None" = None
+
+
+def model_released(model_id: str) -> str:
+    """The date the served weights were published, as `YYYY-MM-DD`, or "" when unknown.
+
+    Read from the Hugging Face repository when the model came from one and it is
+    reachable. A local directory, or an unreachable repository, falls back to the
+    date the weights landed on this machine. A served id carries its engine's
+    prefix (`openai:`, `vision:`, ...), which no repository id can contain, so the
+    repository is what follows the last colon. A remote service's own model names
+    (`typesafe:jev-latest`) name no repository: their date is unknown, and the
+    hosted API's clients require a string.
+    """
+    repo = model_id.rsplit(":", 1)[-1]
+    local = Path(repo)
+    if not local.exists():
+        try:
+            return HfApi().model_info(repo).last_modified.date().isoformat()
+        except (HfHubHTTPError, HFValidationError, OSError, AttributeError):
+            try:
+                local = Path(snapshot_download(repo, allow_patterns=["config.json"], local_files_only=True))
+            except (HFValidationError, LocalEntryNotFoundError):
+                return ""
+    return date.fromtimestamp(local.stat().st_mtime).isoformat()
+
+
+def model_revision(model_id: str) -> str | None:
+    """The Hugging Face commit the weights came from, or None for a local directory."""
+    if Path(model_id).exists():
+        return None
+    # mlx_lm.load has already fetched the weights; only the cached config is consulted here.
+    path = Path(snapshot_download(model_id, allow_patterns=["config.json"], local_files_only=True))
+    return path.name if path.parent.name == "snapshots" else None
+
+
+class Engine(TwoStageChoice):
+    def __init__(
+        self,
+        model_id: str,
+        calibration: Calibration,
+        max_input_tokens: int,
+        rotations: int = 1,
+        prior_debias: bool = True,
+        adapter: Path | None = None,
+        max_branch_tokens: int | None = None,
+    ):
+        if rotations < 1:
+            raise ValueError("rotations must be at least 1")
+        self.model_id = model_id
+        self.max_input_tokens = max_input_tokens
+        self.max_branch_tokens = max_branch_tokens or max_input_tokens
+        self.rotations = rotations
+        self.prior_debias = prior_debias
+        self.adapter = adapter_digest(adapter)
+        self.model, self.tokenizer = self._load(model_id, adapter)
+        # Adapter weights come back lazy, and a request thread has no CPU stream to realize them
+        # on; settle every parameter here, on the thread that loaded the model.
+        mx.eval(self.model.parameters())
+        self.revision = model_revision(model_id)
+        self.codes = AnswerCodes.from_tokenizer(self.tokenizer)
+        self._pad_id = self.tokenizer.pad_token_id
+        if self._pad_id is None:
+            self._pad_id = self.tokenizer.eos_token_id
+        self._system_role = self._template_accepts_system_role()
+        self.max_options = MAX_CHOICE_OPTIONS
+        self.reads_images = False
+        self._shared = _SharedModel()
+        self._shared.lora = json.loads((adapter / ADAPTER_CONFIG).read_text()) if adapter else None
+        self._shared.active = self
+        self._lock = self._shared.lock
+        # The weights to put into the shared LoRA layers before answering. None until another
+        # engine shares this model, since until then nothing can have swapped them out.
+        self._lora_weights: dict[str, mx.array] | None = None
+        self.adapter_path = adapter
+        self._priors: dict[int, np.ndarray] = {}
+        calibration.check(self.provenance)
+        self.calibration = calibration
+
+    def _load(self, model_id: str, adapter: Path | None):
+        return load(model_id, adapter_path=str(adapter) if adapter else None)
+
+    class IncompatibleAdapter(ValueError):
+        """The adapter's LoRA layout differs from the one the shared model was converted to."""
+
+    def variant(self, adapter: Path | None, calibration: Calibration) -> "Engine":
+        """Another engine on this loaded model with its own adapter (or none) and calibration.
+
+        The first variant converts the shared model's layers to LoRA in the adapter's
+        layout; every later adapter must match that layout exactly, since mlx-lm can
+        only load weights into layers of the same rank and depth. Weights are read and
+        realized here, on the loading thread, and merely rebound per request.
+        """
+        config = json.loads((adapter / ADAPTER_CONFIG).read_text()) if adapter else None
+        shared = self._shared
+        if shared.lora is None and config is not None:
+            linear_to_lora_layers(self.model, config["num_layers"], config["lora_parameters"])
+            shared.lora = config
+            # Fresh LoRA layers have zero `lora_b`, so these weights reproduce the plain base.
+            self._lora_weights = lora_parameters(self.model)
+            mx.eval(list(self._lora_weights.values()))
+        elif config is not None and config != shared.lora:
+            raise Engine.IncompatibleAdapter(
+                f"{adapter} has LoRA layout {config}, but the loaded {self.model_id} was converted for {shared.lora}")
+        if self._lora_weights is None and shared.lora is not None:
+            self._lora_weights = mx.load(str(self.adapter_path / ADAPTER_WEIGHTS))
+            mx.eval(list(self._lora_weights.values()))
+        other = type(self).__new__(type(self))
+        other.__dict__.update(self.__dict__)
+        other._priors = {}
+        other.adapter_path = adapter
+        other.adapter = adapter_digest(adapter)
+        if shared.lora is None:
+            other._lora_weights = None
+        elif adapter is None:
+            other._lora_weights = {key: mx.zeros_like(value) for key, value in lora_parameters(self.model).items()}
+            mx.eval(list(other._lora_weights.values()))
+        else:
+            other._lora_weights = mx.load(str(adapter / ADAPTER_WEIGHTS))
+            mx.eval(list(other._lora_weights.values()))
+        calibration.check(other.provenance)
+        other.calibration = calibration
+        return other
+
+    def _activate(self) -> None:
+        """Put this engine's LoRA weights into the shared layers; call with the lock held."""
+        if self._lora_weights is not None and self._shared.active is not self:
+            self.model.load_weights(list(self._lora_weights.items()), strict=False)
+            self._shared.active = self
+
+    @property
+    def provenance(self) -> Provenance:
+        return Provenance(self.model_id, self.revision, self.rotations, self.prior_debias, self.adapter)
+
+    @property
+    def pass_options(self) -> int:
+        return self.codes.max_options
+
+    # ---- prompt assembly -------------------------------------------------
+
+    def split_prompt(self, state: State, thinking: bool = False) -> tuple[str, str]:
+        """Return the chat-templated prefix (through the state) and the template tail.
+
+        With `thinking`, the tail ends where the model starts reasoning instead of answering.
+        """
+        rendered_state = render_state(state)
+        if SPLIT_MARKER in rendered_state:
+            raise ValueError("state contains the reserved split marker")
+        content = self._user_content(state, rendered_state)
+        if self._system_role:
+            messages = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": content}]
+        else:
+            messages = [{"role": "user", "content": _prepend(SYSTEM_PROMPT + "\n\n", content)}]
+        text = self.tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True, enable_thinking=thinking
+        )
+        prefix, tail = text.split(SPLIT_MARKER)
+        return prefix, tail
+
+    def _user_content(self, state: State, rendered_state: str) -> str | list[dict]:
+        """The user turn through the split marker."""
+        return rendered_state + SPLIT_MARKER
+
+    def _template_accepts_system_role(self) -> bool:
+        try:
+            self.tokenizer.apply_chat_template(
+                [{"role": "system", "content": "s"}, {"role": "user", "content": "u"}],
+                tokenize=False, add_generation_prompt=True, **TEMPLATE_KWARGS,
+            )
+        except Exception:  # templates raise jinja TemplateError subclasses; the type varies by model
+            return False
+        return True
+
+    # ---- scoring -----------------------------------------------------------
+
+    def _score_plain(self, state: State, questions: dict[str, Choice | Score | Noul]) -> Scored:
+        prefix = self._prefix(state)
+        prefix_tokens = prefix.tokens
+        passes = [
+            _Pass(qid, rendered.keys, self._encode(rendered.text + prefix.tail))
+            for qid, question in questions.items()
+            for rendered in (render_question(question, self.codes.codes, r) for r in self._rotations_for(question))
+        ]
+        widest = max((p for p in passes), key=lambda p: len(p.tokens), default=None)
+        if widest is not None and len(prefix_tokens) + len(widest.tokens) > self.max_branch_tokens:
+            raise InputTooLong(f"question {widest.question_id!r} makes a branch of "
+                               f"{len(prefix_tokens) + len(widest.tokens)} tokens, over the maximum context length "
+                               f"of {self.max_branch_tokens} for one branch", question_id=widest.question_id)
+        total = len(prefix_tokens) + sum(len(p.tokens) for p in passes)
+        if total > self.max_input_tokens:
+            raise InputTooLong(f"request is {total} tokens, over the maximum context length of "
+                               f"{self.max_input_tokens} for a whole request")
+
+        per_question: dict[str, dict[str, list[float]]] = {qid: {} for qid in questions}
+        with self._lock:
+            self._activate()
+            prefix_cache = self._encode_prefix(prefix)
+            for chunk in self._chunks(passes, len(prefix_tokens)):
+                for one, logits in zip(chunk, self._forward_batch(prefix_cache, chunk, prefix)):
+                    log_probs = log_softmax(logits)
+                    if self.prior_debias:
+                        log_probs = log_probs - self._prior(len(one.keys))
+                    for key, value in zip(one.keys, log_probs):
+                        per_question[one.question_id].setdefault(key, []).append(float(value))
+
+        raw = {}
+        for qid, question in questions.items():
+            keys = option_keys(question)
+            raw[qid] = RawScore(keys=keys, logits=np.array([np.mean(per_question[qid][k]) for k in keys]))
+        return Scored(raw=raw, input_tokens=total)
+
+    def _prior(self, cardinality: int) -> np.ndarray:
+        """Log-probabilities the model assigns to each letter with no content to go on."""
+        if cardinality not in self._priors:
+            prefix = self._prefix(PRIOR_STATE)
+            rendered = render_question(prior_question(cardinality), self.codes.codes)
+            one = _Pass("prior", rendered.keys, self._encode(rendered.text + prefix.tail))
+            self._priors[cardinality] = log_softmax(self._forward_batch(self._encode_prefix(prefix), [one], prefix)[0])
+        return self._priors[cardinality]
+
+    def _chunks(self, passes: list[_Pass], prefix_length: int) -> list[list[_Pass]]:
+        chunks: list[list[_Pass]] = [[]]
+        for one in passes:
+            candidate = chunks[-1] + [one]
+            width = max(len(p.tokens) for p in candidate)
+            if chunks[-1] and len(candidate) * (prefix_length + width) > BATCH_TOKEN_BUDGET:
+                chunks.append([one])
+            else:
+                chunks[-1] = candidate
+        return chunks
+
+    def _forward_batch(self, prefix_cache, passes: list[_Pass], prefix: _Prefix) -> list[np.ndarray]:
+        """One padded forward pass for every suffix, against replicas of the prefix cache."""
+        batch = len(passes)
+        width = max(len(p.tokens) for p in passes)
+        cache = make_prompt_cache(self.model)
+        for replica, original in zip(cache, prefix_cache):
+            state = original.state
+            if isinstance(state, tuple):
+                replica.state = tuple(mx.concatenate([a] * batch, axis=0) for a in state)
+            else:
+                replica.state = [None if a is None else mx.concatenate([a] * batch, axis=0) for a in state]
+        tokens = mx.array([p.tokens + [self._pad_id] * (width - len(p.tokens)) for p in passes])
+        last = mx.array([len(p.tokens) - 1 for p in passes])
+        cardinality = max(len(p.keys) for p in passes)
+        ids = mx.array([self.codes.token_ids[:cardinality] for _ in passes])
+        positions = suffix_positions(prefix, batch, width)
+        logits = np.array(option_logits(self.model, tokens, cache, last, ids, positions), dtype=np.float64)
+        return [row[: len(one.keys)] for row, one in zip(logits, passes)]
+
+    # ---- answers -----------------------------------------------------------
+
+    def evaluate(self, request: SystemOneRequest) -> SystemOneResponse:
+        scored = self.score(request.state, request.questions)
+        answers = {qid: self.answer(request.questions[qid], scored.raw[qid]) for qid in request.questions}
+        return SystemOneResponse(
+            model=self.model_id, answers=answers, usage=Usage(input_tokens=scored.input_tokens)
+        )
+
+    def answer(self, question: Choice | Score | Noul, raw: RawScore) -> ChoiceAnswer | ScoreAnswer | NoulAnswer:
+        return answer(question, raw, self.calibration)
+
+    # ---- helpers -----------------------------------------------------------
+
+    def _rotations_for(self, question: Choice | Score | Noul) -> range:
+        orderings = 2 if isinstance(question, Score) else len(option_keys(question))
+        return range(min(self.rotations, orderings))
+
+    def _encode(self, text: str) -> list[int]:
+        return self.tokenizer.encode(text, add_special_tokens=False)
+
+    def _prefix(self, state: State) -> _Prefix:
+        prefix_text, tail = self.split_prompt(state)
+        return _Prefix(self._encode(prefix_text), tail)
+
+    def _encode_prefix(self, prefix: _Prefix):
+        tokens = prefix.tokens
+        cache = make_prompt_cache(self.model)
+        array = mx.array(tokens)[None]
+        for start in range(0, len(tokens), PREFILL_STEP):
+            self.model(array[:, start : start + PREFILL_STEP], cache=cache)
+            mx.eval([c.state for c in cache])
+        return cache
+
+
+def suffix_positions(prefix: _Prefix, batch: int, width: int) -> mx.array | None:
+    """Multimodal position ids for `batch` suffixes of `width` tokens after `prefix`, or None
+    when the model counts positions from the cache as usual: after images, a token's
+    position on every axis is its index plus the prefix's (negative) delta."""
+    if not prefix.position_delta:
+        return None
+    row = mx.arange(width) + len(prefix.tokens) + prefix.position_delta
+    return mx.broadcast_to(row[None, None, :], (3, batch, width))
+
+
+def _prepend(text: str, content: str | list[dict]) -> str | list[dict]:
+    if isinstance(content, str):
+        return text + content
+    return [{"type": "text", "text": text}, *content]