sharada 0.1.0__py3-none-any.whl

sharada/__init__.py

@@ -0,0 +1,31 @@
+"""Sharada — typed decisions from text in one forward pass.
+
+    from sharada import DecisionModel
+
+    model = DecisionModel.from_pretrained("lenabarretta/sharada-base")
+    d = model.decide("My card hasn't arrived yet",
+                     "Which team should handle this?",
+                     ["billing", "technical", "sales"])
+    d.answer, d.confidence
+
+The model reads the text, the question and every option in one pass and returns a probability for each
+option. Every option is a branch of the sequence starting at the same position, so their order cannot
+change the answer; an option reads the text, the question and itself, so its score does not depend on
+which other options are offered.
+"""
+
+from .calibrate import calibrate, check_passport, fit_temperature, save_passport
+from .data import Example, Request
+from .evaluate import calibration_error, evaluate, latency, risk_coverage
+from .layout import Limits
+from .model import Config, Decision, DecisionModel
+from .policy import Action, Policy, escalation_budget
+from .train import fit
+
+__all__ = [
+    "DecisionModel", "Config", "Decision", "Request", "Example", "Limits",
+    "fit", "calibrate", "fit_temperature", "save_passport", "check_passport",
+    "evaluate", "calibration_error", "risk_coverage", "latency",
+    "Policy", "Action", "escalation_budget",
+]
+__version__ = "0.1.0"

sharada/calibrate.py

@@ -0,0 +1,115 @@
+"""Calibration, and the passport that says where it holds.
+
+A model is not calibrated in general; it is calibrated on a distribution, and it goes stale when the
+traffic moves. So fitting a temperature also writes a passport: what it was fitted on, how far off the
+probabilities were, and what should make you fit it again.
+"""
+
+from __future__ import annotations
+
+import datetime as dt
+import json
+import math
+import pathlib
+
+from .data import Example
+
+VALID_DAYS = 90
+
+
+def _negative_log_likelihood(log_probabilities: list[list[float]], labels: list[int], temperature: float) -> float:
+    total = 0.0
+    for row, label in zip(log_probabilities, labels):
+        scaled = [v / temperature for v in row]
+        top = max(scaled)
+        total += -(scaled[label] - (top + math.log(sum(math.exp(v - top) for v in scaled))))
+    return total / len(labels)
+
+
+def fit_temperature(probabilities: list[list[float]], labels: list[int],
+                    grid: tuple[float, float, int] = (0.25, 8.0, 61)) -> float:
+    """One number per task: divide the scores by it. Chosen to minimise the log loss on held-out
+    answers, which is what a temperature can honestly be chosen on."""
+    low, high, steps = grid
+    log_probabilities = [[math.log(max(p, 1e-12)) for p in row] for row in probabilities]
+    candidates = [math.exp(math.log(low) + (math.log(high) - math.log(low)) * i / (steps - 1)) for i in range(steps)]
+    return min(candidates, key=lambda t: _negative_log_likelihood(log_probabilities, labels, t))
+
+
+def fingerprint(model, examples: list[Example]) -> dict:
+    """What the examples looked like, so a later drift is visible."""
+    out: dict[str, dict] = {}
+    for task in sorted({e.task for e in examples}):
+        subset = [e for e in examples if e.task == task]
+        lengths = sorted(len(model.layout(e.text, e.question, e.options).ids) for e in subset)
+        counts: dict[str, int] = {}
+        for e in subset:
+            counts[e.options[e.label]] = counts.get(e.options[e.label], 0) + 1
+        out[task] = {"n": len(subset), "options": subset[0].options,
+                     "median_tokens": lengths[len(lengths) // 2],
+                     "answers": counts}
+    return out
+
+
+def calibrate(model, examples: list[Example], batch_size: int = 32, valid_days: int = VALID_DAYS) -> dict:
+    """Fit one temperature per task on these examples and return the passport. The temperatures are
+    stored on the model, so `decide` uses them from then on."""
+    from .evaluate import evaluate
+
+    raw = model.probabilities(examples, batch_size=batch_size, calibrated=False)
+    for task in sorted({e.task for e in examples}):
+        index = [i for i, e in enumerate(examples) if e.task == task]
+        model.config.temperatures[task] = fit_temperature([raw[i] for i in index],
+                                                          [examples[i].label for i in index])
+    measured = evaluate(model, examples, batch_size=batch_size, calibrated=True)
+    today = dt.date.today()
+    return {
+        "created": today.isoformat(),
+        "valid_until": (today + dt.timedelta(days=valid_days)).isoformat(),
+        "encoder": model.config.encoder,
+        "examples": len(examples),
+        "temperatures": dict(model.config.temperatures),
+        "fingerprint": fingerprint(model, examples),
+        "calibration_error": measured["calibration_error"],
+        "calibration_error_interval": measured["calibration_error_interval"],
+        "accuracy": measured["accuracy"],
+        "reliability": measured["reliability"],
+        "risk_coverage": measured["risk_coverage"],
+        "recalibrate_if": {
+            "after": "valid_until",
+            "options_change": True,
+            "median_tokens_shift": 0.5,       # half as long, or twice as long
+        },
+    }
+
+
+def save_passport(passport: dict, path: str | pathlib.Path) -> pathlib.Path:
+    path = pathlib.Path(path)
+    path.write_text(json.dumps(passport, indent=1))
+    return path
+
+
+def check_passport(passport: dict, examples: list[Example] | None = None, model=None,
+                   today: dt.date | None = None) -> list[str]:
+    """What is wrong with trusting these probabilities today. An empty list means nothing found —
+    run it in the deployment pipeline and fail the build on anything it returns."""
+    today = today or dt.date.today()
+    problems = []
+    if today.isoformat() > passport["valid_until"]:
+        problems.append(f"the calibration expired on {passport['valid_until']}; fit it again on recent answers")
+    if passport["examples"] < 200:
+        problems.append(f"fitted on {passport['examples']} answers — the interval on the calibration error is wide")
+    if examples and model:
+        now = fingerprint(model, examples)
+        for task, then in passport["fingerprint"].items():
+            if task not in now:
+                continue
+            if now[task]["options"] != then["options"]:
+                problems.append(f"{task}: the options have changed since the calibration")
+            ratio = now[task]["median_tokens"] / max(then["median_tokens"], 1)
+            if ratio < 0.5 or ratio > 2:
+                problems.append(f"{task}: texts are now {ratio:.1f}× the length they were calibrated on")
+        for task in now:
+            if task not in passport["fingerprint"]:
+                problems.append(f"{task}: never calibrated — its probabilities are raw")
+    return problems

sharada/data.py

@@ -0,0 +1,105 @@
+"""Requests, examples and batches."""
+
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass, field
+
+import torch
+
+from .layout import FIRST_OPTION, Laid, Layout
+
+KINDS = ("choice", "scale", "binary")       # unordered labels, an ordered scale, yes or no
+
+
+@dataclass
+class Request:
+    """One question about one text."""
+
+    text: str
+    question: str
+    options: list[str]
+    kind: str = "choice"
+    task: str = "task"                       # a name, so temperatures can be fitted per task
+
+    def __post_init__(self):
+        if self.kind not in KINDS:
+            raise ValueError(f"kind must be one of {KINDS}, got {self.kind!r}")
+
+
+@dataclass
+class Example(Request):
+    """A request whose answer is known: `label` is the index of the right option."""
+
+    label: int = 0
+
+    def __post_init__(self):
+        super().__post_init__()
+        if not 0 <= self.label < len(self.options):
+            raise ValueError(f"label {self.label} is outside the {len(self.options)} options")
+
+
+@dataclass
+class Batch:
+    ids: torch.Tensor
+    positions: torch.Tensor
+    segments: torch.Tensor
+    present: torch.Tensor                    # [batch, options] — which option slots are real
+    kinds: torch.Tensor
+    n_options: torch.Tensor
+    labels: torch.Tensor | None = None
+    tasks: list[str] = field(default_factory=list)
+    width: int = 0                           # the widest number of options in this batch
+
+    def to(self, device) -> "Batch":
+        move = lambda v: v.to(device) if torch.is_tensor(v) else v
+        return Batch(move(self.ids), move(self.positions), move(self.segments), move(self.present),
+                     move(self.kinds), move(self.n_options), move(self.labels), self.tasks, self.width)
+
+
+def collate(laid: list[Laid], requests: list[Request], pad_id: int) -> Batch:
+    length = max(len(l) for l in laid)
+    width = max(l.n_options for l in laid)
+    n = len(laid)
+    ids = torch.full((n, length), pad_id, dtype=torch.long)
+    positions = torch.zeros((n, length), dtype=torch.long)
+    segments = torch.full((n, length), -1, dtype=torch.long)
+    for i, l in enumerate(laid):
+        k = len(l)
+        ids[i, :k] = torch.tensor(l.ids)
+        positions[i, :k] = torch.tensor(l.positions)
+        segments[i, :k] = torch.tensor(l.segments)
+    counts = torch.tensor([l.n_options for l in laid])
+    labels = None
+    if all(isinstance(r, Example) for r in requests):
+        labels = torch.tensor([r.label for r in requests])
+    return Batch(ids=ids, positions=positions, segments=segments,
+                 present=torch.arange(width)[None] < counts[:, None],
+                 kinds=torch.tensor([KINDS.index(r.kind) for r in requests]),
+                 n_options=counts, labels=labels, tasks=[r.task for r in requests], width=width)
+
+
+def batches(requests: list[Request], layout: Layout, batch_size: int, pad_id: int,
+            shuffle: bool = False, seed: int = 0):
+    """Mini-batches. When shuffling, requests of similar length travel together, so a batch of short
+    questions is not padded out to the length of one with sixty options."""
+    order = list(range(len(requests)))
+    laid = [layout(r.text, r.question, r.options) for r in requests]
+    if shuffle:
+        rng = random.Random(seed)
+        rng.shuffle(order)
+        window = batch_size * 50
+        order = [j for w in range(0, len(order), window)
+                 for j in sorted(order[w:w + window], key=lambda j: len(laid[j]))]
+    chunks = [order[i:i + batch_size] for i in range(0, len(order), batch_size)]
+    if shuffle:
+        random.Random(seed + 1).shuffle(chunks)
+    for chunk in chunks:
+        yield chunk, collate([laid[j] for j in chunk], [requests[j] for j in chunk], pad_id)
+
+
+def option_mass(segments: torch.Tensor, states: torch.Tensor, width: int) -> torch.Tensor:
+    """Mean of each option's own final token states -> [batch, options, hidden]."""
+    index = FIRST_OPTION + torch.arange(width, device=states.device)
+    member = (segments[:, :, None] == index[None, None, :]).to(states.dtype)
+    return torch.einsum("blk,bld->bkd", member, states) / member.sum(1).clamp(min=1)[..., None]

sharada/evaluate.py

@@ -0,0 +1,105 @@
+"""Measuring a decision model: the answer, the stated number, and the whole distribution."""
+
+from __future__ import annotations
+
+import math
+import random
+import time
+
+from .data import Example
+
+
+def _bands(confidence: list[float], n: int) -> list[int]:
+    return [min(int(c * n), n - 1) for c in confidence]
+
+
+def calibration_error(confidence: list[float], correct: list[int], bands: int = 15) -> float:
+    """The average gap between the probability stated and how often it turns out right, over equal
+    bands, weighted by how many answers fall in each. Zero means this estimate finds no gap — on a
+    finite sample that is not proof of none."""
+    index = _bands(confidence, bands)
+    total = 0.0
+    for b in range(bands):
+        inside = [i for i, j in enumerate(index) if j == b]
+        if inside:
+            said = sum(confidence[i] for i in inside) / len(inside)
+            was = sum(correct[i] for i in inside) / len(inside)
+            total += len(inside) / len(confidence) * abs(said - was)
+    return total
+
+
+def calibration_interval(confidence: list[float], correct: list[int], bands: int = 15,
+                         draws: int = 200, seed: int = 0) -> tuple[float, float]:
+    """A bootstrap interval for the calibration error — on a few hundred answers it is wide, and
+    saying so is the point."""
+    rng = random.Random(seed)
+    n = len(confidence)
+    values = []
+    for _ in range(draws):
+        pick = [rng.randrange(n) for _ in range(n)]
+        values.append(calibration_error([confidence[i] for i in pick], [correct[i] for i in pick], bands))
+    values.sort()
+    return values[int(0.025 * draws)], values[int(0.975 * draws) - 1]
+
+
+def reliability_curve(confidence: list[float], correct: list[int], bands: int = 10) -> list[dict]:
+    index = _bands(confidence, bands)
+    curve = []
+    for b in range(bands):
+        inside = [i for i, j in enumerate(index) if j == b]
+        curve.append({"from": b / bands, "n": len(inside),
+                      "stated": sum(confidence[i] for i in inside) / len(inside) if inside else None,
+                      "observed": sum(correct[i] for i in inside) / len(inside) if inside else None})
+    return curve
+
+
+def risk_coverage(confidence: list[float], correct: list[int], steps: int = 20) -> list[dict]:
+    """Answer only where the model is surest: what accuracy, at what share of the traffic."""
+    order = sorted(range(len(confidence)), key=lambda i: -confidence[i])
+    out = []
+    for k in range(1, steps + 1):
+        take = max(1, round(len(order) * k / steps))
+        kept = [correct[i] for i in order[:take]]
+        out.append({"coverage": take / len(order), "accuracy": sum(kept) / len(kept),
+                    "threshold": confidence[order[take - 1]]})
+    return out
+
+
+def evaluate(model, examples: list[Example], batch_size: int = 32, calibrated: bool = True) -> dict:
+    """Everything above, on one set of labelled examples."""
+    probabilities = model.probabilities(examples, batch_size=batch_size, calibrated=calibrated)
+    confidence, correct, true_p, squared = [], [], [], []
+    for p, e in zip(probabilities, examples):
+        best = max(range(len(p)), key=p.__getitem__)
+        confidence.append(p[best])
+        correct.append(int(best == e.label))
+        true_p.append(p[e.label])
+        squared.append(sum((q - (1.0 if i == e.label else 0.0)) ** 2 for i, q in enumerate(p)))
+    low, high = calibration_interval(confidence, correct)
+    return {"n": len(examples),
+            "accuracy": sum(correct) / len(correct),
+            "mean_confidence": sum(confidence) / len(confidence),
+            "calibration_error": calibration_error(confidence, correct),
+            "calibration_error_interval": [low, high],
+            "log_loss": sum(-math.log(max(p, 1e-9)) for p in true_p) / len(true_p),
+            "brier": sum(squared) / len(squared),
+            "reliability": reliability_curve(confidence, correct),
+            "risk_coverage": risk_coverage(confidence, correct)}
+
+
+def latency(model, request, batch_size: int = 1, repeats: int = 20) -> dict:
+    """Milliseconds per call and per question, as the model is loaded right now."""
+    import torch
+
+    requests = [request] * batch_size
+    for _ in range(3):
+        model.probabilities(requests, batch_size=batch_size)
+    if next(model.parameters()).is_cuda:
+        torch.cuda.synchronize()
+    start = time.time()
+    for _ in range(repeats):
+        model.probabilities(requests, batch_size=batch_size)
+    if next(model.parameters()).is_cuda:
+        torch.cuda.synchronize()
+    per_call = (time.time() - start) / repeats * 1e3
+    return {"batch": batch_size, "ms_per_call": per_call, "ms_per_question": per_call / batch_size}

sharada/layout.py

@@ -0,0 +1,78 @@
+"""One request, laid out as one sequence.
+
+    position:  0     1 … n      n+1    n+2 … m     m+1  │ m+2 …         │ m+2 …
+    token:     [CLS] the text   [SEP]  question    [SEP]│ option 0 [SEP] │ option 1 [SEP]
+
+Every option is a branch of its own and every branch starts at the same position, the one right after
+the question. The encoder's positions are rotary, so to it no option is earlier or later than another:
+the order of the options cannot change their scores. `masking.who_reads_whom` then decides who may read
+whom; together the two give the properties the model is built on.
+"""
+
+from __future__ import annotations
+
+import functools
+from dataclasses import dataclass, field
+
+TEXT = 0
+QUESTION = 1
+FIRST_OPTION = 2        # option i has segment FIRST_OPTION + i
+
+
+@dataclass(frozen=True)
+class Limits:
+    """How many tokens each part of a request may use."""
+
+    text: int = 256
+    question: int = 48
+    option: int = 12
+
+    def as_dict(self) -> dict:
+        return {"text": self.text, "question": self.question, "option": self.option}
+
+
+@dataclass
+class Laid:
+    """A request turned into token ids, position ids and one segment per token."""
+
+    ids: list[int]
+    positions: list[int]
+    segments: list[int]
+    n_options: int
+    option_start: int = field(default=0)
+
+    def __len__(self) -> int:
+        return len(self.ids)
+
+
+class Layout:
+    """Lays requests out for one tokenizer. Token ids of repeated strings — option names, instructions,
+    the same text asked about twice — are cached, which is most of the cost of preparing a batch."""
+
+    def __init__(self, tokenizer, limits: Limits | None = None):
+        self.tokenizer = tokenizer
+        self.limits = limits or Limits()
+        self._ids = functools.lru_cache(maxsize=100_000)(self._encode)
+
+    def _encode(self, text: str) -> tuple[int, ...]:
+        return tuple(self.tokenizer(text, add_special_tokens=False)["input_ids"])
+
+    def __call__(self, text: str, question: str, options: list[str]) -> Laid:
+        if len(options) < 2:
+            raise ValueError("a decision needs at least two options")
+        cls_id, sep_id = self.tokenizer.cls_token_id, self.tokenizer.sep_token_id
+        body = list(self._ids(text))[: self.limits.text]
+        ask = list(self._ids(question))[: self.limits.question]
+
+        ids = [cls_id] + body + [sep_id] + ask + [sep_id]
+        segments = [TEXT] * (len(body) + 2) + [QUESTION] * (len(ask) + 1)
+        positions = list(range(len(ids)))
+
+        start = len(ids)
+        for i, option in enumerate(options):
+            tokens = list(self._ids(" " + option))[: self.limits.option] + [sep_id]
+            ids += tokens
+            segments += [FIRST_OPTION + i] * len(tokens)
+            positions += range(start, start + len(tokens))      # every option restarts here
+        return Laid(ids=ids, positions=positions, segments=segments, n_options=len(options),
+                    option_start=start)

sharada/masking.py

@@ -0,0 +1,48 @@
+"""Who reads whom.
+
+The text reads only the text, so its states do not depend on the question or the options — a text asked
+several questions could be read once. The question reads the text and itself. An option reads the text,
+the question and its own tokens, and nothing else: with seventy-seven options on offer, each one still
+sees only the message and itself, and its score does not depend on which other options came along.
+
+`read_each_other=True` keeps the other behaviour — options that also read one another — which is what
+this model was first trained with. On label sets of a few options it costs nothing; on seventy-seven it
+cost thirty points of accuracy, which is why it is not the default.
+
+The encoder's local-attention layers only look a fixed distance away, and the distance is measured in
+positions, so an option sees the end of the question the same way wherever it sits in the sequence.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from .layout import QUESTION, TEXT
+
+
+def who_reads_whom(segments: torch.Tensor, positions: torch.Tensor, local_reach: int,
+                   read_each_other: bool = False) -> tuple[torch.Tensor, torch.Tensor]:
+    """-> two [batch, length, length] boolean masks: for the global layers and for the local ones.
+    True at (i, j) means query token i may read key token j."""
+    q, k = segments[:, :, None], segments[:, None, :]
+    real = k >= 0                                     # padding is not a key
+    context = (k == TEXT) | (k == QUESTION)
+    options_see = real if read_each_other else (context | (k == q))
+    full = torch.where(q == TEXT, k == TEXT, torch.where(q == QUESTION, context, options_see)) & real
+    full = full | torch.eye(segments.size(1), dtype=torch.bool, device=segments.device)[None]
+    near = (positions[:, :, None] - positions[:, None, :]).abs() <= local_reach
+    return full, full & near
+
+
+def encoder_masks(segments: torch.Tensor, positions: torch.Tensor, local_reach: int,
+                  attn_implementation: str, dtype: torch.dtype, read_each_other: bool = False) -> dict:
+    """The masks in the form the encoder expects: booleans for fused attention, an additive float mask
+    for the plain implementation."""
+    out = {}
+    names = ("full_attention", "sliding_attention")
+    for name, allowed in zip(names, who_reads_whom(segments, positions, local_reach, read_each_other)):
+        mask = allowed[:, None]
+        if attn_implementation == "eager":
+            mask = torch.zeros(mask.shape, dtype=dtype, device=mask.device).masked_fill(~mask, torch.finfo(dtype).min)
+        out[name] = mask
+    return out

sharada/model.py

@@ -0,0 +1,167 @@
+"""The decision model: an encoder, a small read-out, one number per option."""
+
+from __future__ import annotations
+
+import json
+import pathlib
+from dataclasses import asdict, dataclass, field
+
+import torch
+import torch.nn as nn
+
+from .data import Batch, Example, Request, batches, collate, option_mass
+from .layout import Laid, Layout, Limits
+
+DEFAULT_ENCODER = "answerdotai/ModernBERT-base"
+
+
+@dataclass
+class Config:
+    encoder: str = DEFAULT_ENCODER
+    limits: dict = field(default_factory=lambda: Limits().as_dict())
+    local_reach: int | None = None          # taken from the encoder when left out
+    read_each_other: bool = False           # options reading one another; off by default, see masking.py
+    temperatures: dict = field(default_factory=dict)    # per task, fitted by `calibrate`
+    version: int = 1
+
+
+@dataclass
+class Decision:
+    """What the model answers, and how sure it says it is."""
+
+    answer: str
+    index: int
+    probabilities: dict[str, float]
+    confidence: float
+
+    def __repr__(self) -> str:
+        return f"Decision({self.answer!r}, confidence={self.confidence:.3f})"
+
+
+class DecisionModel(nn.Module):
+    def __init__(self, config: Config | None = None, tokenizer=None, encoder=None):
+        super().__init__()
+        from transformers import AutoModel, AutoTokenizer
+
+        self.config = config or Config()
+        self.tokenizer = tokenizer or AutoTokenizer.from_pretrained(self.config.encoder)
+        self.attn = "eager" if not torch.cuda.is_available() else "sdpa"
+        self.encoder = encoder or AutoModel.from_pretrained(self.config.encoder, attn_implementation=self.attn)
+        hidden = self.encoder.config.hidden_size
+        self.read = nn.Sequential(nn.Linear(2 * hidden, hidden), nn.GELU(), nn.LayerNorm(hidden),
+                                  nn.Linear(hidden, 1))
+        if self.config.local_reach is None:
+            self.config.local_reach = getattr(self.encoder.config, "local_attention", 128) // 2
+        self.layout = Layout(self.tokenizer, Limits(**self.config.limits))
+
+    # ── the forward pass ──────────────────────────────────────────────────────────────────────
+    def forward(self, batch: Batch) -> torch.Tensor:
+        """-> [batch, options] scores; options that do not exist cannot win."""
+        from .masking import encoder_masks
+
+        dtype = self.encoder.embeddings.tok_embeddings.weight.dtype
+        masks = encoder_masks(batch.segments, batch.positions, self.config.local_reach,
+                              self.attn, dtype, self.config.read_each_other)
+        states = self.encoder(input_ids=batch.ids, position_ids=batch.positions,
+                              attention_mask=masks).last_hidden_state
+        options = option_mass(batch.segments, states, batch.width)
+        summary = states[:, :1].expand(-1, batch.width, -1)
+        scores = self.read(torch.cat([options, options * summary], -1)).squeeze(-1).float()
+        return scores.masked_fill(~batch.present, -1e4)
+
+    # ── asking it things ──────────────────────────────────────────────────────────────────────
+    @torch.no_grad()
+    def probabilities(self, requests: list[Request], batch_size: int = 32,
+                      calibrated: bool = True) -> list[list[float]]:
+        """One distribution over the options of each request, in the order they were given."""
+        self.eval()
+        device = next(self.parameters()).device
+        out: list[list[float]] = [[] for _ in requests]
+        for chunk, batch in batches(requests, self.layout, batch_size, self.tokenizer.pad_token_id):
+            scores = self(batch.to(device))
+            for row, j in enumerate(chunk):
+                k = int(batch.n_options[row])
+                z = scores[row, :k].float().cpu()
+                if calibrated:
+                    z = z / self.temperature_for(requests[j])
+                out[j] = torch.softmax(z, -1).tolist()
+        return out
+
+    def decide(self, text: str, question: str, options: list[str], kind: str = "choice",
+               task: str = "task", calibrated: bool = True) -> Decision:
+        """One decision, start to finish."""
+        request = Request(text=text, question=question, options=options, kind=kind, task=task)
+        return self.decide_many([request], calibrated=calibrated)[0]
+
+    def decide_many(self, requests: list[Request], batch_size: int = 32,
+                    calibrated: bool = True) -> list[Decision]:
+        decisions = []
+        for request, probabilities in zip(requests, self.probabilities(requests, batch_size, calibrated)):
+            best = max(range(len(probabilities)), key=probabilities.__getitem__)
+            decisions.append(Decision(answer=request.options[best], index=best,
+                                      probabilities=dict(zip(request.options, probabilities)),
+                                      confidence=probabilities[best]))
+        return decisions
+
+    def temperature_for(self, request: Request) -> float:
+        return float(self.config.temperatures.get(request.task, 1.0))
+
+    # ── saving and loading ────────────────────────────────────────────────────────────────────
+    def save(self, path: str | pathlib.Path) -> pathlib.Path:
+        from safetensors.torch import save_file
+
+        path = pathlib.Path(path)
+        path.mkdir(parents=True, exist_ok=True)
+        save_file({k: v.contiguous() for k, v in self.state_dict().items()}, path / "model.safetensors")
+        (path / "config.json").write_text(json.dumps(asdict(self.config), indent=1))
+        self.tokenizer.save_pretrained(path)
+        return path
+
+    @classmethod
+    def from_pretrained(cls, path: str | pathlib.Path, device: str | None = None) -> "DecisionModel":
+        """A local directory saved by `save`, or a model id on the Hub."""
+        from safetensors.torch import load_file
+        from transformers import AutoConfig, AutoModel, AutoTokenizer
+
+        local = pathlib.Path(path)
+        if not local.is_dir():
+            from huggingface_hub import snapshot_download
+
+            local = pathlib.Path(snapshot_download(str(path)))
+        config = Config(**json.loads((local / "config.json").read_text()))
+        tokenizer = AutoTokenizer.from_pretrained(local)
+        attn = "eager" if not torch.cuda.is_available() else "sdpa"
+        encoder = AutoModel.from_config(AutoConfig.from_pretrained(config.encoder), attn_implementation=attn)
+        model = cls(config=config, tokenizer=tokenizer, encoder=encoder)
+        model.load_state_dict(load_file(local / "model.safetensors"))
+        return model.to(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+
+    def push_to_hub(self, repo_id: str, private: bool = False) -> str:
+        """Upload to the Hub. The token comes from the environment; this never asks for one."""
+        import tempfile
+
+        from huggingface_hub import HfApi
+
+        api = HfApi()
+        api.create_repo(repo_id, private=private, exist_ok=True)
+        with tempfile.TemporaryDirectory() as tmp:
+            self.save(tmp)
+            api.upload_folder(folder_path=tmp, repo_id=repo_id)
+        return f"https://huggingface.co/{repo_id}"
+
+    # ── training lives in train.py, calibration in calibrate.py ───────────────────────────────
+    def fit(self, examples: list[Example], **kwargs):
+        from .train import fit
+
+        return fit(self, examples, **kwargs)
+
+    def calibrate(self, examples: list[Example], **kwargs):
+        from .calibrate import calibrate
+
+        return calibrate(self, examples, **kwargs)
+
+
+def encode_one(model: DecisionModel, request: Request) -> tuple[Laid, Batch]:
+    """Handy in tests: one request, laid out and collated."""
+    laid = model.layout(request.text, request.question, request.options)
+    return laid, collate([laid], [request], model.tokenizer.pad_token_id)