ruling 0.2.0__py3-none-any.whl

ruling/corpus.py

@@ -0,0 +1,202 @@
+"""Openly licensed labeled datasets, rewritten as typed-question records.
+
+Only sources whose licenses permit training and redistribution of derived
+weights are used: WANLI (CC-BY-4.0), SNLI (CC-BY-SA-4.0), CLINC150
+(CC-BY-3.0), BoolQ (CC-BY-SA-3.0), GoEmotions (Apache-2.0), and Civil Comments
+(CC0-1.0). Every download is pinned to a Hugging Face commit. Where a dataset
+records how many raters chose each answer, the rater fraction becomes a soft
+`target`, which is the only ground truth for calibration that does not come
+from a model.
+"""
+
+from __future__ import annotations
+
+import json
+from collections import defaultdict
+from pathlib import Path
+
+import numpy as np
+import pyarrow.parquet as pq
+from huggingface_hub import hf_hub_download
+
+from ruling.evals import Record
+
+NLI_OPTIONS = {
+    "supported": "The evidence establishes the claim",
+    "insufficient": "The evidence does not establish either",
+    "contradicted": "The evidence establishes the opposite",
+}
+NLI_GOLD = {"entailment": "supported", "neutral": "insufficient", "contradiction": "contradicted"}
+SNLI_LABELS = {0: "entailment", 1: "neutral", 2: "contradiction"}
+NONE_OPTION = "none_of_these"
+CLINC_OUT_OF_SCOPE = "oos"
+
+SOURCES = {
+    "wanli": ("alisawuffles/WANLI", "61c95318fd71c55b6ba355d76253254615f387ec", {"train": ["train.jsonl"]}),
+    "snli": ("stanfordnlp/snli", "cdb5c3d5eed6ead6e5a341c8e56e669bb666725b",
+             {"train": ["plain_text/train-00000-of-00001.parquet"], "dev": ["plain_text/validation-00000-of-00001.parquet"]}),
+    "clinc": ("clinc/clinc_oos", "155b9c710419136e17307b80d0a13e68cd46b4ec",
+              {"train": ["plus/train-00000-of-00001.parquet"], "dev": ["plus/validation-00000-of-00001.parquet"]}),
+    "boolq": ("google/boolq", "35b264d03638db9f4ce671b711558bf7ff0f80d5",
+              {"train": ["data/train-00000-of-00001.parquet"], "dev": ["data/validation-00000-of-00001.parquet"]}),
+    "emotions": ("google-research-datasets/go_emotions", "add492243ff905527e67aeb8b80c082af02207c3",
+                 {"train": ["raw/train-00000-of-00001.parquet"]}),
+    "civil": ("google/civil_comments", "f2970eb3a55777454c94069077cc8d9b5866312d",
+              {"train": ["data/train-00000-of-00002.parquet", "data/train-00001-of-00002.parquet"],
+               "dev": ["data/validation-00000-of-00001.parquet"]}),
+}
+"""Pinned source files per split. `dev` uses validation splits only, so test splits stay untouched and no
+evaluation set in `datasets/` shares rows with it; WANLI's test split and GoEmotions have no suitable dev split."""
+
+
+def nli_record(family: str, premise: str, hypothesis: str, gold: str) -> Record:
+    verdict = NLI_GOLD[gold]
+    return Record(
+        family=family,
+        state=premise,
+        questions={
+            "verdict": {"type": "choice", "instructions": "Assess the claim using only the supplied evidence: " + hypothesis,
+                        "criteria": NLI_OPTIONS},
+            "supported": {"type": "noul", "instructions": "The evidence establishes this claim: " + hypothesis},
+        },
+        labels={"verdict": verdict, "supported": verdict == "supported"},
+    )
+
+
+def clinc_record(text: str, gold: str | None, intents: list[str], rng: np.random.Generator) -> Record:
+    """A routing question over a random subset of intents, with an explicit none option.
+
+    A quarter of in-scope rows drop the gold intent from the offered options, so
+    the right answer becomes `none_of_these`; that teaches abstention.
+    """
+    size = int(rng.integers(3, 13))
+    pool = [i for i in intents if i != gold]
+    offered = list(rng.choice(pool, size=size, replace=False))
+    label = NONE_OPTION
+    if gold is not None and rng.random() >= 0.25:
+        offered[int(rng.integers(size))] = gold
+        label = gold
+    criteria = {str(i): str(i).replace("_", " ") for i in offered}
+    criteria[NONE_OPTION] = "The request matches none of the other options"
+    return Record(
+        family="clinc",
+        state=text,
+        questions={"intent": {"type": "choice", "instructions": "Which request type does this user message belong to?",
+                              "criteria": criteria}},
+        labels={"intent": label},
+    )
+
+
+def boolq_record(passage: str, question: str, answer: bool) -> Record:
+    return Record(
+        family="boolq",
+        state=passage,
+        questions={"answer": {"type": "noul", "instructions": question[0].upper() + question[1:] + "?"}},
+        labels={"answer": bool(answer)},
+    )
+
+
+def emotions_records(rows: list[dict], emotions: list[str], rng: np.random.Generator, questions_per_text: int) -> list[Record]:
+    by_text: dict[str, list[dict]] = defaultdict(list)
+    for row in rows:
+        by_text[row["id"]].append(row)
+    records = []
+    for raters in by_text.values():
+        fractions = {e: float(np.mean([r[e] for r in raters])) for e in emotions}
+        present = [e for e in emotions if fractions[e] > 0]
+        absent = [e for e in emotions if fractions[e] == 0]
+        chosen = present[:questions_per_text]
+        if len(chosen) < questions_per_text and absent:
+            chosen += list(rng.choice(absent, size=min(questions_per_text - len(chosen), len(absent)), replace=False))
+        questions = {e: {"type": "noul", "instructions": f"The text expresses {e.replace('_', ' ')}."} for e in chosen}
+        records.append(Record(
+            family="emotions",
+            state=raters[0]["text"],
+            questions=questions,
+            labels={e: fractions[e] >= 0.5 for e in chosen},
+            references={e: {"target": {"yes": fractions[e], "no": 1.0 - fractions[e]}} for e in chosen},
+        ))
+    return records
+
+
+CIVIL_QUESTIONS = {
+    "toxic": ("toxicity", "The comment is rude, disrespectful, or unreasonable."),
+    "insult": ("insult", "The comment insults someone."),
+    "threat": ("threat", "The comment threatens someone."),
+}
+
+
+def civil_record(row: dict) -> Record:
+    return Record(
+        family="civil",
+        state=row["text"],
+        questions={qid: {"type": "noul", "instructions": text} for qid, (_, text) in CIVIL_QUESTIONS.items()},
+        labels={qid: float(row[column]) >= 0.5 for qid, (column, _) in CIVIL_QUESTIONS.items()},
+        references={qid: {"target": {"yes": float(row[column]), "no": 1.0 - float(row[column])}}
+                    for qid, (column, _) in CIVIL_QUESTIONS.items()},
+    )
+
+
+def _rows(name: str, split: str) -> list[dict]:
+    repo, revision, splits = SOURCES[name]
+    rows = []
+    for file in splits[split]:
+        path = hf_hub_download(repo, file, repo_type="dataset", revision=revision)
+        if path.endswith(".jsonl"):
+            rows.extend(json.loads(line) for line in Path(path).read_text().splitlines() if line.strip())
+        else:
+            rows.extend(pq.read_table(path).to_pylist())
+    return rows
+
+
+def build(per_source: int, seed: int, split: str = "train") -> list[Record]:
+    """Up to `per_source` records from each source that has `split`, sampled with a fixed seed."""
+    rng = np.random.default_rng(seed)
+    available = {name for name, (_, _, splits) in SOURCES.items() if split in splits}
+
+    def sample(rows):
+        return [rows[i] for i in rng.permutation(len(rows))[:per_source]]
+
+    records: list[Record] = []
+    if "wanli" in available:
+        records += [nli_record("wanli", r["premise"], r["hypothesis"], r["gold"]) for r in sample(_rows("wanli", split))]
+    records += [nli_record("snli", r["premise"], r["hypothesis"], SNLI_LABELS[r["label"]])
+                for r in sample([r for r in _rows("snli", split) if r["label"] in SNLI_LABELS])]
+
+    names = _clinc_names()
+    in_scope = [name for name in names.values() if name != CLINC_OUT_OF_SCOPE]
+    records += [clinc_record(r["text"], None if names[r["intent"]] == CLINC_OUT_OF_SCOPE else names[r["intent"]], in_scope, rng)
+                for r in sample(_rows("clinc", split))]
+
+    records += [boolq_record(r["passage"], r["question"], r["answer"]) for r in sample(_rows("boolq", split))]
+
+    if "emotions" in available:
+        raw = _rows("emotions", split)
+        emotions = [c for c in raw[0] if c not in {"text", "id", "author", "subreddit", "link_id", "parent_id",
+                                                   "created_utc", "rater_id", "example_very_unclear", "neutral"}]
+        wanted = set(sample(sorted({r["id"] for r in raw})))
+        records += emotions_records([r for r in raw if r["id"] in wanted and not r["example_very_unclear"]], emotions, rng, 3)
+
+    civil = _rows("civil", split)
+    toxic = [r for r in civil if r["toxicity"] >= 0.2]
+    calm = [r for r in civil if r["toxicity"] < 0.2]
+    half = per_source // 2
+    records += [civil_record(r) for r in [toxic[i] for i in rng.permutation(len(toxic))[:half]]
+                + [calm[i] for i in rng.permutation(len(calm))[:per_source - half]]]
+    return records
+
+
+def _clinc_names() -> dict[int, str]:
+    """Intent names for the `plus` configuration, read from the dataset card at the pinned revision."""
+    repo, revision, _ = SOURCES["clinc"]
+    lines = Path(hf_hub_download(repo, "README.md", repo_type="dataset", revision=revision)).read_text().splitlines()
+    start = next(i for i, line in enumerate(lines) if line.strip().removeprefix("- ") == "config_name: plus")
+    names: dict[int, str] = {}
+    for line in lines[next(i for i in range(start, len(lines)) if lines[i].strip() == "names:") + 1:]:
+        key, sep, value = line.strip().partition(":")
+        if not (sep and key.strip("'\"").isdigit()):
+            break
+        names[int(key.strip("'\""))] = value.strip().strip("'\"")
+    if CLINC_OUT_OF_SCOPE not in names.values():
+        raise ValueError("CLINC150 card lacks the out-of-scope intent")
+    return names

ruling/decision.py

@@ -0,0 +1,246 @@
+"""A decision encoder: every typed question answered in one bidirectional pass.
+
+The state and every question go into a single sequence. Each option gets a
+`[MASK]` slot followed by its text, and a linear head scores the slot's hidden
+state; a softmax over one question's slots is that question's distribution.
+Attention is restricted so every question sees the state and its own text,
+every option sees only itself, its question, and the state, and all options of
+a question start at the same position. Answers are therefore identical
+whatever other questions are asked, in any order, with options in any order.
+Nothing is decoded and there are no letter codes.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Callable
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+from mlx.utils import tree_flatten
+
+from ruling.calibration import Calibration, Provenance
+from ruling.engine import InputTooLong, RawScore, Scored, answer
+from ruling.modernbert import Config, Layer, MaskedLM, load
+from ruling.questions import (
+    Choice,
+    MAX_CHOICE_OPTIONS,
+    Noul,
+    option_keys,
+    question_statement,
+    render_text,
+    Score,
+    State,
+    SystemOneRequest,
+    SystemOneResponse,
+    Usage,
+)
+
+CHECKPOINT_FILE = "decision.json"
+WEIGHTS_FILE = "weights.safetensors"
+
+__all__ = ["InputTooLong", "Special", "Packed", "pack", "attention_mask", "DecisionModel", "DecisionEngine"]
+
+
+@dataclass(frozen=True)
+class Special:
+    cls: int
+    sep: int
+    mask: int
+    pad: int
+
+
+@dataclass(frozen=True)
+class Packed:
+    tokens: list[int]
+    positions: list[int]
+    segments: list[tuple[int, int]]
+    """Per token: (0, 0) for the state, (q, 0) for question q's text, (q, o) for its option o."""
+    slots: dict[str, list[tuple[str, int]]]
+    question_index: dict[str, int]
+
+
+def question_text(question: Choice | Score | Noul) -> str:
+    instructions = question_statement(question)
+    if isinstance(question, Choice):
+        return f"Choose the best option. {instructions}"
+    if isinstance(question, Score):
+        return f"Choose the level that fits, from lowest to highest. {instructions}"
+    return f"Is this statement true? {instructions}"
+
+
+def option_text(question: Choice | Score | Noul, key: str) -> str:
+    if isinstance(question, Choice):
+        description = question.criteria[key]
+        return key if description is None else f"{key}: {render_text(description)}"
+    if isinstance(question, Score):
+        index = int(key)
+        return f"level {index + 1} of {len(question.criteria)}: {render_text(question.criteria[index])}"
+    criteria = question.criteria
+    if key == "yes":
+        return "true" if criteria is None or criteria.true is None else f"true: {render_text(criteria.true)}"
+    return "false" if criteria is None or criteria.false is None else f"false: {render_text(criteria.false)}"
+
+
+def pack(state: State, questions: dict[str, Choice | Score | Noul], encode: Callable[[str], list[int]],
+         special: Special, max_tokens: int) -> Packed:
+    body = state if isinstance(state, str) else json.dumps(state, indent=2, ensure_ascii=False)
+    tokens = [special.cls, *encode(body), special.sep]
+    positions = list(range(len(tokens)))
+    segments = [(0, 0)] * len(tokens)
+    state_length = len(tokens)
+    slots: dict[str, list[tuple[str, int]]] = {}
+    index: dict[str, int] = {}
+    for q, (qid, question) in enumerate(questions.items(), start=1):
+        if isinstance(question, Choice) and len(question.criteria) > MAX_CHOICE_OPTIONS:
+            raise ValueError(f"question {qid!r} has more than {MAX_CHOICE_OPTIONS} options")
+        index[qid] = q
+        text = [*encode(question_text(question)), special.sep]
+        tokens += text
+        positions += range(state_length, state_length + len(text))
+        segments += [(q, 0)] * len(text)
+        start = state_length + len(text)
+        slots[qid] = []
+        for o, key in enumerate(option_keys(question), start=1):
+            option = [special.mask, *encode(option_text(question, key)), special.sep]
+            slots[qid].append((key, len(tokens)))
+            tokens += option
+            positions += range(start, start + len(option))
+            segments += [(q, o)] * len(option)
+    if len(tokens) > max_tokens:
+        raise InputTooLong(f"request is {len(tokens)} tokens, over the maximum context length of {max_tokens}")
+    return Packed(tokens, positions, segments, slots, index)
+
+
+def attention_mask(segment_lists: list[list[tuple[int, int]]]) -> np.ndarray:
+    """Boolean [batch, width, width]; padding positions neither attend nor are attended to."""
+    width = max(len(s) for s in segment_lists)
+    mask = np.zeros((len(segment_lists), width, width), dtype=bool)
+    for b, segments in enumerate(segment_lists):
+        seg = np.array(segments)
+        q, o = seg[:, 0], seg[:, 1]
+        sees_state = (q == 0)[None, :]
+        own_question_text = (q[:, None] == q[None, :]) & (o == 0)[None, :]
+        own_option = (q[:, None] == q[None, :]) & (o[:, None] == o[None, :])
+        n = len(segments)
+        mask[b, :n, :n] = sees_state | own_question_text | own_option
+    return mask
+
+
+class DecisionModel(nn.Module):
+    def __init__(self, config: Config):
+        super().__init__()
+        self.lm = MaskedLM(config)
+        self.score = nn.Linear(config.hidden_size, 1)
+        self.score.weight = mx.zeros_like(self.score.weight)
+        self.score.bias = mx.zeros_like(self.score.bias)
+
+    def __call__(self, tokens: mx.array, positions: mx.array, mask: mx.array, rows: mx.array, cols: mx.array) -> mx.array:
+        """One score per requested slot, where slot n is at (rows[n], cols[n])."""
+        hidden = self.lm.encoder(tokens, positions, mask)
+        return self.score(self.lm.transform(hidden[rows, cols]))[:, 0].astype(mx.float32)
+
+
+@dataclass(frozen=True)
+class Checkpoint:
+    base: str
+    revision: str
+
+    def save(self, directory: Path, model: DecisionModel, extra: dict) -> None:
+        directory.mkdir(parents=True, exist_ok=True)
+        mx.save_safetensors(str(directory / WEIGHTS_FILE), dict(tree_flatten(model.parameters())))
+        (directory / CHECKPOINT_FILE).write_text(json.dumps({**asdict(self), **extra}, indent=2) + "\n")
+
+
+def checkpoint_layers(model: DecisionModel) -> None:
+    """Recompute each encoder layer's activations in the backward pass instead of storing them."""
+    call = Layer.__call__
+
+    def checkpointed(layer, *args):
+        def inner(params, *args):
+            layer.update(params)
+            return call(layer, *args)
+
+        return mx.checkpoint(inner)(layer.trainable_parameters(), *args)
+
+    Layer.__call__ = checkpointed
+
+
+def is_checkpoint(model: str) -> bool:
+    return (Path(model) / CHECKPOINT_FILE).exists()
+
+
+def load_model(base: str, revision: str, checkpoint: Path | None):
+    """The base encoder with a fresh zero-initialized head, or a trained checkpoint on top of it."""
+    lm, tokenizer, config = load(base, revision)
+    model = DecisionModel(config)
+    model.lm = lm
+    if checkpoint is not None:
+        model.load_weights(str(checkpoint / WEIGHTS_FILE), strict=True)
+    mx.eval(model.parameters())
+    special = Special(cls=tokenizer.cls_token_id, sep=tokenizer.sep_token_id, mask=tokenizer.mask_token_id, pad=tokenizer.pad_token_id)
+    return model, tokenizer, special
+
+
+def batch_inputs(packs: list[Packed], pad: int) -> tuple[mx.array, mx.array, mx.array]:
+    width = max(len(p.tokens) for p in packs)
+    tokens = mx.array([p.tokens + [pad] * (width - len(p.tokens)) for p in packs])
+    positions = mx.array([p.positions + [0] * (width - len(p.positions)) for p in packs])
+    return tokens, positions, mx.array(attention_mask([p.segments for p in packs]))
+
+
+class DecisionEngine:
+    """Serves a trained decision checkpoint through the same interface as the decoder engine."""
+
+    rotations = 1
+    prior_debias = False
+    max_options = MAX_CHOICE_OPTIONS
+    reads_images = False
+
+    def __init__(self, checkpoint: Path, calibration: Calibration, max_input_tokens: int):
+        meta = json.loads((checkpoint / CHECKPOINT_FILE).read_text())
+        self.model_id = str(checkpoint)
+        self.revision = meta["revision"]
+        self.adapter = hashlib.sha256((checkpoint / WEIGHTS_FILE).read_bytes()).hexdigest()[:16]
+        self.model, self.tokenizer, self.special = load_model(meta["base"], meta["revision"], checkpoint)
+        self.max_input_tokens = min(max_input_tokens, self.model.lm.encoder.config.max_position_embeddings)
+        # The encoder packs state and every question into one sequence, so a branch is the whole input.
+        self.max_branch_tokens = self.max_input_tokens
+        self.model.eval()
+        calibration.check(self.provenance)
+        self.calibration = calibration
+
+    @property
+    def provenance(self) -> Provenance:
+        return Provenance(self.model_id, self.revision, self.rotations, self.prior_debias, self.adapter)
+
+    def encode(self, text: str) -> list[int]:
+        return self.tokenizer.encode(text, add_special_tokens=False)
+
+    def score(self, state: State, questions: dict[str, Choice | Score | Noul]) -> Scored:
+        packed = pack(state, questions, self.encode, self.special, self.max_input_tokens)
+        tokens, positions, mask = batch_inputs([packed], self.special.pad)
+        order = [(qid, key, col) for qid, slots in packed.slots.items() for key, col in slots]
+        scores = self.model(tokens, positions, mask, mx.zeros((len(order),), dtype=mx.int32),
+                            mx.array([col for _, _, col in order]))
+        values = np.array(scores, dtype=np.float64)
+        raw = {}
+        cursor = 0
+        for qid, question in questions.items():
+            count = len(packed.slots[qid])
+            logits = values[cursor:cursor + count]
+            raw[qid] = RawScore(keys=option_keys(question), logits=logits - np.log(np.exp(logits - logits.max()).sum()) - logits.max())
+            cursor += count
+        return Scored(raw=raw, input_tokens=len(packed.tokens))
+
+    def evaluate(self, request: SystemOneRequest) -> SystemOneResponse:
+        scored = self.score(request.state, request.questions)
+        answers = {qid: answer(request.questions[qid], scored.raw[qid], self.calibration) for qid in request.questions}
+        return SystemOneResponse(model=self.model_id, answers=answers, usage=Usage(input_tokens=scored.input_tokens))
+
+    def answer(self, question, raw: RawScore):
+        return answer(question, raw, self.calibration)

ruling/decision_train.py

@@ -0,0 +1,208 @@
+"""Full fine-tuning of a decision encoder on typed-question records.
+
+Each training example is a whole record: its state and every question that
+has a target, packed and scored in one pass exactly as the engine serves it.
+Answers do not depend on option order or on which other questions are asked,
+so there is nothing to augment. The loss is cross-entropy against each
+question's target distribution, a proper scoring rule, averaged over the
+questions in a batch. The new scoring head starts at zero, so the untrained
+model answers uniformly, and learns at a higher rate than the pretrained
+encoder.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import time
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+import mlx.core as mx
+import mlx.nn as nn
+import mlx.optimizers as optim
+import numpy as np
+from mlx.utils import tree_flatten
+
+from ruling.calibration import expected_calibration_error
+from ruling.decision import Checkpoint, InputTooLong, checkpoint_layers, Packed, batch_inputs, load_model, pack
+from ruling.evals import Record, load_dataset
+from ruling.train import holdout_overlap, split_records, target_vector
+
+
+@dataclass(frozen=True)
+class DecisionTrainConfig:
+    base: str
+    revision: str
+    data: list[Path]
+    holdout: list[Path]
+    output: Path
+    valid_fraction: float = 0.02
+    steps: int = 3000
+    batch_tokens: int = 8192
+    learning_rate: float = 3e-5
+    head_learning_rate: float = 1e-3
+    weight_decay: float = 0.01
+    warmup: int = 150
+    max_tokens: int = 2048
+    eval_every: int = 250
+    valid_records: int = 800
+    grad_checkpoint: bool = True
+    seed: int = 0
+
+
+@dataclass(frozen=True)
+class Example:
+    packed: Packed
+    targets: list[list[float]]
+
+
+class DecisionTrainer:
+    def __init__(self, config: DecisionTrainConfig):
+        self.config = config
+        self.rng = np.random.default_rng(config.seed)
+        records = [r for path in config.data for r in load_dataset(path)]
+        held_out = [r for path in config.holdout for r in load_dataset(path)]
+        overlap = holdout_overlap(records, held_out)
+        if overlap:
+            raise ValueError(f"{overlap} training states also appear in held-out evaluation data")
+        train_records, valid_records = split_records(records, config.valid_fraction)
+        mx.random.seed(config.seed)
+        self.model, self.tokenizer, self.special = load_model(config.base, config.revision, None)
+        if config.grad_checkpoint:
+            checkpoint_layers(self.model)
+        self.train = self._examples(train_records)
+        valid = self._examples(valid_records)
+        self.valid = [valid[i] for i in self.rng.permutation(len(valid))[: config.valid_records]]
+        if not self.train or not self.valid:
+            raise ValueError("training and validation both need records with labeled questions")
+
+    def _examples(self, records: list[Record]) -> list[Example]:
+        encode = lambda text: self.tokenizer.encode(text, add_special_tokens=False)
+        examples = []
+        for record in records:
+            questions = {qid: q for qid, q in record.questions.items() if target_vector(record, qid) is not None}
+            if not questions:
+                continue
+            try:
+                packed = pack(record.state, questions, encode, self.special, self.config.max_tokens)
+            except InputTooLong:
+                continue
+            targets = [[target_vector(record, qid)[key] for key, _ in packed.slots[qid]] for qid in questions]
+            examples.append(Example(packed, targets))
+        return examples
+
+    def _batch(self, examples: list[Example]):
+        tokens, positions, mask = batch_inputs([e.packed for e in examples], self.special.pad)
+        rows, cols, slot_index, targets = [], [], [], []
+        width = max(len(t) for e in examples for t in e.targets)
+        for b, example in enumerate(examples):
+            for slots, target in zip(example.packed.slots.values(), example.targets):
+                start = len(rows)
+                for _, col in slots:
+                    rows.append(b)
+                    cols.append(col)
+                slot_index.append([start + i for i in range(len(slots))] + [0] * (width - len(slots)))
+                targets.append(target + [0.0] * (width - len(target)))
+        valid = [[i < len(e_t) for i in range(width)] for e in examples for e_t in e.targets]
+        return (tokens, positions, mask, mx.array(rows), mx.array(cols), mx.array(slot_index),
+                mx.array(valid), mx.array(targets, dtype=mx.float32))
+
+    @staticmethod
+    def _log_probs(model, tokens, positions, mask, rows, cols, slot_index, valid):
+        scores = model(tokens, positions, mask, rows, cols)
+        logits = mx.where(valid, scores[slot_index], -mx.inf)
+        return logits - mx.logsumexp(logits, axis=-1, keepdims=True)
+
+    @classmethod
+    def loss(cls, model, tokens, positions, mask, rows, cols, slot_index, valid, targets):
+        log_probs = cls._log_probs(model, tokens, positions, mask, rows, cols, slot_index, valid)
+        return -mx.where(valid, targets * log_probs, 0.0).sum(axis=-1).mean()
+
+    def _groups(self, examples: list[Example], shuffle: bool) -> list[list[Example]]:
+        order = self.rng.permutation(len(examples)) if shuffle else np.arange(len(examples))
+        groups = []
+        for start in range(0, len(order), 512):
+            chunk = sorted((examples[i] for i in order[start:start + 512]), key=lambda e: len(e.packed.tokens))
+            current: list[Example] = []
+            for example in chunk:
+                width = max([len(e.packed.tokens) for e in current] + [len(example.packed.tokens)])
+                if current and width * (len(current) + 1) > self.config.batch_tokens:
+                    groups.append(current)
+                    current = []
+                current.append(example)
+            if current:
+                groups.append(current)
+        if shuffle:
+            groups = [groups[i] for i in self.rng.permutation(len(groups))]
+        return groups
+
+    def validate(self) -> dict:
+        self.model.eval()
+        losses, top, correct = [], [], []
+        for group in self._groups(self.valid, shuffle=False):
+            batch = self._batch(group)
+            log_probs = self._log_probs(self.model, *batch[:7])
+            valid, targets = np.array(batch[6]), np.array(batch[7])
+            probs = np.exp(np.where(valid, np.array(log_probs), -np.inf))
+            for row, v, target in zip(probs, valid, targets):
+                row, target = row[v], target[v]
+                losses.append(float(-(target * np.log(np.maximum(row, 1e-12))).sum()))
+                top.append(float(row.max()))
+                correct.append(int(row.argmax()) == int(target.argmax()))
+        self.model.train()
+        return {"loss": float(np.mean(losses)), "accuracy": float(np.mean(correct)),
+                "expected_calibration_error": expected_calibration_error(top, correct), "questions": len(losses)}
+
+    def save(self, history: list[dict]) -> None:
+        data = [{"path": str(p), "sha256": hashlib.sha256(p.read_bytes()).hexdigest()} for p in self.config.data]
+        extra = {"training": {**asdict(self.config), "data": data, "output": str(self.config.output),
+                              "holdout": [str(p) for p in self.config.holdout]},
+                 "train_records": len(self.train), "valid_records": len(self.valid), "history": history}
+        Checkpoint(self.config.base, self.config.revision).save(self.config.output, self.model, extra)
+
+    def run(self, log=print) -> list[dict]:
+        config = self.config
+
+        def schedule(peak):
+            return optim.join_schedules([optim.linear_schedule(peak * 0.01, peak, config.warmup),
+                                         optim.cosine_decay(peak, max(config.steps - config.warmup, 1), peak * 0.05)],
+                                        [config.warmup])
+
+        optimizer = optim.MultiOptimizer(
+            [optim.AdamW(learning_rate=schedule(config.head_learning_rate), weight_decay=0.0),
+             optim.AdamW(learning_rate=schedule(config.learning_rate), weight_decay=config.weight_decay)],
+            [lambda path, _: path.startswith("score.")],
+        )
+        self.model.train()
+        step_fn = nn.value_and_grad(self.model, self.loss)
+        log(f"{sum(v.size for _, v in tree_flatten(self.model.trainable_parameters())):,} parameters; "
+            f"{len(self.train):,} train and {len(self.valid):,} validation records")
+        history = [{"step": 0, **self.validate()}]
+        log(json.dumps(history[-1]))
+        best = history[-1]["loss"]
+        step, tokens_seen, started, window = 0, 0, time.perf_counter(), []
+        while step < config.steps:
+            for group in self._groups(self.train, shuffle=True):
+                step += 1
+                batch = self._batch(group)
+                loss, grads = step_fn(self.model, *batch)
+                grads, _ = optim.clip_grad_norm(grads, 1.0)
+                optimizer.update(self.model, grads)
+                mx.eval(self.model.parameters(), optimizer.state, loss)
+                window.append(float(loss))
+                tokens_seen += batch[0].size
+                if step % config.eval_every == 0 or step == config.steps:
+                    entry = {"step": step, "train_loss": float(np.mean(window)),
+                             "tokens_per_second": tokens_seen / (time.perf_counter() - started), **self.validate()}
+                    window = []
+                    history.append(entry)
+                    log(json.dumps(entry))
+                    if entry["loss"] < best:
+                        best = entry["loss"]
+                        self.save(history)
+                if step == config.steps:
+                    break
+        if best == history[0]["loss"]:
+            self.save(history)
+        return history