yev 0.2.0

README.md

@@ -0,0 +1,65 @@
+---
+tags:
+  - yev
+  - decision-model
+---
+
+# yev
+
+The client for [Yev-Decision-1](https://huggingface.co/HallD/Yev-Decision-1), a local **System One**
+decision model. Give it a state, which can be text, images or both, plus named questions. It returns
+a typed probability distribution for each one: a **Choice** between options, a **Score** on an
+ordered scale, or a **Noul** (yes or no). It generates no text and nothing leaves the machine.
+
+## Install
+
+```sh
+uv add "yev[inference] @ git+https://huggingface.co/HallD/yev"
+```
+
+`pip install "yev[inference] @ git+https://huggingface.co/HallD/yev"` works the same way. Python 3.13
+or later.
+
+| Extra | Brings | For |
+| --- | --- | --- |
+| *(none)* | pydantic, huggingface-hub, numpy | Question and answer types only |
+| `inference` | torch, transformers, peft, accelerate, pillow | Answering questions with a model |
+| `quantized` | `inference` plus bitsandbytes | `quantize4bit=True` on a CUDA GPU |
+
+## Use
+
+```python
+from yev import Choice, Image, Noul, Score, YevClient
+
+with YevClient() as client:   # YEV_DEFAULT_MODEL, else HallD/Yev-Decision-1
+    response = client.system_one(
+        state="I was charged twice. Please fix this ASAP.",
+        questions={
+            "billing": Noul(instructions="Is this ticket about billing?"),
+            "tone": Choice(instructions="What is the customer's tone?",
+                           criteria={"calm": None, "frustrated": None, "angry": None}),
+            "urgency": Score(instructions="How urgent is this ticket?",
+                             criteria=["can wait", "this week", "today"]),
+        },
+    )
+print(response.nouls["billing"].noul, response.choices["tone"].probabilities)
+```
+
+The first call downloads the model and its backbone,
+[`Qwen/Qwen3.5-4B-Base`](https://huggingface.co/Qwen/Qwen3.5-4B-Base) at a pinned revision, into the
+Hugging Face cache. After that it runs offline (`HF_HUB_OFFLINE=1`). The backbone runs in bfloat16;
+a GPU with 12 GB holds it.
+
+The call shape follows the TypeSafe SDK (0.7.5): `YevClient` and `AsyncYevClient` take its
+`system_one(state, questions, model=, assets=, response_model=)`, its question objects (TypeSafe's
+own are accepted), and return its answer and response types. It runs in-process, so there are no
+API keys, retries or timeouts; Box questions are refused, and a Score takes at most 10 levels.
+
+`yev --export HallD/Yev-Decision-1 --state "I was charged twice." --noul billing "Is this ticket about billing?"`
+asks from the command line.
+
+## JevBench
+
+`examples/jevbench_adapter.py` is an adapter for the
+[JevBench](https://github.com/fstandhartinger/jevbench) harness: it sends the harness's own question
+to `system_one` and returns the probabilities.

examples/jevbench_adapter.py

@@ -0,0 +1,48 @@
+"""A JevBench adapter for yev: the harness's own question dict goes to ``system_one`` unchanged.
+
+    from examples.jevbench_adapter import YevAdapter
+    from jevbench.adapters.conformance import check_adapter
+    print(check_adapter(YevAdapter("HallD/Yev-Decision-1")).summary())
+"""
+
+import time
+
+from jevbench.adapters.base import DecisionResult, build_question
+
+from yev import Yev
+
+
+class YevAdapter:
+    name = "yev_local"
+    cost_basis = "local_no_provider_tariff"
+    price_input_per_m = None
+    price_output_per_m = None
+
+    def __init__(self, model="HallD/Yev-Decision-1"):
+        self.model, self.client = model, Yev(model, lazy=False)
+
+    def reserve_estimate(self, task):
+        return None
+
+    def run(self, task):
+        result = DecisionResult(adapter=self.name, ok=False, probs_source="native", model=self.model)
+        question = build_question(task)
+        result.request_body = {"state": task.state, "questions": {"decision": question}}
+        started = time.perf_counter()
+        try:
+            response = self.client.system_one(task.state, {"decision": question}, windows="mean")
+        except Exception as error:   # a failed answer is a failed attempt
+            result.latency_s = time.perf_counter() - started
+            result.error = f"{type(error).__name__}: {str(error)[:300]}"
+            return result
+        result.latency_s = time.perf_counter() - started
+        kind = task.question["type"]
+        if kind == "noul":
+            p = float(response.nouls["decision"].noul)
+            result.probs = {"yes": p, "no": 1.0 - p}
+        else:
+            answer = response.choices["decision"] if kind == "choice" else response.scores["decision"]
+            result.probs = {str(k): float(v) for k, v in answer.probabilities.items()}
+        result.ok = True
+        result.raw = {"runtime": {"probability_origin": "native-softmax", "windows": "mean"}}
+        return result

pyproject.toml

@@ -0,0 +1,35 @@
+[project]
+name = "yev"
+version = "0.2.0"
+description = "A local System One decision model: typed Choice, Score and Noul probabilities, no generated text."
+readme = "README.md"
+requires-python = ">=3.13"
+dependencies = [
+    "huggingface-hub>=0.36.2",
+    "numpy>=1.26",
+    "pydantic>=2.9,<3",
+]
+
+[project.optional-dependencies]
+inference = [
+    "torch>=2.6,<3",
+    "transformers>=5.17,<6",
+    "peft>=0.21,<0.22",
+    "accelerate>=1.12,<2",
+    "pillow>=11",
+]
+quantized = [
+    "yev[inference]",
+    "bitsandbytes>=0.50",
+]
+
+[project.scripts]
+yev = "yev.cli:main"
+
+[build-system]
+requires = ["uv_build>=0.12.17,<0.13.0"]
+build-backend = "uv_build"
+
+[tool.uv.build-backend]
+module-name = "yev"
+module-root = "src"

src/yev/__init__.py

@@ -0,0 +1,98 @@
+"""yev: a local System One decision model with the TypeSafe SDK's call shape.
+
+The package holds two layers:
+
+- The client surface, the TypeSafe SDK's (0.7.5) ported (``YevClient`` and ``AsyncYevClient``, the
+  questions, ``Image`` and the answer types), which answers ``system_one(state, questions)`` from a
+  locally exported model.
+- ``yev.training``: the model, its loading and export, which the client reads an export with.
+
+Importing this module does not load PyTorch; the training stack is imported on first use.
+"""
+
+from . import constants as constants
+from .answers import (
+    Answer,
+    BoxAnswer,
+    BoxCoordinates,
+    ChoiceAnswer,
+    ListModelsResponse,
+    ModelMetadata,
+    NoulAnswer,
+    ScoreAnswer,
+    SystemOneResponse,
+    Usage,
+)
+from .client import AsyncModels, AsyncYevClient, Models, Yev, YevClient
+from .errors import YevError, YevModelNotFoundError, YevResponseValidationError
+from .primitives import (
+    CHOICE_LIMITS,
+    SCORE_LIMITS,
+    Asset,
+    Box,
+    BoxModel,
+    Choice,
+    ChoiceModel,
+    Content,
+    ContentValue,
+    Image,
+    JSONContent,
+    JSONValue,
+    Noul,
+    NoulCriteria,
+    NoulModel,
+    Question,
+    QuestionModel,
+    Questions,
+    Score,
+    ScoreModel,
+    normalize_question,
+    normalize_questions,
+)
+
+__version__ = "0.2.0"
+
+__all__ = [
+    "Answer",
+    "Asset",
+    "AsyncModels",
+    "AsyncYevClient",
+    "Box",
+    "BoxAnswer",
+    "BoxCoordinates",
+    "BoxModel",
+    "CHOICE_LIMITS",
+    "Choice",
+    "ChoiceAnswer",
+    "ChoiceModel",
+    "Content",
+    "ContentValue",
+    "Image",
+    "JSONContent",
+    "JSONValue",
+    "ListModelsResponse",
+    "ModelMetadata",
+    "Models",
+    "Noul",
+    "NoulAnswer",
+    "NoulCriteria",
+    "NoulModel",
+    "Question",
+    "QuestionModel",
+    "Questions",
+    "SCORE_LIMITS",
+    "Score",
+    "ScoreAnswer",
+    "ScoreModel",
+    "SystemOneResponse",
+    "Usage",
+    "Yev",
+    "YevClient",
+    "YevError",
+    "YevModelNotFoundError",
+    "YevResponseValidationError",
+    "__version__",
+    "constants",
+    "normalize_question",
+    "normalize_questions",
+]

src/yev/answers.py

@@ -0,0 +1,137 @@
+"""Typed answers and response objects, ported from the TypeSafe SDK's (0.7.5) response types.
+
+The same classes, fields and accessors (``response.nouls``, ``.choices``, ``.scores``, ``.boxes``),
+immutable and tolerant of unknown fields. yev never returns a ``BoxAnswer`` (it has no box head);
+the type is here for ``response_model``s written against the TypeSafe SDK. yev adds ``SystemOneResponse.elapsed_seconds``, the time the
+model took, and its ``Usage.output_tokens`` is always 0: nothing is generated.
+"""
+
+from functools import cached_property
+from typing import Annotated, Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class _Schema(BaseModel):
+    model_config = ConfigDict(extra="ignore", frozen=True)
+
+
+class NoulAnswer(_Schema):
+    """A yes/no answer."""
+
+    type: Literal["noul"] = "noul"
+    noul: float
+    """Probability of a yes answer or a true statement, from 0 to 1."""
+
+
+class ChoiceAnswer(_Schema):
+    """A selected label and its probabilities."""
+
+    type: Literal["choice"] = "choice"
+    choice: str
+    """The name of the choice with the highest probability among the question's criteria."""
+    confidence: float
+    """Confidence in the selected choice, from 0 to 1."""
+    probabilities: dict[str, float]
+    """Probability of each choice in criteria, keyed by choice name; values sum to about 1."""
+
+
+class ScoreAnswer(_Schema):
+    """An expected score with its rubric and probabilities."""
+
+    type: Literal["score"] = "score"
+    score: float
+    """Expected score: the probability-weighted average of the rubric levels."""
+    confidence: float
+    """Confidence in the score, from 0 to 1."""
+    legend: dict[int, str | dict[str, Any] | list[Any]]
+    """Rubric descriptions keyed by integer score."""
+    probabilities: dict[int, float]
+    """Probabilities keyed by integer score."""
+
+
+class BoxCoordinates(_Schema):
+    """A box's edges as fractions of the image's height (``top``, ``bottom``) and width (``left``, ``right``)."""
+
+    top: float
+    bottom: float
+    left: float
+    right: float
+
+
+class BoxAnswer(_Schema):
+    """A located box with its confidence."""
+
+    type: Literal["box"] = "box"
+    box: BoxCoordinates
+    """The box's edges as fractions of the image's height and width, from 0 to 1."""
+    confidence: float | None
+    """Confidence for the box, from 0 to 1, or ``None`` when the model does not provide it."""
+
+
+Answer = Annotated[NoulAnswer | ChoiceAnswer | ScoreAnswer | BoxAnswer, Field(discriminator="type")]
+"""An answer to a single question, identified by its ``type``."""
+
+
+class Usage(_Schema):
+    """Token counts for a request."""
+
+    input_tokens: int | None = None
+    """Number of input tokens the model read."""
+    output_tokens: int | None = 0
+    """Always 0: yev generates no text."""
+
+
+class SystemOneResponse(_Schema):
+    """Answers grouped by question type with model and usage metadata.
+
+    A subclass that declares fields named after questions gets those answers as attributes when it is
+    passed as ``response_model``, as in the TypeSafe SDK.
+    """
+
+    model: str
+    """The model used to answer the request."""
+    usage: Usage
+    """Token usage for the request."""
+    answers: dict[str, Answer] = Field(default_factory=dict)
+    """All answer objects keyed by question name."""
+    elapsed_seconds: float | None = None
+    """Seconds the model took to answer (for a batch, the whole batch)."""
+
+    @cached_property
+    def nouls(self) -> dict[str, NoulAnswer]:
+        """Yes/no answers keyed by question name."""
+        return {name: answer for name, answer in self.answers.items() if isinstance(answer, NoulAnswer)}
+
+    @cached_property
+    def choices(self) -> dict[str, ChoiceAnswer]:
+        """Choice answers keyed by question name."""
+        return {name: answer for name, answer in self.answers.items() if isinstance(answer, ChoiceAnswer)}
+
+    @cached_property
+    def scores(self) -> dict[str, ScoreAnswer]:
+        """Score answers keyed by question name."""
+        return {name: answer for name, answer in self.answers.items() if isinstance(answer, ScoreAnswer)}
+
+    @cached_property
+    def boxes(self) -> dict[str, BoxAnswer]:
+        """Box answers keyed by question name; always empty from yev."""
+        return {name: answer for name, answer in self.answers.items() if isinstance(answer, BoxAnswer)}
+
+
+class ModelMetadata(_Schema):
+    """Metadata describing a single available model."""
+
+    name: str
+    """Model name, as answers report it."""
+    description: str
+    """What the model reads: its backbone, token budget and whether it takes images."""
+    release_date: str
+    """The date the export was written or downloaded, formatted as YYYY-MM-DD."""
+
+
+class ListModelsResponse(_Schema):
+    """The models a client answers with."""
+
+    models: tuple[ModelMetadata, ...]
+    """The available models."""

src/yev/cli.py

@@ -0,0 +1,176 @@
+"""``yev`` command line: ask an exported model Choice, Score and Noul questions.
+
+    yev --export path/to/export --state "I was charged twice." \\
+        --noul billing "Is this ticket about billing?" \\
+        --choice tone "What is the tone?" "calm,frustrated,angry" \\
+        --score urgency "How urgent?" "can wait|this week|today"
+
+    yev --export path/to/export --request request.json --round
+    yev --export path/to/export --interactive
+"""
+
+import argparse
+import json
+from pathlib import Path
+import sys
+
+from .client import Yev
+from .errors import YevError
+from .primitives import Choice, Image, Noul, Score
+
+
+def parse_choice_criteria(text):
+    """``"a,b,c"`` or ``"a=description,b=description"`` into an ordered criteria map."""
+    criteria = {}
+    for part in text.split(","):
+        part = part.strip()
+        if not part:
+            continue
+        label, _, description = part.partition("=")
+        criteria[label.strip()] = description.strip() or None
+    if len(criteria) < 2:
+        raise YevError("A Choice needs at least two comma-separated options")
+    return criteria
+
+
+def parse_score_levels(text):
+    """Pipe-separated level descriptions, lowest first."""
+    levels = [level.strip() for level in text.split("|") if level.strip()]
+    if not 2 <= len(levels) <= 10:
+        raise YevError("A Score needs 2 to 10 levels separated by '|'")
+    return levels
+
+
+def build_request(arguments):
+    """A ``{"state": ..., "questions": {...}}`` request from a JSON file or the question flags."""
+    if arguments.request is not None:
+        stream = sys.stdin if str(arguments.request) == "-" else open(arguments.request, encoding="utf-8")
+        with stream:
+            request = json.load(stream)
+        if not isinstance(request, dict) or "state" not in request or not isinstance(request.get("questions"), dict):
+            raise YevError("Request JSON needs 'state' and a 'questions' object")
+        return request
+    if arguments.state is None and not arguments.image:
+        raise YevError("--state or --image is required unless --request or --interactive is used")
+    questions = {}
+    for identifier, instructions, options in arguments.choice or []:
+        questions[identifier] = Choice(instructions=instructions, criteria=parse_choice_criteria(options))
+    for identifier, instructions, levels in arguments.score or []:
+        questions[identifier] = Score(instructions=instructions, criteria=parse_score_levels(levels))
+    for identifier, instructions in arguments.noul or []:
+        questions[identifier] = Noul(instructions=instructions)
+    if not questions:
+        raise YevError("Add at least one --choice, --score or --noul question")
+    if arguments.image:
+        try:
+            images = [Image(path) for path in arguments.image]
+        except OSError as error:
+            raise YevError(f"Cannot read image: {error}") from error
+        return {"state": [arguments.state, *images] if arguments.state else images[0] if len(images) == 1 else images,
+                "questions": questions}
+    return {"state": arguments.state, "questions": questions}
+
+
+def render(response, digits=None):
+    """The response as a JSON-ready dict, optionally rounded to the API's hundredth grid."""
+    payload = response.model_dump()
+    if digits is None:
+        return payload
+    for answer in payload["answers"].values():
+        for key in ("noul", "score", "confidence"):
+            if key in answer:
+                answer[key] = round(answer[key], digits)
+        if "probabilities" in answer:
+            answer["probabilities"] = {label: round(value, digits) for label, value in answer["probabilities"].items()}
+    return payload
+
+
+def interactive(client, digits):
+    print("Enter a state, then questions as  kind id instructions | criteria  (blank line answers, Ctrl-D quits).")
+    print("  choice tone What is the tone? | calm,frustrated,angry")
+    print("  score urgency How urgent? | can wait|this week|today")
+    print("  noul billing Is this about billing?")
+    while True:
+        try:
+            state = input("\nstate> ").strip()
+        except EOFError:
+            return
+        if not state:
+            continue
+        questions = {}
+        while True:
+            try:
+                line = input("question> ").strip()
+            except EOFError:
+                return
+            if not line:
+                break
+            try:
+                kind, identifier, rest = line.split(None, 2)
+                instructions, _, criteria = rest.partition("|")
+                if kind == "choice":
+                    questions[identifier] = Choice(instructions=instructions.strip(), criteria=parse_choice_criteria(criteria))
+                elif kind == "score":
+                    questions[identifier] = Score(instructions=instructions.strip(), criteria=parse_score_levels(criteria))
+                elif kind == "noul":
+                    questions[identifier] = Noul(instructions=instructions.strip())
+                else:
+                    raise YevError("kind must be choice, score or noul")
+            except (ValueError, YevError) as error:
+                print("  ignored:", error)
+        if questions:
+            print(json.dumps(render(client.system_one(state, questions), digits), indent=2))
+
+
+def parse_arguments(argv=None):
+    parser = argparse.ArgumentParser(prog="yev", description="Ask an exported yev model typed questions locally.")
+    parser.add_argument("--export", required=True, help="Exported model directory, or a Hugging Face 'owner/name'")
+    parser.add_argument("--revision", help="Hub branch, tag or commit (with an 'owner/name' export)")
+    parser.add_argument("--request", type=Path, help="JSON file with 'state' and 'questions' ('-' for stdin)")
+    parser.add_argument("--state", help="State text for flag-built requests")
+    parser.add_argument("--choice", nargs=3, action="append", metavar=("ID", "INSTRUCTIONS", "OPTIONS"),
+                        help="Options as 'a,b,c' or 'a=description,b=description'")
+    parser.add_argument("--score", nargs=3, action="append", metavar=("ID", "INSTRUCTIONS", "LEVELS"),
+                        help="Levels low to high separated by '|'")
+    parser.add_argument("--noul", nargs=2, action="append", metavar=("ID", "INSTRUCTIONS"))
+    parser.add_argument("--image", type=Path, action="append", default=[],
+                        help="An image file to add to the state, after its text (v2 exports); repeatable")
+    parser.add_argument("--windows", choices=("mean", "max"),
+                        help="Score a long state in overlapping windows: average them, or take the most likely Yes (Noul only)")
+    parser.add_argument("--overlap", type=int, help="Tokens shared by consecutive windows (default: a quarter of a window)")
+    parser.add_argument("--round", action="store_true", help="Round probabilities to hundredths, like the hosted API")
+    parser.add_argument("--micro-batch", type=int, default=8)
+    parser.add_argument("--quantize4bit", action="store_true", help="NF4 4-bit backbone (CUDA and bitsandbytes only)")
+    parser.add_argument("--interactive", action="store_true")
+    return parser.parse_args(argv)
+
+
+def run(argv=None):
+    """Execute one invocation and return the response, or ``None`` in interactive mode."""
+    arguments = parse_arguments(argv)
+    digits = 2 if arguments.round else None
+    request = None if arguments.interactive else build_request(arguments)
+    client = Yev(arguments.export, quantize4bit=arguments.quantize4bit, micro_batch=arguments.micro_batch,
+                 revision=arguments.revision)
+    if arguments.interactive:
+        interactive(client, digits)
+        return None
+    response = client.system_one(request["state"], request["questions"], windows=arguments.windows, overlap=arguments.overlap)
+    print(json.dumps(render(response, digits), indent=2))
+    return response
+
+
+def main(argv=None) -> int:
+    """Console-script entry point: prints the response and returns a process exit code."""
+    try:
+        run(argv)
+    except YevError as error:
+        print(f"yev: {error}", file=sys.stderr)
+        return 2
+    except KeyboardInterrupt:
+        return 130
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

src/yev/client.py

@@ -0,0 +1,564 @@
+"""A local System One client: the TypeSafe SDK's clients and call shape, answered by a local model.
+
+``YevClient`` is ``TypeSafeClient`` and ``AsyncYevClient`` is ``AsyncTypeSafeClient``: the same
+constructor ``model`` argument and ``YEV_DEFAULT_MODEL`` environment default, the same
+``system_one(state, questions, *, model=None, response_model=None)``, the same ``models.list()``,
+context managers and ``close()``. ``model`` names a Hugging Face Hub repository or an export
+directory instead of a hosted model. The HTTP settings (API key, base URL, retries, timeouts,
+headers, extra body) have nothing to act on and are not taken. Every question is one forward
+pass; no text is generated and nothing leaves the machine.
+
+Questions in one request never see each other: each is an independent sequence, so a fact
+placed in one question cannot influence another. Answers do move by a few thousandths with
+the batch shape, because bf16 kernels are not batch-invariant; that drift is below the
+hosted API's hundredth display grid and is independent of what the neighbouring questions
+say. Pass ``micro_batch=1`` for bitwise-stable answers at the cost of throughput.
+
+    from yev import Choice, Noul, Score, YevClient
+
+    with YevClient() as client:                    # YEV_DEFAULT_MODEL, else Yev-Decision-1 from the Hub
+        response = client.system_one(
+            state="I was charged twice. Please fix this ASAP.",
+            questions={
+                "billing": Noul(instructions="Is this ticket about billing?"),
+                "tone": Choice(instructions="What is the customer's tone?",
+                               criteria={"calm": None, "frustrated": None, "angry": None}),
+                "urgency": Score(instructions="How urgent is this ticket?",
+                                 criteria=["can wait", "this week", "today"]),
+            },
+        )
+    response.nouls["billing"].noul
+    response.choices["tone"].choice
+    response.scores["urgency"].score
+
+What yev adds:
+
+- ``system_one_batch(states, questions)`` answers the same questions about many states in one
+  pass over the model, which is where a local model's throughput comes from.
+- A state longer than the export's token budget is cut off after its start by default. Pass
+  ``windows="mean"`` to score overlapping windows that cover all of it and average them, or
+  ``windows="max"`` (Noul only) to answer from the window most likely to be Yes, for questions
+  such as "does any part of this contain an instruction aimed at the agent?". Make ``overlap``
+  longer than the longest span you are looking for, so each one sits whole inside some window.
+- Images, as in the TypeSafe SDK (``Image``, ``assets=``, ``attach_assets``), in the state of a request
+  to an export with a vision backbone; yev does not read images in instructions or criteria.
+- A Choice with one option, or a Score with one level, is answered as certain without the model.
+"""
+
+import asyncio
+from collections import Counter
+from collections.abc import Mapping, Sequence
+from contextlib import contextmanager
+from datetime import date
+from functools import cached_property
+import json
+import os
+from pathlib import Path
+import re
+import threading
+import time
+
+from pydantic import ValidationError
+
+from .answers import ChoiceAnswer, ListModelsResponse, ModelMetadata, NoulAnswer, ScoreAnswer, SystemOneResponse, Usage
+from .constants import DEFAULT_MICRO_BATCH, DEFAULT_MODEL, DEFAULT_MODEL_ENV
+from .errors import YevError, YevModelNotFoundError, YevResponseValidationError
+from .primitives import (
+    Choice,
+    Noul,
+    content_parts,
+    content_text,
+    holds_image,
+    is_image,
+    normalize_questions,
+    question_content,
+    question_labels_and_descriptions,
+)
+
+WINDOW_MODES = (None, "mean", "max")
+WINDOW_MARGIN = 8   # decoding a slice of tokens and encoding it again can merge or split a few
+HUB_REPO_ID = re.compile(r"[A-Za-z0-9][\w.-]*/[A-Za-z0-9][\w.-]*")
+
+
+def resolve_export(export, revision=None):
+    """A local export directory, downloading ``owner/name`` from the Hugging Face Hub if needed.
+
+    A path that exists locally always wins. Otherwise an ``owner/name`` string is fetched with
+    ``huggingface_hub.snapshot_download`` into the Hub cache, so later calls load offline;
+    private repositories use the token saved by ``hf auth login``.
+    """
+    local = Path(export)
+    if local.is_dir():
+        return local
+    if isinstance(export, str) and HUB_REPO_ID.fullmatch(export):
+        try:
+            from huggingface_hub import snapshot_download
+            from huggingface_hub.errors import HfHubHTTPError, RepositoryNotFoundError, RevisionNotFoundError
+        except ImportError as error:
+            raise YevError("Loading from the Hugging Face Hub needs the 'inference' extra") from error
+        try:
+            return Path(snapshot_download(repo_id=export, revision=revision, repo_type="model"))
+        except (RepositoryNotFoundError, RevisionNotFoundError) as error:
+            raise YevModelNotFoundError(f"No Hugging Face model repository {export!r} visible to this login") from error
+        except HfHubHTTPError as error:
+            raise YevError(f"Could not download {export!r} from the Hugging Face Hub: {error}") from error
+    raise YevModelNotFoundError(f"No exported model directory at {local}")
+
+
+def state_windows(tokenizer, text, size, overlap=None):
+    """Overlapping slices of at most ``size`` tokens covering all of ``text``; consecutive slices
+    share ``overlap`` tokens, a quarter of a slice by default. Text that fits is returned whole."""
+    if type(size) is not int or size < 1:
+        raise YevError("The question leaves no room for the state in the token budget")
+    shared = size // 4 if overlap is None else overlap
+    if type(shared) is not int or not 0 <= shared < size:
+        raise YevError(f"overlap must be an integer from 0 to {size - 1} tokens for this question")
+    ids = tokenizer.encode(text, add_special_tokens=False)
+    if len(ids) <= size:
+        return [text]
+    starts = [0]
+    while starts[-1] + size < len(ids):
+        starts.append(starts[-1] + size - shared)
+    return [tokenizer.decode(ids[start:start + size]) for start in starts]
+
+
+def _typed_answer(question, labels, probabilities, descriptions):
+    """Build the typed answer for one question from its normalized distribution."""
+    count = len(probabilities)
+    if isinstance(question, Noul):
+        return NoulAnswer(noul=probabilities[1])
+    mode = max(range(count), key=probabilities.__getitem__)
+    if isinstance(question, Choice):
+        confidence = 1.0 if count == 1 else (probabilities[mode] - 1 / count) / (1 - 1 / count)
+        return ChoiceAnswer(
+            choice=labels[mode],
+            probabilities=dict(zip(labels, probabilities)),
+            confidence=max(0.0, confidence),
+        )
+    distance = sum(value * abs(index - mode) for index, value in enumerate(probabilities))
+    uniform = sum(abs(index - (count - 1) / 2) for index in range(count)) / count
+    return ScoreAnswer(
+        score=sum(index * value for index, value in enumerate(probabilities)),
+        probabilities={index: value for index, value in enumerate(probabilities)},
+        legend=dict(enumerate(question.criteria)),
+        confidence=max(0.0, 1 - distance / uniform) if uniform else 1.0,
+    )
+
+
+def requested_model(model=None):
+    """The model to load: the argument, else ``YEV_DEFAULT_MODEL`` unless it is blank, else ``DEFAULT_MODEL``."""
+    return model if model is not None else os.environ.get(DEFAULT_MODEL_ENV, "").strip() or DEFAULT_MODEL
+
+
+def field_path(error):
+    """The first validation error's location as the TypeSafe SDK's dotted ``field_path``."""
+    location = list(error.errors(include_url=False)[0]["loc"])
+    if location[:1] == ["answers"] and len(location) > 2 and location[2] in ("noul", "choice", "score"):
+        del location[2]   # the union member pydantic tried, not a field
+    path = ""
+    for segment in location:
+        if segment != "[key]":
+            path += f"[{segment}]" if isinstance(segment, int) else f".{segment}" if path else str(segment)
+    return path
+
+
+def typed_response(response_model, model, usage, answers, elapsed):
+    """The response as ``response_model`` reads it: a ``SystemOneResponse`` subclass gets the answers it declares as
+    fields lifted to the top level, and any other model validates the plain response body, as in the TypeSafe SDK."""
+    if response_model is None:
+        return SystemOneResponse(model=model, usage=usage, answers=answers, elapsed_seconds=elapsed)
+    body = {"model": model, "usage": usage.model_dump(), "answers": {name: answer.model_dump() for name, answer in answers.items()}}
+    if isinstance(response_model, type) and issubclass(response_model, SystemOneResponse):
+        body["elapsed_seconds"] = elapsed
+        body |= {name: body["answers"][name] for name in set(response_model.model_fields) - set(SystemOneResponse.model_fields)
+                 if name in answers}
+    try:
+        return response_model.model_validate_json(json.dumps(body))
+    except ValidationError as error:
+        raise YevResponseValidationError(field_path(error)) from error
+
+
+def model_metadata(name, export):
+    """What ``models.list()`` reports for an export, from its ``metadata.json``."""
+    path = export / "metadata.json"
+    metadata = json.loads(path.read_text()) if path.exists() else {}
+    parts = [f"yev decision model on {metadata.get('base_model', 'an unknown backbone')}"]
+    if metadata.get("max_length"):
+        parts.append(f"{metadata['max_length']:,} tokens per question")
+    parts.append("text and images" if metadata.get("images") else "text")
+    written = date.fromtimestamp((path if path.exists() else export).stat().st_mtime)
+    return ModelMetadata(name=name, description=", ".join(parts), release_date=written.isoformat())
+
+
+class Models:
+    """The models a client answers with, reached through ``YevClient.models``: the default model and any a call
+    named with ``model=``."""
+
+    def __init__(self, client):
+        self._client = client
+
+    def list(self) -> ListModelsResponse:
+        """List the models this client has: names, descriptions and export dates. Reads files only."""
+        clients = [self._client, *self._client._others.values()]
+        return ListModelsResponse(models=tuple(model_metadata(client.model, client.export) for client in clients))
+
+
+class AttachedAssets:
+    """A client's attached assets, keyed by reference and shared by its calls from any thread: the TypeSafe SDK's."""
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._assets = {}
+        self._holds = Counter()   # overlapping blocks may attach the same asset; it stays until the last one exits
+
+    @contextmanager
+    def attach(self, assets):
+        attached = [assets] if is_image(assets) else list(assets)
+        for asset in attached:
+            if not is_image(asset):
+                raise YevError(f"Attached value {asset!r} must be an asset.")
+        with self._lock:
+            for asset in attached:
+                self._assets[str(asset)] = asset
+                self._holds[str(asset)] += 1
+        try:
+            yield
+        finally:
+            with self._lock:
+                for asset in attached:
+                    self._holds[str(asset)] -= 1
+                    if not self._holds[str(asset)]:
+                        del self._holds[str(asset)], self._assets[str(asset)]
+
+    def snapshot(self):
+        with self._lock:
+            return dict(self._assets)
+
+
+def named_assets(assets):
+    """The ``assets`` a call names, checked as the TypeSafe SDK checks them."""
+    for name, asset in (assets or {}).items():
+        if not isinstance(name, str) or not name.strip():
+            raise YevError("Asset aliases must be non-empty strings.")
+        if not is_image(asset):
+            raise YevError(f"Asset {name!r} must be an asset.")
+    return dict(assets or {})
+
+
+def checked_request(states, questions, windows, overlap, known=()):
+    """The request's questions as yev questions, once ``states``, the window settings and where images are are
+    checked; ``known`` holds the references of the images the request can use."""
+    if isinstance(states, (str, Mapping)) or not isinstance(states, Sequence):
+        raise YevError("states must be a list of states")
+    if any(state is None for state in states):
+        raise YevError("State is required.")
+    questions = normalize_questions(questions)
+    for name, question in questions.items():
+        if any(holds_image(value, known) for value in question_content(question)):
+            raise YevError(f'Question "{name}" holds an image; yev reads images in the state only. Put the image in the '
+                           "state and refer to it by name, as in Is the dog in `photo` wearing a collar?")
+    if windows not in WINDOW_MODES:
+        raise YevError(f"windows must be one of {WINDOW_MODES}")
+    if windows is None and overlap is not None:
+        raise YevError("overlap applies only when windows is set")
+    if windows == "max" and any(not isinstance(question, Noul) for question in questions.values()):
+        raise YevError("windows='max' answers Noul questions only")
+    return questions
+
+
+def question_examples(name, question, state_text, digests, header_text, tokenizer, max_length, windows, overlap):
+    """The sequences one question about one state is answered from: the state, or each window of it. A question with
+    one option has none: it is certain, and the heads need two."""
+    from .training.data import DecisionExample
+    from .training.model import encode_example
+
+    labels, descriptions = question_labels_and_descriptions(question)
+    if len(labels) == 1:
+        return []
+    instructions = content_text(question.instructions)
+
+    def example(text):
+        return DecisionExample(uid=name, group=name, source="yev", split="test", kind=question.type, state=text,
+                               instructions=instructions, criteria=descriptions, target=[1 / len(descriptions)] * len(descriptions),
+                               images=digests)   # empty whenever windows are used
+
+    if windows is None:
+        return [example(header_text + state_text)]
+    try:
+        overhead = len(encode_example(example(header_text), tokenizer, max_length)["input_ids"])
+    except ValueError as error:
+        raise YevError(f"Question {name!r} does not fit the token budget: {error}") from error
+    return [example(header_text + part) for part in state_windows(tokenizer, state_text, max_length - overhead - WINDOW_MARGIN, overlap)]
+
+
+def combined_answer(question, rows, windows):
+    """The typed answer from a question's sequence distributions: the window most likely Yes for ``windows="max"``,
+    their mean otherwise, and certainty when there are none."""
+    labels = question_labels_and_descriptions(question)[0]
+    rows = [row[:len(labels)].tolist() for row in rows] or [[1.0]]
+    if windows == "max":
+        probabilities = max(rows, key=lambda row: row[1])
+    else:
+        probabilities = [sum(column) / len(rows) for column in zip(*rows)]
+    return _typed_answer(question, labels, probabilities, None)
+
+
+class YevClient:
+    """A System One client backed by a local exported model: the TypeSafe SDK's ``TypeSafeClient``.
+
+    Args:
+        model: The model: a Hugging Face Hub id such as ``"owner/name"``, or an export directory
+            holding ``adapter/``, ``heads.pt``, ``config.json``, ``metadata.json`` and the
+            tokenizer. Defaults to ``YEV_DEFAULT_MODEL``, then ``yev.constants.DEFAULT_MODEL``.
+        revision: Hub branch, tag or commit to download; ignored for a local directory.
+        quantize4bit: Load the backbone's language model in NF4 4-bit (bitsandbytes 0.50 or later, the
+            ``quantized`` extra) on a CUDA GPU, a CPU or an Apple GPU; the vision encoder stays in bf16.
+        micro_batch: Question sequences evaluated per forward pass.
+        lazy: Defer loading weights until the first request.
+
+    Raises:
+        YevModelNotFoundError: The directory does not exist, or the Hub has no such repository.
+        YevError: ``micro_batch`` is not a positive integer, or the download failed.
+    """
+
+    def __init__(self, model=None, *, revision=None, quantize4bit=False, micro_batch=DEFAULT_MICRO_BATCH, lazy=True):
+        if type(micro_batch) is not int or micro_batch < 1:
+            raise YevError("micro_batch must be a positive integer")
+        requested = requested_model(model)
+        hub = not Path(requested).is_dir() and isinstance(requested, str) and HUB_REPO_ID.fullmatch(requested)
+        self.export = resolve_export(requested, revision)
+        self.quantize4bit = quantize4bit
+        self.micro_batch = micro_batch
+        resolved = self.export.resolve()
+        default = requested.split("/")[1] if hub else (
+            resolved.parent.name if resolved.name in ("export", "model") else resolved.name)
+        self.model = default or "yev"
+        """The name answers report."""
+        self._requested = str(requested)
+        self._model = None
+        self._tokenizer = None
+        self._others = {}
+        self._attached = AttachedAssets()
+        if not lazy:
+            self._load()
+
+    @cached_property
+    def models(self) -> Models:
+        """An accessor for the models this client answers with: ``client.models.list()``."""
+        return Models(self)
+
+    def _load(self):
+        if self._model is None:
+            from .training.run import reload_export
+
+            self._model, self._tokenizer = reload_export(self.export, quantize4bit=self.quantize4bit)
+        return self._model, self._tokenizer
+
+    def _other(self, model):
+        """The client for a model a call names that is not this one's, loaded once and kept until ``close``."""
+        if model not in self._others:
+            self._others[model] = YevClient(model, quantize4bit=self.quantize4bit, micro_batch=self.micro_batch)
+            self._others[model]._attached = self._attached   # attached assets hold for every model the client calls
+        return self._others[model]
+
+    @property
+    def max_length(self) -> int:
+        """Token budget per question, from the export's training configuration."""
+        return self._load()[0].training_config.max_length
+
+    @property
+    def base_model(self) -> str:
+        """Hugging Face id of the frozen backbone this export adapts."""
+        return self._load()[0].training_config.model_id
+
+    def attach_assets(self, assets):
+        """Attach assets to calls made within a ``with`` block, for referencing them inside strings.
+
+        Until the block exits, a state holding an attached asset's reference, ``str(asset)``, reads the
+        asset in its place. Only references to attached assets, assets passed in ``assets`` and assets
+        placed in the content are recognised; any other is plain text.
+
+        Raises:
+            YevError: An attached value is not an asset.
+        """
+        return self._attached.attach(assets)
+
+    def system_one(self, state, questions, *, model=None, assets=None, response_model=None, windows=None, overlap=None,
+                   header=""):
+        """Answer named questions about a state.
+
+        Args:
+            state: Text, a JSON object or an array to evaluate, which may hold ``Image``s.
+            questions: Nonempty mapping of names to ``Noul``, ``Choice`` or ``Score`` questions,
+                question dictionaries, or TypeSafe SDK question objects.
+            model: Model override, a Hub id or export directory; ``None`` uses this client's.
+                Another model is loaded on first use and kept until ``close``.
+            assets: Images keyed by the reference used for them in the state.
+            response_model: Optional Pydantic model for the response body; a ``SystemOneResponse``
+                subclass gets the answers it declares as fields, by question name.
+            windows: ``None`` answers from the start of a state too long for the token
+                budget; ``"mean"`` averages overlapping windows covering all of it; ``"max"``
+                takes the window most likely to be Yes (Noul questions only).
+            overlap: Tokens shared by consecutive windows; a quarter of a window by default.
+            header: Text placed before the state, and before every window of it, whole.
+
+        Returns:
+            A ``SystemOneResponse`` with answers keyed by question name, or a ``response_model``.
+
+        Raises:
+            YevError: The request is invalid before the model reads it: no state or questions, a
+                question that is malformed, holds an image or has more options than yev answers, an
+                asset that is not an image or not referenced, or window settings that do not fit.
+            YevResponseValidationError: The answers do not fit ``response_model``.
+        """
+        return self.system_one_batch([state], questions, model=model, assets=assets, response_model=response_model,
+                                     windows=windows, overlap=overlap, header=header)[0]
+
+    def system_one_batch(self, states, questions, *, model=None, assets=None, response_model=None, windows=None, overlap=None,
+                         header=""):
+        """Answer the same named questions about each of ``states`` in one pass over the model, a response per state in
+        order; the arguments are ``system_one``'s."""
+        if model not in (None, self.model, self._requested):
+            return self._other(model).system_one_batch(states, questions, assets=assets, response_model=response_model,
+                                                       windows=windows, overlap=overlap, header=header)
+        named = named_assets(assets)
+        known = {**self._attached.snapshot(), **{str(asset): asset for asset in named.values()}, **named}
+        questions = checked_request(states, questions, windows, overlap, known)
+        from .training.model import collator_for
+
+        loaded, tokenizer = self._load()
+        max_length = loaded.training_config.max_length
+        used = set()
+        parsed = [content_parts(state, known, used) for state in states]
+        unreferenced = sorted(named.keys() - used)
+        if unreferenced:
+            raise YevError(f"Asset {unreferenced[0]!r} is not referenced in the state.")
+        vision = self._vision(loaded, parsed, windows)
+        header_text = content_text(header) if header else ""
+        prepared, examples, owners = [], [], []
+        for index, (state_text, images) in enumerate(parsed):
+            digests = [vision.store.put(data) for data in images]
+            for name, question in questions.items():
+                parts = question_examples(name, question, state_text, digests, header_text, tokenizer, max_length, windows, overlap)
+                prepared.append((index, name, question, len(examples), len(parts)))
+                examples.extend(parts)
+                owners.extend([index] * len(parts))
+        started = time.perf_counter()
+        distributions, tokens = self._forward(loaded, collator_for(loaded, tokenizer, max_length, vision), examples, owners,
+                                              len(states))
+        elapsed = round(time.perf_counter() - started, 4)
+        answers = [{} for _ in states]
+        for index, name, question, first, count in prepared:
+            answers[index][name] = combined_answer(question, distributions[first:first + count], windows)
+        return [typed_response(response_model, self.model, Usage(input_tokens=tokens[index], output_tokens=0), answers[index],
+                               elapsed) for index in range(len(states))]
+
+    def _vision(self, loaded, parsed, windows):
+        """The export's vision settings, with an image store of this request's own when a state holds images."""
+        vision = getattr(loaded, "vision", None)
+        if not any(images for _, images in parsed):
+            return vision
+        if vision is None:
+            raise YevError(f"{self.model} reads text only; images need a v2 export (and Pillow installed)")
+        if windows is not None:
+            raise YevError("windows cannot split a state that holds images")
+        from dataclasses import replace
+
+        from .training.images import ImageStore
+
+        return replace(vision, store=ImageStore())
+
+    def _forward(self, loaded, collator, examples, owners, count):
+        """Each example's distribution, in order, and the tokens read for each of ``count`` states."""
+        from .training.run import device_batch, evaluation_mode
+
+        distributions, tokens = [], [0] * count
+        if not examples:
+            return distributions, tokens
+        with evaluation_mode(loaded):
+            for start in range(0, len(examples), self.micro_batch):
+                batch = device_batch(loaded, collator, examples[start:start + self.micro_batch])
+                for owner, read in zip(owners[start:start + self.micro_batch], batch["attention_mask"].sum(-1).tolist()):
+                    tokens[owner] += int(read)
+                distributions.extend(loaded(**batch)["logits"].float().softmax(-1).cpu())
+        return distributions, tokens
+
+    def close(self) -> None:
+        """Release the loaded weights, of every model this client loaded."""
+        for other in self._others.values():
+            other.close()
+        self._others = {}
+        self._model = None
+        self._tokenizer = None
+
+    def __enter__(self):
+        self._load()
+        return self
+
+    def __exit__(self, *exception):
+        self.close()
+        return False
+
+
+Yev = YevClient
+"""The name yev's client had first; the same class."""
+
+
+class AsyncModels:
+    """The models an ``AsyncYevClient`` answers with: ``await client.models.list()``."""
+
+    def __init__(self, client):
+        self._client = client
+
+    async def list(self) -> ListModelsResponse:
+        """List the models this client has: names, descriptions and export dates."""
+        return self._client.client.models.list()
+
+
+class AsyncYevClient:
+    """The TypeSafe SDK's ``AsyncTypeSafeClient`` over a ``YevClient``, with the same arguments. Each call runs in a
+    worker thread, one at a time (a model answers one batch at a time), so the event loop stays free meanwhile."""
+
+    def __init__(self, model=None, *, revision=None, quantize4bit=False, micro_batch=DEFAULT_MICRO_BATCH, lazy=True):
+        self.client = YevClient(model, revision=revision, quantize4bit=quantize4bit, micro_batch=micro_batch, lazy=lazy)
+        """The synchronous client answering for this one."""
+        self._lock = asyncio.Lock()
+
+    @property
+    def model(self) -> str:
+        """The name answers report."""
+        return self.client.model
+
+    @cached_property
+    def models(self) -> AsyncModels:
+        """An accessor for the models this client answers with: ``await client.models.list()``."""
+        return AsyncModels(self)
+
+    def attach_assets(self, assets):
+        """``YevClient.attach_assets``: assets referenced inside strings by calls made within a ``with`` block."""
+        return self.client.attach_assets(assets)
+
+    async def system_one(self, state, questions, *, model=None, assets=None, response_model=None, windows=None, overlap=None,
+                         header=""):
+        """``YevClient.system_one``, awaited."""
+        async with self._lock:
+            return await asyncio.to_thread(self.client.system_one, state, questions, model=model, assets=assets,
+                                           response_model=response_model, windows=windows, overlap=overlap, header=header)
+
+    async def system_one_batch(self, states, questions, *, model=None, assets=None, response_model=None, windows=None,
+                               overlap=None, header=""):
+        """``YevClient.system_one_batch``, awaited."""
+        async with self._lock:
+            return await asyncio.to_thread(self.client.system_one_batch, states, questions, model=model, assets=assets,
+                                           response_model=response_model, windows=windows, overlap=overlap, header=header)
+
+    async def aclose(self) -> None:
+        """Release the loaded weights."""
+        self.client.close()
+
+    async def __aenter__(self):
+        await asyncio.to_thread(self.client._load)
+        return self
+
+    async def __aexit__(self, *exception):
+        await self.aclose()

src/yev/constants.py

@@ -0,0 +1,10 @@
+"""Public environment-variable names and client defaults, as ``typesafe_sdk.constants`` holds them."""
+
+DEFAULT_MODEL_ENV = "YEV_DEFAULT_MODEL"
+"""Environment variable for the default model: a Hugging Face Hub id or an export directory."""
+
+DEFAULT_MODEL = "HallD/Yev-Decision-1"
+"""Default model: Yev-Decision-1, a 4B trained without Jev."""
+
+DEFAULT_MICRO_BATCH = 8
+"""Default number of question sequences scored per forward pass."""

src/yev/errors.py

@@ -0,0 +1,18 @@
+"""Exceptions raised by the yev client: ``YevError`` is the TypeSafe SDK's ``TypeSafeError``."""
+
+
+class YevError(Exception):
+    """Base exception for every error raised by this package."""
+
+
+class YevModelNotFoundError(YevError, FileNotFoundError):
+    """The requested export directory does not exist or is not a model export."""
+
+
+class YevResponseValidationError(YevError):
+    """The answers do not fit the ``response_model`` the call asked for."""
+
+    def __init__(self, field_path: str) -> None:
+        super().__init__(f"Invalid response data at {field_path!r}.")
+        self.field_path = field_path
+        """Dotted path to the offending field, such as ``answers.tone.confidence``."""

src/yev/primitives.py

@@ -0,0 +1,417 @@
+"""Question primitives and images, ported from the TypeSafe SDK (0.7.5): ``Noul``, ``Choice``, ``Score``,
+``Box`` and ``Image``.
+
+The same fields, types, defaults and wire form (an unset optional field is left out of
+``model_dump``), the same question dictionaries (``NoulModel``, ``ChoiceModel``, ``ScoreModel``,
+``BoxModel``) and the same checks, with the same messages. Instructions and criteria accept text,
+JSON objects or arrays; non-string content is rendered as compact JSON before it reaches the
+model. ``normalize_questions`` also takes TypeSafe SDK question objects, and states take TypeSafe
+SDK images, so code written against ``typesafe_sdk`` runs on ``yev``. The differences: yev
+questions are frozen; the number of options is also checked against what yev's heads answer; yev
+reads images in the state only, not in instructions or criteria; and it has no head for ``Box``
+questions, so asking one is an error.
+"""
+
+from base64 import b64decode
+from collections.abc import Mapping, Sequence
+import io
+import json
+import os
+from pathlib import Path
+import re
+from typing import Any, Literal, NotRequired
+from uuid import uuid4
+
+from pydantic import BaseModel, ConfigDict, Field, model_serializer
+from pydantic.functional_serializers import SerializerFunctionWrapHandler
+from pydantic_core import core_schema
+from typing_extensions import TypedDict   # closed=True, which typing's TypedDict takes from Python 3.15
+
+from .errors import YevError
+
+type JSONValue = str | int | float | bool | Sequence[JSONValue | None] | Mapping[str, JSONValue | None]
+"""A JSON-like value. May be nested and contain ``None``."""
+
+type JSONContent = str | Mapping[str, JSONValue | None] | Sequence[JSONValue | None]
+"""Either a plain string or a mapping/sequence of ``JSONValue`` entries."""
+
+CHOICE_LIMITS = (1, 255)   # options yev answers; a single option is certain and needs no model
+SCORE_LIMITS = (1, 10)     # levels yev's heads were trained on, from zero
+REFERENCE = re.compile(r"<image:[0-9a-f]{32}>")   # str(image), in yev and in the TypeSafe SDK alike
+
+
+class Image:
+    """An image asset: the TypeSafe SDK's ``Image``.
+
+    Made from image file bytes or any other buffer (``bytearray``, ``memoryview``, ``mmap``), a path,
+    base64 text (``Image.from_base64``) or a Pillow image (``Image.from_pil``); yev also takes a path
+    given as a string, a Pillow image directly and a TypeSafe SDK ``Image``. Put it in a state on its
+    own, in a list with text (``["Is this paid?", Image(Path("receipt.png"))]``) or as a value in a
+    JSON object; ``str(image)`` is its reference, for placing it inside text (see
+    ``YevClient.attach_assets``). Only an export with a vision backbone reads it.
+    """
+
+    __slots__ = ("_reference", "data")
+    type = "image"
+    """Asset type discriminator, always ``"image"``."""
+
+    def __init__(self, data):
+        if isinstance(data, Image):
+            data = data.data
+        elif isinstance(data, (str, os.PathLike)):
+            data = Path(data).read_bytes()
+        elif callable(getattr(data, "save", None)):   # a Pillow image
+            buffer = io.BytesIO()
+            data.save(buffer, format="PNG")
+            data = buffer.getvalue()
+        elif getattr(data, "type", None) == "image" and hasattr(data, "data"):   # a TypeSafe SDK Image
+            data = Image(data.data).data
+        else:
+            try:
+                data = memoryview(data).tobytes()
+            except TypeError as error:
+                raise TypeError(f"An Image is made from a path, bytes or a Pillow image, not {type(data).__name__}") from error
+        if not data:
+            raise ValueError("An Image needs non-empty data")
+        self.data = data
+        """The image file bytes."""
+        self._reference = f"<image:{uuid4().hex}>"
+
+    @classmethod
+    def from_base64(cls, data):
+        """An ``Image`` from base64 text or an ASCII buffer; whitespace, such as wrapped lines, is ignored.
+
+        Raises:
+            YevError: The data isn't valid base64 text.
+        """
+        text = data if isinstance(data, str) else memoryview(data).tobytes().decode("latin-1")
+        text = "".join(text.split()) if text.isascii() else ""
+        if not text:
+            raise YevError("Image data must contain non-empty ASCII base64 text.")
+        try:
+            return cls(b64decode(text, validate=True))
+        except ValueError as error:
+            raise YevError("Image data must contain valid base64 text.") from error
+
+    @classmethod
+    def from_pil(cls, image, format="jpeg"):   # Pillow's own name for it
+        """An ``Image`` from a Pillow image, saved with Pillow's defaults except JPEG quality 95.
+
+        Raises:
+            YevError: Pillow can't encode the image in this format.
+        """
+        buffer = io.BytesIO()
+        try:
+            image.save(buffer, format=format, quality=95)
+        except (KeyError, OSError, ValueError) as error:   # Pillow raises KeyError for a format it doesn't know
+            raise YevError(f"Pillow image could not be encoded as {format}: {error}") from error
+        return cls(buffer.getvalue())
+
+    def __str__(self):
+        """The unique reference to this image, for using it inside text, such as an f-string."""
+        return self._reference
+
+    def __repr__(self):
+        return f"Image({len(self.data)} bytes)"
+
+    @classmethod
+    def __get_pydantic_core_schema__(cls, source, handler):
+        return core_schema.is_instance_schema(cls)   # question objects hold an Image as it is
+
+
+type Asset = Image
+"""An asset referenced from request content; only images are supported."""
+
+type ContentValue = str | int | float | bool | Sequence[ContentValue | None] | Mapping[str, ContentValue | None] | Image
+"""A JSON-like value that may also hold images. May be nested and contain ``None``."""
+
+type Content = str | Mapping[str, ContentValue | None] | Sequence[ContentValue | None] | Image
+"""Text, an image, or a mapping/sequence of ``ContentValue`` entries."""
+
+
+def is_image(value) -> bool:
+    """A yev ``Image`` or a TypeSafe SDK one."""
+    return isinstance(value, Image) or (not isinstance(value, Mapping) and getattr(value, "type", None) == "image"
+                                        and hasattr(value, "data"))
+
+
+class NoulCriteria(TypedDict, total=False, closed=True):
+    """Optional descriptions of the yes and no outcomes."""
+
+    true: Content | None
+    """Description of the yes outcome as text, an asset, a JSON object, or an array; ``None`` leaves it undescribed."""
+    false: Content | None
+    """Description of the no outcome as text, an asset, a JSON object, or an array; ``None`` leaves it undescribed."""
+
+
+class NoulModel(TypedDict, closed=True):
+    """A yes/no question dictionary with ``type="noul"``."""
+
+    type: Literal["noul"]
+    instructions: NotRequired[Content | None]
+    """The question to ask, as text, an asset, a JSON object, or an array; optional when ``criteria`` describes an outcome."""
+    criteria: NotRequired[NoulCriteria | None]
+    """Optional descriptions of the yes and no outcomes; required when ``instructions`` is ``None``."""
+
+
+class ChoiceModel(TypedDict, closed=True):
+    """A choice question dictionary with ``type="choice"``."""
+
+    type: Literal["choice"]
+    instructions: NotRequired[Content | None]
+    """The question to ask, expressed as text, an asset, a JSON object, or an array; optional."""
+    criteria: Mapping[str, Content | None]
+    """A nonempty mapping of labels to text, asset, object, or array descriptions, or ``None`` for undescribed labels."""
+
+
+class ScoreModel(TypedDict, closed=True):
+    """A score question dictionary with ``type="score"``."""
+
+    type: Literal["score"]
+    instructions: NotRequired[Content | None]
+    """The question to ask, expressed as text, an asset, a JSON object, or an array; optional."""
+    criteria: Sequence[Content]
+    """A nonempty, ordered list of text, asset, object, or array descriptions, one per score from zero."""
+
+
+class BoxModel(TypedDict, closed=True):
+    """A box question dictionary with ``type="box"``."""
+
+    type: Literal["box"]
+    instructions: NotRequired[Content | None]
+    """What to locate in the request's image, expressed as text, an asset, a JSON object, or an array; optional."""
+
+
+class _Question(BaseModel):
+    """Reject unknown fields, freeze, and omit optional fields left at their default from the wire form."""
+
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    @model_serializer(mode="wrap")
+    def _omit_none(self, handler: SerializerFunctionWrapHandler) -> dict[str, Any]:
+        # An unset optional field is left off, while ``None`` values nested inside criteria or instructions are kept.
+        return {key: value for key, value in handler(self).items() if value is not None}
+
+
+class Noul(_Question):
+    """A yes/no question with optional descriptions for either outcome."""
+
+    type: Literal["noul"] = "noul"
+    instructions: Content | None = None
+    """The question to ask, as text, an asset, a JSON object, or an array; optional when ``criteria`` describes an outcome."""
+    criteria: NoulCriteria | None = None
+    """Optional descriptions of the yes and no outcomes; required when ``instructions`` is ``None``."""
+
+
+class Choice(_Question):
+    """A question that selects between named alternatives."""
+
+    type: Literal["choice"] = "choice"
+    criteria: Mapping[str, Content | None]
+    """A nonempty mapping of labels to text, asset, object, or array descriptions, or ``None`` for undescribed labels."""
+    instructions: Content | None = None
+    """The question to ask, expressed as text, an asset, a JSON object, or an array; optional."""
+
+
+class Score(_Question):
+    """A question that assigns a score using an ordered rubric."""
+
+    type: Literal["score"] = "score"
+    criteria: Sequence[Content] = Field(min_length=1)
+    """A nonempty, ordered list of text, asset, object, or array descriptions, one per score from zero."""
+    instructions: Content | None = None
+    """The question to ask, expressed as text, an asset, a JSON object, or an array; optional."""
+
+
+class Box(_Question):
+    """A question that locates what its instructions describe in the request's image. yev has no head for it: asking
+    one raises ``YevError``. It is here so TypeSafe SDK code that builds one runs until it asks."""
+
+    type: Literal["box"] = "box"
+    instructions: Content | None = None
+    """What to locate in the request's image, expressed as text, an asset, a JSON object, or an array; optional."""
+
+
+type QuestionModel = NoulModel | ChoiceModel | ScoreModel | BoxModel
+"""A question dictionary identified by its ``type`` key."""
+type Question = Noul | Choice | Score | Box | QuestionModel
+"""A question object or question dictionary."""
+type Questions = Mapping[str, Question]
+"""Question inputs keyed by the names used to identify their answers."""
+QUESTION_TYPES = {"noul": Noul, "choice": Choice, "score": Score, "box": Box}
+
+
+class ContentReader:
+    """Reads content as the model's text, replacing each image with the image placeholder; see ``content_parts``."""
+
+    def __init__(self, known, used):
+        self.known, self.used, self.images = known, used, []
+        ordered = sorted(known, key=len, reverse=True)   # at one position, the longest reference wins
+        self.pattern = re.compile("|".join(map(re.escape, ordered))) if known else None
+
+    def image(self, image, references=()):
+        if self.used is not None:
+            self.used.update(references)
+        self.images.append(Image(image).data)
+        from .training.images import IMAGE_PLACEHOLDER
+
+        return IMAGE_PLACEHOLDER
+
+    def reference(self, match):
+        return self.image(self.known[match[0]], [match[0]])
+
+    def text(self, value):
+        if self.pattern is None:
+            return value
+        return self.image(self.known[value], [value]) if value in self.known else self.pattern.sub(self.reference, value)
+
+    def swap(self, value):
+        if is_image(value):
+            return self.image(value, [reference for reference, image in self.known.items() if image is value])
+        if isinstance(value, str):
+            return self.text(value)
+        if isinstance(value, Mapping):
+            return {key: self.swap(item) for key, item in value.items()}
+        if isinstance(value, (list, tuple)):
+            return [self.swap(item) for item in value]
+        return value
+
+
+def content_parts(content, assets=None, used=None) -> tuple[str, list[bytes]]:
+    """The model's text for a state, with each image replaced by the image placeholder, and the images' bytes in the
+    order their placeholders appear. Text is kept as is, an image alone is just its placeholder, a list of text and
+    images becomes lines, and any other object or array becomes compact JSON with the placeholders as string values.
+
+    ``assets`` maps references (``str(image)``, or the aliases a call names) to images. As the TypeSafe API reads them,
+    a string that is a reference becomes that image, longer strings have references replaced left to right, and
+    object keys stay text. Images placed in the content are known by their own references too. The references that
+    matched are added to ``used``, and so are the names of an image placed as itself.
+    """
+    reader = ContentReader({**dict(assets or {}), **{str(image): image for image in placed(content)}}, used)
+    if is_image(content):
+        return reader.swap(content), reader.images
+    if (isinstance(content, (list, tuple)) and any(is_image(item) for item in content)
+            and all(isinstance(item, str) or is_image(item) for item in content)):
+        return "\n".join(reader.swap(item) for item in content), reader.images
+    return content_text(reader.swap(content)), reader.images
+
+
+def placed(content):
+    """The images placed in content as image objects, in order."""
+    if is_image(content):
+        yield content
+    elif isinstance(content, Mapping):
+        for value in content.values():
+            yield from placed(value)
+    elif isinstance(content, (list, tuple)):
+        for value in content:
+            yield from placed(value)
+
+
+def holds_image(content, references=()) -> bool:
+    """Whether content holds an image: an image object, a TypeSafe-style reference, or one of ``references``."""
+    if is_image(content):
+        return True
+    if isinstance(content, str):
+        return bool(REFERENCE.search(content)) or any(reference in content for reference in references)
+    if isinstance(content, Mapping):
+        return any(holds_image(value, references) for value in content.values())
+    if isinstance(content, (list, tuple)):
+        return any(holds_image(value, references) for value in content)
+    return False
+
+
+def content_text(content) -> str:
+    """Render text as is and JSON objects or arrays as compact JSON."""
+    if content is None:
+        return ""
+    if isinstance(content, str):
+        return content
+    return json.dumps(content, ensure_ascii=False, separators=(", ", ": "), sort_keys=False)
+
+
+def normalize_question(question) -> Noul | Choice | Score | Box:
+    """Coerce a yev question, a TypeSafe SDK question or a ``{"type": ...}`` dict into a yev question."""
+    if isinstance(question, (Noul, Choice, Score, Box)):
+        return question
+    if isinstance(question, Mapping):
+        payload = dict(question)
+    else:
+        dump = getattr(question, "model_dump", None)
+        if dump is None:
+            raise TypeError(f"Unsupported question object: {type(question).__name__}")
+        payload = dump()
+    kind = payload.get("type")
+    if kind not in QUESTION_TYPES:
+        raise ValueError(f"Unknown question type: {kind!r}")
+    return QUESTION_TYPES[kind].model_validate(payload)
+
+
+def checked_dictionary(name, question):
+    """The TypeSafe SDK's checks on a question dictionary, before it is read."""
+    kind = question.get("type")
+    if not isinstance(kind, str) or not kind:
+        raise YevError(f'Question "{name}" must be a question object or a dictionary with a nonempty string "type".')
+    if kind in ("choice", "score") and "criteria" not in question:
+        raise YevError(f'Question "{name}" requires "criteria".')
+    if kind in ("choice", "score") and not question["criteria"]:
+        raise YevError(f'{kind.title()} question "{name}" has no criteria; at least one {kind} is required.')
+    if kind == "noul" and question.get("criteria") is not None and not isinstance(question["criteria"], Mapping):
+        raise YevError(f'Noul question "{name}" criteria must be an object.')
+
+
+def checked_question(name, question) -> Noul | Choice | Score:
+    """One question as a yev question, checked as the TypeSafe SDK checks it, then against what yev answers;
+    ``YevError`` names it."""
+    if isinstance(question, Mapping):
+        checked_dictionary(name, question)
+    try:
+        question = normalize_question(question)
+    except (TypeError, ValueError) as error:
+        raise YevError(f'Question "{name}" is not valid: {error}') from error
+    if isinstance(question, Box):
+        raise YevError(f'Box question "{name}" cannot be answered: yev has Noul, Choice and Score heads, and no box head.')
+    if isinstance(question, Noul):
+        if question.instructions is None and not any(value is not None for value in (question.criteria or {}).values()):
+            raise YevError(f'Noul question "{name}" has neither instructions nor criteria; at least one is required.')
+        return question
+    count = len(question.criteria)
+    low, high = CHOICE_LIMITS if question.type == "choice" else SCORE_LIMITS
+    if count < low:
+        raise YevError(f'{question.type.title()} question "{name}" has no criteria; at least one {question.type} is required.')
+    if count > high:
+        raise YevError(f'{question.type.title()} question "{name}" has {count} criteria; yev answers at most {high}.')
+    return question
+
+
+def normalize_questions(questions: Questions) -> dict[str, Noul | Choice | Score]:
+    """Every question as a yev question, checked by ``checked_question``."""
+    if not questions:
+        raise YevError("At least one question is required.")
+    return {name: checked_question(name, question) for name, question in questions.items()}
+
+
+def question_content(question):
+    """A question's instructions and criteria values: everything an image could hide in."""
+    criteria = getattr(question, "criteria", None) or {}
+    return [question.instructions, *(criteria.values() if isinstance(criteria, Mapping) else criteria)]
+
+
+def question_labels_and_descriptions(question: Noul | Choice | Score) -> tuple[list[str], list[str]]:
+    """Labels the answer will be keyed by, and the criterion text shown to the model, in order."""
+    if isinstance(question, Noul):
+        criteria = question.criteria or {}
+        descriptions = [
+            content_text(criteria.get("false")) or "false",
+            content_text(criteria.get("true")) or "true",
+        ]
+        return ["false", "true"], descriptions
+    if isinstance(question, Choice):
+        labels = list(question.criteria)
+        descriptions = [
+            f"{label}: {content_text(description)}" if description not in (None, "") else label
+            for label, description in question.criteria.items()
+        ]
+        return labels, descriptions
+    descriptions = [content_text(level) for level in question.criteria]
+    return [str(index) for index in range(len(descriptions))], descriptions

src/yev/training/comparison.py

@@ -0,0 +1,111 @@
+"""Student-versus-teacher distribution comparison on held-out Jev-labelled questions.
+
+Pure metrics plus a batched scorer. Nothing here calls the Jev API or downloads data.
+"""
+
+from collections import defaultdict
+import math
+import random
+
+from .model import PRIMITIVES, collator_for
+from .run import device_batch, evaluation_mode
+
+
+def stratified_sample(examples, limit, seed=42):
+    """Deterministically pick up to ``limit`` examples with primitives balanced."""
+    if type(limit) is not int or limit < 1:
+        raise ValueError("limit must be a positive integer")
+    pools = {kind: [example for example in examples if example.kind == kind] for kind in PRIMITIVES}
+    for kind in PRIMITIVES:
+        random.Random(f"{seed}:compare:{kind}").shuffle(pools[kind])
+    selected = []
+    cursor = 0
+    while len(selected) < limit and any(cursor < len(pool) for pool in pools.values()):
+        for kind in PRIMITIVES:
+            if cursor < len(pools[kind]) and len(selected) < limit:
+                selected.append(pools[kind][cursor])
+        cursor += 1
+    return selected
+
+
+def comparison_metrics(kind, student, teacher):
+    """Per-question distances between two normalized distributions."""
+    if len(student) != len(teacher) or not student:
+        raise ValueError("Distributions must have the same nonzero length")
+    if not all(math.isfinite(value) and value >= 0 for value in student + teacher):
+        raise ValueError("Distributions must be finite and nonnegative")
+    student_mode = max(range(len(student)), key=student.__getitem__)
+    teacher_mode = max(range(len(teacher)), key=teacher.__getitem__)
+    metrics = {
+        "tvd": 0.5 * sum(abs(p - q) for p, q in zip(student, teacher)),
+        "agree": int(student_mode == teacher_mode),
+        "log_loss": -sum(q * math.log(max(p, 1e-12)) for p, q in zip(student, teacher)),
+        "brier": sum((p - q) ** 2 for p, q in zip(student, teacher)),
+    }
+    if kind == "score":
+        metrics["mean_abs_error"] = abs(sum(index * p for index, p in enumerate(student)) - sum(index * q for index, q in enumerate(teacher)))
+    if kind == "noul":
+        metrics["probability_abs_error"] = abs(student[1] - teacher[1])
+    return metrics
+
+
+def aggregate(rows):
+    """Mean of every metric over rows of ``{"metrics": {...}}``."""
+    if not rows:
+        return {"count": 0}
+    totals = defaultdict(float)
+    counts = defaultdict(int)
+    for row in rows:
+        for name, value in row["metrics"].items():
+            totals[name] += value
+            counts[name] += 1
+    return {"count": len(rows), **{name: totals[name] / counts[name] for name in sorted(totals)}}
+
+
+def domain_of(example):
+    return example.source.split(":", 1)[1] if ":" in example.source else example.source
+
+
+def student_distributions(model, tokenizer, examples, max_length, micro_batch=8):
+    collator = collator_for(model, tokenizer, max_length)
+    results = []
+    with evaluation_mode(model):
+        for start in range(0, len(examples), micro_batch):
+            items = examples[start:start + micro_batch]
+            batch = device_batch(model, collator, items)
+            probabilities = model(**batch)["logits"].float().softmax(-1).cpu()
+            for index, item in enumerate(items):
+                results.append(probabilities[index, :len(item.target)].tolist())
+    return results
+
+
+def offline_report(model, tokenizer, examples, max_length, micro_batch=8):
+    predictions = student_distributions(model, tokenizer, examples, max_length, micro_batch)
+    rows, baseline_rows = [], []
+    for example, student in zip(examples, predictions):
+        uniform = [1 / len(example.target)] * len(example.target)
+        rows.append({"uid": example.uid, "kind": example.kind, "domain": domain_of(example), "student": student,
+                     "teacher": example.target, "metrics": comparison_metrics(example.kind, student, example.target)})
+        baseline_rows.append({"kind": example.kind, "domain": domain_of(example), "metrics": comparison_metrics(example.kind, uniform, example.target)})
+    by = lambda rows, key: {value: aggregate([row for row in rows if row[key] == value]) for value in sorted({row[key] for row in rows})}  # noqa: E731
+    return {
+        "student": {"overall": aggregate(rows), "by_primitive": by(rows, "kind"), "by_domain": by(rows, "domain")},
+        "uniform_baseline": {"overall": aggregate(baseline_rows), "by_primitive": by(baseline_rows, "kind")},
+        "rows": rows,
+    }
+
+
+def format_table(report):
+    lines = [f"{'slice':28} {'n':>6} {'TVD':>7} {'agree':>7} {'logloss':>8} {'brier':>7}"]
+    student, baseline = report["student"], report["uniform_baseline"]
+    def line(name, stats):  # noqa: E306
+        return f"{name:28} {stats['count']:>6} {stats.get('tvd', float('nan')):>7.3f} {stats.get('agree', float('nan')):>7.3f} {stats.get('log_loss', float('nan')):>8.3f} {stats.get('brier', float('nan')):>7.3f}"
+    lines.append(line("student / overall", student["overall"]))
+    lines.append(line("uniform / overall", baseline["overall"]))
+    for kind in PRIMITIVES:
+        if kind in student["by_primitive"]:
+            lines.append(line(f"student / {kind}", student["by_primitive"][kind]))
+            lines.append(line(f"uniform / {kind}", baseline["by_primitive"][kind]))
+    for domain, stats in student["by_domain"].items():
+        lines.append(line(f"student / {domain}", stats))
+    return "\n".join(lines)

src/yev/training/data.py

@@ -0,0 +1,649 @@
+from collections import Counter
+import csv
+from dataclasses import asdict, dataclass, field
+import hashlib
+import io
+import json
+import math
+from pathlib import Path
+import random
+import re
+from urllib.request import Request, urlopen
+
+from .images import IMAGE_PLACEHOLDER
+
+
+SPLITS = ("train", "dev", "calibration", "test")
+BANKING_REPOSITORY = "https://github.com/PolyAI-LDN/task-specific-datasets"
+BANKING_REVISION_URL = "https://api.github.com/repos/PolyAI-LDN/task-specific-datasets/commits?path=banking_data&per_page=1"
+SEVERITY_CRITERIA = ["Routine", "Low", "Moderate", "High", "Critical"]
+PRIORITY_CRITERIA = ["P4 Low", "P3 Normal", "P2 High", "P1 Urgent"]
+PRIORITY_INSTRUCTIONS = (
+    "Apply the highest matching ticket priority rule. P1 Urgent: a confirmed outage, or at least 100 customers affected. "
+    "P2 High: at least 20 customers affected, or the ticket has been open for at least 48 hours. "
+    "P3 Normal: at least 5 customers affected, or open for at least 8 hours. Otherwise P4 Low. "
+    "The criteria are ordered from lowest to highest priority."
+)
+BINOMIAL_INSTRUCTIONS = (
+    "How many of the drawn beads will be red? Preserve the actual mathematical outcome probabilities. "
+    "The criteria are ordered by the number of red beads drawn."
+)
+SEVERITY_INSTRUCTIONS = (
+    "Apply the highest matching incident severity rule. Critical: data loss, downtime >= 120 minutes, "
+    "or affected users >= 75%. High: downtime >= 60 minutes or affected users >= 40%. "
+    "Moderate: downtime >= 15 minutes or affected users >= 10%. "
+    "Low: downtime >= 5 minutes or affected users >= 1%. Otherwise Routine. "
+    "The criteria are ordered from least to most severe."
+)
+
+
+@dataclass
+class DecisionExample:
+    uid: str
+    group: str
+    source: str
+    split: str
+    kind: str
+    state: str
+    instructions: str
+    criteria: list[str]
+    target: list[float]
+    images: list[str] = field(default_factory=list)   # digests of the images the state's placeholders stand for, in order
+
+    def __post_init__(self):
+        limits = {"choice": (2, 255), "score": (2, 10), "noul": (2, 2)}
+        if self.kind not in limits or self.split not in {"train", "dev", "calibration", "test"}:
+            raise ValueError("Unknown primitive or split")
+        minimum, maximum = limits[self.kind]
+        if not minimum <= len(self.criteria) <= maximum or len(self.target) != len(self.criteria):
+            raise ValueError("Criteria and targets must match the primitive's size limits")
+        if any(not math.isfinite(value) or value < 0 for value in self.target):
+            raise ValueError("Targets must contain finite nonnegative probabilities")
+        if not math.isclose(sum(self.target), 1.0, abs_tol=1e-6):
+            raise ValueError("Target probabilities must sum to one")
+        self.images = list(self.images)
+        if self.state.count(IMAGE_PLACEHOLDER) != len(self.images):
+            raise ValueError(f"The state must hold one {IMAGE_PLACEHOLDER} for each image")
+
+    def to_dict(self):
+        """The record as stored; a text-only example has no ``images`` key, so text corpora and
+        their fingerprints are unchanged."""
+        record = asdict(self)
+        if not self.images:
+            record.pop("images")
+        return record
+
+
+def is_certain(target):
+    """Whether one answer holds all of the target's probability."""
+    return max(target) > 0.999
+
+
+def top_answer(target):
+    """The target as its single most likely answer."""
+    top = target.index(max(target))
+    return [float(index == top) for index in range(len(target))]
+
+
+def document_group(source, state):
+    normalized = " ".join(state.lower().split())
+    digest = hashlib.sha256(normalized.encode()).hexdigest()
+    return f"{source}:{digest}"
+
+
+def split_for_group(group, original_split):
+    if original_split == "test":
+        return "test"
+    bucket = int(hashlib.sha256(group.encode()).hexdigest()[:8], 16) % 100
+    if bucket < 80:
+        return "train"
+    return "dev" if bucket < 90 else "calibration"
+
+
+def probability_examples(seed=42, count=1000):
+    generator = random.Random(seed)
+    examples = []
+    seen = set()
+    if not 1 <= count <= 10000:
+        raise ValueError("Synthetic count must be between 1 and 10000")
+    remaining = synthetic_split_quotas(count)
+    while len(seen) < count:
+        total = generator.randint(2, 1000)
+        red = generator.randint(0, total)
+        state = f"A bag contains {red} red beads and {total - red} blue beads. One bead is drawn uniformly at random."
+        group = document_group("synthetic_probability", state)
+        split = synthetic_split(group)
+        if group in seen or remaining[split] == 0:
+            continue
+        seen.add(group)
+        remaining[split] -= 1
+        probability = red / total
+        for kind in ("choice", "score", "noul"):
+            examples.append(DecisionExample(
+                uid=f"{group}:{kind}", group=group, source="synthetic_probability", split=split,
+                kind=kind, state=state,
+                instructions="Will the drawn bead be red? Preserve the actual mathematical outcome probabilities.",
+                criteria=["false", "true"] if kind == "noul" else ["The drawn bead is blue (no).", "The drawn bead is red (yes)."],
+                target=[1 - probability, probability],
+            ))
+    return examples
+
+
+def synthetic_split(group):
+    bucket = int(hashlib.sha256(group.encode()).hexdigest()[:8], 16) % 100
+    if bucket < 70:
+        return "train"
+    if bucket < 80:
+        return "dev"
+    return "calibration" if bucket < 90 else "test"
+
+
+def synthetic_split_quotas(count):
+    if count < 4:
+        return {split: int(index < count) for index, split in enumerate(SPLITS)}
+    heldout = max(1, count // 10)
+    return dict(zip(SPLITS, (count - 3 * heldout, heldout, heldout, heldout)))
+
+
+def download_bytes(url, max_bytes=8 * 1024 * 1024):
+    request = Request(url, headers={"User-Agent": "jev-investigation-data-preparation"})
+    with urlopen(request, timeout=30) as response:
+        content = response.read(max_bytes + 1)
+    if len(content) > max_bytes:
+        raise ValueError(f"Source exceeds the {max_bytes}-byte download limit: {url}")
+    return content
+
+
+def parse_banking_csv(content, categories):
+    reader = csv.DictReader(io.StringIO(content.decode("utf-8-sig"), newline=""))
+    if reader.fieldnames != ["text", "category"]:
+        raise ValueError("BANKING77 CSV must have the original text,category header")
+    records = list(reader)
+    for row in records:
+        if set(row) != {"text", "category"} or not row["text"] or not row["text"].strip():
+            raise ValueError("BANKING77 CSV contains an empty or malformed document")
+        if row["category"] not in categories:
+            raise ValueError(f"Unknown BANKING77 category: {row['category']!r}")
+    return records
+
+
+def load_banking77():
+    revision_content = download_bytes(BANKING_REVISION_URL, max_bytes=256 * 1024)
+    revision = json.loads(revision_content)[0]["sha"]
+    if not isinstance(revision, str) or re.fullmatch(r"[0-9a-f]{40}", revision) is None:
+        raise ValueError("GitHub did not return a valid BANKING77 commit SHA")
+    base_url = f"https://raw.githubusercontent.com/PolyAI-LDN/task-specific-datasets/{revision}/banking_data"
+    contents = {}
+    files = {}
+    for filename in ("categories.json", "train.csv", "test.csv"):
+        url = f"{base_url}/{filename}"
+        content = download_bytes(url)
+        contents[filename] = content
+        files[filename] = {"url": url, "sha256": hashlib.sha256(content).hexdigest(), "bytes": len(content)}
+    categories = json.loads(contents["categories.json"])
+    if (
+        not isinstance(categories, list) or len(categories) != 77
+        or any(not isinstance(category, str) or not category for category in categories)
+        or len(set(categories)) != 77
+    ):
+        raise ValueError("BANKING77 categories.json must contain 77 unique category names")
+    records = {split: parse_banking_csv(contents[f"{split}.csv"], categories) for split in ("train", "test")}
+    provenance = {
+        "url": BANKING_REPOSITORY,
+        "revision": revision,
+        "revision_url": BANKING_REVISION_URL,
+        "revision_response_sha256": hashlib.sha256(revision_content).hexdigest(),
+        "files": files,
+        "license": "CC-BY-4.0",
+        "license_url": "https://creativecommons.org/licenses/by/4.0/",
+        "license_evidence_url": "https://huggingface.co/datasets/PolyAI/banking77#licensing-information",
+        "attribution": "BANKING77 by PolyAI; Casanueva et al., Efficient Intent Detection with Dual Sentence Encoders (2020).",
+        "modifications": "Normalized-document deduplication; seeded subsampling; subset-choice and binary intent questions.",
+        "category_count": len(categories),
+        "raw_records": {split: len(rows) for split, rows in records.items()},
+        "target_type": "Observed single-label annotations, not calibrated soft probabilities.",
+    }
+    return records, categories, provenance
+
+
+def banking_group_policy(records, categories):
+    if len(categories) < 2 or len(set(categories)) != len(categories):
+        raise ValueError("At least two distinct intent categories are required")
+    grouped_labels = {}
+    occurrences = Counter()
+    for original_split in ("train", "test"):
+        for row in records[original_split]:
+            if not row["text"].strip() or row["category"] not in categories:
+                raise ValueError("BANKING77 records require nonempty text and a known category")
+            group = document_group("banking77", row["text"])
+            grouped_labels.setdefault(group, {}).setdefault(original_split, set()).add(row["category"])
+            occurrences[group] += 1
+    owners = {group: "test" if "test" in labels else "train" for group, labels in grouped_labels.items()}
+    ambiguous = {group for group, labels in grouped_labels.items() if len(labels[owners[group]]) > 1}
+    audit = {
+        "raw_records": {split: len(records[split]) for split in ("train", "test")},
+        "unique_source_groups": len(grouped_labels),
+        "duplicate_groups": sum(count > 1 for count in occurrences.values()),
+        "conflicting_groups": sum(len(set().union(*labels.values())) > 1 for labels in grouped_labels.values()),
+        "ambiguous_owner_groups_dropped": len(ambiguous),
+        "test_priority_groups": sum(len(labels) == 2 for labels in grouped_labels.values()),
+        "sampled_raw_records": {},
+        "retained_documents": {},
+        "excluded_sampled_records": {},
+    }
+    return owners, ambiguous, audit
+
+
+def banking_questions(state, group, split, label, categories, seed):
+    generator = random.Random(f"{seed}:{group}")
+    alternatives = [category for category in categories if category != label]
+    candidate_count = generator.randint(2, min(8, len(categories)))
+    candidates = generator.sample(alternatives, candidate_count - 1) + [label]
+    generator.shuffle(candidates)
+    examples = [DecisionExample(
+        uid=f"{group}:choice", group=group, source="banking77", split=split, kind="choice",
+        state=state, instructions="Select the customer's intent from the offered candidates.",
+        criteria=candidates, target=[float(candidate == label) for candidate in candidates],
+    )]
+    negative = generator.choice(alternatives)
+    for predicate, truth in ((label, True), (negative, False)):
+        examples.append(DecisionExample(
+            uid=f"{group}:noul:{str(truth).lower()}", group=group, source="banking77", split=split,
+            kind="noul", state=state,
+            instructions=f"Does the customer's request express the intent '{predicate}'?",
+            criteria=["false", "true"], target=[float(not truth), float(truth)],
+        ))
+    return examples
+
+
+def banking_examples(records, categories, seed=42, bank_limit=3000):
+    if type(bank_limit) is not int or bank_limit < 0:
+        raise ValueError("bank_limit must be a nonnegative integer")
+    owners, ambiguous, audit = banking_group_policy(records, categories)
+    generator = random.Random(seed)
+    examples = []
+    seen = set()
+    for original_split in ("train", "test"):
+        shuffled = list(records[original_split])
+        generator.shuffle(shuffled)
+        sampled = shuffled[:bank_limit]
+        audit["sampled_raw_records"][original_split] = len(sampled)
+        retained = 0
+        excluded = Counter()
+        for row in sampled:
+            state = row["text"].strip()
+            group = document_group("banking77", state)
+            if owners[group] != original_split:
+                excluded["test_priority"] += 1
+                continue
+            if group in ambiguous:
+                excluded["ambiguous_label"] += 1
+                continue
+            if group in seen:
+                excluded["duplicate"] += 1
+                continue
+            seen.add(group)
+            retained += 1
+            split = split_for_group(group, original_split)
+            examples.extend(banking_questions(state, group, split, row["category"], categories, seed))
+        audit["retained_documents"][original_split] = retained
+        audit["excluded_sampled_records"][original_split] = dict(excluded)
+    return examples, audit
+
+
+def severity_label(downtime_minutes, affected_percent, data_loss):
+    if data_loss or downtime_minutes >= 120 or affected_percent >= 75:
+        return 4
+    if downtime_minutes >= 60 or affected_percent >= 40:
+        return 3
+    if downtime_minutes >= 15 or affected_percent >= 10:
+        return 2
+    if downtime_minutes >= 5 or affected_percent >= 1:
+        return 1
+    return 0
+
+
+def severity_examples(seed=42, count=200):
+    if type(count) is not int or not 1 <= count <= 200:
+        raise ValueError("Severity count must be between 1 and 200")
+    generator = random.Random(seed)
+    remaining = synthetic_split_quotas(count)
+    templates = {
+        "train": "Service incident: downtime_minutes={downtime}; affected_percent={affected}; data_loss={loss}.",
+        "dev": "Monitoring summary: {affected}% of users affected; downtime {downtime} minutes; data loss {loss}.",
+        "calibration": "Support audit records data loss {loss}, {downtime} minutes of downtime and {affected}% affected users.",
+        "test": "An operations review measured downtime of {downtime} minutes. The affected share was {affected}% of users. Data loss: {loss}.",
+    }
+    bands = ((0, 4, 0, 0), (5, 14, 1, 9), (15, 59, 10, 39), (60, 119, 40, 74), (120, 240, 75, 100))
+    examples = []
+    seen = set()
+    while len(examples) < count:
+        minimum_time, maximum_time, minimum_impact, maximum_impact = generator.choice(bands)
+        downtime = generator.randint(minimum_time, maximum_time)
+        affected = generator.randint(minimum_impact, maximum_impact)
+        loss = generator.random() < 0.1
+        scenario = json.dumps([downtime, affected, loss])
+        group = document_group("synthetic_severity", scenario)
+        split = synthetic_split(group)
+        if group in seen or remaining[split] == 0:
+            continue
+        seen.add(group)
+        remaining[split] -= 1
+        label = severity_label(downtime, affected, loss)
+        examples.append(DecisionExample(
+            uid=f"{group}:score", group=group, source="synthetic_severity", split=split, kind="score",
+            state=templates[split].format(downtime=downtime, affected=affected, loss=str(loss).lower()),
+            instructions=SEVERITY_INSTRUCTIONS, criteria=list(SEVERITY_CRITERIA),
+            target=[float(index == label) for index in range(len(SEVERITY_CRITERIA))],
+        ))
+    return examples
+
+
+def binomial_examples(seed=42, count=500):
+    """Ordered Score questions with exact binomial targets over 3 to 6 levels.
+
+    A bag with known red/blue counts is sampled with replacement ``draws`` times; level ``i``
+    is "exactly i red beads". The target is Binomial(draws, red / total), so calibration sees
+    genuinely spread ordinal distributions rather than one-hot rules. Templates are split-exclusive.
+    """
+    if type(count) is not int or not 1 <= count <= 10000:
+        raise ValueError("Binomial count must be between 1 and 10000")
+    generator = random.Random(seed)
+    remaining = synthetic_split_quotas(count)
+    templates = {
+        "train": "A bag contains {red} red beads and {blue} blue beads. {draws} beads are drawn one at a time, each returned to the bag before the next draw.",
+        "dev": "There are {blue} blue and {red} red beads in a bag. A bead is drawn uniformly at random and put back; this is repeated {draws} times in total.",
+        "calibration": "Bag contents: {red} red, {blue} blue. Procedure: draw one bead at random, note its colour, replace it; perform {draws} draws.",
+        "test": "{draws} independent draws with replacement are made from a bag holding {red} red beads and {blue} blue beads.",
+    }
+    examples = []
+    seen = set()
+    while len(examples) < count:
+        total = generator.randint(2, 60)
+        red = generator.randint(0, total)
+        draws = generator.randint(2, 5)
+        scenario = json.dumps([red, total - red, draws])
+        group = document_group("synthetic_binomial", scenario)
+        split = synthetic_split(group)
+        if group in seen or remaining[split] == 0:
+            continue
+        seen.add(group)
+        remaining[split] -= 1
+        probability = red / total
+        target = [math.comb(draws, hits) * probability ** hits * (1 - probability) ** (draws - hits) for hits in range(draws + 1)]
+        examples.append(DecisionExample(
+            uid=f"{group}:score", group=group, source="synthetic_binomial", split=split, kind="score",
+            state=templates[split].format(red=red, blue=total - red, draws=draws),
+            instructions=BINOMIAL_INSTRUCTIONS,
+            criteria=[f"Exactly {hits} of the {draws} drawn beads are red." for hits in range(draws + 1)],
+            target=target,
+        ))
+    return examples
+
+
+def priority_label(customers_affected, hours_open, outage):
+    if outage or customers_affected >= 100:
+        return 3
+    if customers_affected >= 20 or hours_open >= 48:
+        return 2
+    if customers_affected >= 5 or hours_open >= 8:
+        return 1
+    return 0
+
+
+def priority_examples(seed=42, count=400):
+    """Four-level one-hot Score questions from executable ticket-priority rules with held-out wording."""
+    if type(count) is not int or not 1 <= count <= 400:
+        raise ValueError("Priority count must be between 1 and 400")
+    generator = random.Random(seed)
+    remaining = synthetic_split_quotas(count)
+    templates = {
+        "train": "Support ticket: customers_affected={customers}; hours_open={hours}; confirmed_outage={outage}.",
+        "dev": "Ticket status: open for {hours} hours, {customers} customers affected, confirmed outage {outage}.",
+        "calibration": "The queue report lists a confirmed outage flag of {outage}, {customers} affected customers and {hours} hours since the ticket was opened.",
+        "test": "A ticket has been open {hours} hours. It affects {customers} customers. Outage confirmed: {outage}.",
+    }
+    customer_bands = ((0, 4), (5, 19), (20, 99), (100, 500))
+    hour_bands = ((0, 7), (8, 47), (48, 200))
+    examples = []
+    seen = set()
+    while len(examples) < count:
+        customers = generator.randint(*generator.choice(customer_bands))
+        hours = generator.randint(*generator.choice(hour_bands))
+        outage = generator.random() < 0.1
+        scenario = json.dumps([customers, hours, outage])
+        group = document_group("synthetic_priority", scenario)
+        split = synthetic_split(group)
+        if group in seen or remaining[split] == 0:
+            continue
+        seen.add(group)
+        remaining[split] -= 1
+        label = priority_label(customers, hours, outage)
+        examples.append(DecisionExample(
+            uid=f"{group}:score", group=group, source="synthetic_priority", split=split, kind="score",
+            state=templates[split].format(customers=customers, hours=hours, outage=str(outage).lower()),
+            instructions=PRIORITY_INSTRUCTIONS, criteria=list(PRIORITY_CRITERIA),
+            target=[float(index == label) for index in range(len(PRIORITY_CRITERIA))],
+        ))
+    return examples
+
+
+def jsonl_content(examples):
+    return "".join(json.dumps(example.to_dict(), sort_keys=True, ensure_ascii=True) + "\n" for example in examples).encode("utf-8")
+
+
+def partition_examples(examples, seed, require_all_primitives):
+    splits = {split: [] for split in SPLITS}
+    group_splits = {}
+    uids = set()
+    for example in examples:
+        if group_splits.setdefault(example.group, example.split) != example.split or example.uid in uids:
+            raise ValueError("Duplicate example UID or source-document leakage across splits")
+        uids.add(example.uid)
+        splits[example.split].append(example)
+    for split, examples in splits.items():
+        random.Random(f"{seed}:{split}").shuffle(examples)
+        if require_all_primitives and {example.kind for example in examples} != {"choice", "score", "noul"}:
+            raise ValueError(f"Missing primitive in {split}")
+    return splits
+
+
+def prepare_data(output_dir, seed=42, bank_limit=3000, synthetic_count=500, include_sst5=False) -> dict[str, list[DecisionExample]]:
+    if include_sst5:
+        raise ValueError("SST5 is deferred: its dataset card has no verified license metadata; include_sst5=True is unsupported")
+    if type(seed) is not int:
+        raise ValueError("seed must be an integer")
+    if type(bank_limit) is not int or bank_limit < 0:
+        raise ValueError("bank_limit must be a nonnegative integer")
+    if type(synthetic_count) is not int or not 1 <= synthetic_count <= 10000:
+        raise ValueError("synthetic_count must be between 1 and 10000")
+    records, categories, provenance = load_banking77()
+    bank, audit = banking_examples(records, categories, seed=seed, bank_limit=bank_limit)
+    probabilities = probability_examples(seed=seed, count=synthetic_count)
+    severity = severity_examples(seed=seed, count=min(200, synthetic_count))
+    binomial = binomial_examples(seed=seed, count=synthetic_count)
+    priority = priority_examples(seed=seed, count=min(400, synthetic_count))
+    splits = partition_examples(bank + probabilities + severity + binomial + priority, seed, require_all_primitives=synthetic_count >= 4)
+    sources = {"banking77": {**provenance, "audit": audit}}
+    target_types = {
+        "synthetic_probability": "Exact mathematical probabilities",
+        "synthetic_severity": "Executable severity_label rule, one-hot targets",
+        "synthetic_binomial": "Exact binomial probabilities over 3-6 ordered levels",
+        "synthetic_priority": "Executable priority_label rule, one-hot targets over 4 levels",
+    }
+    for name, examples, generator_revision in (
+        ("synthetic_probability", probabilities, "uniform-bags-v1"),
+        ("synthetic_severity", severity, "ordered-incident-rules-v1"),
+        ("synthetic_binomial", binomial, "binomial-draws-v1"),
+        ("synthetic_priority", priority, "ordered-ticket-priority-rules-v1"),
+    ):
+        sources[name] = {
+            "urls": [], "origin": "Locally generated; no external dataset or teacher model",
+            "license": "Not specified for locally authored synthetic data",
+            "generator_revision": generator_revision,
+            "sha256": hashlib.sha256(jsonl_content(examples)).hexdigest(),
+            "documents": len({example.group for example in examples}), "examples": len(examples),
+            "target_type": target_types[name],
+        }
+    manifest = {
+        "format_version": 1,
+        "config": {"seed": seed, "bank_limit": bank_limit, "synthetic_count": synthetic_count, "include_sst5": include_sst5},
+        "sources": sources,
+        "counts": {
+            "examples": {split: len(examples) for split, examples in splits.items()},
+            "groups": {split: len({example.group for example in examples}) for split, examples in splits.items()},
+            "by_source": {split: dict(Counter(example.source for example in examples)) for split, examples in splits.items()},
+            "by_primitive": {split: dict(Counter(example.kind for example in examples)) for split, examples in splits.items()},
+        },
+        "policies": {
+            "banking_task": "Subset-choice classification with 2-8 candidates, NOT a full-77 benchmark; one positive and one negative Noul per document.",
+            "groups": "Source plus SHA256 of lowercase, whitespace-collapsed document; no punctuation removal or semantic deduplication.",
+            "banking_splits": "Original test groups reserved before sampling. Original train groups hash to train/dev/calibration at 80/10/10.",
+            "duplicates": "Test labels win cross-split conflicts; conflicting labels within the owning split drop the whole group; no refill after raw-record caps.",
+            "synthetic_splits": "Canonical groups hash to 70/10/10/10; rejection sampling meets quotas without moving or copying groups between splits.",
+            "severity_holdout": "Four split-exclusive rendering templates, shared rules; canonical numeric scenarios grouped independently of template.",
+            "gold_score": "Score calibration/test draw on binary bags, exact binomial draws (3-6 levels), and two executable rule tasks (4 and 5 levels); no human ratings.",
+            "small_counts": "Every primitive in every split guaranteed for synthetic_count >= 4; fewer groups can leave splits empty.",
+            "sst5": "Deferred, license unverified; enabling it raises ValueError.",
+        },
+        "outputs": {},
+    }
+    destination = Path(output_dir)
+    destination.mkdir(parents=True, exist_ok=True)
+    for split, examples in splits.items():
+        filename = f"{split}.jsonl"
+        content = jsonl_content(examples)
+        (destination / filename).write_bytes(content)
+        manifest["outputs"][filename] = {"sha256": hashlib.sha256(content).hexdigest(), "bytes": len(content)}
+    (destination / "manifest.json").write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8")
+    return splits
+
+
+def validate_split_integrity(splits):
+    group_splits = {}
+    uids = set()
+    for split, examples in splits.items():
+        if split not in SPLITS:
+            raise ValueError(f"Unknown split: {split}")
+        for example in examples:
+            if example.split != split:
+                raise ValueError(f"Example {example.uid} belongs to {example.split}, not {split}")
+            if not isinstance(example.uid, str) or not example.uid.strip():
+                raise ValueError("Example UID must be a nonempty string")
+            if not isinstance(example.group, str) or not example.group.strip():
+                raise ValueError("Example group must be a nonempty string")
+            if example.uid in uids:
+                raise ValueError(f"Duplicate example UID: {example.uid}")
+            if group_splits.setdefault(example.group, split) != split:
+                raise ValueError(f"Semantic group leakage across splits: {example.group}")
+            uids.add(example.uid)
+
+
+# Where a label came from: Jev's distribution, or people (dataset labels and annotator vote shares).
+LABEL_SOURCES = ("jev_teacher:", "human:", "logit_teacher:")
+# Manifest ``model`` values name the labellers a dataset's targets come from, joined by "+": Jev,
+# people, and generative models read through their answer-token logits ("logit:<model>").
+LABELLERS = ("jev-1.13.0", "human")
+LOGIT_LABELLER = "logit:"
+
+
+def labeller_parts(model):
+    """The labellers in a manifest ``model`` value, or None when any of them is unrecognised."""
+    if not isinstance(model, str) or not model:
+        return None
+    parts = model.split("+")
+    if any(part not in LABELLERS and not (part.startswith(LOGIT_LABELLER) and len(part) > len(LOGIT_LABELLER))
+           for part in parts):
+        return None
+    return parts
+
+
+def join_labellers(parts):
+    """One manifest ``model`` value for several labellers: Jev, then people, then logit teachers."""
+    unique = set(parts)
+    return "+".join([name for name in LABELLERS if name in unique] + sorted(unique - set(LABELLERS)))
+
+
+def label_source(source):
+    """The labeller prefix of a source string, or None when it has no recognised prefix."""
+    return next((prefix for prefix in LABEL_SOURCES if isinstance(source, str) and source.startswith(prefix)), None)
+
+
+def validate_teacher_splits(splits):
+    if set(splits) != {"train", "dev"}:
+        raise ValueError("Teacher data must contain only train and dev splits")
+    validate_split_integrity(splits)
+    for examples in splits.values():
+        for example in examples:
+            prefix = label_source(example.source)
+            if prefix is None or not example.source.removeprefix(prefix).strip():
+                raise ValueError(f"Invalid teacher source for {example.uid}: {example.source!r}")
+
+
+def load_teacher_data(directory) -> dict[str, list[DecisionExample]]:
+    directory = Path(directory)
+    manifest = json.loads((directory / "manifest.json").read_text(encoding="utf-8"))
+    if not isinstance(manifest, dict) or manifest.get("complete") is not True:
+        raise ValueError("Teacher dataset must be complete before import")
+    if type(manifest.get("format_version")) is not int or manifest["format_version"] != 1:
+        raise ValueError("Unsupported teacher manifest format_version")
+    if labeller_parts(manifest.get("model")) is None:
+        raise ValueError(f"Teacher manifest model must name {' or '.join(LABELLERS)} or {LOGIT_LABELLER}<model>, joined by '+'")
+    by_split = manifest.get("by_split")
+    outputs = manifest.get("outputs")
+    if not isinstance(by_split, dict) or set(by_split) != {"train", "dev"} or not isinstance(outputs, dict):
+        raise ValueError("Teacher manifest requires train/dev counts and outputs")
+    splits = {}
+    for split in ("train", "dev"):
+        filename = f"{split}.jsonl"
+        integrity = outputs.get(filename)
+        if not isinstance(integrity, dict):
+            raise ValueError(f"Missing teacher output integrity: {filename}")
+        content = (directory / filename).read_bytes()
+        if hashlib.sha256(content).hexdigest() != integrity.get("sha256"):
+            raise ValueError(f"Teacher output SHA256 mismatch: {filename}")
+        examples = [DecisionExample(**json.loads(line)) for line in content.decode("utf-8").splitlines()]
+        for count in (by_split[split], integrity.get("examples")):
+            if type(count) is not int or count != len(examples):
+                raise ValueError(f"Teacher example count mismatch: {filename}")
+        splits[split] = examples
+    accepted = manifest.get("accepted_examples")
+    if type(accepted) is not int or accepted != sum(map(len, splits.values())):
+        raise ValueError("Teacher manifest accepted_examples mismatch")
+    validate_teacher_splits(splits)
+    return splits
+
+
+def write_teacher_data(splits, destination, manifest):
+    """Write ``splits`` as a teacher dataset ``load_teacher_data`` reads: ``manifest`` (the source's,
+    with any notes of what was filtered) gets the counts and hashes of what is written."""
+    destination = Path(destination)
+    destination.mkdir(parents=True, exist_ok=True)
+    outputs = {}
+    for split in ("train", "dev"):
+        content = "".join(json.dumps(example.to_dict(), ensure_ascii=True, allow_nan=False) + "\n" for example in splits[split])
+        (destination / f"{split}.jsonl").write_text(content, encoding="utf-8")
+        outputs[f"{split}.jsonl"] = {"sha256": hashlib.sha256(content.encode("utf-8")).hexdigest(), "examples": len(splits[split])}
+    manifest = {**manifest, "complete": True, "outputs": outputs,
+                "accepted_examples": sum(len(splits[split]) for split in ("train", "dev")),
+                "by_split": {split: len(splits[split]) for split in ("train", "dev")},
+                "by_primitive": dict(Counter(example.kind for split in ("train", "dev") for example in splits[split]))}
+    (destination / "manifest.json").write_text(json.dumps(manifest, indent=2, ensure_ascii=True, allow_nan=False) + "\n",
+                                               encoding="utf-8")
+    load_teacher_data(destination)   # a written directory must be a loadable one
+    return manifest
+
+
+def merge_teacher_data(base_splits, teacher_splits, seed=42) -> dict[str, list[DecisionExample]]:
+    if set(base_splits) != set(SPLITS):
+        raise ValueError("Base data must contain train, dev, calibration and test splits")
+    if type(seed) is not int:
+        raise ValueError("seed must be an integer")
+    validate_teacher_splits(teacher_splits)
+    merged = {split: list(base_splits[split]) for split in SPLITS}
+    for split in ("train", "dev"):
+        merged[split].extend(teacher_splits[split])
+    validate_split_integrity(merged)
+    for split in ("train", "dev"):
+        random.Random(f"{seed}:{split}").shuffle(merged[split])
+    return merged

src/yev/training/images.py

@@ -0,0 +1,206 @@
+"""Images in System One states (v2 plan, Phase 1).
+
+A state carries its images as ``IMAGE_PLACEHOLDER`` markers in its text, one for each entry of the
+example's ``images``: content hashes (the SHA-256 of the encoded file) that an ``ImageStore``
+resolves, so a dataset keeps its images beside its JSONL rather than inside it, in an ``images/``
+folder of files named by digest or of uncompressed zip shards (``images-00000.zip``...) holding
+such files, which keeps a Hub repository's file count small.
+
+A Qwen3.5 backbone reads each image as ``<|vision_start|>``, one ``<|image_pad|>`` per merged
+patch, then ``<|vision_end|>``. The image processor scales every image to between
+``MIN_IMAGE_TOKENS`` and ``max_image_tokens`` of those, so images have their own budget and never
+push the question or its criteria out of the sequence; ``max_length`` still bounds the text.
+Pillow is imported only when an image is actually read.
+"""
+
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass, field
+import hashlib
+import io
+from pathlib import Path
+import zipfile
+
+IMAGE_PLACEHOLDER = "<|image|>"
+DEFAULT_MAX_IMAGE_TOKENS = 256
+MIN_IMAGE_TOKENS = 16
+
+
+def image_digest(data):
+    return hashlib.sha256(data).hexdigest()
+
+
+def save_image(directory, data):
+    """Write encoded image bytes to ``directory`` under their digest; returns the digest."""
+    digest = image_digest(data)
+    path = Path(directory) / digest
+    if not path.exists():
+        path.parent.mkdir(parents=True, exist_ok=True)
+        temporary = path.with_suffix(".tmp")
+        temporary.write_bytes(data)
+        temporary.replace(path)
+    return digest
+
+
+class ImageShards:
+    """Writes encoded images into ``images-NNNNN.zip`` shards of at most ``shard_bytes`` each,
+    stored uncompressed (images are compressed already) under their digest."""
+
+    def __init__(self, directory, shard_bytes=1 << 30):
+        self.directory, self.shard_bytes = Path(directory), shard_bytes
+        self.directory.mkdir(parents=True, exist_ok=True)
+        self.written, self.archive, self.size = set(), None, 0
+
+    def add(self, data):
+        digest = image_digest(data)
+        if digest in self.written:
+            return digest
+        if self.archive is None or self.size + len(data) > self.shard_bytes:
+            self.close()
+            index = len(list(self.directory.glob("images-*.zip")))
+            self.archive, self.size = zipfile.ZipFile(self.directory / f"images-{index:05d}.zip", "w", zipfile.ZIP_STORED), 0
+        self.archive.writestr(digest, data)
+        self.written.add(digest)
+        self.size += len(data)
+        return digest
+
+    def close(self):
+        if self.archive is not None:
+            self.archive.close()
+            self.archive = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exception):
+        self.close()
+
+
+class ImageStore:
+    """Encoded images by digest: added in memory, or read from registered directories holding files
+    named by their digest and zip shards of such files (a dataset's ``images/`` folder)."""
+
+    def __init__(self, *directories):
+        self.directories = []
+        self.memory = {}
+        self.shards = {}    # digest -> zip path
+        self.open = {}      # zip path -> open ZipFile
+        for directory in directories:
+            self.add_directory(directory)
+
+    def add_directory(self, directory):
+        directory = Path(directory)
+        if directory not in self.directories:
+            self.directories.append(directory)
+            for archive in sorted(directory.glob("images-*.zip")) if directory.is_dir() else []:
+                with zipfile.ZipFile(archive) as opened:
+                    self.shards.update(dict.fromkeys(opened.namelist(), archive))
+
+    def put(self, data):
+        digest = image_digest(data)
+        self.memory[digest] = data
+        return digest
+
+    def read(self, digest):
+        if digest in self.memory:
+            return self.memory[digest]
+        archive = self.shards.get(digest)
+        if archive is not None:
+            if archive not in self.open:
+                self.open[archive] = zipfile.ZipFile(archive)
+            return self.open[archive].read(digest)
+        for directory in self.directories:
+            path = directory / digest
+            if path.is_file():
+                return path.read_bytes()
+        raise KeyError(f"No image {digest} in memory or in {[str(directory) for directory in self.directories]}")
+
+    def load(self, digest):
+        from PIL import Image
+
+        with Image.open(io.BytesIO(self.read(digest))) as image:
+            return image.convert("RGB")
+
+
+def image_processor(source, max_image_tokens=DEFAULT_MAX_IMAGE_TOKENS, revision=None):
+    """The backbone's image processor (PIL backend on every platform, so outputs never depend on
+    whether torchvision is installed), sized to ``MIN_IMAGE_TOKENS``..``max_image_tokens`` per image."""
+    from transformers.models.qwen2_vl.image_processing_pil_qwen2_vl import Qwen2VLImageProcessorPil
+
+    processor = Qwen2VLImageProcessorPil.from_pretrained(source, revision=revision)
+    pixels_per_token = (processor.patch_size * processor.merge_size) ** 2
+    return Qwen2VLImageProcessorPil.from_pretrained(
+        source, revision=revision,
+        size={"shortest_edge": MIN_IMAGE_TOKENS * pixels_per_token, "longest_edge": max_image_tokens * pixels_per_token})
+
+
+@dataclass
+class VisionInputs:
+    """What the collator needs to turn an example's images into backbone inputs."""
+    processor: object
+    start_id: int
+    pad_id: int
+    end_id: int
+    max_image_tokens: int = DEFAULT_MAX_IMAGE_TOKENS
+    store: ImageStore = field(default_factory=ImageStore)
+
+    @classmethod
+    def for_backbone(cls, backbone_config, processor, max_image_tokens=DEFAULT_MAX_IMAGE_TOKENS, store=None):
+        return cls(processor, backbone_config.vision_start_token_id, backbone_config.image_token_id,
+                   backbone_config.vision_end_token_id, max_image_tokens, store or ImageStore())
+
+    def prepare(self, examples):
+        """Per example, the token block of each of its images; and the batch's pixel values and grids."""
+        digests = [digest for example in examples for digest in example.images]
+        if not digests:
+            return [[] for _ in examples], None
+        encoded = self.processor(images=[self.store.load(digest) for digest in digests], return_tensors="pt")
+        counts = (encoded["image_grid_thw"].prod(-1) // self.processor.merge_size ** 2).tolist()
+        blocks, cursor = [], 0
+        for example in examples:
+            blocks.append([[self.start_id] + [self.pad_id] * count + [self.end_id]
+                           for count in counts[cursor:cursor + len(example.images)]])
+            cursor += len(example.images)
+        return blocks, (encoded["pixel_values"], encoded["image_grid_thw"])
+
+
+class ImagePrefetcher:
+    """Prepares the images of upcoming micro-batches in background threads while the model trains on
+    the current one, and stands in for the ``VisionInputs`` it wraps.
+
+    Only image work runs in the threads: reading from the store, decoding, resizing and patching,
+    which Pillow and NumPy do with the GIL released. Tokenizing stays on the caller's thread. A
+    micro-batch nobody submitted is prepared on the spot, so the result never depends on what was
+    prefetched.
+    """
+
+    def __init__(self, vision, workers=2):
+        self.vision = vision
+        self.pool = ThreadPoolExecutor(workers, thread_name_prefix="yev-images")
+        self.pending = {}
+
+    def __getattr__(self, name):
+        return getattr(self.vision, name)
+
+    @staticmethod
+    def key(examples):
+        return tuple(example.uid for example in examples) if any(example.images for example in examples) else None
+
+    def submit(self, examples):
+        key = self.key(examples)
+        if key is not None and key not in self.pending:
+            self.pending[key] = self.pool.submit(self.vision.prepare, list(examples))
+
+    def submit_update(self, train, indices, micro_batch):
+        """The update's image micro-batches, grouped as ``run.micro_batches`` groups them."""
+        pictures = [index for index in indices if train[index].images]
+        for start in range(0, len(pictures), micro_batch):
+            self.submit([train[index] for index in pictures[start:start + micro_batch]])
+
+    def prepare(self, examples):
+        future = self.pending.pop(self.key(examples), None)
+        return future.result() if future is not None else self.vision.prepare(examples)
+
+    def close(self):
+        self.pool.shutdown(wait=False, cancel_futures=True)
+        self.pending.clear()
+

src/yev/training/model.py

@@ -0,0 +1,309 @@
+from dataclasses import replace
+import math
+
+import torch
+from torch import nn
+
+from .data import DecisionExample
+from .images import IMAGE_PLACEHOLDER
+
+
+PRIMITIVES = ("choice", "score", "noul")
+# "mlp": v1's per-primitive MLPs over single token positions. "bilinear": MoJev's readout, one
+# rank-r dot product between pooled context-plus-question and each pooled candidate span.
+HEADS = ("mlp", "bilinear")
+# The bilinear head's optional last-token readout (``last_token``), one projection per role.
+LAST_TOKEN_HEADS = ("option_last", "question_last")
+
+
+def supports_native_bfloat16():
+    return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8 and torch.cuda.is_bf16_supported()
+
+
+def candidate_order(example, canonical=False):
+    """The order criteria are laid out in: sorted by text for a canonical Choice, otherwise as given.
+
+    Candidates share one causal sequence, so a candidate's score moves with the rivals placed
+    before it. Sorting makes the sequence a function of the candidate set, so a Choice answer
+    cannot depend on the order the caller listed its criteria in (MoJev's ``sort_candidates``).
+    Score keeps its scale order and Noul its (false, true) order: there the order is the question.
+    """
+    count = len(example.criteria)
+    if canonical and example.kind == "choice":
+        return sorted(range(count), key=example.criteria.__getitem__)
+    return list(range(count))
+
+
+def fits_context(example, tokenizer, max_length):
+    """Whether a question's instructions and criteria fit ``max_length``. A longer state is cut to
+    fit, but answer options never are, so a question that fails this cannot be asked at that length."""
+    try:
+        encode_example(replace(example, state=".", images=[]), tokenizer, max_length)
+    except ValueError:
+        return False
+    return True
+
+
+def encode_example(example, tokenizer, max_length, canonical=False, image_blocks=None):
+    """Token layout of one question. ``image_blocks`` holds the token block of each of the
+    example's images, which replaces its placeholder; images are outside ``max_length``, which
+    bounds the text, and when the text is cut the images all stay."""
+    order = candidate_order(example, canonical)
+    prefix = tokenizer.encode(f"Task: {example.kind}\nState:\n", add_special_tokens=False)
+    if example.images:
+        if image_blocks is None or len(image_blocks) != len(example.images):
+            raise ValueError("A state with images needs a backbone with a vision encoder")
+        state_parts = [tokenizer.encode(part, add_special_tokens=False) for part in example.state.split(IMAGE_PLACEHOLDER)]
+    else:
+        state_parts, image_blocks = [tokenizer.encode(example.state, add_special_tokens=False)], []
+    question = tokenizer.encode(f"\nQuestion: {example.instructions}\nCriteria:", add_special_tokens=False)
+    criteria = [tokenizer.encode(f"\n{index}: {example.criteria[original]}\n", add_special_tokens=False) for index, original in enumerate(order)]
+    suffix = tokenizer.encode("\nDecision probabilities:", add_special_tokens=False)
+    state_budget = max_length - len(prefix) - len(question) - sum(map(len, criteria)) - len(suffix)
+    if state_budget < 1:
+        raise ValueError("Instructions and criteria exceed max_length; increase it instead of truncating answer options")
+    truncated = max(0, sum(map(len, state_parts)) - state_budget)
+    input_ids, room = list(prefix), state_budget
+    for index, part in enumerate(state_parts):
+        input_ids.extend(part[:room])
+        room = max(0, room - len(part))
+        if index < len(image_blocks):
+            input_ids.extend(image_blocks[index])
+    context_end = len(input_ids)
+    input_ids.extend(question)
+    question_span = (context_end, len(input_ids))
+    positions, spans = [], []
+    for criterion in criteria:
+        start = len(input_ids)
+        input_ids.extend(criterion)
+        positions.append(len(input_ids) - 1)
+        spans.append((start, len(input_ids)))
+    input_ids.extend(suffix)
+    return {
+        "input_ids": input_ids,
+        "criterion_positions": positions,
+        "criterion_spans": spans,
+        "criterion_order": order,
+        "context_end": context_end,
+        "question_span": question_span,
+        "decision_position": len(input_ids) - 1,
+        "kind": PRIMITIVES.index(example.kind),
+        "target": example.target,
+        "truncated_state_tokens": truncated,
+    }
+
+
+class DecisionCollator:
+    """Pads encoded examples into a batch. Criteria arrive in layout order; ``criterion_order`` maps
+    each layout slot back to the caller's index, and targets stay in the caller's order."""
+
+    def __init__(self, tokenizer, max_length=512, canonical=False, vision=None):
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.canonical = canonical
+        self.vision = vision
+
+    def __call__(self, examples):
+        blocks, pixels = ([None] * len(examples), None)
+        if self.vision is not None and any(example.images for example in examples):
+            blocks, pixels = self.vision.prepare(examples)
+        encoded = [encode_example(example, self.tokenizer, self.max_length, self.canonical, image_blocks)
+                   for example, image_blocks in zip(examples, blocks)]
+        length = max(len(item["input_ids"]) for item in encoded)
+        count = max(len(item["target"]) for item in encoded)
+        batch_size = len(encoded)
+        batch = {
+            "input_ids": torch.full((batch_size, length), self.tokenizer.pad_token_id, dtype=torch.long),
+            "attention_mask": torch.zeros(batch_size, length, dtype=torch.long),
+            "criterion_positions": torch.zeros(batch_size, count, dtype=torch.long),
+            "criterion_mask": torch.zeros(batch_size, count, dtype=torch.bool),
+            "criterion_spans": torch.zeros(batch_size, count, 2, dtype=torch.long),
+            "criterion_order": torch.arange(count).repeat(batch_size, 1),
+            "context_ends": torch.tensor([item["context_end"] for item in encoded]),
+            "question_spans": torch.tensor([item["question_span"] for item in encoded]),
+            "decision_positions": torch.tensor([item["decision_position"] for item in encoded]),
+            "kinds": torch.tensor([item["kind"] for item in encoded]),
+            "targets": torch.zeros(batch_size, count, dtype=torch.float32),
+        }
+        for index, item in enumerate(encoded):
+            token_count = len(item["input_ids"])
+            option_count = len(item["target"])
+            batch["input_ids"][index, :token_count] = torch.tensor(item["input_ids"])
+            batch["attention_mask"][index, :token_count] = 1
+            batch["criterion_positions"][index, :option_count] = torch.tensor(item["criterion_positions"])
+            batch["criterion_spans"][index, :option_count] = torch.tensor(item["criterion_spans"])
+            batch["criterion_order"][index, :option_count] = torch.tensor(item["criterion_order"])
+            batch["criterion_mask"][index, :option_count] = True
+            batch["targets"][index, :option_count] = torch.tensor(item["target"])
+        if pixels is not None:   # a text-only batch carries no image keys and runs exactly as before
+            batch["pixel_values"], batch["image_grid_thw"] = pixels
+            batch["mm_token_type_ids"] = (batch["input_ids"] == self.vision.pad_id).int()
+        return batch
+
+
+def collator_for(model, tokenizer, max_length, vision=None):
+    """The collator matching ``model``: its candidate layout, and its vision inputs when it reads images."""
+    return DecisionCollator(tokenizer, max_length, model.canonical_order, vision or getattr(model, "vision", None))
+
+
+def distribution_loss(logits, targets, kinds, criterion_mask, ordinal_weight=0.1, brier_weight=0.0):
+    """Plackett-Luce top-1 (cross-entropy against the target distribution), plus the multi-class
+    Brier score at ``brier_weight`` (RLCD's calibration term; MoJev trains at 1.0), plus the
+    ordinal CDF term on Score. Absent criteria have probability and target 0, so add nothing."""
+    cross_entropy = -(targets * logits.log_softmax(-1)).sum(-1)
+    probabilities = logits.softmax(-1)
+    brier = (probabilities - targets).square().sum(-1)
+    counts = criterion_mask.sum(-1)
+    cdf_mask = torch.arange(logits.shape[-1], device=logits.device)[None, :] < (counts - 1)[:, None]
+    cdf_error = ((probabilities.cumsum(-1) - targets.cumsum(-1)).square() * cdf_mask).sum(-1)
+    ordinal = cdf_error / (counts - 1).clamp_min(1)
+    return (cross_entropy + brier_weight * brier + ordinal_weight * ordinal * (kinds == PRIMITIVES.index("score"))).mean()
+
+
+def span_mask(starts, ends, length):
+    positions = torch.arange(length, device=starts.device)
+    return ((positions >= starts[..., None]) & (positions < ends[..., None])).float()
+
+
+def mean_pool(hidden, mask):
+    """Mean hidden state over each span of ``mask`` (B, S, L); an empty span pools to zero."""
+    return torch.matmul(mask, hidden) / mask.sum(-1, keepdim=True).clamp_min(1)
+
+
+class DecisionModel(nn.Module):
+    def __init__(self, backbone, head_width=128, ordinal_weight=0.1, head="mlp", brier_weight=0.0, last_token=False):
+        super().__init__()
+        if head not in HEADS:
+            raise ValueError(f"head must be one of {HEADS}")
+        if last_token and head != "bilinear":
+            raise ValueError("The last-token readout extends the bilinear head")
+        self.backbone = backbone
+        self.head = head
+        self.head_width = head_width
+        self.ordinal_weight = ordinal_weight
+        self.brier_weight = brier_weight
+        self.last_token = last_token
+        # The bilinear head reads candidate text, so its layout is canonical; v1's MLP heads keep
+        # the order they were trained with.
+        self.canonical_order = head == "bilinear"
+        dimension = backbone.config.get_text_config().hidden_size   # multimodal configs nest it
+        if head == "bilinear":
+            # No per-primitive parameters: the primitive is in the text ("Task: noul"), and a Noul
+            # is a two-candidate preference scored exactly like a Choice.
+            self.heads = nn.ModuleDict({
+                "norm": nn.LayerNorm(dimension),
+                "context": nn.Linear(dimension, head_width, bias=False),
+                "option": nn.Linear(dimension, head_width, bias=False),
+            })
+            if last_token:
+                # A span's mean blurs texts that differ by one word ("The food is a taco, meaning
+                # ..." against "The food is not a taco, meaning it does not ..."); the span's last
+                # token has read all of it and tells them apart. Zero-initialised, so adding it to a
+                # trained head leaves the model's answers unchanged until it is trained.
+                for name in LAST_TOKEN_HEADS:
+                    self.heads[name] = nn.Linear(dimension, head_width, bias=False)
+                    nn.init.zeros_(self.heads[name].weight)
+        else:
+            self.heads = nn.ModuleDict({
+                "choice": nn.Sequential(nn.Linear(3 * dimension, head_width), nn.SiLU(), nn.Linear(head_width, 1)),
+                "score": nn.Sequential(nn.Linear(3 * dimension, head_width), nn.SiLU(), nn.Linear(head_width, 1)),
+                "noul": nn.Sequential(nn.Linear(dimension, head_width), nn.SiLU(), nn.Linear(head_width, 1)),
+            })
+        self.register_buffer("temperatures", torch.ones(len(PRIMITIVES)))
+
+    def bilinear_logits(self, hidden, context_ends, question_spans, criterion_spans):
+        """MoJev's readout: LayerNorm, mean-pool the context, question and candidate spans, then
+        score each candidate as (proj(context) + proj(question)) . proj(candidate) / sqrt(rank).
+        With ``last_token``, the question's and each candidate's last token add their own projections."""
+        hidden = self.heads["norm"](hidden.float())
+        length = hidden.shape[1]
+        context = mean_pool(hidden, span_mask(torch.zeros_like(context_ends), context_ends, length)[:, None]).squeeze(1)
+        question = mean_pool(hidden, span_mask(question_spans[:, :1], question_spans[:, 1:], length)).squeeze(1)
+        options = mean_pool(hidden, span_mask(criterion_spans[..., 0], criterion_spans[..., 1], length))
+        query = self.heads["context"](context) + self.heads["context"](question)
+        keys = self.heads["option"](options)
+        if self.last_token:
+            rows = torch.arange(hidden.shape[0], device=hidden.device)
+            query = query + self.heads["question_last"](hidden[rows, question_spans[:, 1] - 1])
+            ends = (criterion_spans[..., 1] - 1).clamp_min(0)   # an absent candidate's (0, 0) span reads token 0, then is masked
+            keys = keys + self.heads["option_last"](hidden[rows[:, None], ends])
+        return (query[:, None, :] * keys).sum(-1) / self.head_width ** 0.5
+
+    def forward(self, input_ids, attention_mask, criterion_positions, criterion_mask, decision_positions, kinds,
+                context_ends, question_spans, criterion_spans, criterion_order, targets=None,
+                pixel_values=None, image_grid_thw=None, mm_token_type_ids=None):
+        dtype = torch.bfloat16 if supports_native_bfloat16() else torch.float16
+        images = {} if pixel_values is None else {"pixel_values": pixel_values, "image_grid_thw": image_grid_thw,
+                                                  "mm_token_type_ids": mm_token_type_ids}
+        with torch.autocast("cuda", dtype=dtype, enabled=input_ids.is_cuda):
+            hidden = self.backbone(input_ids=input_ids, attention_mask=attention_mask, use_cache=False, return_dict=True,
+                                   **images).last_hidden_state
+        if self.head == "bilinear":
+            logits = self.bilinear_logits(hidden, context_ends, question_spans, criterion_spans)
+        else:
+            logits = self.mlp_logits(hidden, criterion_positions, decision_positions, kinds)
+        logits = torch.zeros_like(logits).scatter(1, criterion_order, logits)   # layout slots -> caller's order
+        logits = (logits / self.temperatures[kinds, None]).masked_fill(~criterion_mask, torch.finfo(torch.float32).min)
+        result = {"logits": logits}
+        if targets is not None:
+            result["loss"] = distribution_loss(logits, targets, kinds, criterion_mask, self.ordinal_weight, self.brier_weight)
+        return result
+
+    def mlp_logits(self, hidden, criterion_positions, decision_positions, kinds):
+        rows = torch.arange(hidden.shape[0], device=hidden.device)
+        question = hidden[rows, decision_positions].float()
+        options = hidden[rows[:, None], criterion_positions].float()
+        shared = question[:, None, :].expand_as(options)
+        features = torch.cat([shared, options, shared * options], dim=-1)
+        choice_logits = self.heads["choice"](features).squeeze(-1)
+        score_logits = self.heads["score"](features).squeeze(-1)
+        logits = torch.where((kinds == 1)[:, None], score_logits, choice_logits)
+        binary_logits = torch.zeros_like(logits)
+        binary_logits[:, 1] = self.heads["noul"](question).squeeze(-1)
+        return torch.where((kinds == 2)[:, None], binary_logits, logits)
+
+
+def typed_answer(kind, probabilities, labels):
+    if len(probabilities) != len(labels) or not probabilities:
+        raise ValueError("One label is required for each probability")
+    if not all(math.isfinite(value) and value >= 0 for value in probabilities) or not math.isclose(sum(probabilities), 1, abs_tol=1e-5):
+        raise ValueError("Expected a normalized finite probability vector")
+    count = len(probabilities)
+    if kind == "noul":
+        if count != 2:
+            raise ValueError("Noul requires false and true outcomes")
+        return {"type": "noul", "noul": probabilities[1]}
+    mode = max(range(count), key=probabilities.__getitem__)
+    if kind == "choice":
+        confidence = 1.0 if count == 1 else (probabilities[mode] - 1 / count) / (1 - 1 / count)
+        return {"type": kind, "choice": labels[mode], "probabilities": dict(zip(labels, probabilities)), "confidence": confidence}
+    if kind != "score" or count < 2:
+        raise ValueError("Expected Choice, Score or Noul")
+    distance = sum(value * abs(index - mode) for index, value in enumerate(probabilities))
+    uniform = sum(abs(index - (count - 1) / 2) for index in range(count)) / count
+    return {
+        "type": kind,
+        "score": sum(index * value for index, value in enumerate(probabilities)),
+        "probabilities": {str(index): value for index, value in enumerate(probabilities)},
+        "legend": {str(index): label for index, label in enumerate(labels)},
+        "confidence": max(0.0, 1 - distance / uniform),
+    }
+
+
+@torch.no_grad()
+def predict(model, tokenizer, kind, state, instructions, criteria, max_length=512):
+    if isinstance(criteria, dict):
+        labels = list(criteria)
+        descriptions = [f"{label}: {description}" if description else label for label, description in criteria.items()]
+    else:
+        descriptions = list(criteria)
+        labels = descriptions
+    if kind == "choice" and len(labels) == 1:
+        return typed_answer(kind, [1.0], labels)
+    example = DecisionExample("inference", "inference", "user", "test", kind, state, instructions, descriptions, [1 / len(labels)] * len(labels))
+    batch = collator_for(model, tokenizer, max_length)([example])
+    device = next(model.heads.parameters()).device
+    batch = {name: value.to(device) for name, value in batch.items()}
+    model.eval()
+    probabilities = model(**batch)["logits"].softmax(-1)[0].tolist()
+    return typed_answer(kind, probabilities, labels)

src/yev/training/run.py

@@ -0,0 +1,782 @@
+from contextlib import contextmanager
+from dataclasses import asdict, dataclass
+import hashlib
+import importlib.util
+from importlib.metadata import PackageNotFoundError, version
+import json
+import math
+import os
+from pathlib import Path
+import random
+import tempfile
+from time import perf_counter
+
+import numpy as np
+import torch
+
+from .images import DEFAULT_MAX_IMAGE_TOKENS, IMAGE_PLACEHOLDER, MIN_IMAGE_TOKENS, ImagePrefetcher, VisionInputs, image_processor
+from .model import HEADS, LAST_TOKEN_HEADS, DecisionModel, PRIMITIVES, collator_for, encode_example, supports_native_bfloat16
+
+
+STAGES = ("head_warmup", "decision_adaptation", "distribution_adaptation")
+LORA_TARGETS = ("q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj")
+# Gated DeltaNet (linear-attention) projections in hybrid backbones such as Qwen3.5. The vision
+# tower uses other names (qkv, proj, linear_fc*), so these adapt the language model only.
+LINEAR_ATTENTION_TARGETS = ("in_proj_qkv", "in_proj_z", "in_proj_a", "in_proj_b", "out_proj")
+
+
+def lora_targets(backbone):
+    """Attention and MLP projections, plus linear-attention projections when the backbone has them."""
+    names = {name.rsplit(".", 1)[-1] for name, _ in backbone.named_modules()}
+    return list(LORA_TARGETS) + [name for name in LINEAR_ATTENTION_TARGETS if name in names]
+
+
+@dataclass
+class TrainingConfig:
+    model_id: str = "Qwen/Qwen3-4B-Base"
+    model_revision: str | None = None
+    output_dir: str = "artifacts/training"
+    max_length: int = 512
+    micro_batch: int = 1
+    gradient_accumulation: int = 16
+    seed: int = 42
+    head_width: int = 128
+    head: str = "mlp"
+    head_last_token: bool = False   # the bilinear head's last-token readout (model.LAST_TOKEN_HEADS)
+    brier_weight: float = 0.0
+    lora_rank: int = 16
+    eval_limit: int = 128
+    eval_every: int = 50
+    save_every: int = 50
+    quantize4bit: bool = True
+    gradient_checkpointing: bool = True
+    stage_steps: tuple[int, int, int] = (25, 200, 100)
+    stream_start: int = 0   # distribution-stage updates of the data stream to skip before the first one
+    human_eval_limit: int = 0   # human-labelled dev examples per primitive scored to select best_human.pt; 0 is off
+    max_image_tokens: int = DEFAULT_MAX_IMAGE_TOKENS   # tokens per image, outside max_length (vision backbones only)
+    snapshot_every: int = 0   # also keep the adapters and heads every this many updates under snapshots/; 0 is off
+    dataset_fingerprint: str | None = None
+    head_learning_rate: float = 5e-4
+    learning_rate: float = 1e-4
+    weight_decay: float = 0.01
+    max_grad_norm: float = 1.0
+
+    def __post_init__(self):
+        self.output_dir = str(self.output_dir)
+        self.stage_steps = tuple(self.stage_steps)
+        for name in ("max_length", "micro_batch", "gradient_accumulation", "head_width",
+                     "lora_rank", "eval_limit", "eval_every", "save_every"):
+            if type(getattr(self, name)) is not int or getattr(self, name) < 1:
+                raise ValueError(f"{name} must be a positive integer")
+        if len(self.stage_steps) != 3 or any(type(steps) is not int or steps < 0 for steps in self.stage_steps):
+            raise ValueError("stage_steps must contain three nonnegative optimizer-step limits")
+        for name in ("head_learning_rate", "learning_rate", "max_grad_norm"):
+            if not math.isfinite(getattr(self, name)) or getattr(self, name) <= 0:
+                raise ValueError(f"{name} must be finite and positive")
+        if not math.isfinite(self.weight_decay) or self.weight_decay < 0:
+            raise ValueError("weight_decay must be finite and nonnegative")
+        if self.head not in HEADS:
+            raise ValueError(f"head must be one of {HEADS}")
+        if self.head_last_token and self.head != "bilinear":
+            raise ValueError("head_last_token extends the bilinear head")
+        if not math.isfinite(self.brier_weight) or self.brier_weight < 0:
+            raise ValueError("brier_weight must be finite and nonnegative")
+        if type(self.seed) is not int or not 0 <= self.seed < 2**32:
+            raise ValueError("seed must be an integer in [0, 2**32)")
+        if type(self.stream_start) is not int or self.stream_start < 0:
+            raise ValueError("stream_start must be a nonnegative integer")
+        if type(self.human_eval_limit) is not int or self.human_eval_limit < 0:
+            raise ValueError("human_eval_limit must be a nonnegative integer")
+        if type(self.snapshot_every) is not int or self.snapshot_every < 0:
+            raise ValueError("snapshot_every must be a nonnegative integer")
+        if type(self.max_image_tokens) is not int or self.max_image_tokens < MIN_IMAGE_TOKENS:
+            raise ValueError(f"max_image_tokens must be an integer of at least {MIN_IMAGE_TOKENS}")
+
+
+def seed_rng(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def select_device():
+    """CUDA when present, then Apple MPS, otherwise CPU."""
+    if torch.cuda.is_available():
+        return torch.device("cuda", torch.cuda.current_device())
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def compute_dtype(device_type):
+    if device_type == "mps":
+        return torch.bfloat16
+    if device_type != "cuda":
+        return torch.float32
+    return torch.bfloat16 if supports_native_bfloat16() else torch.float16
+
+
+def prepare_backbone(backbone, config, trainable):
+    checkpointing = trainable and config.gradient_checkpointing
+    if config.quantize4bit and trainable:   # for training only: it casts every bf16 weight left unquantized to fp32
+        from peft import prepare_model_for_kbit_training
+
+        return prepare_model_for_kbit_training(
+            backbone, use_gradient_checkpointing=checkpointing,
+            gradient_checkpointing_kwargs={"use_reentrant": False},
+        )
+    if checkpointing:
+        backbone.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+    return backbone
+
+
+def load_model(config, adapter_path=None, trainable=True, processor_source=None):
+    """Load a revision-pinned feature backbone; quantized loads require CUDA. A backbone with a
+    vision encoder also gets ``model.vision``, its image processor read from ``processor_source``
+    (an export) or the base model."""
+    if config.quantize4bit and trainable and not torch.cuda.is_available():   # serving in 4-bit also runs on CPU and Apple GPUs
+        raise RuntimeError("4-bit training requires CUDA; select a Colab GPU before loading weights")
+    from peft import LoraConfig, PeftModel, TaskType, get_peft_model
+    from transformers import AutoConfig, AutoModel, AutoTokenizer, BitsAndBytesConfig
+
+    seed_rng(config.seed)
+    device = select_device()
+    cuda = device.type == "cuda"
+    dtype = compute_dtype(device.type)
+    if cuda:
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+    base_config = AutoConfig.from_pretrained(config.model_id, revision=config.model_revision)
+    config.model_revision = getattr(base_config, "_commit_hash", None) or config.model_revision
+    if config.model_revision is None and not Path(config.model_id).is_dir():
+        raise ValueError("Could not resolve the base model revision")
+    base_config.use_cache = False
+    tokenizer = AutoTokenizer.from_pretrained(config.model_id, revision=config.model_revision)
+    if tokenizer.pad_token_id is None:
+        if tokenizer.eos_token_id is None:
+            raise ValueError("Tokenizer needs a padding or EOS token")
+        tokenizer.pad_token = tokenizer.eos_token
+    options = {"config": base_config, "revision": config.model_revision, "dtype": dtype}
+    if cuda:
+        options["device_map"] = {"": device}
+    if config.quantize4bit:
+        options["quantization_config"] = BitsAndBytesConfig(
+            load_in_4bit=True, bnb_4bit_quant_type="nf4",
+            bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=dtype,
+            llm_int8_skip_modules=["visual"],   # the vision encoder has no adapters to learn around its rounding
+        )
+    backbone = AutoModel.from_pretrained(config.model_id, **options)
+    if not cuda and device.type != "cpu":
+        backbone = backbone.to(device)   # device_map loading onto Apple MPS segfaults under transformers 5
+    backbone = prepare_backbone(backbone, config, trainable)
+    if adapter_path is not None:
+        backbone = PeftModel.from_pretrained(backbone, str(adapter_path), is_trainable=trainable)
+    else:
+        backbone = get_peft_model(backbone, LoraConfig(
+            task_type=TaskType.FEATURE_EXTRACTION, r=config.lora_rank,
+            lora_alpha=2 * config.lora_rank, lora_dropout=0.05,
+            target_modules=lora_targets(backbone), inference_mode=not trainable,
+        ))
+    model = DecisionModel(backbone, head_width=config.head_width, head=config.head, brier_weight=config.brier_weight,
+                          last_token=config.head_last_token)
+    device = backbone.get_input_embeddings().weight.device
+    model.heads.to(device=device, dtype=torch.float32)
+    model.temperatures = model.temperatures.to(device)
+    model.training_config = config
+    if (getattr(base_config, "vision_config", None) is not None and getattr(base_config, "image_token_id", None) is not None
+            and importlib.util.find_spec("PIL") is not None):   # without Pillow the model still answers text
+        processor = image_processor(processor_source or config.model_id, config.max_image_tokens,
+                                    None if processor_source else config.model_revision)
+        model.vision = VisionInputs.for_backbone(base_config, processor, config.max_image_tokens)
+    if not trainable:
+        model.requires_grad_(False)
+    model.train(trainable)
+    return model, tokenizer
+
+
+def forward_batch(model, batch):
+    return model(**batch)
+
+
+def device_batch(model, collator, examples):
+    device = next(model.heads.parameters()).device
+    return {name: value.to(device) for name, value in collator(examples).items()}
+
+
+@contextmanager
+def evaluation_mode(model):
+    modes = [(module, module.training) for module in model.modules()]
+    model.eval()
+    try:
+        with torch.no_grad():
+            yield
+    finally:
+        for module, mode in modes:
+            module.training = mode
+
+
+def one_hot(target):
+    return sum(value == 1 for value in target) == 1 and all(value in (0, 1) for value in target)
+
+
+def summarize_metrics(rows):
+    hard = [row[2] for row in rows if row[2] is not None]
+    return {
+        "count": len(rows),
+        "soft_log_loss": sum(row[0] for row in rows) / len(rows) if rows else None,
+        "brier": sum(row[1] for row in rows) / len(rows) if rows else None,
+        "hard_count": len(hard),
+        "hard_accuracy": sum(hard) / len(hard) if hard else None,
+    }
+
+
+def select_per_primitive(examples, limit):
+    """The first ``limit`` examples of each primitive, in their original order."""
+    if type(limit) is not int or limit < 1:
+        raise ValueError("Evaluation limit must be a positive integer")
+    return [item for kind in PRIMITIVES for item in [item for item in examples if item.kind == kind][:limit]]
+
+
+def human_dev_sample(dev, limit):
+    """``limit`` human-labelled dev examples per primitive, spread over their datasets by a stable hash of the uid."""
+    human = sorted((item for item in dev if item.source.startswith("human:")),
+                   key=lambda item: hashlib.sha256(item.uid.encode("utf-8")).hexdigest())
+    return select_per_primitive(human, limit) if human else []
+
+
+def scored(items, batch, logits):
+    """(kind, (log loss, Brier, correct or None)) for each question of a scored micro-batch."""
+    losses = -(batch["targets"] * logits.log_softmax(-1)).sum(-1)
+    briers = (logits.softmax(-1) - batch["targets"]).square().sum(-1)
+    predictions = logits.argmax(-1).tolist()
+    for index, item in enumerate(items):
+        loss, brier = float(losses[index]), float(briers[index])
+        if not math.isfinite(loss) or not math.isfinite(brier):
+            raise FloatingPointError("Nonfinite evaluation scores")
+        yield item.kind, (loss, brier, int(predictions[index] == item.target.index(1)) if one_hot(item.target) else None)
+
+
+def eval_model(model, tokenizer, examples, config, limit=None):
+    """Mean proper scores and one-hot-only accuracy, overall and per primitive.
+
+    ``limit`` caps the sample **per primitive**, so a split dominated by one primitive still
+    yields a meaningful Score or Noul estimate.
+    """
+    selected = select_per_primitive(examples, config.eval_limit if limit is None else limit)
+    collator = collator_for(model, tokenizer, config.max_length)
+    rows = {kind: [] for kind in PRIMITIVES}
+    groups = micro_batches(selected, range(len(selected)), config.micro_batch, config.micro_batch * config.max_length,
+                           token_lengths(collator))
+    with evaluation_mode(model):
+        for group in groups:
+            items = [selected[index] for index in group]
+            batch = device_batch(model, collator, items)
+            for kind, row in scored(items, batch, forward_batch(model, batch)["logits"].float()):
+                rows[kind].append(row)
+    result = summarize_metrics([row for values in rows.values() for row in values])
+    result["per_primitive"] = {kind: summarize_metrics(values) for kind, values in rows.items()}
+    return result
+
+
+def best_temperature(rows, candidates):
+    """The candidate with the lowest multi-class Brier score over (logits, target) rows."""
+    losses = [sum(float(((logits / candidate).softmax(-1) - target).square().sum()) for logits, target in rows)
+              for candidate in candidates]
+    if not all(math.isfinite(loss) for loss in losses):
+        raise FloatingPointError("Nonfinite calibration logits")
+    return float(candidates[min(range(len(losses)), key=losses.__getitem__)])
+
+
+def calibrate_model(model, tokenizer, examples, config, split="calibration", tied=()):
+    """Fit temperatures on at most eval_limit examples per primitive from one named split.
+
+    Each primitive's temperature minimises the multi-class Brier score of the tempered
+    distributions (RLCD's calibration step). The default fits on gold calibration data. Passing ``split="dev"`` with teacher-labelled
+    development examples fits temperatures for teacher imitation instead; the caller is
+    responsible for keeping those examples out of any later scoring sample. The primitives in ``tied`` share
+    one temperature fitted on their examples together, so a Noul and its two-option Choice stay equivalent.
+    """
+    if split not in ("calibration", "dev"):
+        raise ValueError("Temperatures may be fitted on the calibration split or on held-out dev examples only")
+    if not examples or any(item.split != split for item in examples):
+        raise ValueError(f"Temperature fitting requires {split} examples ONLY")
+    collator = collator_for(model, tokenizer, config.max_length)
+    collected = {kind: [] for kind in PRIMITIVES}
+    previous = model.temperatures.detach().clone()
+    try:
+        with evaluation_mode(model):
+            model.temperatures.fill_(1)
+            selected = select_per_primitive(examples, config.eval_limit)
+            for start in range(0, len(selected), config.micro_batch):
+                items = selected[start:start + config.micro_batch]
+                batch = device_batch(model, collator, items)
+                logits = forward_batch(model, batch)["logits"].cpu().double()
+                for index, item in enumerate(items):
+                    collected[item.kind].append((logits[index, :len(item.target)], torch.tensor(item.target, dtype=torch.float64)))
+            candidates = [1.0] + [2 ** (index / 8) for index in range(-24, 25) if index != 0]   # 0.125 to 8.0
+            temperatures = {}
+            groups = ([tuple(tied)] if tied else []) + [(kind,) for kind in PRIMITIVES if kind not in tied]
+            for group in groups:
+                temperatures.update(dict.fromkeys(group, best_temperature([row for kind in group for row in collected[kind]], candidates)))
+            model.temperatures.copy_(model.temperatures.new_tensor([temperatures[kind] for kind in PRIMITIVES]))
+    except BaseException:
+        model.temperatures.copy_(previous)
+        raise
+    return temperatures
+
+
+def fingerprint(value):
+    content = json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False)
+    return hashlib.sha256(content.encode()).hexdigest()
+
+
+def data_fingerprint(train, dev):
+    digest = hashlib.sha256()
+    for split, examples in (("train", train), ("dev", dev)):
+        digest.update(split.encode())
+        for item in examples:
+            digest.update(fingerprint(item.to_dict()).encode())
+    return digest.hexdigest()
+
+
+def config_fingerprint(config, model):
+    settings = asdict(config)
+    settings.pop("output_dir")
+    for name in ("stream_start", "human_eval_limit", "snapshot_every"):
+        if not settings[name]:   # runs from before these settings existed keep their fingerprint
+            settings.pop(name)
+    if settings["max_image_tokens"] == DEFAULT_MAX_IMAGE_TOKENS:
+        settings.pop("max_image_tokens")
+    settings["ordinal_weight"] = model.ordinal_weight
+    return fingerprint(settings)
+
+
+def stage_pool_size(train, stage):
+    """Presentations in one epoch of a stage's stream, mirroring ``stage_indices``."""
+    hard = sum(1 for item in train if one_hot(item.target))
+    soft = len(train) - hard
+    if stage < 2 or not hard or not soft:
+        return hard or soft
+    return 2 * max(hard, soft)
+
+
+def plan_stage_steps(train, config, head_steps=25, decision_epochs=1, distribution_epochs=1):
+    """Derive stage limits that correspond to whole passes over the label pools.
+
+    Stage 1 warms the frozen-backbone heads, stage 2 adapts on hard labels only, and
+    stage 3 replays the balanced hard/soft stream that ``stage_indices`` builds, so one
+    distribution epoch costs ``2 * max(hard, soft)`` presentations. Counters describe
+    presentations rather than unique examples, and the caller must still measure real
+    throughput: this is arithmetic, not a runtime prediction.
+    """
+    for name, value in (("head_steps", head_steps), ("decision_epochs", decision_epochs),
+                        ("distribution_epochs", distribution_epochs)):
+        if type(value) is not int or value < 0:
+            raise ValueError(f"{name} must be a nonnegative integer")
+    if not train:
+        raise ValueError("A nonempty training split is required to plan a schedule")
+    batch = config.micro_batch * config.gradient_accumulation
+    hard = sum(1 for item in train if one_hot(item.target))
+    soft = len(train) - hard
+    decision_updates = -(-stage_pool_size(train, 1) // batch)
+    distribution_updates = -(-stage_pool_size(train, 2) // batch)
+    return {
+        "stage_steps": (head_steps, decision_updates * decision_epochs, distribution_updates * distribution_epochs),
+        "effective_batch": batch,
+        "hard_examples": hard,
+        "soft_examples": soft,
+        "head_warmup_presentations": head_steps * batch,
+        "decision_epochs": decision_epochs,
+        "decision_updates_per_epoch": decision_updates,
+        "distribution_epochs": distribution_epochs,
+        "distribution_updates_per_epoch": distribution_updates,
+        "total_updates": head_steps + decision_updates * decision_epochs + distribution_updates * distribution_epochs,
+    }
+
+
+def stage_indices(train, stage, seed, epoch):
+    generator = random.Random(f"{seed}:{stage}:{epoch}")
+    hard = [index for index, item in enumerate(train) if one_hot(item.target)]
+    soft = [index for index, item in enumerate(train) if not one_hot(item.target)]
+    generator.shuffle(hard)
+    generator.shuffle(soft)
+    if stage < 2:
+        return hard or soft
+    if not hard or not soft:
+        return hard or soft
+    return [pool[index % len(pool)] for index in range(max(len(hard), len(soft))) for pool in (soft, hard)]
+
+
+def adapter_parameters(model):
+    return {name: parameter for name, parameter in model.named_parameters() if "lora_" in name}
+
+
+def checkpoint_tensors(model):
+    return {**adapter_parameters(model),
+            **{f"heads.{name}": value for name, value in model.heads.state_dict().items()},
+            "temperatures": model.temperatures}
+
+
+def tensor_state(model):
+    return {name: value.detach().cpu().clone() for name, value in checkpoint_tensors(model).items()}
+
+
+def with_new_heads(model, tensors):
+    """Checkpoint ``tensors`` completed with the model's zero-initialised last-token readout where
+    the checkpoint predates it, so a continuation can add the readout and start where it left off."""
+    added = {name: value for name, value in tensor_state(model).items()
+             if name not in tensors and name.startswith("heads.") and name.split(".")[1] in LAST_TOKEN_HEADS}
+    return {**tensors, **added}
+
+
+def restore_tensors(model, tensors):
+    expected = checkpoint_tensors(model)
+    if tensors.keys() != expected.keys() or any(tensors[name].shape != value.shape for name, value in expected.items()):
+        raise ValueError("Checkpoint adapter/head layout does not match the model")
+    with torch.no_grad():
+        for name, value in expected.items():
+            value.copy_(tensors[name])
+
+
+def capture_rng():
+    numpy_state = np.random.get_state()
+    return {
+        "python": random.getstate(), "torch": torch.get_rng_state(),
+        "numpy": (numpy_state[0], numpy_state[1].tolist(), *numpy_state[2:]),
+        "cuda": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else [],
+    }
+
+
+def restore_rng(states):
+    random.setstate(states["python"])
+    torch.set_rng_state(states["torch"])
+    numpy_state = states["numpy"]
+    np.random.set_state((numpy_state[0], np.array(numpy_state[1], dtype=np.uint32), *numpy_state[2:]))
+    if states["cuda"]:
+        if not torch.cuda.is_available() or len(states["cuda"]) != torch.cuda.device_count():
+            raise ValueError("Resume requires the same CUDA device count")
+        torch.cuda.set_rng_state_all(states["cuda"])
+
+
+def atomic_save(payload, path):
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = None
+    try:
+        with tempfile.NamedTemporaryFile(dir=path.parent, prefix=f".{path.name}.", delete=False) as stream:
+            temporary = Path(stream.name)
+            torch.save(payload, stream)
+            stream.flush()
+            os.fsync(stream.fileno())
+        os.replace(temporary, path)
+    finally:
+        if temporary is not None:
+            temporary.unlink(missing_ok=True)
+
+
+def stage_optimizer(model, config, steps, stage=0):
+    head_rate = config.head_learning_rate * (1.0, 0.6, 0.2)[stage]
+    adapter_rate = config.learning_rate * (1.0, 1.0, 0.3)[stage]
+    groups = [
+        {"params": [parameter for name, parameter in model.heads.named_parameters() if not name.endswith("bias")], "lr": head_rate},
+        {"params": [parameter for name, parameter in model.heads.named_parameters() if name.endswith("bias")], "lr": head_rate, "weight_decay": 0.0},
+    ]
+    adapters = list(adapter_parameters(model).values())
+    if adapters:
+        groups.append({"params": adapters, "lr": adapter_rate})
+    optimizer = torch.optim.AdamW(groups, lr=config.learning_rate, weight_decay=config.weight_decay)
+    warmup = max(1, math.ceil(0.03 * steps))
+
+    def schedule(step):
+        if step < warmup:
+            return (step + 1) / warmup
+        progress = min(1.0, (step - warmup) / max(1, steps - warmup))
+        return 0.5 * (1 + math.cos(math.pi * progress))
+
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, schedule)
+    return optimizer, scheduler
+
+
+def validate_training(model, splits, config):
+    train, dev = splits["train"], splits["dev"]
+    if not train or not dev:
+        raise ValueError("Nonempty train and dev splits are required")
+    if any(item.split != "train" for item in train) or any(item.split != "dev" for item in dev):
+        raise ValueError("Training and dev selection must use their own splits")
+    if {item.group for item in train} & {item.group for item in dev}:
+        raise ValueError("Train and dev document groups must be disjoint")
+    actual_fingerprint = data_fingerprint(train, dev)
+    if config.dataset_fingerprint not in (None, actual_fingerprint):
+        raise ValueError("Dataset fingerprint mismatch")
+    config.dataset_fingerprint = actual_fingerprint
+    if model.head_width != config.head_width:
+        raise ValueError("Config head_width does not match the model")
+    if model.head != config.head or model.brier_weight != config.brier_weight:
+        raise ValueError("Config head and brier_weight must match the model")
+    if model.last_token != config.head_last_token:
+        raise ValueError("Config head_last_token must match the model")
+    if any(config.stage_steps[1:]) and not adapter_parameters(model):
+        raise ValueError("Adaptation stages require a LoRA backbone")
+    return train, dev
+
+
+def token_lengths(collator):
+    """Each text question's encoded length, counted once."""
+    lengths = {}
+
+    def length(example):
+        if example.uid not in lengths:
+            lengths[example.uid] = len(encode_example(example, collator.tokenizer, collator.max_length, collator.canonical)["input_ids"])
+        return lengths[example.uid]
+    return length
+
+
+def micro_batches(train, indices, micro_batch, budget, length):
+    """An update's micro-batches. Image questions go in fixed groups of ``micro_batch``, the groups the
+    image prefetcher prepares; text questions are sorted by length and packed while a group's padded
+    size (its longest question times its count) stays within ``budget`` tokens, so short questions no
+    longer run one or two to a forward pass. The update's gradient is the same however it is split."""
+    pictures = [index for index in indices if train[index].images]
+    groups = [pictures[start:start + micro_batch] for start in range(0, len(pictures), micro_batch)]
+    group = []
+    for index in sorted((index for index in indices if not train[index].images), key=lambda index: length(train[index])):
+        if group and length(train[index]) * (len(group) + 1) > budget:
+            groups.append(group)
+            group = []
+        group.append(index)
+    return groups + ([group] if group else [])
+
+
+def optimizer_update(model, collator, train, indices, config, optimizer, scaler, length):
+    optimizer.zero_grad(set_to_none=True)
+    update_loss = 0.0
+    for micro_indices in micro_batches(train, indices, config.micro_batch, config.micro_batch * config.max_length, length):
+        batch = device_batch(model, collator, [train[index] for index in micro_indices])
+        loss = forward_batch(model, batch)["loss"]
+        if not torch.isfinite(loss):
+            raise FloatingPointError("Nonfinite training loss; last.pt remains at its saved optimizer boundary")
+        weighted_loss = loss * (len(micro_indices) / len(indices))
+        scaler.scale(weighted_loss).backward()
+        update_loss += float(weighted_loss.detach())
+    scaler.unscale_(optimizer)
+    torch.nn.utils.clip_grad_norm_(
+        [parameter for parameter in model.parameters() if parameter.requires_grad],
+        config.max_grad_norm, error_if_nonfinite=not scaler.is_enabled(),
+    )
+    previous_scale = scaler.get_scale()
+    scaler.step(optimizer)
+    scaler.update()
+    optimizer.zero_grad(set_to_none=True)
+    return None if scaler.get_scale() < previous_scale else update_loss
+
+
+def run_training(model, tokenizer, splits, config, resume=True):
+    """Bounded optimizer steps; last.pt resumes only through its last saved update.
+
+    Warmup and decision adaptation use hard labels when available. Distribution
+    adaptation balances soft targets and hard replay, repeating the smaller pool.
+    Each stage has a fresh AdamW optimizer, 3% warmup and cosine decay. Changed schedules
+    are refused on resume. ``stream_start`` begins distribution adaptation that many updates
+    into its data stream, so continuing a run on the same data and seed skips what it trained on.
+    With ``human_eval_limit`` each evaluation also scores a fixed sample of human-labelled dev
+    examples and keeps best_human.pt, the update with their lowest log loss. best.pt is selected by dev soft log loss, while the
+    returned model remains at the last update, not automatically at the best.
+    """
+    train, dev = validate_training(model, splits, config)
+    actual_fingerprint = config.dataset_fingerprint
+    adapters = adapter_parameters(model)
+    device = next(model.heads.parameters()).device
+    model.heads.to(dtype=torch.float32)
+    model.training_config = config
+    signature = config_fingerprint(config, model)
+    directory = Path(config.output_dir)
+    directory.mkdir(parents=True, exist_ok=True)
+    last_path = directory / "last.pt"
+    saved = torch.load(last_path, map_location="cpu", weights_only=True) if resume and last_path.exists() else None
+    if saved is not None:
+        if saved["format_version"] != 1 or saved["config_fingerprint"] != signature:
+            raise ValueError("Checkpoint config/schedule fingerprint mismatch; restart explicitly with resume=False")
+        if saved["data_fingerprint"] != actual_fingerprint:
+            raise ValueError("Checkpoint dataset fingerprint mismatch")
+        if saved["device_type"] != device.type:
+            raise ValueError("Exact continuation requires the same device type")
+        restore_tensors(model, saved["tensors"])
+        history, best, best_human = saved["history"], saved["best"], saved.get("best_human")
+        history.setdefault("best_human_dev_log_loss", None)
+        history.setdefault("best_human_step", None)
+        if best is not None:
+            atomic_save(best, directory / "best.pt")
+        if best_human is not None:
+            atomic_save(best_human, directory / "best_human.pt")
+    else:
+        seed_rng(config.seed)
+        model.temperatures.fill_(1)
+        history = {"completed_steps": 0, "best_dev_log_loss": None, "best_step": None,
+                   "best_human_dev_log_loss": None, "best_human_step": None, "updates": []}
+        best = best_human = None
+    human_dev = human_dev_sample(dev, config.human_eval_limit) if config.human_eval_limit else []
+    # Image questions are prepared one update ahead in background threads, so decoding and resizing
+    # overlap the GPU's work on the current update; text-only training is unchanged.
+    prefetcher = (ImagePrefetcher(model.vision) if getattr(model, "vision", None) is not None
+                  and any(example.images for example in train) else None)
+    collator = collator_for(model, tokenizer, config.max_length, prefetcher)
+    length = token_lengths(collator)
+    update_size = config.micro_batch * config.gradient_accumulation
+    global_step = history["completed_steps"]
+    start_stage = saved["stage"] if saved is not None else 0
+    for stage in range(start_stage, len(STAGES)):
+        steps = config.stage_steps[stage]
+        if steps == 0:
+            continue
+        for name, parameter in model.named_parameters():
+            parameter.requires_grad_(name.startswith("heads.") or (stage > 0 and name in adapters))
+        model.train()
+        if stage == 0:
+            model.backbone.eval()
+        print(f"{STAGES[stage]}: {steps} optimizer steps", flush=True)
+        optimizer, scheduler = stage_optimizer(model, config, steps, stage)
+        scaler = torch.amp.GradScaler("cuda", enabled=device.type == "cuda" and not supports_native_bfloat16())
+        stage_step, epoch, offset = 0, 0, 0
+        resumed = saved is not None and stage == saved["stage"]
+        if resumed:
+            optimizer.load_state_dict(saved["optimizer"])
+            scheduler.load_state_dict(saved["scheduler"])
+            scaler.load_state_dict(saved["scaler"])
+            stage_step, epoch, offset = saved["stage_step"], saved["epoch"], saved["offset"]
+            restore_rng(saved["rng"])
+            saved = None
+        order = stage_indices(train, stage, config.seed, epoch)
+        if stage == 2 and not resumed:
+            for _ in range(config.stream_start):   # advance as that many updates did, short epoch-end batches included
+                offset += min(config.micro_batch * config.gradient_accumulation, len(order) - offset)
+                if offset == len(order):
+                    epoch, offset = epoch + 1, 0
+                    order = stage_indices(train, stage, config.seed, epoch)
+        overflows = 0
+        stage_clock, stage_start_step = perf_counter(), stage_step
+        while stage_step < steps:
+            indices = order[offset:offset + update_size]
+            if prefetcher is not None:
+                following = offset + len(indices)
+                upcoming = (order[following:following + update_size] if following < len(order)
+                            else stage_indices(train, stage, config.seed, epoch + 1)[:update_size])
+                prefetcher.submit_update(train, indices, config.micro_batch)
+                prefetcher.submit_update(train, upcoming, config.micro_batch)
+            update_loss = optimizer_update(model, collator, train, indices, config, optimizer, scaler, length)
+            if update_loss is None:
+                overflows += 1
+                if overflows > 8:
+                    raise FloatingPointError("Repeated fp16 gradient overflow; resume from last.pt")
+                continue
+            overflows = 0
+            scheduler.step()
+            stage_step += 1
+            global_step += 1
+            offset += len(indices)
+            if offset == len(order):
+                epoch += 1
+                offset = 0
+                order = stage_indices(train, stage, config.seed, epoch)
+            record = {"stage": STAGES[stage], "stage_step": stage_step, "step": global_step,
+                      "loss": update_loss, "examples": len(indices)}
+            history["completed_steps"] = global_step
+            history["updates"].append(record)
+            if stage_step % config.eval_every == 0 or stage_step == steps:
+                metrics = eval_model(model, tokenizer, dev, config)
+                record["dev"] = metrics
+                progress = {"stage": STAGES[stage], "step": global_step, "train_loss": update_loss, "dev_log_loss": metrics["soft_log_loss"]}
+                human = eval_model(model, tokenizer, human_dev, config, limit=config.human_eval_limit) if human_dev else None
+                if human is not None:
+                    record["human_dev"] = human
+                    progress.update(human_dev_log_loss=human["soft_log_loss"], human_dev_accuracy=human["hard_accuracy"])
+                progress["seconds_per_update"] = round((perf_counter() - stage_clock) / max(1, stage_step - stage_start_step), 3)
+                if device.type == "cuda":
+                    progress["max_memory_gib"] = round(torch.cuda.max_memory_allocated(device) / 2**30, 2)
+                print(json.dumps(progress), flush=True)
+                if history["best_dev_log_loss"] is None or metrics["soft_log_loss"] < history["best_dev_log_loss"]:
+                    history["best_dev_log_loss"] = metrics["soft_log_loss"]
+                    history["best_step"] = global_step
+                    best = {"format_version": 1, "tensors": tensor_state(model), "config": asdict(config),
+                            "config_fingerprint": signature, "data_fingerprint": actual_fingerprint,
+                            "step": global_step, "dev": metrics}
+                    atomic_save(best, directory / "best.pt")
+                if human is not None and (history["best_human_dev_log_loss"] is None
+                                          or human["soft_log_loss"] < history["best_human_dev_log_loss"]):
+                    history["best_human_dev_log_loss"] = human["soft_log_loss"]
+                    history["best_human_step"] = global_step
+                    best_human = {"format_version": 1, "tensors": tensor_state(model), "config": asdict(config),
+                                  "config_fingerprint": signature, "data_fingerprint": actual_fingerprint,
+                                  "step": global_step, "dev": metrics, "human_dev": human}
+                    atomic_save(best_human, directory / "best_human.pt")
+            if config.snapshot_every and global_step % config.snapshot_every == 0:
+                # Same layout as best.pt, so any snapshot can be exported, and snapshots can be blended
+                atomic_save({"format_version": 1, "tensors": tensor_state(model), "config": asdict(config),
+                             "config_fingerprint": signature, "data_fingerprint": actual_fingerprint,
+                             "step": global_step}, directory / "snapshots" / f"update-{global_step:06d}.pt")
+            if stage_step % config.save_every == 0 or stage_step == steps:
+                checkpoint = {
+                    "format_version": 1, "tensors": tensor_state(model), "config": asdict(config),
+                    "config_fingerprint": signature, "data_fingerprint": actual_fingerprint,
+                    "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict(),
+                    "scaler": scaler.state_dict(), "stage": stage, "stage_step": stage_step,
+                    "global_step": global_step, "epoch": epoch, "offset": offset,
+                    "rng": capture_rng(), "device_type": device.type, "history": history, "best": best,
+                    "best_human": best_human,
+                }
+                atomic_save(checkpoint, last_path)
+    if prefetcher is not None:
+        prefetcher.close()
+    return history
+
+
+def export_model(model, tokenizer, config, destination):
+    """Export only PEFT weights, float32 heads, temperatures, tokenizer and provenance."""
+    from peft import PeftModel
+
+    if not isinstance(model.backbone, PeftModel):
+        raise ValueError("Export requires a PEFT backbone, not full frozen model weights")
+    if config.head_width != model.head_width:
+        raise ValueError("Config head_width does not match the exported heads")
+    destination = Path(destination)
+    destination.mkdir(parents=True, exist_ok=True)
+    model.backbone.save_pretrained(destination / "adapter", safe_serialization=True, save_embedding_layers=False)
+    tokenizer.save_pretrained(destination)
+    atomic_save({"heads": {name: value.detach().float().cpu() for name, value in model.heads.state_dict().items()},
+                 "temperatures": model.temperatures.detach().float().cpu(),
+                 "ordinal_weight": model.ordinal_weight}, destination / "heads.pt")
+    dependencies = {}
+    for package in ("torch", "transformers", "peft", "accelerate", "bitsandbytes", "numpy"):
+        try:
+            dependencies[package] = version(package)
+        except PackageNotFoundError:
+            dependencies[package] = None
+    vision = getattr(model, "vision", None)
+    if vision is not None:   # export format 2: the image processor and its token budget travel with the weights
+        vision.processor.save_pretrained(destination)
+    metadata = {"format_version": 2 if vision is not None else 1, "base_model": config.model_id, "base_revision": config.model_revision,
+                "head": model.head, "head_width": model.head_width, "head_last_token": model.last_token,
+                "brier_weight": model.brier_weight,
+                "max_length": config.max_length,
+                "dataset_fingerprint": config.dataset_fingerprint, "dependencies": dependencies,
+                "temperatures": dict(zip(PRIMITIVES, model.temperatures.tolist()))}
+    if vision is not None:
+        metadata["images"] = {"max_image_tokens": vision.max_image_tokens, "placeholder": IMAGE_PLACEHOLDER}
+    (destination / "config.json").write_text(json.dumps(asdict(config), indent=2, allow_nan=False) + "\n")
+    (destination / "metadata.json").write_text(json.dumps(metadata, indent=2, allow_nan=False) + "\n")
+
+
+def reload_export(destination, quantize4bit=True):
+    """Reload inference weights; saved config is available as model.training_config."""
+    from transformers import AutoTokenizer
+
+    destination = Path(destination)
+    config = TrainingConfig(**json.loads((destination / "config.json").read_text()))
+    config.quantize4bit = quantize4bit
+    processor = destination if (destination / "preprocessor_config.json").exists() else None
+    model, tokenizer = load_model(config, adapter_path=destination / "adapter", trainable=False, processor_source=processor)
+    tokenizer = AutoTokenizer.from_pretrained(destination)
+    payload = torch.load(destination / "heads.pt", map_location="cpu", weights_only=True)
+    model.heads.load_state_dict(payload["heads"])
+    model.temperatures.copy_(payload["temperatures"].to(model.temperatures.device))
+    model.ordinal_weight = payload["ordinal_weight"]
+    return model, tokenizer