Card: compare with the leaderboard's specialist rows; Jev and Decider 1 are zero-shot generalists there. LF line endings.

README.md

@@ -1,116 +1,116 @@
----
-license: apache-2.0
-base_model: google/diffusiongemma-26B-A4B-it
-language: [en]
-library_name: transformers
-tags: [diffusion, hydra, state-space-model, system-one, decision-model, calibrated, jev, typed-decisions]
-datasets: [LocalLLaMA/typed-decisions, Praveenrajus/jev-bench]
----
-
-# Hev: DiffusionGemma with a Hydra decision head
-
-By Dwain Barnes. First published 2026-09-24.
-
-## What it does
-
-Give Hev a lump of text (a support ticket, an agent log, a security alert, whatever your program is looking at) and a list of typed questions, and it gives you back probabilities rather than prose. "Which team should handle this?" comes back as `billing 0.88, technical 0.12, sales 0.00`. "Is this urgent?" comes back as `0.67`. That is the job TypeSafe's Jev does; Hev does it with open weights you can run yourself.
-
-There are two parts:
-
-1. DiffusionGemma, Google's open text-diffusion model, frozen. I never fine-tuned it. One pass of its encoder reads your text and one pass of its bidirectional denoiser reads a small JSON answer canvas.
-2. A Hydra head I trained on top. Hydra is a bidirectional state-space model, the two-way cousin of Mamba. It reads every token of your text and every answer slot together, in linear time, and turns that into calibrated probabilities. It is an ensemble of 3 heads (seeds 42, 43, 44) whose probabilities are averaged, 44,476,866 parameters in total, and it trains in minutes.
-
-## What is new here
-
-As far as I can tell (checked 2026-09-24), nobody had done any of these three things before:
-
-- put Hydra on top of a diffusion language model's representations;
-- used a state-space model anywhere inside a System-One or Jev-style decision model. Every open clone I found is an autoregressive Qwen or Gemma with LoRA, or a BERT-style encoder with a linear head;
-- published an open-weights entry that scores above Jev on the typed-decisions benchmark (specialist setting, fitted on the train split): clearly on Brier, and by a margin within the noise on accuracy, as the results below spell out.
-
-The idea, the name and the direction are mine. If you build on it, please cite this repo; the commit history is the timestamp.
-
-## Results
-
-Scored on the `test` split of [LocalLLaMA/typed-decisions](https://huggingface.co/datasets/LocalLLaMA/typed-decisions) (400 cases, 2,000 decisions) in the specialist setting (the head is fitted on the train split), using the metrics defined on that dataset's card (see the evaluation notes under Limits). The "zero-shot canvas read" row is DiffusionGemma with no head, reading its own option probabilities off the canvas, which is what the vLLM structured-diffusion mode does.
-
-| Model | n | acc | soft_acc | macro_f1 | kl | tv | brier | nll | ece | score_mae | within_1 |
-|---|---|---|---|---|---|---|---|---|---|---|---|
-| DiffusionGemma zero-shot canvas read | 2000 | 0.708 | 0.573 | 0.551 | 1.333 | 0.362 | 0.304 | 2.100 | 0.213 | 0.411 | 0.934 |
-| Hev | 2000 | 0.731 | 0.584 | 0.626 | 0.185 | 0.213 | 0.094 | 0.952 | 0.162 | 0.303 | 0.976 |
-
-For reference, the dataset's own leaderboard lists TypeSafe Jev 1.13.0 in the specialist setting (fitted on the train split) at accuracy 0.727 and Brier 0.148. The dataset card says specialist and generalist numbers are not comparable, so compare Hev only with specialist entries. That accuracy lead is 0.4 points on 2,000 decisions, less than one standard error (1.0 points), so it is a point estimate ahead of Jev rather than a settled result; the Brier lead of 0.054 is well outside noise. Hev's head lifts the frozen backbone's zero-shot accuracy from 0.708 to 0.731 and reaches Brier 0.094.
-
-## How to run it
-
-You need a GPU with 80 GB of memory for the bf16 backbone. The head itself is tiny and runs anywhere; the backbone is the heavy part. Loading with `load_in_4bit=True` is untested: DiffusionGemma's fused expert parameters are not `nn.Linear`, so bitsandbytes leaves most weights in bf16 and the memory saving is small.
-
-Install:
-
-```bash
-pip install "transformers>=5.17" accelerate datasets safetensors
-```
-
-```python
-from transformers import AutoModel
-
-model = AutoModel.from_pretrained("EryriLabs/DiffusionGemma-26B-A4B-Hev", trust_remote_code=True)
-answer = model.decide(
-    state="Help! My payouts have been failing for 3 days.",
-    questions={
-        "department": {"type": "choice", "instructions": "Which team should handle this?",
-                        "criteria": {"billing": "Payments, invoicing, refunds",
-                                     "technical": "Bugs, outages, integrations",
-                                     "sales": "Pricing, upgrades, new accounts"}},
-        "urgent": {"type": "noul", "instructions": "This needs attention today."},
-        "severity": {"type": "score", "instructions": "How bad is the customer impact?",
-                      "criteria": ["none", "minor", "major", "critical"]},
-    },
-)
-print(answer["answers"])
-```
-
-The request and response shapes match Jev's `/v1/systemone` API, so existing client code needs no changes. To serve that endpoint locally, `pip install fastapi uvicorn` and run:
-
-```python
-import uvicorn
-from transformers import AutoModel
-from transformers.dynamic_module_utils import get_class_from_dynamic_module
-
-model = AutoModel.from_pretrained("EryriLabs/DiffusionGemma-26B-A4B-Hev", trust_remote_code=True)
-build_app = get_class_from_dynamic_module("client.build_app", "EryriLabs/DiffusionGemma-26B-A4B-Hev")
-uvicorn.run(build_app(model), host="127.0.0.1", port=8080)
-```
-
-The `.py` files in this repo are a flattened copy of the `hev` package and use relative imports, so they cannot be run as scripts.
-
-## How it was trained
-
-- Features: one forward pass per example through the frozen backbone, canvas length 64, slot seed mode `pad`, hidden states from the final encoder and decoder norms, plus the unit-RMS outputs of layers 14 and 22 of the encoder for the prompt rows and of the decoder for the answer rows, so each row carries 8448 features.
-- Head: linear projection 8448 to 512, 3 Hydra blocks, pointer readout (each answer slot scored against each option's description span), plus a learned residual on the backbone's own option log-probabilities, plus fitted temperatures per question type for each member (seed 42: [1.339, 1.776, 1.816]; seed 43: [1.401, 1.899, 1.832]; seed 44: [1.39, 1.616, 1.888]).
-- Loss: KL to the gold distribution plus 0.1 times Brier. AdamW, cosine schedule, early stopping on validation log loss.
-- Data: typed-decisions train split plus samples from 19 jev-bench sources. Dataset licences are as declared by each source on the Hub.
-- Hydra is implemented in plain PyTorch in `hydra_torch.py`. Check against the official CUDA kernels: skipped (session 4 reuses session 2's kernel check).
-
-## Limits
-
-- Accuracy on typed-decisions measures agreement with a synthetic teacher, not ground truth. The teacher's own self-agreement is 0.735.
-- Version one supports 2 to 64 options per choice question, 2 to 10 score levels, and prompts up to 1,024 tokens. English only. Text only.
-- Request size: every question takes a few canvas tokens (its id plus the JSON around the answer slot), so the 64-token answer canvas holds about 8 questions with short ids. A longer request fails with a clear error (HTTP 422 from the server).
-- Evaluation notes: `ece` in the table is top-label expected calibration error with 15 equal-width bins, and is not comparable with the ECE column on the dataset card. `acc` counts a hit when the top option matches the card's discrete `label` for each typed-decisions question, so ties in the gold distribution are broken the way the dataset breaks them; `soft_acc` is the gold probability of the top option.
-- No GGUF yet. llama.cpp's DiffusionGemma support is still a draft branch and cannot expose the hidden states the head needs, and there is no Hydra kernel in ggml. If that changes I will port it. Until then this ships as safetensors.
-- On fastino/fast-decisions, 17 domains the head never saw, it scores 0.575 exact match on the 2,600 single-label tasks against 0.611 for the backbone's own zero-shot read and 0.670 for GLiNER2.5-Decide run the same way, so as a generalist the head is a step back from the backbone; it is a specialist for the schema it was trained on.
-- The backbone is Google's, unchanged, Apache-2.0. Everything else here is Apache-2.0 too.
-
-## Citation
-
-```bibtex
-@misc{barnes2026hev,
-  title  = {Hev: A Hydra decision head on frozen DiffusionGemma for calibrated System-One decisions},
-  author = {Barnes, Dwain},
-  year   = {2026},
-  url    = {https://huggingface.co/EryriLabs/DiffusionGemma-26B-A4B-Hev}
-}
-```
-
+---
+license: apache-2.0
+base_model: google/diffusiongemma-26B-A4B-it
+language: [en]
+library_name: transformers
+tags: [diffusion, hydra, state-space-model, system-one, decision-model, calibrated, jev, typed-decisions]
+datasets: [LocalLLaMA/typed-decisions, Praveenrajus/jev-bench]
+---
+
+# Hev: DiffusionGemma with a Hydra decision head
+
+By Dwain Barnes. First published 2026-09-24.
+
+## What it does
+
+Give Hev a lump of text (a support ticket, an agent log, a security alert, whatever your program is looking at) and a list of typed questions, and it gives you back probabilities rather than prose. "Which team should handle this?" comes back as `billing 0.88, technical 0.12, sales 0.00`. "Is this urgent?" comes back as `0.67`. That is the job TypeSafe's Jev does; Hev does it with open weights you can run yourself.
+
+There are two parts:
+
+1. DiffusionGemma, Google's open text-diffusion model, frozen. I never fine-tuned it. One pass of its encoder reads your text and one pass of its bidirectional denoiser reads a small JSON answer canvas.
+2. A Hydra head I trained on top. Hydra is a bidirectional state-space model, the two-way cousin of Mamba. It reads every token of your text and every answer slot together, in linear time, and turns that into calibrated probabilities. It is an ensemble of 3 heads (seeds 42, 43, 44) whose probabilities are averaged, 44,476,866 parameters in total, and it trains in minutes.
+
+## What is new here
+
+As far as I can tell (checked 2026-09-24), nobody had done any of these three things before:
+
+- put Hydra on top of a diffusion language model's representations;
+- used a state-space model anywhere inside a System-One or Jev-style decision model. Every open clone I found is an autoregressive Qwen or Gemma with LoRA, or a BERT-style encoder with a linear head;
+- published an open-weights specialist that leads the specialist rows of the typed-decisions leaderboard (fitted on the train split, as they are). How it stands next to the zero-shot generalists, Jev included, is spelled out under Results, with the dataset's own warning that the two kinds of number are not comparable.
+
+The idea, the name and the direction are mine. If you build on it, please cite this repo; the commit history is the timestamp.
+
+## Results
+
+Scored on the `test` split of [LocalLLaMA/typed-decisions](https://huggingface.co/datasets/LocalLLaMA/typed-decisions) (400 cases, 2,000 decisions) in the specialist setting (the head is fitted on the train split), using the metrics defined on that dataset's card (see the evaluation notes under Limits). The "zero-shot canvas read" row is DiffusionGemma with no head, reading its own option probabilities off the canvas, which is what the vLLM structured-diffusion mode does.
+
+| Model | n | acc | soft_acc | macro_f1 | kl | tv | brier | nll | ece | score_mae | within_1 |
+|---|---|---|---|---|---|---|---|---|---|---|---|
+| DiffusionGemma zero-shot canvas read | 2000 | 0.708 | 0.573 | 0.551 | 1.333 | 0.362 | 0.304 | 2.100 | 0.213 | 0.411 | 0.934 |
+| Hev | 2000 | 0.731 | 0.584 | 0.626 | 0.185 | 0.213 | 0.094 | 0.952 | 0.162 | 0.303 | 0.976 |
+
+Hev is a specialist: its head was fitted on the train split of the same four workflows it is scored on, and it cannot answer a question schema it was not fitted for. On the dataset's leaderboard (read 2026-09-25) the comparable rows are the specialists fitted per workflow: ModernBERT-base at accuracy 0.646 and Brier 0.119, and MiniLM-L6 at 0.587 and 0.143. Hev is 8.5 points above the best of them on accuracy and its Brier is 0.025 lower. The generalist rows, scored zero-shot on workflows they never saw, are TypeSafe Jev 1.13.0 at 0.727 and 0.148 and meraGPT Decider 1 at 0.768 and 0.052. Hev's accuracy is 0.4 points above Jev's (less than one standard error, 1.0 points, on 2,000 decisions) and its Brier is 0.054 lower, but the dataset card is explicit that a specialist number next to a generalist one is not a ranking, and Decider 1 leads both on accuracy and on Brier. Hev's head lifts the frozen backbone's zero-shot accuracy from 0.708 to 0.731 and reaches Brier 0.094.
+
+## How to run it
+
+You need a GPU with 80 GB of memory for the bf16 backbone. The head itself is tiny and runs anywhere; the backbone is the heavy part. Loading with `load_in_4bit=True` is untested: DiffusionGemma's fused expert parameters are not `nn.Linear`, so bitsandbytes leaves most weights in bf16 and the memory saving is small.
+
+Install:
+
+```bash
+pip install "transformers>=5.17" accelerate datasets safetensors
+```
+
+```python
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained("EryriLabs/DiffusionGemma-26B-A4B-Hev", trust_remote_code=True)
+answer = model.decide(
+    state="Help! My payouts have been failing for 3 days.",
+    questions={
+        "department": {"type": "choice", "instructions": "Which team should handle this?",
+                        "criteria": {"billing": "Payments, invoicing, refunds",
+                                     "technical": "Bugs, outages, integrations",
+                                     "sales": "Pricing, upgrades, new accounts"}},
+        "urgent": {"type": "noul", "instructions": "This needs attention today."},
+        "severity": {"type": "score", "instructions": "How bad is the customer impact?",
+                      "criteria": ["none", "minor", "major", "critical"]},
+    },
+)
+print(answer["answers"])
+```
+
+The request and response shapes match Jev's `/v1/systemone` API, so existing client code needs no changes. To serve that endpoint locally, `pip install fastapi uvicorn` and run:
+
+```python
+import uvicorn
+from transformers import AutoModel
+from transformers.dynamic_module_utils import get_class_from_dynamic_module
+
+model = AutoModel.from_pretrained("EryriLabs/DiffusionGemma-26B-A4B-Hev", trust_remote_code=True)
+build_app = get_class_from_dynamic_module("client.build_app", "EryriLabs/DiffusionGemma-26B-A4B-Hev")
+uvicorn.run(build_app(model), host="127.0.0.1", port=8080)
+```
+
+The `.py` files in this repo are a flattened copy of the `hev` package and use relative imports, so they cannot be run as scripts.
+
+## How it was trained
+
+- Features: one forward pass per example through the frozen backbone, canvas length 64, slot seed mode `pad`, hidden states from the final encoder and decoder norms, plus the unit-RMS outputs of layers 14 and 22 of the encoder for the prompt rows and of the decoder for the answer rows, so each row carries 8448 features.
+- Head: linear projection 8448 to 512, 3 Hydra blocks, pointer readout (each answer slot scored against each option's description span), plus a learned residual on the backbone's own option log-probabilities, plus fitted temperatures per question type for each member (seed 42: [1.339, 1.776, 1.816]; seed 43: [1.401, 1.899, 1.832]; seed 44: [1.39, 1.616, 1.888]).
+- Loss: KL to the gold distribution plus 0.1 times Brier. AdamW, cosine schedule, early stopping on validation log loss.
+- Data: typed-decisions train split plus samples from 19 jev-bench sources. Dataset licences are as declared by each source on the Hub.
+- Hydra is implemented in plain PyTorch in `hydra_torch.py`. Check against the official CUDA kernels: skipped (session 4 reuses session 2's kernel check).
+
+## Limits
+
+- Accuracy on typed-decisions measures agreement with a synthetic teacher, not ground truth. The teacher's own self-agreement is 0.735.
+- Version one supports 2 to 64 options per choice question, 2 to 10 score levels, and prompts up to 1,024 tokens. English only. Text only.
+- Request size: every question takes a few canvas tokens (its id plus the JSON around the answer slot), so the 64-token answer canvas holds about 8 questions with short ids. A longer request fails with a clear error (HTTP 422 from the server).
+- Evaluation notes: `ece` in the table is top-label expected calibration error with 15 equal-width bins, and is not comparable with the ECE column on the dataset card. `acc` counts a hit when the top option matches the card's discrete `label` for each typed-decisions question, so ties in the gold distribution are broken the way the dataset breaks them; `soft_acc` is the gold probability of the top option.
+- No GGUF yet. llama.cpp's DiffusionGemma support is still a draft branch and cannot expose the hidden states the head needs, and there is no Hydra kernel in ggml. If that changes I will port it. Until then this ships as safetensors.
+- On fastino/fast-decisions, 17 domains the head never saw, it scores 0.575 exact match on the 2,600 single-label tasks against 0.611 for the backbone's own zero-shot read and 0.670 for GLiNER2.5-Decide run the same way, so as a generalist the head is a step back from the backbone; it is a specialist for the schema it was trained on.
+- The backbone is Google's, unchanged, Apache-2.0. Everything else here is Apache-2.0 too.
+
+## Citation
+
+```bibtex
+@misc{barnes2026hev,
+  title  = {Hev: A Hydra decision head on frozen DiffusionGemma for calibrated System-One decisions},
+  author = {Barnes, Dwain},
+  year   = {2026},
+  url    = {https://huggingface.co/EryriLabs/DiffusionGemma-26B-A4B-Hev}
+}
+```
+
 Thanks to Google DeepMind for DiffusionGemma, to Hwang, Lahoti, Dao and Gu for Hydra, to Matt Mastracci for showing that a diffusion model's canvas is a decision engine, and to the LocalLLaMA community for the typed-decisions benchmark.

cache.py

@@ -1,185 +1,185 @@
-"""On-disk feature cache: one safetensors file per example plus a JSONL manifest."""
-from __future__ import annotations
-
-import json
-from dataclasses import asdict, dataclass, field
-from pathlib import Path
-from typing import Optional
-
-import torch
-from safetensors.torch import load_file, save_file
-
-from .schema import TYPE_INDEX
-
-
-@dataclass
-class Decision:
-    qid: str
-    type: str
-    slot: int
-    spans: list[list[int]]
-    alias: list[int]
-    prior: list[float]
-    gold: Optional[list[float]]
-
-
-@dataclass
-class Record:
-    id: str
-    source: str
-    path: str
-    n_prompt: int
-    n_canvas: int
-    decisions: list[Decision] = field(default_factory=list)
-
-    def to_json(self) -> str:
-        return json.dumps(asdict(self))
-
-    @classmethod
-    def from_json(cls, line: str) -> "Record":
-        d = json.loads(line)
-        d["decisions"] = [Decision(**x) for x in d["decisions"]]
-        return cls(**d)
-
-
-class CacheWriter:
-    """Writes one safetensors file per example and a JSONL manifest.
-
-    A writer always starts a fresh manifest, so re-running into the same folder does not duplicate records.
-    """
-
-    def __init__(self, out_dir):
-        self.dir = Path(out_dir)
-        self.dir.mkdir(parents=True, exist_ok=True)
-        self._f = (self.dir / "manifest.jsonl").open("w", encoding="utf-8")
-
-    def write(self, record: Record, feats: torch.Tensor) -> None:
-        save_file({"feats": feats.detach().to("cpu", torch.float16).contiguous()}, str(self.dir / record.path))
-        self._f.write(record.to_json() + "\n")
-        self._f.flush()
-
-    def close(self) -> None:
-        self._f.close()
-
-
-class FeatureCache:
-    def __init__(self, cache_dir):
-        self.dir = Path(cache_dir)
-        with (self.dir / "manifest.jsonl").open("r", encoding="utf-8") as f:
-            self.records = [Record.from_json(line) for line in f if line.strip()]
-
-    def __len__(self) -> int:
-        return len(self.records)
-
-    def __getitem__(self, i: int) -> tuple[Record, torch.Tensor]:
-        rec = self.records[i]
-        return rec, load_file(str(self.dir / rec.path))["feats"]
-
-
-@dataclass
-class Batch:
-    """A padded feature batch plus two views of its decisions.
-
-    `decisions` is the per-decision list (metadata, spans, prior, gold) that evaluation, serving and the tests
-    read. The flat tensors are what the head and the losses compute on: options are numbered decision by
-    decision in that same order, so one gather, one index_add and one padded log-softmax cover the whole batch
-    instead of a Python loop issuing hundreds of tiny ops per decision (session 3's CPU-bound trainer).
-    """
-    feats: torch.Tensor
-    seg: torch.Tensor
-    mask: torch.Tensor
-    decisions: list[dict]
-    dec_b: torch.Tensor  # (n_dec,) batch row of each decision
-    dec_slot: torch.Tensor  # (n_dec,) answer-slot row of each decision
-    dec_type: torch.Tensor  # (n_dec,) question type id
-    dec_has_gold: torch.Tensor  # (n_dec,) bool
-    opt_dec: torch.Tensor  # (n_opt,) decision index of each option
-    opt_pos: torch.Tensor  # (n_opt,) position of each option within its decision
-    opt_prior: torch.Tensor  # (n_opt,) backbone log-prior per option
-    opt_gold: torch.Tensor  # (n_opt,) gold probability per option, 0 where the decision has no gold
-    span_rows: torch.Tensor  # (n_span_rows,) row index into the padded (B * L) feature rows, option by option
-    span_opt: torch.Tensor  # (n_span_rows,) option index of each span row
-    span_len: torch.Tensor  # (n_opt,) span length per option, float
-    counts: list[int]  # options per decision; plain Python so torch.split never syncs on it
-    n_gold: int  # decisions with gold; plain Python so the losses never sync on it
-
-    _FLAT = ("dec_b", "dec_slot", "dec_type", "dec_has_gold", "opt_dec", "opt_pos", "opt_prior", "opt_gold",
-             "span_rows", "span_opt", "span_len")
-
-    def to(self, device) -> "Batch":
-        # spans stay on the CPU: consumers read them back as Python ints, so a device copy would only add a
-        # device-to-host sync per decision.
-        return Batch(self.feats.to(device), self.seg.to(device), self.mask.to(device),
-                     [{**d, "prior": d["prior"].to(device),
-                       "gold": None if d["gold"] is None else d["gold"].to(device)} for d in self.decisions],
-                     *(getattr(self, name).to(device) for name in self._FLAT), self.counts, self.n_gold)
-
-
-def _check_record(rec: Record, n_rows: int) -> None:
-    """Raise a ValueError naming the record and question when a record does not fit its feature rows.
-
-    Rows [0, n_prompt) hold the prompt and [n_prompt, n_prompt + n_canvas) the canvas. Without this check a bad
-    record would be read silently from padding or from the wrong segment.
-    """
-    if n_rows != rec.n_prompt + rec.n_canvas:
-        raise ValueError(f"record {rec.id!r}: {n_rows} feature rows, but n_prompt {rec.n_prompt} + n_canvas "
-                         f"{rec.n_canvas} = {rec.n_prompt + rec.n_canvas}")
-    for d in rec.decisions:
-        where = f"record {rec.id!r}, question {d.qid!r}"
-        if not rec.n_prompt <= d.slot < n_rows:
-            raise ValueError(f"{where}: slot {d.slot} is outside the canvas rows [{rec.n_prompt}, {n_rows})")
-        for s, e in d.spans:
-            if not 0 <= s < e <= rec.n_prompt:
-                raise ValueError(f"{where}: option span [{s}, {e}) is empty or outside the prompt rows "
-                                 f"[0, {rec.n_prompt})")
-        if len(d.prior) != len(d.spans):
-            raise ValueError(f"{where}: {len(d.spans)} option spans but {len(d.prior)} prior values")
-        if d.gold is not None and len(d.gold) != len(d.spans):
-            raise ValueError(f"{where}: {len(d.spans)} option spans but {len(d.gold)} gold values")
-
-
-def collate(items: list[tuple[Record, torch.Tensor]]) -> Batch:
-    lmax = max(f.shape[0] for _, f in items)
-    dim = items[0][1].shape[1]
-    feats = torch.zeros(len(items), lmax, dim, dtype=torch.float16)
-    seg = torch.zeros(len(items), lmax, dtype=torch.long)
-    mask = torch.zeros(len(items), lmax, dtype=torch.bool)
-    decisions = []
-    dec_b, dec_slot, dec_type, dec_has_gold, counts = [], [], [], [], []
-    opt_dec, opt_pos, opt_prior, opt_gold, span_rows, span_opt, span_len = [], [], [], [], [], [], []
-    for b, (rec, f) in enumerate(items):
-        n = f.shape[0]
-        _check_record(rec, n)
-        feats[b, :n] = f.to(torch.float16)
-        mask[b, :n] = True
-        seg[b, rec.n_prompt:n] = 1
-        for d in rec.decisions:
-            i = len(decisions)
-            decisions.append({
-                "b": b, "slot": d.slot, "type_id": TYPE_INDEX[d.type], "qid": d.qid, "type": d.type,
-                "spans": torch.tensor(d.spans, dtype=torch.long),
-                "prior": torch.tensor(d.prior, dtype=torch.float32),
-                "gold": None if d.gold is None else torch.tensor(d.gold, dtype=torch.float32),
-            })
-            dec_b.append(b)
-            dec_slot.append(d.slot)
-            dec_type.append(TYPE_INDEX[d.type])
-            dec_has_gold.append(d.gold is not None)
-            counts.append(len(d.spans))
-            for j, (s, e) in enumerate(d.spans):
-                o = len(opt_dec)
-                opt_dec.append(i)
-                opt_pos.append(j)
-                opt_prior.append(d.prior[j])
-                opt_gold.append(0.0 if d.gold is None else d.gold[j])
-                span_rows.extend(range(b * lmax + s, b * lmax + e))
-                span_opt.extend([o] * (e - s))
-                span_len.append(float(e - s))
-    long, f32 = torch.long, torch.float32
-    return Batch(feats=feats, seg=seg, mask=mask, decisions=decisions,
-                 dec_b=torch.tensor(dec_b, dtype=long), dec_slot=torch.tensor(dec_slot, dtype=long),
-                 dec_type=torch.tensor(dec_type, dtype=long), dec_has_gold=torch.tensor(dec_has_gold, dtype=torch.bool),
-                 opt_dec=torch.tensor(opt_dec, dtype=long), opt_pos=torch.tensor(opt_pos, dtype=long),
-                 opt_prior=torch.tensor(opt_prior, dtype=f32), opt_gold=torch.tensor(opt_gold, dtype=f32),
-                 span_rows=torch.tensor(span_rows, dtype=long), span_opt=torch.tensor(span_opt, dtype=long),
-                 span_len=torch.tensor(span_len, dtype=f32), counts=counts, n_gold=sum(dec_has_gold))
+"""On-disk feature cache: one safetensors file per example plus a JSONL manifest."""
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Optional
+
+import torch
+from safetensors.torch import load_file, save_file
+
+from .schema import TYPE_INDEX
+
+
+@dataclass
+class Decision:
+    qid: str
+    type: str
+    slot: int
+    spans: list[list[int]]
+    alias: list[int]
+    prior: list[float]
+    gold: Optional[list[float]]
+
+
+@dataclass
+class Record:
+    id: str
+    source: str
+    path: str
+    n_prompt: int
+    n_canvas: int
+    decisions: list[Decision] = field(default_factory=list)
+
+    def to_json(self) -> str:
+        return json.dumps(asdict(self))
+
+    @classmethod
+    def from_json(cls, line: str) -> "Record":
+        d = json.loads(line)
+        d["decisions"] = [Decision(**x) for x in d["decisions"]]
+        return cls(**d)
+
+
+class CacheWriter:
+    """Writes one safetensors file per example and a JSONL manifest.
+
+    A writer always starts a fresh manifest, so re-running into the same folder does not duplicate records.
+    """
+
+    def __init__(self, out_dir):
+        self.dir = Path(out_dir)
+        self.dir.mkdir(parents=True, exist_ok=True)
+        self._f = (self.dir / "manifest.jsonl").open("w", encoding="utf-8")
+
+    def write(self, record: Record, feats: torch.Tensor) -> None:
+        save_file({"feats": feats.detach().to("cpu", torch.float16).contiguous()}, str(self.dir / record.path))
+        self._f.write(record.to_json() + "\n")
+        self._f.flush()
+
+    def close(self) -> None:
+        self._f.close()
+
+
+class FeatureCache:
+    def __init__(self, cache_dir):
+        self.dir = Path(cache_dir)
+        with (self.dir / "manifest.jsonl").open("r", encoding="utf-8") as f:
+            self.records = [Record.from_json(line) for line in f if line.strip()]
+
+    def __len__(self) -> int:
+        return len(self.records)
+
+    def __getitem__(self, i: int) -> tuple[Record, torch.Tensor]:
+        rec = self.records[i]
+        return rec, load_file(str(self.dir / rec.path))["feats"]
+
+
+@dataclass
+class Batch:
+    """A padded feature batch plus two views of its decisions.
+
+    `decisions` is the per-decision list (metadata, spans, prior, gold) that evaluation, serving and the tests
+    read. The flat tensors are what the head and the losses compute on: options are numbered decision by
+    decision in that same order, so one gather, one index_add and one padded log-softmax cover the whole batch
+    instead of a Python loop issuing hundreds of tiny ops per decision (session 3's CPU-bound trainer).
+    """
+    feats: torch.Tensor
+    seg: torch.Tensor
+    mask: torch.Tensor
+    decisions: list[dict]
+    dec_b: torch.Tensor  # (n_dec,) batch row of each decision
+    dec_slot: torch.Tensor  # (n_dec,) answer-slot row of each decision
+    dec_type: torch.Tensor  # (n_dec,) question type id
+    dec_has_gold: torch.Tensor  # (n_dec,) bool
+    opt_dec: torch.Tensor  # (n_opt,) decision index of each option
+    opt_pos: torch.Tensor  # (n_opt,) position of each option within its decision
+    opt_prior: torch.Tensor  # (n_opt,) backbone log-prior per option
+    opt_gold: torch.Tensor  # (n_opt,) gold probability per option, 0 where the decision has no gold
+    span_rows: torch.Tensor  # (n_span_rows,) row index into the padded (B * L) feature rows, option by option
+    span_opt: torch.Tensor  # (n_span_rows,) option index of each span row
+    span_len: torch.Tensor  # (n_opt,) span length per option, float
+    counts: list[int]  # options per decision; plain Python so torch.split never syncs on it
+    n_gold: int  # decisions with gold; plain Python so the losses never sync on it
+
+    _FLAT = ("dec_b", "dec_slot", "dec_type", "dec_has_gold", "opt_dec", "opt_pos", "opt_prior", "opt_gold",
+             "span_rows", "span_opt", "span_len")
+
+    def to(self, device) -> "Batch":
+        # spans stay on the CPU: consumers read them back as Python ints, so a device copy would only add a
+        # device-to-host sync per decision.
+        return Batch(self.feats.to(device), self.seg.to(device), self.mask.to(device),
+                     [{**d, "prior": d["prior"].to(device),
+                       "gold": None if d["gold"] is None else d["gold"].to(device)} for d in self.decisions],
+                     *(getattr(self, name).to(device) for name in self._FLAT), self.counts, self.n_gold)
+
+
+def _check_record(rec: Record, n_rows: int) -> None:
+    """Raise a ValueError naming the record and question when a record does not fit its feature rows.
+
+    Rows [0, n_prompt) hold the prompt and [n_prompt, n_prompt + n_canvas) the canvas. Without this check a bad
+    record would be read silently from padding or from the wrong segment.
+    """
+    if n_rows != rec.n_prompt + rec.n_canvas:
+        raise ValueError(f"record {rec.id!r}: {n_rows} feature rows, but n_prompt {rec.n_prompt} + n_canvas "
+                         f"{rec.n_canvas} = {rec.n_prompt + rec.n_canvas}")
+    for d in rec.decisions:
+        where = f"record {rec.id!r}, question {d.qid!r}"
+        if not rec.n_prompt <= d.slot < n_rows:
+            raise ValueError(f"{where}: slot {d.slot} is outside the canvas rows [{rec.n_prompt}, {n_rows})")
+        for s, e in d.spans:
+            if not 0 <= s < e <= rec.n_prompt:
+                raise ValueError(f"{where}: option span [{s}, {e}) is empty or outside the prompt rows "
+                                 f"[0, {rec.n_prompt})")
+        if len(d.prior) != len(d.spans):
+            raise ValueError(f"{where}: {len(d.spans)} option spans but {len(d.prior)} prior values")
+        if d.gold is not None and len(d.gold) != len(d.spans):
+            raise ValueError(f"{where}: {len(d.spans)} option spans but {len(d.gold)} gold values")
+
+
+def collate(items: list[tuple[Record, torch.Tensor]]) -> Batch:
+    lmax = max(f.shape[0] for _, f in items)
+    dim = items[0][1].shape[1]
+    feats = torch.zeros(len(items), lmax, dim, dtype=torch.float16)
+    seg = torch.zeros(len(items), lmax, dtype=torch.long)
+    mask = torch.zeros(len(items), lmax, dtype=torch.bool)
+    decisions = []
+    dec_b, dec_slot, dec_type, dec_has_gold, counts = [], [], [], [], []
+    opt_dec, opt_pos, opt_prior, opt_gold, span_rows, span_opt, span_len = [], [], [], [], [], [], []
+    for b, (rec, f) in enumerate(items):
+        n = f.shape[0]
+        _check_record(rec, n)
+        feats[b, :n] = f.to(torch.float16)
+        mask[b, :n] = True
+        seg[b, rec.n_prompt:n] = 1
+        for d in rec.decisions:
+            i = len(decisions)
+            decisions.append({
+                "b": b, "slot": d.slot, "type_id": TYPE_INDEX[d.type], "qid": d.qid, "type": d.type,
+                "spans": torch.tensor(d.spans, dtype=torch.long),
+                "prior": torch.tensor(d.prior, dtype=torch.float32),
+                "gold": None if d.gold is None else torch.tensor(d.gold, dtype=torch.float32),
+            })
+            dec_b.append(b)
+            dec_slot.append(d.slot)
+            dec_type.append(TYPE_INDEX[d.type])
+            dec_has_gold.append(d.gold is not None)
+            counts.append(len(d.spans))
+            for j, (s, e) in enumerate(d.spans):
+                o = len(opt_dec)
+                opt_dec.append(i)
+                opt_pos.append(j)
+                opt_prior.append(d.prior[j])
+                opt_gold.append(0.0 if d.gold is None else d.gold[j])
+                span_rows.extend(range(b * lmax + s, b * lmax + e))
+                span_opt.extend([o] * (e - s))
+                span_len.append(float(e - s))
+    long, f32 = torch.long, torch.float32
+    return Batch(feats=feats, seg=seg, mask=mask, decisions=decisions,
+                 dec_b=torch.tensor(dec_b, dtype=long), dec_slot=torch.tensor(dec_slot, dtype=long),
+                 dec_type=torch.tensor(dec_type, dtype=long), dec_has_gold=torch.tensor(dec_has_gold, dtype=torch.bool),
+                 opt_dec=torch.tensor(opt_dec, dtype=long), opt_pos=torch.tensor(opt_pos, dtype=long),
+                 opt_prior=torch.tensor(opt_prior, dtype=f32), opt_gold=torch.tensor(opt_gold, dtype=f32),
+                 span_rows=torch.tensor(span_rows, dtype=long), span_opt=torch.tensor(span_opt, dtype=long),
+                 span_len=torch.tensor(span_len, dtype=f32), counts=counts, n_gold=sum(dec_has_gold))

client.py

@@ -1,104 +1,104 @@
-"""Jev-shaped request/response wrapper around HevForDecisions, plus an optional FastAPI server."""
-from __future__ import annotations
-
-import argparse
-import json
-import threading
-
-from .schema import Question
-
-MAX_QUESTIONS = 32
-# A question's canvas entry (its id, the JSON around it and the answer slot) takes 6 or more DiffusionGemma tokens,
-# so 4 is a safe floor for the early check below.
-MIN_CANVAS_TOKENS_PER_QUESTION = 4
-
-
-def max_questions(model) -> int:
-    """Upper bound on questions per request for `model`'s answer canvas (32 when the model has no canvas_length).
-
-    Every question needs more than MIN_CANVAS_TOKENS_PER_QUESTION canvas tokens, so this never rejects a request that
-    would fit. With short ids a 64-token canvas holds about 8 questions (9 one-letter ids, 7 word ids); the exact
-    check happens when the canvas is built and raises ValueError, which the server turns into a 422.
-    """
-    canvas = getattr(getattr(model, "config", None), "canvas_length", None)
-    if not isinstance(canvas, int) or isinstance(canvas, bool) or canvas <= 0:
-        return MAX_QUESTIONS
-    return max(1, min(MAX_QUESTIONS, canvas // MIN_CANVAS_TOKENS_PER_QUESTION))
-
-
-def validate_request(request: dict, max_q: int = MAX_QUESTIONS) -> None:
-    if not isinstance(request, dict) or "state" not in request:
-        raise ValueError("request needs a 'state' (string, object or array)")
-    qs = request.get("questions")
-    if not isinstance(qs, dict) or not qs:
-        raise ValueError("request needs a non-empty 'questions' mapping")
-    for qid, spec in qs.items():
-        if not isinstance(spec, dict):
-            raise ValueError(f"question {qid!r} must be an object")
-        try:
-            Question.from_jev(qid, spec)
-        except ValueError as err:
-            raise ValueError(f"question {qid!r}: {err}") from err
-    n = len(qs)
-    if n > max_q:
-        raise ValueError(f"at most {max_q} questions per request fit this model's answer canvas, got {n}")
-
-
-def _state_text(state) -> str:
-    return state if isinstance(state, str) else json.dumps(state, ensure_ascii=False)
-
-
-def decide(model, request: dict) -> dict:
-    validate_request(request, max_questions(model))
-    resp = model.decide(_state_text(request["state"]), request["questions"])
-    resp["model"] = request.get("model", "hev")
-    return resp
-
-
-def build_app(model):
-    """FastAPI app serving `model` at POST /v1/systemone and GET /health.
-
-    FastAPI runs plain `def` endpoints in a thread pool, and HevForDecisions.decide() is not
-    thread-safe (it shares one hook-capture buffer on the backbone), so calls to the model are
-    serialised with a lock: concurrent requests queue and are answered one at a time.
-    """
-    from fastapi import FastAPI, HTTPException
-
-    app = FastAPI(title="Hev", version="0.1.0")
-    model_lock = threading.Lock()
-
-    @app.get("/health")
-    def health():
-        return {"ok": True}
-
-    @app.post("/v1/systemone")
-    def systemone(request: dict):
-        try:
-            with model_lock:
-                return decide(model, request)
-        except ValueError as err:
-            raise HTTPException(status_code=422, detail=str(err)) from err
-
-    return app
-
-
-def main():
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--repo", required=True)
-    ap.add_argument("--port", type=int, default=8080)
-    ap.add_argument("--host", default="127.0.0.1")
-    ap.add_argument("--load-in-4bit", action="store_true")
-    ap.add_argument("--request", default=None, help="JSON file: answer once and exit")
-    args = ap.parse_args()
-    from .modeling_hev import HevForDecisions
-    model = HevForDecisions.from_pretrained(args.repo, load_in_4bit=args.load_in_4bit)
-    if args.request:
-        with open(args.request, "r", encoding="utf-8") as f:
-            print(json.dumps(decide(model, json.load(f)), indent=2))
-        return
-    import uvicorn
-    uvicorn.run(build_app(model), host=args.host, port=args.port)
-
-
-if __name__ == "__main__":
-    main()
+"""Jev-shaped request/response wrapper around HevForDecisions, plus an optional FastAPI server."""
+from __future__ import annotations
+
+import argparse
+import json
+import threading
+
+from .schema import Question
+
+MAX_QUESTIONS = 32
+# A question's canvas entry (its id, the JSON around it and the answer slot) takes 6 or more DiffusionGemma tokens,
+# so 4 is a safe floor for the early check below.
+MIN_CANVAS_TOKENS_PER_QUESTION = 4
+
+
+def max_questions(model) -> int:
+    """Upper bound on questions per request for `model`'s answer canvas (32 when the model has no canvas_length).
+
+    Every question needs more than MIN_CANVAS_TOKENS_PER_QUESTION canvas tokens, so this never rejects a request that
+    would fit. With short ids a 64-token canvas holds about 8 questions (9 one-letter ids, 7 word ids); the exact
+    check happens when the canvas is built and raises ValueError, which the server turns into a 422.
+    """
+    canvas = getattr(getattr(model, "config", None), "canvas_length", None)
+    if not isinstance(canvas, int) or isinstance(canvas, bool) or canvas <= 0:
+        return MAX_QUESTIONS
+    return max(1, min(MAX_QUESTIONS, canvas // MIN_CANVAS_TOKENS_PER_QUESTION))
+
+
+def validate_request(request: dict, max_q: int = MAX_QUESTIONS) -> None:
+    if not isinstance(request, dict) or "state" not in request:
+        raise ValueError("request needs a 'state' (string, object or array)")
+    qs = request.get("questions")
+    if not isinstance(qs, dict) or not qs:
+        raise ValueError("request needs a non-empty 'questions' mapping")
+    for qid, spec in qs.items():
+        if not isinstance(spec, dict):
+            raise ValueError(f"question {qid!r} must be an object")
+        try:
+            Question.from_jev(qid, spec)
+        except ValueError as err:
+            raise ValueError(f"question {qid!r}: {err}") from err
+    n = len(qs)
+    if n > max_q:
+        raise ValueError(f"at most {max_q} questions per request fit this model's answer canvas, got {n}")
+
+
+def _state_text(state) -> str:
+    return state if isinstance(state, str) else json.dumps(state, ensure_ascii=False)
+
+
+def decide(model, request: dict) -> dict:
+    validate_request(request, max_questions(model))
+    resp = model.decide(_state_text(request["state"]), request["questions"])
+    resp["model"] = request.get("model", "hev")
+    return resp
+
+
+def build_app(model):
+    """FastAPI app serving `model` at POST /v1/systemone and GET /health.
+
+    FastAPI runs plain `def` endpoints in a thread pool, and HevForDecisions.decide() is not
+    thread-safe (it shares one hook-capture buffer on the backbone), so calls to the model are
+    serialised with a lock: concurrent requests queue and are answered one at a time.
+    """
+    from fastapi import FastAPI, HTTPException
+
+    app = FastAPI(title="Hev", version="0.1.0")
+    model_lock = threading.Lock()
+
+    @app.get("/health")
+    def health():
+        return {"ok": True}
+
+    @app.post("/v1/systemone")
+    def systemone(request: dict):
+        try:
+            with model_lock:
+                return decide(model, request)
+        except ValueError as err:
+            raise HTTPException(status_code=422, detail=str(err)) from err
+
+    return app
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--repo", required=True)
+    ap.add_argument("--port", type=int, default=8080)
+    ap.add_argument("--host", default="127.0.0.1")
+    ap.add_argument("--load-in-4bit", action="store_true")
+    ap.add_argument("--request", default=None, help="JSON file: answer once and exit")
+    args = ap.parse_args()
+    from .modeling_hev import HevForDecisions
+    model = HevForDecisions.from_pretrained(args.repo, load_in_4bit=args.load_in_4bit)
+    if args.request:
+        with open(args.request, "r", encoding="utf-8") as f:
+            print(json.dumps(decide(model, json.load(f)), indent=2))
+        return
+    import uvicorn
+    uvicorn.run(build_app(model), host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+    main()

config.json

@@ -1,312 +1,312 @@
-{
-  "aliases": [
-    "A",
-    "B",
-    "C",
-    "D",
-    "E",
-    "F",
-    "G",
-    "H",
-    "I",
-    "J",
-    "K",
-    "L",
-    "M",
-    "N",
-    "O",
-    "P",
-    "Q",
-    "R",
-    "S",
-    "T",
-    "U",
-    "V",
-    "W",
-    "X",
-    "Y",
-    "Z",
-    "a",
-    "b",
-    "c",
-    "d",
-    "e",
-    "f",
-    "g",
-    "h",
-    "i",
-    "j",
-    "k",
-    "l",
-    "m",
-    "n",
-    "o",
-    "p",
-    "q",
-    "r",
-    "s",
-    "t",
-    "u",
-    "v",
-    "w",
-    "x",
-    "y",
-    "z",
-    "AA",
-    "AB",
-    "AC",
-    "AD",
-    "AE",
-    "AF",
-    "AG",
-    "AH",
-    "AI",
-    "AJ",
-    "AK",
-    "AL"
-  ],
-  "auto_map": {
-    "AutoConfig": "configuration_hev.HevConfig",
-    "AutoModel": "modeling_hev.HevForDecisions"
-  },
-  "backbone_config": {
-    "_name_or_path": "google/diffusiongemma-26B-A4B-it",
-    "architectures": [
-      "DiffusionGemmaForBlockDiffusion"
-    ],
-    "boi_token_id": 255999,
-    "canvas_length": 256,
-    "chunk_size_feed_forward": 0,
-    "dtype": "bfloat16",
-    "eoi_token_id": 258882,
-    "eos_token_id": [
-      1,
-      106
-    ],
-    "id2label": {
-      "0": "LABEL_0",
-      "1": "LABEL_1"
-    },
-    "image_token_id": 258880,
-    "initializer_range": 0.02,
-    "is_encoder_decoder": false,
-    "label2id": {
-      "LABEL_0": 0,
-      "LABEL_1": 1
-    },
-    "model_type": "diffusion_gemma",
-    "output_attentions": false,
-    "output_hidden_states": false,
-    "problem_type": null,
-    "return_dict": true,
-    "text_config": {
-      "_name_or_path": "",
-      "architectures": null,
-      "attention_bias": false,
-      "attention_dropout": 0.0,
-      "bos_token_id": 2,
-      "chunk_size_feed_forward": 0,
-      "dtype": "bfloat16",
-      "eos_token_id": 1,
-      "final_logit_softcapping": 30.0,
-      "head_dim": 256,
-      "hidden_activation": "gelu_pytorch_tanh",
-      "hidden_size": 2816,
-      "id2label": {
-        "0": "LABEL_0",
-        "1": "LABEL_1"
-      },
-      "initializer_range": 0.02,
-      "intermediate_size": 2112,
-      "is_encoder_decoder": false,
-      "label2id": {
-        "LABEL_0": 0,
-        "LABEL_1": 1
-      },
-      "layer_types": [
-        "sliding_attention",
-        "sliding_attention",
-        "sliding_attention",
-        "sliding_attention",
-        "sliding_attention",
-        "full_attention",
-        "sliding_attention",
-        "sliding_attention",
-        "sliding_attention",
-        "sliding_attention",
-        "sliding_attention",
-        "full_attention",
-        "sliding_attention",
-        "sliding_attention",
-        "sliding_attention",
-        "sliding_attention",
-        "sliding_attention",
-        "full_attention",
-        "sliding_attention",
-        "sliding_attention",
-        "sliding_attention",
-        "sliding_attention",
-        "sliding_attention",
-        "full_attention",
-        "sliding_attention",
-        "sliding_attention",
-        "sliding_attention",
-        "sliding_attention",
-        "sliding_attention",
-        "full_attention"
-      ],
-      "max_position_embeddings": 262144,
-      "model_type": "diffusion_gemma_text",
-      "moe_intermediate_size": 704,
-      "num_attention_heads": 16,
-      "num_experts": 128,
-      "num_hidden_layers": 30,
-      "num_key_value_heads": 8,
-      "output_attentions": false,
-      "output_hidden_states": false,
-      "pad_token_id": 0,
-      "per_layer_config": {
-        "05": {
-          "head_dim": 512,
-          "num_key_value_heads": 2
-        },
-        "11": {
-          "head_dim": 512,
-          "num_key_value_heads": 2
-        },
-        "17": {
-          "head_dim": 512,
-          "num_key_value_heads": 2
-        },
-        "23": {
-          "head_dim": 512,
-          "num_key_value_heads": 2
-        },
-        "29": {
-          "head_dim": 512,
-          "num_key_value_heads": 2
-        }
-      },
-      "problem_type": null,
-      "return_dict": true,
-      "rms_norm_eps": 1e-06,
-      "rope_parameters": {
-        "full_attention": {
-          "partial_rotary_factor": 0.25,
-          "rope_theta": 1000000.0,
-          "rope_type": "proportional"
-        },
-        "sliding_attention": {
-          "rope_theta": 10000.0,
-          "rope_type": "default"
-        }
-      },
-      "sliding_window": 1024,
-      "tie_word_embeddings": true,
-      "top_k_experts": 8,
-      "use_bidirectional_attention": "vision",
-      "vocab_size": 262144
-    },
-    "tie_word_embeddings": true,
-    "transformers_version": "5.17.0",
-    "vision_config": {
-      "_name_or_path": "",
-      "architectures": null,
-      "attention_bias": false,
-      "attention_dropout": 0.0,
-      "chunk_size_feed_forward": 0,
-      "default_output_length": 280,
-      "dtype": "bfloat16",
-      "global_head_dim": 72,
-      "head_dim": 72,
-      "hidden_activation": "gelu_pytorch_tanh",
-      "hidden_size": 1152,
-      "id2label": {
-        "0": "LABEL_0",
-        "1": "LABEL_1"
-      },
-      "initializer_range": 0.02,
-      "intermediate_size": 4304,
-      "is_encoder_decoder": false,
-      "label2id": {
-        "LABEL_0": 0,
-        "LABEL_1": 1
-      },
-      "max_position_embeddings": 131072,
-      "model_type": "gemma4_vision",
-      "num_attention_heads": 16,
-      "num_hidden_layers": 27,
-      "num_key_value_heads": 16,
-      "output_attentions": false,
-      "output_hidden_states": false,
-      "patch_size": 16,
-      "pooling_kernel_size": 3,
-      "position_embedding_size": 10240,
-      "problem_type": null,
-      "return_dict": true,
-      "rms_norm_eps": 1e-06,
-      "rope_parameters": {
-        "rope_theta": 100.0,
-        "rope_type": "axial"
-      },
-      "standardize": true,
-      "use_clipped_linears": false
-    },
-    "vision_soft_tokens_per_image": 280
-  },
-  "backbone_id": "google/diffusiongemma-26B-A4B-it",
-  "canvas_length": 64,
-  "decoder_norm": "model.decoder.norm",
-  "encoder_norm": "model.encoder.language_model.norm",
-  "extra_layers": {
-    "decoder": [
-      "model.decoder.layers.14",
-      "model.decoder.layers.22"
-    ],
-    "encoder": [
-      "model.encoder.language_model.layers.14",
-      "model.encoder.language_model.layers.22"
-    ]
-  },
-  "head_config": {
-    "chunk_size": 64,
-    "d_backbone": 8448,
-    "d_conv": 7,
-    "d_model": 512,
-    "d_state": 64,
-    "expand": 2,
-    "headdim": 64,
-    "mlp_mult": 2,
-    "n_blocks": 3,
-    "n_types": 3
-  },
-  "max_prompt_tokens": 1024,
-  "member_temperatures": [
-    [
-      1.338762879371643,
-      1.776159405708313,
-      1.816263198852539
-    ],
-    [
-      1.4009417295455933,
-      1.899003267288208,
-      1.8320508003234863
-    ],
-    [
-      1.3898975849151611,
-      1.61630380153656,
-      1.888205885887146
-    ]
-  ],
-  "model_type": "hev",
-  "n_heads": 3,
-  "prompt_format_version": 1,
-  "seed_mode": "pad",
-  "slot_token_id": 0,
-  "temperatures": [
-    1.338762879371643,
-    1.776159405708313,
-    1.816263198852539
-  ],
-  "transformers_version": "5.17.0"
-}
+{
+  "aliases": [
+    "A",
+    "B",
+    "C",
+    "D",
+    "E",
+    "F",
+    "G",
+    "H",
+    "I",
+    "J",
+    "K",
+    "L",
+    "M",
+    "N",
+    "O",
+    "P",
+    "Q",
+    "R",
+    "S",
+    "T",
+    "U",
+    "V",
+    "W",
+    "X",
+    "Y",
+    "Z",
+    "a",
+    "b",
+    "c",
+    "d",
+    "e",
+    "f",
+    "g",
+    "h",
+    "i",
+    "j",
+    "k",
+    "l",
+    "m",
+    "n",
+    "o",
+    "p",
+    "q",
+    "r",
+    "s",
+    "t",
+    "u",
+    "v",
+    "w",
+    "x",
+    "y",
+    "z",
+    "AA",
+    "AB",
+    "AC",
+    "AD",
+    "AE",
+    "AF",
+    "AG",
+    "AH",
+    "AI",
+    "AJ",
+    "AK",
+    "AL"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_hev.HevConfig",
+    "AutoModel": "modeling_hev.HevForDecisions"
+  },
+  "backbone_config": {
+    "_name_or_path": "google/diffusiongemma-26B-A4B-it",
+    "architectures": [
+      "DiffusionGemmaForBlockDiffusion"
+    ],
+    "boi_token_id": 255999,
+    "canvas_length": 256,
+    "chunk_size_feed_forward": 0,
+    "dtype": "bfloat16",
+    "eoi_token_id": 258882,
+    "eos_token_id": [
+      1,
+      106
+    ],
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "image_token_id": 258880,
+    "initializer_range": 0.02,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "model_type": "diffusion_gemma",
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "problem_type": null,
+    "return_dict": true,
+    "text_config": {
+      "_name_or_path": "",
+      "architectures": null,
+      "attention_bias": false,
+      "attention_dropout": 0.0,
+      "bos_token_id": 2,
+      "chunk_size_feed_forward": 0,
+      "dtype": "bfloat16",
+      "eos_token_id": 1,
+      "final_logit_softcapping": 30.0,
+      "head_dim": 256,
+      "hidden_activation": "gelu_pytorch_tanh",
+      "hidden_size": 2816,
+      "id2label": {
+        "0": "LABEL_0",
+        "1": "LABEL_1"
+      },
+      "initializer_range": 0.02,
+      "intermediate_size": 2112,
+      "is_encoder_decoder": false,
+      "label2id": {
+        "LABEL_0": 0,
+        "LABEL_1": 1
+      },
+      "layer_types": [
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention"
+      ],
+      "max_position_embeddings": 262144,
+      "model_type": "diffusion_gemma_text",
+      "moe_intermediate_size": 704,
+      "num_attention_heads": 16,
+      "num_experts": 128,
+      "num_hidden_layers": 30,
+      "num_key_value_heads": 8,
+      "output_attentions": false,
+      "output_hidden_states": false,
+      "pad_token_id": 0,
+      "per_layer_config": {
+        "05": {
+          "head_dim": 512,
+          "num_key_value_heads": 2
+        },
+        "11": {
+          "head_dim": 512,
+          "num_key_value_heads": 2
+        },
+        "17": {
+          "head_dim": 512,
+          "num_key_value_heads": 2
+        },
+        "23": {
+          "head_dim": 512,
+          "num_key_value_heads": 2
+        },
+        "29": {
+          "head_dim": 512,
+          "num_key_value_heads": 2
+        }
+      },
+      "problem_type": null,
+      "return_dict": true,
+      "rms_norm_eps": 1e-06,
+      "rope_parameters": {
+        "full_attention": {
+          "partial_rotary_factor": 0.25,
+          "rope_theta": 1000000.0,
+          "rope_type": "proportional"
+        },
+        "sliding_attention": {
+          "rope_theta": 10000.0,
+          "rope_type": "default"
+        }
+      },
+      "sliding_window": 1024,
+      "tie_word_embeddings": true,
+      "top_k_experts": 8,
+      "use_bidirectional_attention": "vision",
+      "vocab_size": 262144
+    },
+    "tie_word_embeddings": true,
+    "transformers_version": "5.17.0",
+    "vision_config": {
+      "_name_or_path": "",
+      "architectures": null,
+      "attention_bias": false,
+      "attention_dropout": 0.0,
+      "chunk_size_feed_forward": 0,
+      "default_output_length": 280,
+      "dtype": "bfloat16",
+      "global_head_dim": 72,
+      "head_dim": 72,
+      "hidden_activation": "gelu_pytorch_tanh",
+      "hidden_size": 1152,
+      "id2label": {
+        "0": "LABEL_0",
+        "1": "LABEL_1"
+      },
+      "initializer_range": 0.02,
+      "intermediate_size": 4304,
+      "is_encoder_decoder": false,
+      "label2id": {
+        "LABEL_0": 0,
+        "LABEL_1": 1
+      },
+      "max_position_embeddings": 131072,
+      "model_type": "gemma4_vision",
+      "num_attention_heads": 16,
+      "num_hidden_layers": 27,
+      "num_key_value_heads": 16,
+      "output_attentions": false,
+      "output_hidden_states": false,
+      "patch_size": 16,
+      "pooling_kernel_size": 3,
+      "position_embedding_size": 10240,
+      "problem_type": null,
+      "return_dict": true,
+      "rms_norm_eps": 1e-06,
+      "rope_parameters": {
+        "rope_theta": 100.0,
+        "rope_type": "axial"
+      },
+      "standardize": true,
+      "use_clipped_linears": false
+    },
+    "vision_soft_tokens_per_image": 280
+  },
+  "backbone_id": "google/diffusiongemma-26B-A4B-it",
+  "canvas_length": 64,
+  "decoder_norm": "model.decoder.norm",
+  "encoder_norm": "model.encoder.language_model.norm",
+  "extra_layers": {
+    "decoder": [
+      "model.decoder.layers.14",
+      "model.decoder.layers.22"
+    ],
+    "encoder": [
+      "model.encoder.language_model.layers.14",
+      "model.encoder.language_model.layers.22"
+    ]
+  },
+  "head_config": {
+    "chunk_size": 64,
+    "d_backbone": 8448,
+    "d_conv": 7,
+    "d_model": 512,
+    "d_state": 64,
+    "expand": 2,
+    "headdim": 64,
+    "mlp_mult": 2,
+    "n_blocks": 3,
+    "n_types": 3
+  },
+  "max_prompt_tokens": 1024,
+  "member_temperatures": [
+    [
+      1.338762879371643,
+      1.776159405708313,
+      1.816263198852539
+    ],
+    [
+      1.4009417295455933,
+      1.899003267288208,
+      1.8320508003234863
+    ],
+    [
+      1.3898975849151611,
+      1.61630380153656,
+      1.888205885887146
+    ]
+  ],
+  "model_type": "hev",
+  "n_heads": 3,
+  "prompt_format_version": 1,
+  "seed_mode": "pad",
+  "slot_token_id": 0,
+  "temperatures": [
+    1.338762879371643,
+    1.776159405708313,
+    1.816263198852539
+  ],
+  "transformers_version": "5.17.0"
+}

configuration_hev.py

@@ -1,47 +1,47 @@
-"""HevConfig: everything needed to rebuild the frozen backbone, the head, and the prompt protocol."""
-from __future__ import annotations
-
-try:
-    from transformers import PreTrainedConfig
-except ImportError:  # transformers < 5
-    from transformers import PretrainedConfig as PreTrainedConfig
-
-from .head import HevHeadConfig
-from .schema import ALIASES
-
-
-class HevConfig(PreTrainedConfig):
-    model_type = "hev"
-
-    def __init__(self, backbone_id: str = "google/diffusiongemma-26B-A4B-it", backbone_config: dict | None = None,
-                 head_config: dict | None = None, canvas_length: int = 64, seed_mode: str = "pad",
-                 slot_token_id: int = 0, encoder_norm: str = "", decoder_norm: str = "",
-                 max_prompt_tokens: int = 1024, temperatures: list[float] | None = None,
-                 n_heads: int = 1, member_temperatures: list[list[float]] | None = None,
-                 aliases: list[str] | None = None, prompt_format_version: int = 1,
-                 extra_layers: dict | None = None, **kwargs):
-        self.backbone_id = backbone_id
-        # Module names of the intermediate encoder/decoder layers whose unit-RMS outputs the features append to
-        # the final-norm block (hev.features, --layers). Empty for a head trained on the final norms alone.
-        extra_layers = extra_layers or {}
-        self.extra_layers = {"encoder": list(extra_layers.get("encoder", [])),
-                             "decoder": list(extra_layers.get("decoder", []))}
-        self.backbone_config = backbone_config or {}
-        self.head_config = head_config or HevHeadConfig().to_dict()
-        self.canvas_length = canvas_length
-        self.seed_mode = seed_mode
-        self.slot_token_id = slot_token_id
-        self.encoder_norm = encoder_norm
-        self.decoder_norm = decoder_norm
-        self.max_prompt_tokens = max_prompt_tokens
-        self.temperatures = temperatures or [1.0, 1.0, 1.0]
-        self.n_heads = int(n_heads)
-        # One temperature triple per member; a single-head repo written before this field existed has none, so
-        # its `temperatures` triple serves member 0.
-        self.member_temperatures = member_temperatures or [list(self.temperatures)]
-        self.aliases = aliases or list(ALIASES)
-        # The hev.prompting format the head was trained with. Configs written before this field existed used 1.
-        self.prompt_format_version = prompt_format_version
-        kwargs.setdefault("auto_map", {"AutoConfig": "configuration_hev.HevConfig",
-                                       "AutoModel": "modeling_hev.HevForDecisions"})
-        super().__init__(**kwargs)
+"""HevConfig: everything needed to rebuild the frozen backbone, the head, and the prompt protocol."""
+from __future__ import annotations
+
+try:
+    from transformers import PreTrainedConfig
+except ImportError:  # transformers < 5
+    from transformers import PretrainedConfig as PreTrainedConfig
+
+from .head import HevHeadConfig
+from .schema import ALIASES
+
+
+class HevConfig(PreTrainedConfig):
+    model_type = "hev"
+
+    def __init__(self, backbone_id: str = "google/diffusiongemma-26B-A4B-it", backbone_config: dict | None = None,
+                 head_config: dict | None = None, canvas_length: int = 64, seed_mode: str = "pad",
+                 slot_token_id: int = 0, encoder_norm: str = "", decoder_norm: str = "",
+                 max_prompt_tokens: int = 1024, temperatures: list[float] | None = None,
+                 n_heads: int = 1, member_temperatures: list[list[float]] | None = None,
+                 aliases: list[str] | None = None, prompt_format_version: int = 1,
+                 extra_layers: dict | None = None, **kwargs):
+        self.backbone_id = backbone_id
+        # Module names of the intermediate encoder/decoder layers whose unit-RMS outputs the features append to
+        # the final-norm block (hev.features, --layers). Empty for a head trained on the final norms alone.
+        extra_layers = extra_layers or {}
+        self.extra_layers = {"encoder": list(extra_layers.get("encoder", [])),
+                             "decoder": list(extra_layers.get("decoder", []))}
+        self.backbone_config = backbone_config or {}
+        self.head_config = head_config or HevHeadConfig().to_dict()
+        self.canvas_length = canvas_length
+        self.seed_mode = seed_mode
+        self.slot_token_id = slot_token_id
+        self.encoder_norm = encoder_norm
+        self.decoder_norm = decoder_norm
+        self.max_prompt_tokens = max_prompt_tokens
+        self.temperatures = temperatures or [1.0, 1.0, 1.0]
+        self.n_heads = int(n_heads)
+        # One temperature triple per member; a single-head repo written before this field existed has none, so
+        # its `temperatures` triple serves member 0.
+        self.member_temperatures = member_temperatures or [list(self.temperatures)]
+        self.aliases = aliases or list(ALIASES)
+        # The hev.prompting format the head was trained with. Configs written before this field existed used 1.
+        self.prompt_format_version = prompt_format_version
+        kwargs.setdefault("auto_map", {"AutoConfig": "configuration_hev.HevConfig",
+                                       "AutoModel": "modeling_hev.HevForDecisions"})
+        super().__init__(**kwargs)

data.py

@@ -1,216 +1,216 @@
-"""Loaders for LocalLLaMA/typed-decisions and Praveenrajus/jev-bench, plus deterministic manifests."""
-from __future__ import annotations
-
-import hashlib
-import json
-import math
-import random
-import re
-from collections import defaultdict
-from pathlib import Path
-
-from .schema import Example, Question, align_distribution
-
-TYPED_DECISIONS = "LocalLLaMA/typed-decisions"
-JEV_BENCH = "Praveenrajus/jev-bench"
-JEV_BENCH_CONFIGS = [
-    "arc_challenge", "boolq", "chaosnli", "civil_comments", "fever_evidence", "go_emotions",
-    "helpsteer2_helpfulness", "helpsteer2_verbosity", "massive", "measuring_hate_speech", "mmlu",
-    "mnli", "paws", "sms_spam", "sst5", "strategyqa_closed", "strategyqa_grounded", "stsb", "yelp5",
-]
-TRUE_LABELS = {"1", "true", "yes"}
-FALSE_LABELS = {"0", "false", "no"}
-
-
-def _questions_from_map(qmap: dict) -> list[Question]:
-    return [Question.from_jev(qid, spec) for qid, spec in qmap.items()]
-
-
-def _label_option(label, q: Question) -> str | None:
-    """The row's discrete gold label as one of q's option names, or None when it is missing or not an option.
-
-    A boolean names the noul option "true" or "false", and a number its string form (a score level such as "2").
-    """
-    if isinstance(label, bool):
-        label = "true" if label else "false"
-    if label is None or str(label) not in q.options:
-        return None
-    return str(label)
-
-
-def typed_decisions_example(row: dict) -> Example:
-    questions = _questions_from_map(json.loads(row["questions"]))
-    gold_raw = json.loads(row["gold"])
-    gold, labels = {}, {}
-    for q in questions:
-        entry = gold_raw[q.qid]
-        gold[q.qid] = align_distribution(entry["probabilities"], q.options)
-        # The card's discrete label breaks ties in the distribution; hev.evaluate scores accuracy against it.
-        label = _label_option(entry.get("label"), q)
-        if label is not None:
-            labels[q.qid] = label
-    meta = {"workflow": row["workflow"]}
-    if labels:
-        meta["labels"] = labels
-    return Example(id=row["id"], source=f"typed-decisions/{row['workflow']}",
-                   state=row["state"], questions=questions, gold=gold, meta=meta)
-
-
-def _positional_distribution(soft, k: int) -> list[float] | None:
-    if len(soft) != k:
-        return None
-    if not all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in soft):
-        return None
-    vals = [float(v) for v in soft]
-    if any(v < 0 or not math.isfinite(v) for v in vals):
-        return None
-    total = sum(vals)
-    if total <= 0:
-        return None
-    return [v / total for v in vals]
-
-
-def jev_bench_example(row: dict, config: str) -> Example:
-    q = Question.from_jev("q", json.loads(row["question"]))
-    soft = row.get("soft_label")
-    if isinstance(soft, str):
-        soft = json.loads(soft)
-    dist = None
-    if isinstance(soft, dict) and set(soft) & set(q.options):
-        dist = align_distribution(soft, q.options)
-    elif (isinstance(soft, (int, float)) and not isinstance(soft, bool) and math.isfinite(soft)
-          and q.type == "noul"):
-        p = min(1.0, max(0.0, float(soft)))
-        dist = [p if o == "true" else 1.0 - p for o in q.options]
-    elif isinstance(soft, (list, tuple)) and q.type == "score":
-        dist = _positional_distribution(soft, q.k)
-    if dist is None:
-        label = str(row["label"])
-        if q.type == "noul":
-            norm = label.strip().lower()
-            if norm in TRUE_LABELS:
-                label = "true"
-            elif norm in FALSE_LABELS:
-                label = "false"
-            else:
-                raise ValueError(f"{row['id']}: unrecognised noul label {row['label']!r}")
-        if label not in q.options:
-            raise ValueError(f"{row['id']}: label {label!r} not in options {q.options}")
-        dist = [1.0 if o == label else 0.0 for o in q.options]
-    meta = row.get("meta") or "{}"
-    meta = json.loads(meta) if isinstance(meta, str) else dict(meta)
-    meta["config"] = config
-    return Example(id=row["id"], source=f"jev-bench/{config}", state=row["state"],
-                   questions=[q], gold={"q": dist}, meta=meta)
-
-
-def split_validation(examples: list[Example], frac: float = 0.1, seed: int = 42):
-    by_source = defaultdict(list)
-    for ex in examples:
-        by_source[ex.source].append(ex)
-    train, val = [], []
-    for source in sorted(by_source):
-        group = sorted(by_source[source], key=lambda e: e.id)
-        rng = random.Random(f"{seed}:{source}")
-        rng.shuffle(group)
-        n_val = max(1, round(len(group) * frac))
-        val.extend(group[:n_val])
-        train.extend(group[n_val:])
-    return train, val
-
-
-def write_manifest(examples: list[Example], path) -> None:
-    path = Path(path)
-    path.parent.mkdir(parents=True, exist_ok=True)
-    with path.open("w", encoding="utf-8") as f:
-        for ex in examples:
-            f.write(json.dumps(ex.to_dict(), ensure_ascii=False) + "\n")
-
-
-def read_manifest(path) -> list[Example]:
-    with Path(path).open("r", encoding="utf-8") as f:
-        return [Example.from_dict(json.loads(line)) for line in f if line.strip()]
-
-
-def manifest_checksum(path) -> str:
-    return hashlib.sha256(Path(path).read_bytes()).hexdigest()[:16]
-
-
-def _is_missing_split(err: Exception, split: str) -> bool:
-    """True for the error `datasets` raises when a config has no such split (the only load error that is skipped)."""
-    msg = str(err)
-    return ("Unknown split" in msg or "Bad split" in msg) and split in msg
-
-
-def _dataset_revision(ds) -> str | None:
-    """The Hub commit a loaded split came from, read off its cache path (.../<config>/<version>/<sha>/...)."""
-    for f in getattr(ds, "cache_files", None) or []:
-        name = Path(str(f.get("filename", ""))).parent.name
-        if re.fullmatch(r"[0-9a-f]{40}", name):
-            return name
-    return None
-
-
-def _sample(rows, k: int, seed: int, key: str):
-    idx = list(range(len(rows)))
-    random.Random(f"{seed}:{key}").shuffle(idx)
-    return [rows[i] for i in sorted(idx[:k])]
-
-
-def build_manifests(out_dir, seed: int = 42, per_config: dict | None = None) -> dict:
-    # Imported here, inside try/except, so neither the tests nor Hub users of the flattened data.py need `datasets`:
-    # transformers' check_imports skips imports inside try blocks.
-    try:
-        from datasets import load_dataset
-    except ImportError as err:
-        raise ImportError('building the manifests needs the `datasets` package: pip install datasets '
-                          '(or pip install -e ".[data]")') from err
-
-    per_config = per_config or {"train": 300, "validation": 100, "test": 200}
-    out = Path(out_dir)
-    revisions = defaultdict(set)
-
-    def load(repo, config, split):
-        ds = load_dataset(repo, config, split=split)
-        sha = _dataset_revision(ds)
-        if sha:
-            revisions[repo].add(sha)
-        return ds
-
-    td_train = [typed_decisions_example(r) for r in load(TYPED_DECISIONS, "all", "train")]
-    td_test = [typed_decisions_example(r) for r in load(TYPED_DECISIONS, "all", "test")]
-    train, val = split_validation(td_train, 0.1, seed)
-    jb_test = []
-    for config in JEV_BENCH_CONFIGS:
-        for split, target in (("train", train), ("validation", val), ("test", jb_test)):
-            try:
-                rows = load(JEV_BENCH, config, split)
-            except ValueError as err:
-                if not _is_missing_split(err, split):
-                    raise
-                print(f"skip {config}/{split}: {err}", flush=True)
-                continue
-            rows = [rows[i] for i in range(len(rows))]
-            for r in _sample(rows, per_config[split], seed, f"{config}/{split}"):
-                target.append(jev_bench_example(r, config))
-    write_manifest(train, out / "train.jsonl")
-    write_manifest(val, out / "validation.jsonl")
-    write_manifest(td_test, out / "test_td.jsonl")
-    write_manifest(jb_test, out / "test_jb.jsonl")
-    counts = {"train": len(train), "validation": len(val), "test_td": len(td_test), "test_jb": len(jb_test)}
-    for name in ("train", "validation", "test_td", "test_jb"):
-        counts[f"{name}_sha"] = manifest_checksum(out / f"{name}.jsonl")
-    if revisions:  # the dataset commit each repo was loaded from (a sorted list if the cache mixed several)
-        counts["revisions"] = {repo: next(iter(shas)) if len(shas) == 1 else sorted(shas)
-                               for repo, shas in sorted(revisions.items())}
-    (out / "counts.json").write_text(json.dumps(counts, indent=2))
-    return counts
-
-
-if __name__ == "__main__":
-    import argparse
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--out", default="data")
-    ap.add_argument("--seed", type=int, default=42)
-    args = ap.parse_args()
-    print(json.dumps(build_manifests(args.out, args.seed), indent=2))
+"""Loaders for LocalLLaMA/typed-decisions and Praveenrajus/jev-bench, plus deterministic manifests."""
+from __future__ import annotations
+
+import hashlib
+import json
+import math
+import random
+import re
+from collections import defaultdict
+from pathlib import Path
+
+from .schema import Example, Question, align_distribution
+
+TYPED_DECISIONS = "LocalLLaMA/typed-decisions"
+JEV_BENCH = "Praveenrajus/jev-bench"
+JEV_BENCH_CONFIGS = [
+    "arc_challenge", "boolq", "chaosnli", "civil_comments", "fever_evidence", "go_emotions",
+    "helpsteer2_helpfulness", "helpsteer2_verbosity", "massive", "measuring_hate_speech", "mmlu",
+    "mnli", "paws", "sms_spam", "sst5", "strategyqa_closed", "strategyqa_grounded", "stsb", "yelp5",
+]
+TRUE_LABELS = {"1", "true", "yes"}
+FALSE_LABELS = {"0", "false", "no"}
+
+
+def _questions_from_map(qmap: dict) -> list[Question]:
+    return [Question.from_jev(qid, spec) for qid, spec in qmap.items()]
+
+
+def _label_option(label, q: Question) -> str | None:
+    """The row's discrete gold label as one of q's option names, or None when it is missing or not an option.
+
+    A boolean names the noul option "true" or "false", and a number its string form (a score level such as "2").
+    """
+    if isinstance(label, bool):
+        label = "true" if label else "false"
+    if label is None or str(label) not in q.options:
+        return None
+    return str(label)
+
+
+def typed_decisions_example(row: dict) -> Example:
+    questions = _questions_from_map(json.loads(row["questions"]))
+    gold_raw = json.loads(row["gold"])
+    gold, labels = {}, {}
+    for q in questions:
+        entry = gold_raw[q.qid]
+        gold[q.qid] = align_distribution(entry["probabilities"], q.options)
+        # The card's discrete label breaks ties in the distribution; hev.evaluate scores accuracy against it.
+        label = _label_option(entry.get("label"), q)
+        if label is not None:
+            labels[q.qid] = label
+    meta = {"workflow": row["workflow"]}
+    if labels:
+        meta["labels"] = labels
+    return Example(id=row["id"], source=f"typed-decisions/{row['workflow']}",
+                   state=row["state"], questions=questions, gold=gold, meta=meta)
+
+
+def _positional_distribution(soft, k: int) -> list[float] | None:
+    if len(soft) != k:
+        return None
+    if not all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in soft):
+        return None
+    vals = [float(v) for v in soft]
+    if any(v < 0 or not math.isfinite(v) for v in vals):
+        return None
+    total = sum(vals)
+    if total <= 0:
+        return None
+    return [v / total for v in vals]
+
+
+def jev_bench_example(row: dict, config: str) -> Example:
+    q = Question.from_jev("q", json.loads(row["question"]))
+    soft = row.get("soft_label")
+    if isinstance(soft, str):
+        soft = json.loads(soft)
+    dist = None
+    if isinstance(soft, dict) and set(soft) & set(q.options):
+        dist = align_distribution(soft, q.options)
+    elif (isinstance(soft, (int, float)) and not isinstance(soft, bool) and math.isfinite(soft)
+          and q.type == "noul"):
+        p = min(1.0, max(0.0, float(soft)))
+        dist = [p if o == "true" else 1.0 - p for o in q.options]
+    elif isinstance(soft, (list, tuple)) and q.type == "score":
+        dist = _positional_distribution(soft, q.k)
+    if dist is None:
+        label = str(row["label"])
+        if q.type == "noul":
+            norm = label.strip().lower()
+            if norm in TRUE_LABELS:
+                label = "true"
+            elif norm in FALSE_LABELS:
+                label = "false"
+            else:
+                raise ValueError(f"{row['id']}: unrecognised noul label {row['label']!r}")
+        if label not in q.options:
+            raise ValueError(f"{row['id']}: label {label!r} not in options {q.options}")
+        dist = [1.0 if o == label else 0.0 for o in q.options]
+    meta = row.get("meta") or "{}"
+    meta = json.loads(meta) if isinstance(meta, str) else dict(meta)
+    meta["config"] = config
+    return Example(id=row["id"], source=f"jev-bench/{config}", state=row["state"],
+                   questions=[q], gold={"q": dist}, meta=meta)
+
+
+def split_validation(examples: list[Example], frac: float = 0.1, seed: int = 42):
+    by_source = defaultdict(list)
+    for ex in examples:
+        by_source[ex.source].append(ex)
+    train, val = [], []
+    for source in sorted(by_source):
+        group = sorted(by_source[source], key=lambda e: e.id)
+        rng = random.Random(f"{seed}:{source}")
+        rng.shuffle(group)
+        n_val = max(1, round(len(group) * frac))
+        val.extend(group[:n_val])
+        train.extend(group[n_val:])
+    return train, val
+
+
+def write_manifest(examples: list[Example], path) -> None:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as f:
+        for ex in examples:
+            f.write(json.dumps(ex.to_dict(), ensure_ascii=False) + "\n")
+
+
+def read_manifest(path) -> list[Example]:
+    with Path(path).open("r", encoding="utf-8") as f:
+        return [Example.from_dict(json.loads(line)) for line in f if line.strip()]
+
+
+def manifest_checksum(path) -> str:
+    return hashlib.sha256(Path(path).read_bytes()).hexdigest()[:16]
+
+
+def _is_missing_split(err: Exception, split: str) -> bool:
+    """True for the error `datasets` raises when a config has no such split (the only load error that is skipped)."""
+    msg = str(err)
+    return ("Unknown split" in msg or "Bad split" in msg) and split in msg
+
+
+def _dataset_revision(ds) -> str | None:
+    """The Hub commit a loaded split came from, read off its cache path (.../<config>/<version>/<sha>/...)."""
+    for f in getattr(ds, "cache_files", None) or []:
+        name = Path(str(f.get("filename", ""))).parent.name
+        if re.fullmatch(r"[0-9a-f]{40}", name):
+            return name
+    return None
+
+
+def _sample(rows, k: int, seed: int, key: str):
+    idx = list(range(len(rows)))
+    random.Random(f"{seed}:{key}").shuffle(idx)
+    return [rows[i] for i in sorted(idx[:k])]
+
+
+def build_manifests(out_dir, seed: int = 42, per_config: dict | None = None) -> dict:
+    # Imported here, inside try/except, so neither the tests nor Hub users of the flattened data.py need `datasets`:
+    # transformers' check_imports skips imports inside try blocks.
+    try:
+        from datasets import load_dataset
+    except ImportError as err:
+        raise ImportError('building the manifests needs the `datasets` package: pip install datasets '
+                          '(or pip install -e ".[data]")') from err
+
+    per_config = per_config or {"train": 300, "validation": 100, "test": 200}
+    out = Path(out_dir)
+    revisions = defaultdict(set)
+
+    def load(repo, config, split):
+        ds = load_dataset(repo, config, split=split)
+        sha = _dataset_revision(ds)
+        if sha:
+            revisions[repo].add(sha)
+        return ds
+
+    td_train = [typed_decisions_example(r) for r in load(TYPED_DECISIONS, "all", "train")]
+    td_test = [typed_decisions_example(r) for r in load(TYPED_DECISIONS, "all", "test")]
+    train, val = split_validation(td_train, 0.1, seed)
+    jb_test = []
+    for config in JEV_BENCH_CONFIGS:
+        for split, target in (("train", train), ("validation", val), ("test", jb_test)):
+            try:
+                rows = load(JEV_BENCH, config, split)
+            except ValueError as err:
+                if not _is_missing_split(err, split):
+                    raise
+                print(f"skip {config}/{split}: {err}", flush=True)
+                continue
+            rows = [rows[i] for i in range(len(rows))]
+            for r in _sample(rows, per_config[split], seed, f"{config}/{split}"):
+                target.append(jev_bench_example(r, config))
+    write_manifest(train, out / "train.jsonl")
+    write_manifest(val, out / "validation.jsonl")
+    write_manifest(td_test, out / "test_td.jsonl")
+    write_manifest(jb_test, out / "test_jb.jsonl")
+    counts = {"train": len(train), "validation": len(val), "test_td": len(td_test), "test_jb": len(jb_test)}
+    for name in ("train", "validation", "test_td", "test_jb"):
+        counts[f"{name}_sha"] = manifest_checksum(out / f"{name}.jsonl")
+    if revisions:  # the dataset commit each repo was loaded from (a sorted list if the cache mixed several)
+        counts["revisions"] = {repo: next(iter(shas)) if len(shas) == 1 else sorted(shas)
+                               for repo, shas in sorted(revisions.items())}
+    (out / "counts.json").write_text(json.dumps(counts, indent=2))
+    return counts
+
+
+if __name__ == "__main__":
+    import argparse
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--out", default="data")
+    ap.add_argument("--seed", type=int, default=42)
+    args = ap.parse_args()
+    print(json.dumps(build_manifests(args.out, args.seed), indent=2))

eval/results.json

@@ -1,223 +1,223 @@
-{
-  "preds": "pod_artifacts/yl3x22ceoi4s80/unpacked/eval/test_td.preds.json",
-  "manifest": "data/test_td.jsonl",
-  "labelled": 2000,
-  "rows": {
-    "DiffusionGemma zero-shot canvas read": {
-      "n": 2000,
-      "acc": 0.7075,
-      "soft_acc": 0.5731857140189666,
-      "macro_f1": 0.550702638008923,
-      "kl": 1.3327719231855613,
-      "tv": 0.3621118097383339,
-      "brier": 0.30395114563065456,
-      "nll": 2.100130478809889,
-      "ece": 0.21345150114595893,
-      "score_mae": 0.41135865707582425,
-      "within_1": 0.93375,
-      "by_type": {
-        "choice": {
-          "n": 600,
-          "acc": 0.6583333333333333,
-          "soft_acc": 0.529550382418425,
-          "macro_f1": 0.4580138346822841,
-          "kl": 1.8344121477390998,
-          "tv": 0.4134490126261477,
-          "brier": 0.3653080227268783,
-          "nll": 2.698206744585924,
-          "ece": 0.27046276137232783,
-          "score_mae": null,
-          "within_1": null
-        },
-        "noul": {
-          "n": 600,
-          "acc": 0.7716666666666666,
-          "soft_acc": 0.6983805466666665,
-          "macro_f1": 0.6222345422638039,
-          "kl": 1.0794531386915138,
-          "tv": 0.27516267522582993,
-          "brier": 0.25676828758281806,
-          "nll": 1.5331753820440792,
-          "ece": 0.19590681572755178,
-          "score_mae": null,
-          "within_1": null
-        },
-        "score": {
-          "n": 800,
-          "acc": 0.69625,
-          "soft_acc": 0.5120160882335887,
-          "macro_f1": 0.5665703123127417,
-          "kl": 1.1465308431409578,
-          "tv": 0.38882075845685227,
-          "brier": 0.2933206313443617,
-          "nll": 2.076789602052242,
-          "ece": 0.1857137256860733,
-          "score_mae": 0.41135865707582425,
-          "within_1": 0.93375
-        }
-      },
-      "by_source": {
-        "typed-decisions/agent_trace_observability": {
-          "n": 500,
-          "acc": 0.604,
-          "soft_acc": 0.4776453329467184,
-          "macro_f1": 0.5448100563542999,
-          "kl": 1.3346726647537777,
-          "tv": 0.4248764365872663,
-          "brier": 0.36639518814697086,
-          "nll": 2.2131429376990006,
-          "ece": 0.2710903854966163,
-          "score_mae": 0.41654339421197734,
-          "within_1": 0.965
-        },
-        "typed-decisions/customer_service": {
-          "n": 500,
-          "acc": 0.772,
-          "soft_acc": 0.6172040574027722,
-          "macro_f1": 0.6487457943630335,
-          "kl": 1.5586844768592238,
-          "tv": 0.34467774901520876,
-          "brier": 0.28435340437718454,
-          "nll": 2.2958391148687096,
-          "ece": 0.1874517269730567,
-          "score_mae": 0.3006309953052926,
-          "within_1": 0.99
-        },
-        "typed-decisions/invoice_processing": {
-          "n": 500,
-          "acc": 0.758,
-          "soft_acc": 0.6795901281829838,
-          "macro_f1": 0.571010468782984,
-          "kl": 0.7263074853218779,
-          "tv": 0.24453685503895994,
-          "brier": 0.20045868313962484,
-          "nll": 1.301254960986774,
-          "ece": 0.1528367287516594,
-          "score_mae": 0.3736864058152724,
-          "within_1": 0.91
-        },
-        "typed-decisions/security_incidents": {
-          "n": 500,
-          "acc": 0.696,
-          "soft_acc": 0.51830333754338,
-          "macro_f1": 0.43824423253537487,
-          "kl": 1.7114230658073863,
-          "tv": 0.4343561983119038,
-          "brier": 0.36459730685883546,
-          "nll": 2.590284901685105,
-          "ece": 0.24981171822547912,
-          "score_mae": 0.5545738329707549,
-          "within_1": 0.87
-        }
-      }
-    },
-    "Hev": {
-      "n": 2000,
-      "acc": 0.731,
-      "soft_acc": 0.5841690466990694,
-      "macro_f1": 0.6263593100482722,
-      "kl": 0.18476683805325994,
-      "tv": 0.21320750403794955,
-      "brier": 0.09350640922665171,
-      "nll": 0.9521253933876698,
-      "ece": 0.16195184084028008,
-      "score_mae": 0.30327667244408074,
-      "within_1": 0.97625,
-      "by_type": {
-        "choice": {
-          "n": 600,
-          "acc": 0.6883333333333334,
-          "soft_acc": 0.5400114920828492,
-          "macro_f1": 0.5819579445698401,
-          "kl": 0.2682545787140505,
-          "tv": 0.2583168463366746,
-          "brier": 0.12431620966551991,
-          "nll": 1.1320491751871322,
-          "ece": 0.15876342798272766,
-          "score_mae": null,
-          "within_1": null
-        },
-        "noul": {
-          "n": 600,
-          "acc": 0.8133333333333334,
-          "soft_acc": 0.7151805511737863,
-          "macro_f1": 0.7030793844723765,
-          "kl": 0.10688763325700651,
-          "tv": 0.1591590390598867,
-          "brier": 0.07606820441416512,
-          "nll": 0.5606098758757383,
-          "ece": 0.08742988765239716,
-          "score_mae": null,
-          "within_1": null
-        },
-        "score": {
-          "n": 800,
-          "acc": 0.70125,
-          "soft_acc": 0.519028584305197,
-          "macro_f1": 0.602120278339018,
-          "kl": 0.1805604361548552,
-          "tv": 0.21991184604745284,
-          "brier": 0.08347771250686478,
-          "nll": 1.1108191951720112,
-          "ece": 0.22084334811195733,
-          "score_mae": 0.30327667244408074,
-          "within_1": 0.97625
-        }
-      },
-      "by_source": {
-        "typed-decisions/agent_trace_observability": {
-          "n": 500,
-          "acc": 0.654,
-          "soft_acc": 0.49528532839939,
-          "macro_f1": 0.6164904940650149,
-          "kl": 0.19085320215634394,
-          "tv": 0.21603136127034667,
-          "brier": 0.09737799980224329,
-          "nll": 1.0693234745297415,
-          "ece": 0.15192090582847595,
-          "score_mae": 0.2841385978687322,
-          "within_1": 0.99
-        },
-        "typed-decisions/customer_service": {
-          "n": 500,
-          "acc": 0.76,
-          "soft_acc": 0.6186040637716651,
-          "macro_f1": 0.696363079487989,
-          "kl": 0.18797709494129541,
-          "tv": 0.2113354651791742,
-          "brier": 0.09376893105456086,
-          "nll": 0.9251317321598946,
-          "ece": 0.14841172212362289,
-          "score_mae": 0.24546261819079518,
-          "within_1": 1.0
-        },
-        "typed-decisions/invoice_processing": {
-          "n": 500,
-          "acc": 0.786,
-          "soft_acc": 0.6944667954742908,
-          "macro_f1": 0.5980108397703063,
-          "kl": 0.22106904117845244,
-          "tv": 0.23180358876640095,
-          "brier": 0.11166511418841754,
-          "nll": 0.7960165171770665,
-          "ece": 0.1501800780892372,
-          "score_mae": 0.39022797410623755,
-          "within_1": 0.92
-        },
-        "typed-decisions/security_incidents": {
-          "n": 500,
-          "acc": 0.724,
-          "soft_acc": 0.5283199991509319,
-          "macro_f1": 0.5945728268697784,
-          "kl": 0.13916801393694356,
-          "tv": 0.19365960093587636,
-          "brier": 0.07121359186138425,
-          "nll": 1.0180298496839644,
-          "ece": 0.2086083108484745,
-          "score_mae": 0.29327749961055816,
-          "within_1": 0.995
-        }
-      }
-    }
-  }
+{
+  "preds": "pod_artifacts/yl3x22ceoi4s80/unpacked/eval/test_td.preds.json",
+  "manifest": "data/test_td.jsonl",
+  "labelled": 2000,
+  "rows": {
+    "DiffusionGemma zero-shot canvas read": {
+      "n": 2000,
+      "acc": 0.7075,
+      "soft_acc": 0.5731857140189666,
+      "macro_f1": 0.550702638008923,
+      "kl": 1.3327719231855613,
+      "tv": 0.3621118097383339,
+      "brier": 0.30395114563065456,
+      "nll": 2.100130478809889,
+      "ece": 0.21345150114595893,
+      "score_mae": 0.41135865707582425,
+      "within_1": 0.93375,
+      "by_type": {
+        "choice": {
+          "n": 600,
+          "acc": 0.6583333333333333,
+          "soft_acc": 0.529550382418425,
+          "macro_f1": 0.4580138346822841,
+          "kl": 1.8344121477390998,
+          "tv": 0.4134490126261477,
+          "brier": 0.3653080227268783,
+          "nll": 2.698206744585924,
+          "ece": 0.27046276137232783,
+          "score_mae": null,
+          "within_1": null
+        },
+        "noul": {
+          "n": 600,
+          "acc": 0.7716666666666666,
+          "soft_acc": 0.6983805466666665,
+          "macro_f1": 0.6222345422638039,
+          "kl": 1.0794531386915138,
+          "tv": 0.27516267522582993,
+          "brier": 0.25676828758281806,
+          "nll": 1.5331753820440792,
+          "ece": 0.19590681572755178,
+          "score_mae": null,
+          "within_1": null
+        },
+        "score": {
+          "n": 800,
+          "acc": 0.69625,
+          "soft_acc": 0.5120160882335887,
+          "macro_f1": 0.5665703123127417,
+          "kl": 1.1465308431409578,
+          "tv": 0.38882075845685227,
+          "brier": 0.2933206313443617,
+          "nll": 2.076789602052242,
+          "ece": 0.1857137256860733,
+          "score_mae": 0.41135865707582425,
+          "within_1": 0.93375
+        }
+      },
+      "by_source": {
+        "typed-decisions/agent_trace_observability": {
+          "n": 500,
+          "acc": 0.604,
+          "soft_acc": 0.4776453329467184,
+          "macro_f1": 0.5448100563542999,
+          "kl": 1.3346726647537777,
+          "tv": 0.4248764365872663,
+          "brier": 0.36639518814697086,
+          "nll": 2.2131429376990006,
+          "ece": 0.2710903854966163,
+          "score_mae": 0.41654339421197734,
+          "within_1": 0.965
+        },
+        "typed-decisions/customer_service": {
+          "n": 500,
+          "acc": 0.772,
+          "soft_acc": 0.6172040574027722,
+          "macro_f1": 0.6487457943630335,
+          "kl": 1.5586844768592238,
+          "tv": 0.34467774901520876,
+          "brier": 0.28435340437718454,
+          "nll": 2.2958391148687096,
+          "ece": 0.1874517269730567,
+          "score_mae": 0.3006309953052926,
+          "within_1": 0.99
+        },
+        "typed-decisions/invoice_processing": {
+          "n": 500,
+          "acc": 0.758,
+          "soft_acc": 0.6795901281829838,
+          "macro_f1": 0.571010468782984,
+          "kl": 0.7263074853218779,
+          "tv": 0.24453685503895994,
+          "brier": 0.20045868313962484,
+          "nll": 1.301254960986774,
+          "ece": 0.1528367287516594,
+          "score_mae": 0.3736864058152724,
+          "within_1": 0.91
+        },
+        "typed-decisions/security_incidents": {
+          "n": 500,
+          "acc": 0.696,
+          "soft_acc": 0.51830333754338,
+          "macro_f1": 0.43824423253537487,
+          "kl": 1.7114230658073863,
+          "tv": 0.4343561983119038,
+          "brier": 0.36459730685883546,
+          "nll": 2.590284901685105,
+          "ece": 0.24981171822547912,
+          "score_mae": 0.5545738329707549,
+          "within_1": 0.87
+        }
+      }
+    },
+    "Hev": {
+      "n": 2000,
+      "acc": 0.731,
+      "soft_acc": 0.5841690466990694,
+      "macro_f1": 0.6263593100482722,
+      "kl": 0.18476683805325994,
+      "tv": 0.21320750403794955,
+      "brier": 0.09350640922665171,
+      "nll": 0.9521253933876698,
+      "ece": 0.16195184084028008,
+      "score_mae": 0.30327667244408074,
+      "within_1": 0.97625,
+      "by_type": {
+        "choice": {
+          "n": 600,
+          "acc": 0.6883333333333334,
+          "soft_acc": 0.5400114920828492,
+          "macro_f1": 0.5819579445698401,
+          "kl": 0.2682545787140505,
+          "tv": 0.2583168463366746,
+          "brier": 0.12431620966551991,
+          "nll": 1.1320491751871322,
+          "ece": 0.15876342798272766,
+          "score_mae": null,
+          "within_1": null
+        },
+        "noul": {
+          "n": 600,
+          "acc": 0.8133333333333334,
+          "soft_acc": 0.7151805511737863,
+          "macro_f1": 0.7030793844723765,
+          "kl": 0.10688763325700651,
+          "tv": 0.1591590390598867,
+          "brier": 0.07606820441416512,
+          "nll": 0.5606098758757383,
+          "ece": 0.08742988765239716,
+          "score_mae": null,
+          "within_1": null
+        },
+        "score": {
+          "n": 800,
+          "acc": 0.70125,
+          "soft_acc": 0.519028584305197,
+          "macro_f1": 0.602120278339018,
+          "kl": 0.1805604361548552,
+          "tv": 0.21991184604745284,
+          "brier": 0.08347771250686478,
+          "nll": 1.1108191951720112,
+          "ece": 0.22084334811195733,
+          "score_mae": 0.30327667244408074,
+          "within_1": 0.97625
+        }
+      },
+      "by_source": {
+        "typed-decisions/agent_trace_observability": {
+          "n": 500,
+          "acc": 0.654,
+          "soft_acc": 0.49528532839939,
+          "macro_f1": 0.6164904940650149,
+          "kl": 0.19085320215634394,
+          "tv": 0.21603136127034667,
+          "brier": 0.09737799980224329,
+          "nll": 1.0693234745297415,
+          "ece": 0.15192090582847595,
+          "score_mae": 0.2841385978687322,
+          "within_1": 0.99
+        },
+        "typed-decisions/customer_service": {
+          "n": 500,
+          "acc": 0.76,
+          "soft_acc": 0.6186040637716651,
+          "macro_f1": 0.696363079487989,
+          "kl": 0.18797709494129541,
+          "tv": 0.2113354651791742,
+          "brier": 0.09376893105456086,
+          "nll": 0.9251317321598946,
+          "ece": 0.14841172212362289,
+          "score_mae": 0.24546261819079518,
+          "within_1": 1.0
+        },
+        "typed-decisions/invoice_processing": {
+          "n": 500,
+          "acc": 0.786,
+          "soft_acc": 0.6944667954742908,
+          "macro_f1": 0.5980108397703063,
+          "kl": 0.22106904117845244,
+          "tv": 0.23180358876640095,
+          "brier": 0.11166511418841754,
+          "nll": 0.7960165171770665,
+          "ece": 0.1501800780892372,
+          "score_mae": 0.39022797410623755,
+          "within_1": 0.92
+        },
+        "typed-decisions/security_incidents": {
+          "n": 500,
+          "acc": 0.724,
+          "soft_acc": 0.5283199991509319,
+          "macro_f1": 0.5945728268697784,
+          "kl": 0.13916801393694356,
+          "tv": 0.19365960093587636,
+          "brier": 0.07121359186138425,
+          "nll": 1.0180298496839644,
+          "ece": 0.2086083108484745,
+          "score_mae": 0.29327749961055816,
+          "within_1": 0.995
+        }
+      }
+    }
+  }
 }

features.py

@@ -1,345 +1,345 @@
-"""Backbone feature extraction: one forward pass per example, hidden states captured by hooks."""
-from __future__ import annotations
-
-import argparse
-import json
-import math
-import random
-import re
-import time
-from pathlib import Path
-
-import torch
-
-from .cache import CacheWriter, Decision, FeatureCache, Record
-from .data import read_manifest
-from .prompting import alias_token_ids, canvas_ids, shuffled_order, tokenize_prompt
-from .schema import Example, argmax
-
-MODEL_ID = "google/diffusiongemma-26B-A4B-it"
-PROBE_STATE = "A short probe state."
-RANDOM_DRAWS = 4
-SEED_MODES = ("pad", "alias", "random")
-ZERO_SHOT_LOG = "configs/session1_zero_shot.jsonl"
-
-
-class NormCapture:
-    """Registers forward hooks on final-norm modules (and any modules named in `extra`) and keeps their latest
-    outputs. With no names given every norm outside the layer stacks is hooked (discovery, see probe)."""
-
-    def __init__(self, model, encoder_name: str | None = None, decoder_name: str | None = None,
-                 extra: tuple[str, ...] | list[str] = ()):
-        self.outputs: dict[str, torch.Tensor] = {}
-        self.handles = []
-        extra = set(extra)
-        for name, mod in model.named_modules():
-            if name in extra:
-                self.handles.append(mod.register_forward_hook(self._hook(name)))
-                continue
-            if ".layers." in f".{name}." or "Norm" not in type(mod).__name__:
-                continue
-            if encoder_name is not None and name not in (encoder_name, decoder_name):
-                continue
-            self.handles.append(mod.register_forward_hook(self._hook(name)))
-
-    def _hook(self, name):
-        def fn(_m, _i, out):
-            self.outputs[name] = out[0] if isinstance(out, tuple) else out
-        return fn
-
-    def remove(self):
-        for h in self.handles:
-            h.remove()
-
-
-def unit_rms(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
-    """Scale every row to unit RMS (a parameter-free RMSNorm). Intermediate layers are residual-stream states
-    of arbitrary magnitude; this keeps them on the final norm's scale and inside float16 for the cache."""
-    return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
-
-
-_LAYER_RE = re.compile(r"(.*\blayers)\.(\d+)")
-
-
-def layer_stacks(model) -> dict[str, int]:
-    """Layer-stack prefixes and their depths, from module names of the form `<prefix>.layers.<n>`."""
-    stacks: dict[str, int] = {}
-    for name, _ in model.named_modules():
-        m = _LAYER_RE.fullmatch(name)
-        if m:
-            stacks[m.group(1)] = max(stacks.get(m.group(1), 0), int(m.group(2)) + 1)
-    return stacks
-
-
-def select_layers(n_layers: int, fracs) -> list[int]:
-    """Zero-based layer indices at the given depth fractions: 0.5 of 30 layers is layer 14, 0.75 is 22."""
-    out = set()
-    for f in fracs:
-        if not 0.0 < float(f) <= 1.0:
-            raise ValueError(f"layer fraction must be in (0, 1], got {f}")
-        out.add(max(0, min(n_layers - 1, math.ceil(float(f) * n_layers) - 1)))
-    return sorted(out)
-
-
-def resolve_extra_layers(model, encoder_norm: str, decoder_norm: str, fracs) -> dict[str, list[str]]:
-    """Module names of the extra layers to capture: for each final norm, the layer stack beside it (same parent,
-    else the shallowest stack below that parent) at `fracs` of its depth."""
-    stacks = layer_stacks(model)
-
-    def parent(name: str) -> str:
-        return name.rsplit(".", 1)[0] if "." in name else ""
-
-    def stack_for(norm: str) -> str:
-        p = parent(norm)
-        beside = [s for s in stacks if parent(s) == p]
-        below = [s for s in stacks if p and s.startswith(p + ".")]
-        cands = beside or below
-        if not cands:
-            raise RuntimeError(f"no layer stack found beside {norm!r}; stacks: {sorted(stacks)}")
-        return min(cands, key=len)
-
-    enc, dec = stack_for(encoder_norm), stack_for(decoder_norm)
-    if enc == dec:
-        raise RuntimeError(f"encoder and decoder norms resolve to the same layer stack {enc!r}")
-    return {"encoder": [f"{enc}.{i}" for i in select_layers(stacks[enc], fracs)],
-            "decoder": [f"{dec}.{i}" for i in select_layers(stacks[dec], fracs)]}
-
-
-def extra_layer_names(info: dict) -> list[str]:
-    extra = info.get("extra_layers") or {}
-    return list(extra.get("encoder", [])) + list(extra.get("decoder", []))
-
-
-def _probe_example() -> Example:
-    from .schema import Question
-    q = Question.from_jev("ok", {"type": "noul", "instructions": "Is the probe fine?"})
-    return Example(id="probe", source="probe", state=PROBE_STATE, questions=[q], gold={"ok": None})
-
-
-@torch.no_grad()
-def _forward(model, prompt_ids: list[int], canvas: list[int]):
-    dev = model.device
-    input_ids = torch.tensor([prompt_ids], dtype=torch.long, device=dev)
-    return model(input_ids=input_ids, attention_mask=torch.ones_like(input_ids),
-                 decoder_input_ids=torch.tensor([canvas], dtype=torch.long, device=dev))
-
-
-def probe(model, tok, hidden_size: int) -> dict:
-    ex = _probe_example()
-    tp = tokenize_prompt(ex, tok)
-    for canvas_length in (64, 256):
-        canvas, _ = canvas_ids(ex, tok, slot_token_id=0, canvas_length=canvas_length)
-        cap = NormCapture(model)
-        try:
-            _forward(model, tp.ids, canvas)
-        except Exception as err:  # noqa: BLE001 - the whole point is to learn which length works
-            cap.remove()
-            print(f"canvas_length {canvas_length} failed: {type(err).__name__}: {err}")
-            continue
-        cap.remove()
-        enc = [n for n, o in cap.outputs.items() if o.shape[-2] == len(tp.ids) and o.shape[-1] == hidden_size]
-        dec = [n for n, o in cap.outputs.items() if o.shape[-2] == canvas_length and o.shape[-1] == hidden_size]
-        if not enc or not dec:
-            shapes = {k: tuple(v.shape) for k, v in cap.outputs.items()}
-            raise RuntimeError(f"could not find final norms; captured {shapes}")
-        if enc[-1] == dec[-1]:
-            raise RuntimeError(f"encoder and decoder norms resolved to the same module: {enc[-1]}")
-        if len(enc) > 1 or len(dec) > 1:
-            print(f"warning: several final-norm candidates, using the last; encoder {enc}, decoder {dec}")
-        return {"canvas_length": canvas_length, "encoder_norm": enc[-1], "decoder_norm": dec[-1],
-                "hidden_size": hidden_size}
-    raise RuntimeError("neither canvas length 64 nor 256 ran")
-
-
-def _slot_token(seed_mode: str, alias_ids: list[int], rng: random.Random, vocab: int) -> list[int]:
-    if seed_mode == "pad":
-        return [0]
-    if seed_mode == "alias":
-        return [alias_ids[0]]
-    if seed_mode == "random":
-        return [rng.randrange(10, vocab) for _ in range(RANDOM_DRAWS)]
-    raise ValueError(f"unknown seed_mode {seed_mode}")
-
-
-def _rows(cap: "NormCapture", name: str, n_rows: int, what: str) -> torch.Tensor:
-    """A captured layer's rows for this forward (batch element 0), checked against the expected row count."""
-    out = cap.outputs[name][0].float()
-    if out.shape[0] != n_rows:
-        raise RuntimeError(f"{name} output has shape {tuple(out.shape)}, expected {n_rows} {what} rows")
-    return out
-
-
-@torch.no_grad()
-def encode_example(model, tok, ex: Example, cap: "NormCapture", info: dict, *, slot_tokens: list[int],
-                   alias_ids: list[int], max_tokens: int = 1024, option_order=None) -> tuple[torch.Tensor, Record]:
-    """One example through the backbone: features averaged over slot tokens, plus a Record with priors.
-
-    The Record's path is left empty; the caller names the file.
-    """
-    canvas_length = info["canvas_length"]
-    tp = tokenize_prompt(ex, tok, max_tokens=max_tokens, option_order=option_order)
-    feats_sum, prob_sum, slots = None, None, None
-    for slot_tok in slot_tokens:
-        canvas, slots = canvas_ids(ex, tok, slot_token_id=slot_tok, canvas_length=canvas_length)
-        cap.outputs.clear()
-        out = _forward(model, tp.ids, canvas)
-        if out.logits.shape[1] != canvas_length:
-            raise RuntimeError(f"logits cover {out.logits.shape[1]} positions, expected canvas_length {canvas_length}")
-        extra = info.get("extra_layers") or {}
-        wanted = [info["encoder_norm"], info["decoder_norm"], *extra_layer_names(info)]
-        missing = [n for n in wanted if n not in cap.outputs]
-        if missing:
-            raise RuntimeError(f"norm outputs missing after forward: {missing}; captured {sorted(cap.outputs)}")
-        hp = cap.outputs[info["encoder_norm"]][0].float()
-        hc = cap.outputs[info["decoder_norm"]][0].float()
-        if hp.shape[0] != len(tp.ids):
-            raise RuntimeError(f"encoder norm output has shape {tuple(hp.shape)}, expected {len(tp.ids)} prompt rows")
-        if hc.shape[0] != canvas_length:
-            raise RuntimeError(f"decoder norm output has shape {tuple(hc.shape)}, expected canvas_length {canvas_length}")
-        enc_extra = [unit_rms(_rows(cap, name, len(tp.ids), "prompt")) for name in extra.get("encoder", [])]
-        dec_extra = [unit_rms(_rows(cap, name, canvas_length, "canvas")) for name in extra.get("decoder", [])]
-        feats = torch.cat([torch.cat([hp, *enc_extra], dim=1), torch.cat([hc, *dec_extra], dim=1)], dim=0)
-        probs = torch.softmax(out.logits[0].float(), dim=-1)[slots]  # (n_q, vocab)
-        feats_sum = feats if feats_sum is None else feats_sum + feats
-        prob_sum = probs if prob_sum is None else prob_sum + probs
-    feats = feats_sum / len(slot_tokens)
-    logp = (prob_sum / len(slot_tokens)).clamp_min(1e-12).log()
-    decisions = []
-    for qi, q in enumerate(ex.questions):
-        ids = torch.tensor([alias_ids[a] for a in tp.option_alias[qi]], device=logp.device)
-        prior = logp[qi, ids]
-        prior = (prior - torch.logsumexp(prior, dim=0)).tolist()
-        decisions.append(Decision(qid=q.qid, type=q.type, slot=len(tp.ids) + slots[qi],
-                                  spans=[list(s) for s in tp.option_spans[qi]],
-                                  alias=list(tp.option_alias[qi]), prior=prior, gold=ex.gold.get(q.qid)))
-    rec = Record(id=ex.id, source=ex.source, path="", n_prompt=len(tp.ids),
-                 n_canvas=canvas_length, decisions=decisions)
-    return feats, rec
-
-
-def extract(model, tok, examples: list[Example], out_dir, *, seed_mode: str, canvas_length: int,
-            alias_ids: list[int], hidden_size: int, max_tokens: int = 1024, shuffle_rng=None,
-            limit: int | None = None, vocab_size: int = 262144, log_every: int = 50,
-            info: dict | None = None, layer_fracs=()) -> dict:
-    """Write a feature cache. With `layer_fracs` (depth fractions, e.g. (0.5, 0.75)) every row also carries the
-    unit-RMS outputs of the encoder (prompt rows) or decoder (canvas rows) layers at those depths, appended to the
-    final-norm block, so the feature dim is hidden_size * (1 + len(layers))."""
-    info = dict(probe(model, tok, hidden_size) if info is None else info)
-    if info["canvas_length"] != canvas_length:
-        raise ValueError(f"requested canvas {canvas_length} but probe says {info['canvas_length']}")
-    if layer_fracs:
-        info["extra_layers"] = resolve_extra_layers(model, info["encoder_norm"], info["decoder_norm"], layer_fracs)
-    writer = CacheWriter(out_dir)  # opened before the hooks go on, so a folder that cannot be written leaks nothing
-    cap = NormCapture(model, info["encoder_norm"], info["decoder_norm"], extra=extra_layer_names(info))
-    rng = random.Random(1234)
-    n_dec, t0 = 0, time.time()
-    examples = examples[:limit] if limit else examples
-    try:
-        for i, ex in enumerate(examples):
-            order = shuffled_order(ex, shuffle_rng) if shuffle_rng is not None else None
-            draws = _slot_token(seed_mode, alias_ids, rng, vocab_size)
-            feats, rec = encode_example(model, tok, ex, cap, info, slot_tokens=draws, alias_ids=alias_ids,
-                                        max_tokens=max_tokens, option_order=order)
-            rec.path = f"{ex.id.replace('/', '_')}.safetensors"
-            writer.write(rec, feats)
-            n_dec += len(rec.decisions)
-            if (i + 1) % log_every == 0:
-                print(f"{i + 1}/{len(examples)} examples, {time.time() - t0:.0f}s", flush=True)
-    finally:  # close the manifest and drop the hooks even when an example raises
-        writer.close()
-        cap.remove()
-    n_extra = len((info.get("extra_layers") or {}).get("encoder", []))
-    summary = {"examples": len(examples), "decisions": n_dec, "seed_mode": seed_mode,
-               "canvas_length": canvas_length, "seconds": round(time.time() - t0, 1),
-               "feature_dim": hidden_size * (1 + n_extra), **info}
-    Path(out_dir, "summary.json").write_text(json.dumps(summary, indent=2))
-    return summary
-
-
-def zero_shot_accuracy(cache_dir) -> float:
-    cache = FeatureCache(cache_dir)
-    hits = total = 0
-    for rec in cache.records:
-        for d in rec.decisions:
-            if d.gold is None:
-                continue
-            total += 1
-            hits += int(argmax(d.prior) == argmax(d.gold))
-    return hits / max(total, 1)
-
-
-def load_backbone(model_id: str = MODEL_ID, load_in_4bit: bool = False):
-    from transformers import AutoTokenizer, DiffusionGemmaForBlockDiffusion
-    kwargs = {"device_map": "auto", "dtype": torch.bfloat16}
-    if load_in_4bit:
-        from transformers import BitsAndBytesConfig
-        kwargs["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)
-    model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, **kwargs).eval()
-    return model, AutoTokenizer.from_pretrained(model_id)
-
-
-def main():
-    ap = argparse.ArgumentParser()
-    sub = ap.add_subparsers(dest="cmd", required=True)
-    p = sub.add_parser("probe")
-    p.add_argument("--model", default=MODEL_ID)
-    e = sub.add_parser("extract")
-    e.add_argument("--model", default=MODEL_ID)
-    e.add_argument("--manifest", required=True)
-    e.add_argument("--out", required=True)
-    e.add_argument("--seed-mode", default="pad",
-                   help="pad, alias or random, or a comma-separated list of them; a list writes <out>_<mode> per mode")
-    e.add_argument("--canvas", type=int, default=64)
-    e.add_argument("--limit", type=int, default=None)
-    e.add_argument("--shuffle-seed", type=int, default=None, help="shuffle choice options (training augmentation)")
-    e.add_argument("--load-in-4bit", action="store_true")
-    e.add_argument("--zero-shot-log", action="store_true", help=f"append each cache's zero-shot accuracy to {ZERO_SHOT_LOG}")
-    for sp in (p, e):
-        sp.add_argument("--layers", default=None,
-                        help="depth fractions of extra encoder/decoder layers to append to the features, e.g. 0.5,0.75 "
-                             "(feature dim becomes hidden_size * (1 + count)); probe records their module names")
-    z = sub.add_parser("zero-shot")
-    z.add_argument("--cache", required=True)
-    args = ap.parse_args()
-    fracs = tuple(float(x) for x in args.layers.split(",")) if getattr(args, "layers", None) else ()
-    if args.cmd == "zero-shot":
-        print(json.dumps({"cache": args.cache, "zero_shot_accuracy": zero_shot_accuracy(args.cache)}))
-        return
-    if args.cmd == "extract":
-        modes = args.seed_mode.split(",")
-        if any(m not in SEED_MODES for m in modes) or len(set(modes)) != len(modes):
-            e.error(f"--seed-mode takes distinct values from {', '.join(SEED_MODES)}; got {args.seed_mode!r}")
-    model, tok = load_backbone(args.model, getattr(args, "load_in_4bit", False))
-    hidden = model.config.text_config.hidden_size
-    if args.cmd == "probe":
-        info = probe(model, tok, hidden)
-        if fracs:
-            info["extra_layers"] = resolve_extra_layers(model, info["encoder_norm"], info["decoder_norm"], fracs)
-            info["feature_dim"] = hidden * (1 + len(info["extra_layers"]["encoder"]))
-        Path("configs").mkdir(exist_ok=True)
-        Path("configs/backbone.json").write_text(json.dumps({**info, "model_id": args.model}, indent=2))
-        print(json.dumps(info, indent=2))
-        return
-    alias = alias_token_ids(tok, 64)
-    vocab = getattr(model.config.text_config, "vocab_size", 262144)
-    examples = read_manifest(args.manifest)
-    info = probe(model, tok, hidden) if len(modes) > 1 else None  # one probe shared by every mode
-    for mode in modes:
-        out = f"{args.out}_{mode}" if len(modes) > 1 else args.out
-        # a fresh rng per mode, so each cache matches what a separate run of that mode would write
-        rng = random.Random(args.shuffle_seed) if args.shuffle_seed is not None else None
-        summary = extract(model, tok, examples, out, seed_mode=mode, canvas_length=args.canvas, alias_ids=alias,
-                          hidden_size=hidden, shuffle_rng=rng, limit=args.limit, vocab_size=vocab, info=info,
-                          layer_fracs=fracs)
-        print(json.dumps(summary, indent=2))
-        acc = zero_shot_accuracy(out)
-        print(f"zero_shot {mode} {acc:.4f}", flush=True)
-        if args.zero_shot_log:
-            log = Path(ZERO_SHOT_LOG)
-            log.parent.mkdir(parents=True, exist_ok=True)
-            with log.open("a", encoding="utf-8") as f:
-                f.write(json.dumps({"cache": out, "seed_mode": mode, "zero_shot_accuracy": acc}) + "\n")
-
-
-if __name__ == "__main__":
-    main()
+"""Backbone feature extraction: one forward pass per example, hidden states captured by hooks."""
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import random
+import re
+import time
+from pathlib import Path
+
+import torch
+
+from .cache import CacheWriter, Decision, FeatureCache, Record
+from .data import read_manifest
+from .prompting import alias_token_ids, canvas_ids, shuffled_order, tokenize_prompt
+from .schema import Example, argmax
+
+MODEL_ID = "google/diffusiongemma-26B-A4B-it"
+PROBE_STATE = "A short probe state."
+RANDOM_DRAWS = 4
+SEED_MODES = ("pad", "alias", "random")
+ZERO_SHOT_LOG = "configs/session1_zero_shot.jsonl"
+
+
+class NormCapture:
+    """Registers forward hooks on final-norm modules (and any modules named in `extra`) and keeps their latest
+    outputs. With no names given every norm outside the layer stacks is hooked (discovery, see probe)."""
+
+    def __init__(self, model, encoder_name: str | None = None, decoder_name: str | None = None,
+                 extra: tuple[str, ...] | list[str] = ()):
+        self.outputs: dict[str, torch.Tensor] = {}
+        self.handles = []
+        extra = set(extra)
+        for name, mod in model.named_modules():
+            if name in extra:
+                self.handles.append(mod.register_forward_hook(self._hook(name)))
+                continue
+            if ".layers." in f".{name}." or "Norm" not in type(mod).__name__:
+                continue
+            if encoder_name is not None and name not in (encoder_name, decoder_name):
+                continue
+            self.handles.append(mod.register_forward_hook(self._hook(name)))
+
+    def _hook(self, name):
+        def fn(_m, _i, out):
+            self.outputs[name] = out[0] if isinstance(out, tuple) else out
+        return fn
+
+    def remove(self):
+        for h in self.handles:
+            h.remove()
+
+
+def unit_rms(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
+    """Scale every row to unit RMS (a parameter-free RMSNorm). Intermediate layers are residual-stream states
+    of arbitrary magnitude; this keeps them on the final norm's scale and inside float16 for the cache."""
+    return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
+
+
+_LAYER_RE = re.compile(r"(.*\blayers)\.(\d+)")
+
+
+def layer_stacks(model) -> dict[str, int]:
+    """Layer-stack prefixes and their depths, from module names of the form `<prefix>.layers.<n>`."""
+    stacks: dict[str, int] = {}
+    for name, _ in model.named_modules():
+        m = _LAYER_RE.fullmatch(name)
+        if m:
+            stacks[m.group(1)] = max(stacks.get(m.group(1), 0), int(m.group(2)) + 1)
+    return stacks
+
+
+def select_layers(n_layers: int, fracs) -> list[int]:
+    """Zero-based layer indices at the given depth fractions: 0.5 of 30 layers is layer 14, 0.75 is 22."""
+    out = set()
+    for f in fracs:
+        if not 0.0 < float(f) <= 1.0:
+            raise ValueError(f"layer fraction must be in (0, 1], got {f}")
+        out.add(max(0, min(n_layers - 1, math.ceil(float(f) * n_layers) - 1)))
+    return sorted(out)
+
+
+def resolve_extra_layers(model, encoder_norm: str, decoder_norm: str, fracs) -> dict[str, list[str]]:
+    """Module names of the extra layers to capture: for each final norm, the layer stack beside it (same parent,
+    else the shallowest stack below that parent) at `fracs` of its depth."""
+    stacks = layer_stacks(model)
+
+    def parent(name: str) -> str:
+        return name.rsplit(".", 1)[0] if "." in name else ""
+
+    def stack_for(norm: str) -> str:
+        p = parent(norm)
+        beside = [s for s in stacks if parent(s) == p]
+        below = [s for s in stacks if p and s.startswith(p + ".")]
+        cands = beside or below
+        if not cands:
+            raise RuntimeError(f"no layer stack found beside {norm!r}; stacks: {sorted(stacks)}")
+        return min(cands, key=len)
+
+    enc, dec = stack_for(encoder_norm), stack_for(decoder_norm)
+    if enc == dec:
+        raise RuntimeError(f"encoder and decoder norms resolve to the same layer stack {enc!r}")
+    return {"encoder": [f"{enc}.{i}" for i in select_layers(stacks[enc], fracs)],
+            "decoder": [f"{dec}.{i}" for i in select_layers(stacks[dec], fracs)]}
+
+
+def extra_layer_names(info: dict) -> list[str]:
+    extra = info.get("extra_layers") or {}
+    return list(extra.get("encoder", [])) + list(extra.get("decoder", []))
+
+
+def _probe_example() -> Example:
+    from .schema import Question
+    q = Question.from_jev("ok", {"type": "noul", "instructions": "Is the probe fine?"})
+    return Example(id="probe", source="probe", state=PROBE_STATE, questions=[q], gold={"ok": None})
+
+
+@torch.no_grad()
+def _forward(model, prompt_ids: list[int], canvas: list[int]):
+    dev = model.device
+    input_ids = torch.tensor([prompt_ids], dtype=torch.long, device=dev)
+    return model(input_ids=input_ids, attention_mask=torch.ones_like(input_ids),
+                 decoder_input_ids=torch.tensor([canvas], dtype=torch.long, device=dev))
+
+
+def probe(model, tok, hidden_size: int) -> dict:
+    ex = _probe_example()
+    tp = tokenize_prompt(ex, tok)
+    for canvas_length in (64, 256):
+        canvas, _ = canvas_ids(ex, tok, slot_token_id=0, canvas_length=canvas_length)
+        cap = NormCapture(model)
+        try:
+            _forward(model, tp.ids, canvas)
+        except Exception as err:  # noqa: BLE001 - the whole point is to learn which length works
+            cap.remove()
+            print(f"canvas_length {canvas_length} failed: {type(err).__name__}: {err}")
+            continue
+        cap.remove()
+        enc = [n for n, o in cap.outputs.items() if o.shape[-2] == len(tp.ids) and o.shape[-1] == hidden_size]
+        dec = [n for n, o in cap.outputs.items() if o.shape[-2] == canvas_length and o.shape[-1] == hidden_size]
+        if not enc or not dec:
+            shapes = {k: tuple(v.shape) for k, v in cap.outputs.items()}
+            raise RuntimeError(f"could not find final norms; captured {shapes}")
+        if enc[-1] == dec[-1]:
+            raise RuntimeError(f"encoder and decoder norms resolved to the same module: {enc[-1]}")
+        if len(enc) > 1 or len(dec) > 1:
+            print(f"warning: several final-norm candidates, using the last; encoder {enc}, decoder {dec}")
+        return {"canvas_length": canvas_length, "encoder_norm": enc[-1], "decoder_norm": dec[-1],
+                "hidden_size": hidden_size}
+    raise RuntimeError("neither canvas length 64 nor 256 ran")
+
+
+def _slot_token(seed_mode: str, alias_ids: list[int], rng: random.Random, vocab: int) -> list[int]:
+    if seed_mode == "pad":
+        return [0]
+    if seed_mode == "alias":
+        return [alias_ids[0]]
+    if seed_mode == "random":
+        return [rng.randrange(10, vocab) for _ in range(RANDOM_DRAWS)]
+    raise ValueError(f"unknown seed_mode {seed_mode}")
+
+
+def _rows(cap: "NormCapture", name: str, n_rows: int, what: str) -> torch.Tensor:
+    """A captured layer's rows for this forward (batch element 0), checked against the expected row count."""
+    out = cap.outputs[name][0].float()
+    if out.shape[0] != n_rows:
+        raise RuntimeError(f"{name} output has shape {tuple(out.shape)}, expected {n_rows} {what} rows")
+    return out
+
+
+@torch.no_grad()
+def encode_example(model, tok, ex: Example, cap: "NormCapture", info: dict, *, slot_tokens: list[int],
+                   alias_ids: list[int], max_tokens: int = 1024, option_order=None) -> tuple[torch.Tensor, Record]:
+    """One example through the backbone: features averaged over slot tokens, plus a Record with priors.
+
+    The Record's path is left empty; the caller names the file.
+    """
+    canvas_length = info["canvas_length"]
+    tp = tokenize_prompt(ex, tok, max_tokens=max_tokens, option_order=option_order)
+    feats_sum, prob_sum, slots = None, None, None
+    for slot_tok in slot_tokens:
+        canvas, slots = canvas_ids(ex, tok, slot_token_id=slot_tok, canvas_length=canvas_length)
+        cap.outputs.clear()
+        out = _forward(model, tp.ids, canvas)
+        if out.logits.shape[1] != canvas_length:
+            raise RuntimeError(f"logits cover {out.logits.shape[1]} positions, expected canvas_length {canvas_length}")
+        extra = info.get("extra_layers") or {}
+        wanted = [info["encoder_norm"], info["decoder_norm"], *extra_layer_names(info)]
+        missing = [n for n in wanted if n not in cap.outputs]
+        if missing:
+            raise RuntimeError(f"norm outputs missing after forward: {missing}; captured {sorted(cap.outputs)}")
+        hp = cap.outputs[info["encoder_norm"]][0].float()
+        hc = cap.outputs[info["decoder_norm"]][0].float()
+        if hp.shape[0] != len(tp.ids):
+            raise RuntimeError(f"encoder norm output has shape {tuple(hp.shape)}, expected {len(tp.ids)} prompt rows")
+        if hc.shape[0] != canvas_length:
+            raise RuntimeError(f"decoder norm output has shape {tuple(hc.shape)}, expected canvas_length {canvas_length}")
+        enc_extra = [unit_rms(_rows(cap, name, len(tp.ids), "prompt")) for name in extra.get("encoder", [])]
+        dec_extra = [unit_rms(_rows(cap, name, canvas_length, "canvas")) for name in extra.get("decoder", [])]
+        feats = torch.cat([torch.cat([hp, *enc_extra], dim=1), torch.cat([hc, *dec_extra], dim=1)], dim=0)
+        probs = torch.softmax(out.logits[0].float(), dim=-1)[slots]  # (n_q, vocab)
+        feats_sum = feats if feats_sum is None else feats_sum + feats
+        prob_sum = probs if prob_sum is None else prob_sum + probs
+    feats = feats_sum / len(slot_tokens)
+    logp = (prob_sum / len(slot_tokens)).clamp_min(1e-12).log()
+    decisions = []
+    for qi, q in enumerate(ex.questions):
+        ids = torch.tensor([alias_ids[a] for a in tp.option_alias[qi]], device=logp.device)
+        prior = logp[qi, ids]
+        prior = (prior - torch.logsumexp(prior, dim=0)).tolist()
+        decisions.append(Decision(qid=q.qid, type=q.type, slot=len(tp.ids) + slots[qi],
+                                  spans=[list(s) for s in tp.option_spans[qi]],
+                                  alias=list(tp.option_alias[qi]), prior=prior, gold=ex.gold.get(q.qid)))
+    rec = Record(id=ex.id, source=ex.source, path="", n_prompt=len(tp.ids),
+                 n_canvas=canvas_length, decisions=decisions)
+    return feats, rec
+
+
+def extract(model, tok, examples: list[Example], out_dir, *, seed_mode: str, canvas_length: int,
+            alias_ids: list[int], hidden_size: int, max_tokens: int = 1024, shuffle_rng=None,
+            limit: int | None = None, vocab_size: int = 262144, log_every: int = 50,
+            info: dict | None = None, layer_fracs=()) -> dict:
+    """Write a feature cache. With `layer_fracs` (depth fractions, e.g. (0.5, 0.75)) every row also carries the
+    unit-RMS outputs of the encoder (prompt rows) or decoder (canvas rows) layers at those depths, appended to the
+    final-norm block, so the feature dim is hidden_size * (1 + len(layers))."""
+    info = dict(probe(model, tok, hidden_size) if info is None else info)
+    if info["canvas_length"] != canvas_length:
+        raise ValueError(f"requested canvas {canvas_length} but probe says {info['canvas_length']}")
+    if layer_fracs:
+        info["extra_layers"] = resolve_extra_layers(model, info["encoder_norm"], info["decoder_norm"], layer_fracs)
+    writer = CacheWriter(out_dir)  # opened before the hooks go on, so a folder that cannot be written leaks nothing
+    cap = NormCapture(model, info["encoder_norm"], info["decoder_norm"], extra=extra_layer_names(info))
+    rng = random.Random(1234)
+    n_dec, t0 = 0, time.time()
+    examples = examples[:limit] if limit else examples
+    try:
+        for i, ex in enumerate(examples):
+            order = shuffled_order(ex, shuffle_rng) if shuffle_rng is not None else None
+            draws = _slot_token(seed_mode, alias_ids, rng, vocab_size)
+            feats, rec = encode_example(model, tok, ex, cap, info, slot_tokens=draws, alias_ids=alias_ids,
+                                        max_tokens=max_tokens, option_order=order)
+            rec.path = f"{ex.id.replace('/', '_')}.safetensors"
+            writer.write(rec, feats)
+            n_dec += len(rec.decisions)
+            if (i + 1) % log_every == 0:
+                print(f"{i + 1}/{len(examples)} examples, {time.time() - t0:.0f}s", flush=True)
+    finally:  # close the manifest and drop the hooks even when an example raises
+        writer.close()
+        cap.remove()
+    n_extra = len((info.get("extra_layers") or {}).get("encoder", []))
+    summary = {"examples": len(examples), "decisions": n_dec, "seed_mode": seed_mode,
+               "canvas_length": canvas_length, "seconds": round(time.time() - t0, 1),
+               "feature_dim": hidden_size * (1 + n_extra), **info}
+    Path(out_dir, "summary.json").write_text(json.dumps(summary, indent=2))
+    return summary
+
+
+def zero_shot_accuracy(cache_dir) -> float:
+    cache = FeatureCache(cache_dir)
+    hits = total = 0
+    for rec in cache.records:
+        for d in rec.decisions:
+            if d.gold is None:
+                continue
+            total += 1
+            hits += int(argmax(d.prior) == argmax(d.gold))
+    return hits / max(total, 1)
+
+
+def load_backbone(model_id: str = MODEL_ID, load_in_4bit: bool = False):
+    from transformers import AutoTokenizer, DiffusionGemmaForBlockDiffusion
+    kwargs = {"device_map": "auto", "dtype": torch.bfloat16}
+    if load_in_4bit:
+        from transformers import BitsAndBytesConfig
+        kwargs["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)
+    model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, **kwargs).eval()
+    return model, AutoTokenizer.from_pretrained(model_id)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    sub = ap.add_subparsers(dest="cmd", required=True)
+    p = sub.add_parser("probe")
+    p.add_argument("--model", default=MODEL_ID)
+    e = sub.add_parser("extract")
+    e.add_argument("--model", default=MODEL_ID)
+    e.add_argument("--manifest", required=True)
+    e.add_argument("--out", required=True)
+    e.add_argument("--seed-mode", default="pad",
+                   help="pad, alias or random, or a comma-separated list of them; a list writes <out>_<mode> per mode")
+    e.add_argument("--canvas", type=int, default=64)
+    e.add_argument("--limit", type=int, default=None)
+    e.add_argument("--shuffle-seed", type=int, default=None, help="shuffle choice options (training augmentation)")
+    e.add_argument("--load-in-4bit", action="store_true")
+    e.add_argument("--zero-shot-log", action="store_true", help=f"append each cache's zero-shot accuracy to {ZERO_SHOT_LOG}")
+    for sp in (p, e):
+        sp.add_argument("--layers", default=None,
+                        help="depth fractions of extra encoder/decoder layers to append to the features, e.g. 0.5,0.75 "
+                             "(feature dim becomes hidden_size * (1 + count)); probe records their module names")
+    z = sub.add_parser("zero-shot")
+    z.add_argument("--cache", required=True)
+    args = ap.parse_args()
+    fracs = tuple(float(x) for x in args.layers.split(",")) if getattr(args, "layers", None) else ()
+    if args.cmd == "zero-shot":
+        print(json.dumps({"cache": args.cache, "zero_shot_accuracy": zero_shot_accuracy(args.cache)}))
+        return
+    if args.cmd == "extract":
+        modes = args.seed_mode.split(",")
+        if any(m not in SEED_MODES for m in modes) or len(set(modes)) != len(modes):
+            e.error(f"--seed-mode takes distinct values from {', '.join(SEED_MODES)}; got {args.seed_mode!r}")
+    model, tok = load_backbone(args.model, getattr(args, "load_in_4bit", False))
+    hidden = model.config.text_config.hidden_size
+    if args.cmd == "probe":
+        info = probe(model, tok, hidden)
+        if fracs:
+            info["extra_layers"] = resolve_extra_layers(model, info["encoder_norm"], info["decoder_norm"], fracs)
+            info["feature_dim"] = hidden * (1 + len(info["extra_layers"]["encoder"]))
+        Path("configs").mkdir(exist_ok=True)
+        Path("configs/backbone.json").write_text(json.dumps({**info, "model_id": args.model}, indent=2))
+        print(json.dumps(info, indent=2))
+        return
+    alias = alias_token_ids(tok, 64)
+    vocab = getattr(model.config.text_config, "vocab_size", 262144)
+    examples = read_manifest(args.manifest)
+    info = probe(model, tok, hidden) if len(modes) > 1 else None  # one probe shared by every mode
+    for mode in modes:
+        out = f"{args.out}_{mode}" if len(modes) > 1 else args.out
+        # a fresh rng per mode, so each cache matches what a separate run of that mode would write
+        rng = random.Random(args.shuffle_seed) if args.shuffle_seed is not None else None
+        summary = extract(model, tok, examples, out, seed_mode=mode, canvas_length=args.canvas, alias_ids=alias,
+                          hidden_size=hidden, shuffle_rng=rng, limit=args.limit, vocab_size=vocab, info=info,
+                          layer_fracs=fracs)
+        print(json.dumps(summary, indent=2))
+        acc = zero_shot_accuracy(out)
+        print(f"zero_shot {mode} {acc:.4f}", flush=True)
+        if args.zero_shot_log:
+            log = Path(ZERO_SHOT_LOG)
+            log.parent.mkdir(parents=True, exist_ok=True)
+            with log.open("a", encoding="utf-8") as f:
+                f.write(json.dumps({"cache": out, "seed_mode": mode, "zero_shot_accuracy": acc}) + "\n")
+
+
+if __name__ == "__main__":
+    main()

head.py

@@ -1,119 +1,119 @@
-"""Hev decision head: projection, Hydra stack, pointer readout over option spans, backbone prior, temperatures."""
-from __future__ import annotations
-
-import math
-from dataclasses import asdict, dataclass
-
-import torch
-from torch import nn
-
-from .cache import Batch
-from .hydra_torch import HydraBlock, RMSNormGated
-
-
-@dataclass
-class HevHeadConfig:
-    d_backbone: int = 2816
-    d_model: int = 512
-    n_blocks: int = 3
-    d_state: int = 64
-    expand: int = 2
-    headdim: int = 64
-    d_conv: int = 7
-    mlp_mult: int = 2
-    chunk_size: int = 64
-    n_types: int = 3
-
-    def to_dict(self) -> dict:
-        return asdict(self)
-
-    @classmethod
-    def from_dict(cls, d: dict) -> "HevHeadConfig":
-        return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})
-
-
-class HevHead(nn.Module):
-    def __init__(self, cfg: HevHeadConfig):
-        super().__init__()
-        self.cfg = cfg
-        self.in_proj = nn.Linear(cfg.d_backbone, cfg.d_model)
-        self.seg_emb = nn.Embedding(2, cfg.d_model)
-        self.in_norm = RMSNormGated(cfg.d_model)
-        self.blocks = nn.ModuleList([
-            HydraBlock(cfg.d_model, mlp_mult=cfg.mlp_mult, d_state=cfg.d_state, d_conv=cfg.d_conv,
-                       expand=cfg.expand, headdim=cfg.headdim, chunk_size=cfg.chunk_size)
-            for _ in range(cfg.n_blocks)
-        ])
-        self.final_norm = RMSNormGated(cfg.d_model)
-        self.w_q = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
-        self.w_o = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
-        nn.init.normal_(self.w_q.weight, std=0.02)
-        nn.init.zeros_(self.w_o.weight)  # pointer term is exactly 0 at init, so a fresh head starts at the backbone prior
-        self.beta = nn.Parameter(torch.ones(cfg.n_types))
-        self.log_temp = nn.Parameter(torch.zeros(cfg.n_types))
-
-    def num_parameters(self) -> int:
-        return sum(p.numel() for p in self.parameters() if p.requires_grad)
-
-    def encode(self, feats: torch.Tensor, seg: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
-        h = self.in_proj(feats.float()) + self.seg_emb(seg)
-        h = self.in_norm(h) * mask[..., None].to(h.dtype)
-        for blk in self.blocks:
-            h = blk(h, mask)
-        return self.final_norm(h)
-
-    def forward_flat(self, batch: Batch, use_temperature: bool = False) -> torch.Tensor:
-        """Logits for every option of every decision in the batch as one flat tensor, in the batch's option order.
-
-        This is the whole readout in a fixed number of tensor ops: a gather of the span rows, one index_add for
-        the span means, one gather for the slot queries, and elementwise arithmetic on the flat option axis.
-        The per-decision loops it replaced issued ~360 ops per decision and left a GPU idle (session 3).
-        """
-        h = self.encode(batch.feats, batch.seg, batch.mask)
-        b, l, d = h.shape
-        if bool((batch.span_len <= 0).any()):  # one host sync per batch; collate already rejects empty spans
-            bad = int((batch.span_len <= 0).nonzero()[0])
-            raise ValueError(f"empty span for decision {batch.decisions[int(batch.opt_dec[bad])]['qid']!r}")
-        hf = h.reshape(b * l, d).float()
-        sums = hf.new_zeros(batch.span_len.shape[0], d).index_add(0, batch.span_opt, hf.index_select(0, batch.span_rows))
-        reps = sums / batch.span_len[:, None]
-        q = self.w_q(h[batch.dec_b, batch.dec_slot].float())  # (n_dec, d_model)
-        scale = 1.0 / math.sqrt(self.cfg.d_model)
-        scores = (self.w_o(reps) * q[batch.opt_dec]).sum(-1).float() * scale
-        logits = scores + self.beta[batch.dec_type][batch.opt_dec] * batch.opt_prior
-        if use_temperature:
-            logits = logits / torch.exp(self.log_temp)[batch.dec_type][batch.opt_dec]
-        return logits
-
-    def forward(self, batch: Batch, use_temperature: bool = False) -> list[torch.Tensor]:
-        """One logit vector per decision: views into forward_flat's output."""
-        return list(torch.split(self.forward_flat(batch, use_temperature), batch.counts))
-
-
-class HevEnsemble(nn.Module):
-    """One or more HevHeads with the same config, averaged in probability space.
-
-    forward() returns, per decision, the log of the mean of the members' softmaxes, so every consumer that
-    applies softmax to the head's output (evaluate.predict, HevForDecisions.decide) gets exactly the averaged
-    distribution. Members keep their own fitted temperatures; the ensemble is not re-calibrated."""
-
-    def __init__(self, members: list[HevHead]):
-        super().__init__()
-        if not members:
-            raise ValueError("an ensemble needs at least one member")
-        cfg = members[0].cfg
-        if any(m.cfg != cfg for m in members):
-            raise ValueError("every ensemble member must share one HevHeadConfig")
-        self.members = nn.ModuleList(members)
-        self.cfg = cfg
-
-    def num_parameters(self) -> int:
-        return sum(m.num_parameters() for m in self.members)
-
-    def forward(self, batch: Batch, use_temperature: bool = False) -> list[torch.Tensor]:
-        outs = [m(batch, use_temperature=use_temperature) for m in self.members]
-        result = []
-        for i in range(len(outs[0])):
-            probs = torch.stack([torch.softmax(o[i].float(), 0) for o in outs]).mean(0)
-            result.append(torch.log(probs.clamp_min(1e-12)))
-        return result
+"""Hev decision head: projection, Hydra stack, pointer readout over option spans, backbone prior, temperatures."""
+from __future__ import annotations
+
+import math
+from dataclasses import asdict, dataclass
+
+import torch
+from torch import nn
+
+from .cache import Batch
+from .hydra_torch import HydraBlock, RMSNormGated
+
+
+@dataclass
+class HevHeadConfig:
+    d_backbone: int = 2816
+    d_model: int = 512
+    n_blocks: int = 3
+    d_state: int = 64
+    expand: int = 2
+    headdim: int = 64
+    d_conv: int = 7
+    mlp_mult: int = 2
+    chunk_size: int = 64
+    n_types: int = 3
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "HevHeadConfig":
+        return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})
+
+
+class HevHead(nn.Module):
+    def __init__(self, cfg: HevHeadConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.in_proj = nn.Linear(cfg.d_backbone, cfg.d_model)
+        self.seg_emb = nn.Embedding(2, cfg.d_model)
+        self.in_norm = RMSNormGated(cfg.d_model)
+        self.blocks = nn.ModuleList([
+            HydraBlock(cfg.d_model, mlp_mult=cfg.mlp_mult, d_state=cfg.d_state, d_conv=cfg.d_conv,
+                       expand=cfg.expand, headdim=cfg.headdim, chunk_size=cfg.chunk_size)
+            for _ in range(cfg.n_blocks)
+        ])
+        self.final_norm = RMSNormGated(cfg.d_model)
+        self.w_q = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
+        self.w_o = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
+        nn.init.normal_(self.w_q.weight, std=0.02)
+        nn.init.zeros_(self.w_o.weight)  # pointer term is exactly 0 at init, so a fresh head starts at the backbone prior
+        self.beta = nn.Parameter(torch.ones(cfg.n_types))
+        self.log_temp = nn.Parameter(torch.zeros(cfg.n_types))
+
+    def num_parameters(self) -> int:
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+
+    def encode(self, feats: torch.Tensor, seg: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        h = self.in_proj(feats.float()) + self.seg_emb(seg)
+        h = self.in_norm(h) * mask[..., None].to(h.dtype)
+        for blk in self.blocks:
+            h = blk(h, mask)
+        return self.final_norm(h)
+
+    def forward_flat(self, batch: Batch, use_temperature: bool = False) -> torch.Tensor:
+        """Logits for every option of every decision in the batch as one flat tensor, in the batch's option order.
+
+        This is the whole readout in a fixed number of tensor ops: a gather of the span rows, one index_add for
+        the span means, one gather for the slot queries, and elementwise arithmetic on the flat option axis.
+        The per-decision loops it replaced issued ~360 ops per decision and left a GPU idle (session 3).
+        """
+        h = self.encode(batch.feats, batch.seg, batch.mask)
+        b, l, d = h.shape
+        if bool((batch.span_len <= 0).any()):  # one host sync per batch; collate already rejects empty spans
+            bad = int((batch.span_len <= 0).nonzero()[0])
+            raise ValueError(f"empty span for decision {batch.decisions[int(batch.opt_dec[bad])]['qid']!r}")
+        hf = h.reshape(b * l, d).float()
+        sums = hf.new_zeros(batch.span_len.shape[0], d).index_add(0, batch.span_opt, hf.index_select(0, batch.span_rows))
+        reps = sums / batch.span_len[:, None]
+        q = self.w_q(h[batch.dec_b, batch.dec_slot].float())  # (n_dec, d_model)
+        scale = 1.0 / math.sqrt(self.cfg.d_model)
+        scores = (self.w_o(reps) * q[batch.opt_dec]).sum(-1).float() * scale
+        logits = scores + self.beta[batch.dec_type][batch.opt_dec] * batch.opt_prior
+        if use_temperature:
+            logits = logits / torch.exp(self.log_temp)[batch.dec_type][batch.opt_dec]
+        return logits
+
+    def forward(self, batch: Batch, use_temperature: bool = False) -> list[torch.Tensor]:
+        """One logit vector per decision: views into forward_flat's output."""
+        return list(torch.split(self.forward_flat(batch, use_temperature), batch.counts))
+
+
+class HevEnsemble(nn.Module):
+    """One or more HevHeads with the same config, averaged in probability space.
+
+    forward() returns, per decision, the log of the mean of the members' softmaxes, so every consumer that
+    applies softmax to the head's output (evaluate.predict, HevForDecisions.decide) gets exactly the averaged
+    distribution. Members keep their own fitted temperatures; the ensemble is not re-calibrated."""
+
+    def __init__(self, members: list[HevHead]):
+        super().__init__()
+        if not members:
+            raise ValueError("an ensemble needs at least one member")
+        cfg = members[0].cfg
+        if any(m.cfg != cfg for m in members):
+            raise ValueError("every ensemble member must share one HevHeadConfig")
+        self.members = nn.ModuleList(members)
+        self.cfg = cfg
+
+    def num_parameters(self) -> int:
+        return sum(m.num_parameters() for m in self.members)
+
+    def forward(self, batch: Batch, use_temperature: bool = False) -> list[torch.Tensor]:
+        outs = [m(batch, use_temperature=use_temperature) for m in self.members]
+        result = []
+        for i in range(len(outs[0])):
+            probs = torch.stack([torch.softmax(o[i].float(), 0) for o in outs]).mean(0)
+            result.append(torch.log(probs.clamp_min(1e-12)))
+        return result

hydra_torch.py

@@ -1,156 +1,156 @@
-"""Pure-PyTorch Hydra (Hwang, Lahoti, Dao, Gu 2024): a quasiseparable bidirectional SSM mixer.
-
-Mirrors goombalab/hydra's HydraMixer parameter layout exactly (same names and shapes), so weights are
-interchangeable with the CUDA implementation. The SSD scan is the chunked algorithm from the Mamba-2
-paper (Listing 1), padded to a chunk multiple. No mamba_ssm, no einops, no CUDA kernels.
-"""
-from __future__ import annotations
-
-import math
-
-import torch
-import torch.nn.functional as F
-from torch import nn
-
-
-def segsum(x: torch.Tensor) -> torch.Tensor:
-    """out[..., i, j] = sum_{k=j+1..i} x[..., k] for i >= j, else -inf."""
-    T = x.shape[-1]
-    x = x.unsqueeze(-1).expand(*x.shape, T)
-    lower = torch.tril(torch.ones(T, T, dtype=torch.bool, device=x.device), diagonal=-1)
-    x = x.masked_fill(~lower, 0)
-    x_segsum = torch.cumsum(x, dim=-2)
-    keep = torch.tril(torch.ones(T, T, dtype=torch.bool, device=x.device), diagonal=0)
-    return x_segsum.masked_fill(~keep, float("-inf"))
-
-
-def ssd_scan(X: torch.Tensor, A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, chunk: int = 64) -> torch.Tensor:
-    b, l, h, p = X.shape
-    pad = (-l) % chunk
-    if pad:
-        X = F.pad(X, (0, 0, 0, 0, 0, pad))
-        A = F.pad(A, (0, 0, 0, pad))
-        B = F.pad(B, (0, 0, 0, 0, 0, pad))
-        C = F.pad(C, (0, 0, 0, 0, 0, pad))
-    L = X.shape[1]
-    c = L // chunk
-    X = X.reshape(b, c, chunk, h, p)
-    B = B.reshape(b, c, chunk, h, -1)
-    C = C.reshape(b, c, chunk, h, -1)
-    A = A.reshape(b, c, chunk, h).permute(0, 3, 1, 2)  # (b, h, c, l)
-    A_cumsum = torch.cumsum(A, dim=-1)
-    Lmat = torch.exp(segsum(A))  # (b, h, c, l, s)
-    Y_diag = torch.einsum("bclhn,bcshn,bhcls,bcshp->bclhp", C, B, Lmat, X)
-    decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum)
-    states = torch.einsum("bclhn,bhcl,bclhp->bchpn", B, decay_states, X)
-    states = torch.cat([torch.zeros_like(states[:, :1]), states], dim=1)
-    decay_chunk = torch.exp(segsum(F.pad(A_cumsum[:, :, :, -1], (1, 0))))
-    new_states = torch.einsum("bhzc,bchpn->bzhpn", decay_chunk, states)
-    states = new_states[:, :-1]
-    Y_off = torch.einsum("bclhn,bchpn,bhcl->bclhp", C, states, torch.exp(A_cumsum))
-    Y = (Y_diag + Y_off).reshape(b, L, h, p)
-    return Y[:, :l]
-
-
-def ssd_naive(X: torch.Tensor, A: torch.Tensor, B: torch.Tensor, C: torch.Tensor) -> torch.Tensor:
-    b, l, h, p = X.shape
-    n = B.shape[-1]
-    state = X.new_zeros(b, h, p, n)
-    ys = []
-    for t in range(l):
-        state = state * torch.exp(A[:, t])[:, :, None, None] + X[:, t][..., None] * B[:, t][:, :, None, :]
-        ys.append(torch.einsum("bhpn,bhn->bhp", state, C[:, t]))
-    return torch.stack(ys, dim=1)
-
-
-class RMSNormGated(nn.Module):
-    """RMSNorm followed by an optional SiLU gate (norm_before_gate=True in mamba_ssm)."""
-
-    def __init__(self, d: int, eps: float = 1e-5):
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(d))
-
-    def forward(self, x: torch.Tensor, z: torch.Tensor | None = None) -> torch.Tensor:
-        xf = x.float()
-        y = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight.float()
-        if z is not None:
-            y = y * F.silu(z.float())
-        return y.to(x.dtype)
-
-
-class HydraMixer(nn.Module):
-    def __init__(self, d_model: int, d_state: int = 64, d_conv: int = 7, expand: int = 2, headdim: int = 64,
-                 ngroups: int = 1, chunk_size: int = 64, dt_min: float = 0.001, dt_max: float = 0.1,
-                 dt_init_floor: float = 1e-4):
-        super().__init__()
-        self.d_model, self.d_state, self.headdim, self.ngroups, self.chunk_size = d_model, d_state, headdim, ngroups, chunk_size
-        self.d_inner = expand * d_model
-        if self.d_inner % headdim != 0:
-            raise ValueError(f"d_inner (expand * d_model = {self.d_inner}) must be divisible by headdim ({headdim})")
-        self.nheads = self.d_inner // headdim
-        if self.nheads % ngroups != 0:
-            raise ValueError(f"nheads (d_inner / headdim = {self.nheads}) must be divisible by ngroups ({ngroups})")
-        if d_conv % 2 != 1:
-            raise ValueError(f"d_conv must be odd so the non-causal conv keeps the sequence length, got {d_conv}")
-        d_in_proj = 2 * self.d_inner + 2 * (2 * ngroups * d_state) + 2 * self.nheads
-        self.in_proj = nn.Linear(d_model, d_in_proj, bias=False)
-        conv_dim = self.d_inner + 2 * (2 * ngroups * d_state)
-        self.conv1d = nn.Conv1d(conv_dim, conv_dim, kernel_size=d_conv, groups=conv_dim, padding=d_conv // 2, bias=True)
-        dt = torch.exp(torch.rand(self.nheads) * (math.log(dt_max) - math.log(dt_min)) + math.log(dt_min))
-        dt = torch.clamp(dt, min=dt_init_floor)
-        self.dt_bias = nn.Parameter(dt + torch.log(-torch.expm1(-dt)))
-        self.dt_bias._no_weight_decay = True
-        self.A_log = nn.Parameter(torch.zeros(self.nheads))  # A = -exp(0) = -1, as in the official init
-        self.A_log._no_weight_decay = True
-        self.D = nn.Parameter(torch.ones(self.nheads))
-        self.D._no_weight_decay = True
-        self.fc_D = nn.Linear(self.d_inner, self.nheads, bias=False)
-        self.norm = RMSNormGated(self.d_inner, eps=1e-5)
-        self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)
-
-    def forward(self, u: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
-        b, l, _ = u.shape
-        H, P, G, N = self.nheads, self.headdim, self.ngroups, self.d_state
-        if mask is not None:
-            u = u * mask[..., None].to(u.dtype)
-        z, xBC, dt = torch.split(self.in_proj(u), [self.d_inner, self.d_inner + 4 * G * N, 2 * H], dim=-1)
-        xBC = F.silu(self.conv1d(xBC.transpose(1, 2)).transpose(1, 2))
-        dt2 = F.softplus(torch.cat([dt[..., :H], dt[..., H:].flip(1)], dim=0) + self.dt_bias)  # (2b, l, H)
-        if mask is not None:
-            m = mask[..., None].to(u.dtype)
-            xBC = xBC * m
-            dt2 = dt2 * torch.cat([m, m.flip(1)], dim=0)
-        x, BC = torch.split(xBC, [self.d_inner, 4 * G * N], dim=-1)
-        x_og = x
-        x2 = torch.cat([x, x.flip(1)], dim=0)
-        BC2 = torch.cat([BC[..., :2 * G * N], BC[..., 2 * G * N:].flip(1)], dim=0)
-        Bm, Cm = torch.split(BC2, [G * N, G * N], dim=-1)
-        A = -torch.exp(self.A_log.float())
-        Xh = x2.reshape(2 * b, l, H, P).float() * dt2.float()[..., None]
-        Ah = A[None, None, :] * dt2.float()
-        Bh = Bm.reshape(2 * b, l, G, N).float().repeat_interleave(H // G, dim=2)
-        Ch = Cm.reshape(2 * b, l, G, N).float().repeat_interleave(H // G, dim=2)
-        y = ssd_scan(Xh, Ah, Bh, Ch, self.chunk_size).reshape(2 * b, l, self.d_inner).to(u.dtype)
-        y = torch.cat([torch.zeros_like(y[:, :1]), y[:, :-1]], dim=1)  # shift by one, first position zero
-        y_fw, y_bw = y[:b], y[b:].flip(1)
-        d_term = F.linear(x_og, self.fc_D.weight, self.D).repeat_interleave(P, dim=-1)
-        y = y_fw + y_bw + x_og * d_term
-        y = self.norm(y, z)
-        return self.out_proj(y)
-
-
-class HydraBlock(nn.Module):
-    def __init__(self, d_model: int, mlp_mult: int = 2, **mixer_kwargs):
-        super().__init__()
-        self.norm1 = RMSNormGated(d_model)
-        self.mixer = HydraMixer(d_model, **mixer_kwargs)
-        self.norm2 = RMSNormGated(d_model)
-        hidden = mlp_mult * d_model
-        self.fc1 = nn.Linear(d_model, 2 * hidden)
-        self.fc2 = nn.Linear(hidden, d_model)
-
-    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
-        x = x + self.mixer(self.norm1(x), mask)
-        h, g = self.fc1(self.norm2(x)).chunk(2, dim=-1)
-        return x + self.fc2(F.silu(g) * h)
+"""Pure-PyTorch Hydra (Hwang, Lahoti, Dao, Gu 2024): a quasiseparable bidirectional SSM mixer.
+
+Mirrors goombalab/hydra's HydraMixer parameter layout exactly (same names and shapes), so weights are
+interchangeable with the CUDA implementation. The SSD scan is the chunked algorithm from the Mamba-2
+paper (Listing 1), padded to a chunk multiple. No mamba_ssm, no einops, no CUDA kernels.
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+def segsum(x: torch.Tensor) -> torch.Tensor:
+    """out[..., i, j] = sum_{k=j+1..i} x[..., k] for i >= j, else -inf."""
+    T = x.shape[-1]
+    x = x.unsqueeze(-1).expand(*x.shape, T)
+    lower = torch.tril(torch.ones(T, T, dtype=torch.bool, device=x.device), diagonal=-1)
+    x = x.masked_fill(~lower, 0)
+    x_segsum = torch.cumsum(x, dim=-2)
+    keep = torch.tril(torch.ones(T, T, dtype=torch.bool, device=x.device), diagonal=0)
+    return x_segsum.masked_fill(~keep, float("-inf"))
+
+
+def ssd_scan(X: torch.Tensor, A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, chunk: int = 64) -> torch.Tensor:
+    b, l, h, p = X.shape
+    pad = (-l) % chunk
+    if pad:
+        X = F.pad(X, (0, 0, 0, 0, 0, pad))
+        A = F.pad(A, (0, 0, 0, pad))
+        B = F.pad(B, (0, 0, 0, 0, 0, pad))
+        C = F.pad(C, (0, 0, 0, 0, 0, pad))
+    L = X.shape[1]
+    c = L // chunk
+    X = X.reshape(b, c, chunk, h, p)
+    B = B.reshape(b, c, chunk, h, -1)
+    C = C.reshape(b, c, chunk, h, -1)
+    A = A.reshape(b, c, chunk, h).permute(0, 3, 1, 2)  # (b, h, c, l)
+    A_cumsum = torch.cumsum(A, dim=-1)
+    Lmat = torch.exp(segsum(A))  # (b, h, c, l, s)
+    Y_diag = torch.einsum("bclhn,bcshn,bhcls,bcshp->bclhp", C, B, Lmat, X)
+    decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum)
+    states = torch.einsum("bclhn,bhcl,bclhp->bchpn", B, decay_states, X)
+    states = torch.cat([torch.zeros_like(states[:, :1]), states], dim=1)
+    decay_chunk = torch.exp(segsum(F.pad(A_cumsum[:, :, :, -1], (1, 0))))
+    new_states = torch.einsum("bhzc,bchpn->bzhpn", decay_chunk, states)
+    states = new_states[:, :-1]
+    Y_off = torch.einsum("bclhn,bchpn,bhcl->bclhp", C, states, torch.exp(A_cumsum))
+    Y = (Y_diag + Y_off).reshape(b, L, h, p)
+    return Y[:, :l]
+
+
+def ssd_naive(X: torch.Tensor, A: torch.Tensor, B: torch.Tensor, C: torch.Tensor) -> torch.Tensor:
+    b, l, h, p = X.shape
+    n = B.shape[-1]
+    state = X.new_zeros(b, h, p, n)
+    ys = []
+    for t in range(l):
+        state = state * torch.exp(A[:, t])[:, :, None, None] + X[:, t][..., None] * B[:, t][:, :, None, :]
+        ys.append(torch.einsum("bhpn,bhn->bhp", state, C[:, t]))
+    return torch.stack(ys, dim=1)
+
+
+class RMSNormGated(nn.Module):
+    """RMSNorm followed by an optional SiLU gate (norm_before_gate=True in mamba_ssm)."""
+
+    def __init__(self, d: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(d))
+
+    def forward(self, x: torch.Tensor, z: torch.Tensor | None = None) -> torch.Tensor:
+        xf = x.float()
+        y = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight.float()
+        if z is not None:
+            y = y * F.silu(z.float())
+        return y.to(x.dtype)
+
+
+class HydraMixer(nn.Module):
+    def __init__(self, d_model: int, d_state: int = 64, d_conv: int = 7, expand: int = 2, headdim: int = 64,
+                 ngroups: int = 1, chunk_size: int = 64, dt_min: float = 0.001, dt_max: float = 0.1,
+                 dt_init_floor: float = 1e-4):
+        super().__init__()
+        self.d_model, self.d_state, self.headdim, self.ngroups, self.chunk_size = d_model, d_state, headdim, ngroups, chunk_size
+        self.d_inner = expand * d_model
+        if self.d_inner % headdim != 0:
+            raise ValueError(f"d_inner (expand * d_model = {self.d_inner}) must be divisible by headdim ({headdim})")
+        self.nheads = self.d_inner // headdim
+        if self.nheads % ngroups != 0:
+            raise ValueError(f"nheads (d_inner / headdim = {self.nheads}) must be divisible by ngroups ({ngroups})")
+        if d_conv % 2 != 1:
+            raise ValueError(f"d_conv must be odd so the non-causal conv keeps the sequence length, got {d_conv}")
+        d_in_proj = 2 * self.d_inner + 2 * (2 * ngroups * d_state) + 2 * self.nheads
+        self.in_proj = nn.Linear(d_model, d_in_proj, bias=False)
+        conv_dim = self.d_inner + 2 * (2 * ngroups * d_state)
+        self.conv1d = nn.Conv1d(conv_dim, conv_dim, kernel_size=d_conv, groups=conv_dim, padding=d_conv // 2, bias=True)
+        dt = torch.exp(torch.rand(self.nheads) * (math.log(dt_max) - math.log(dt_min)) + math.log(dt_min))
+        dt = torch.clamp(dt, min=dt_init_floor)
+        self.dt_bias = nn.Parameter(dt + torch.log(-torch.expm1(-dt)))
+        self.dt_bias._no_weight_decay = True
+        self.A_log = nn.Parameter(torch.zeros(self.nheads))  # A = -exp(0) = -1, as in the official init
+        self.A_log._no_weight_decay = True
+        self.D = nn.Parameter(torch.ones(self.nheads))
+        self.D._no_weight_decay = True
+        self.fc_D = nn.Linear(self.d_inner, self.nheads, bias=False)
+        self.norm = RMSNormGated(self.d_inner, eps=1e-5)
+        self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)
+
+    def forward(self, u: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
+        b, l, _ = u.shape
+        H, P, G, N = self.nheads, self.headdim, self.ngroups, self.d_state
+        if mask is not None:
+            u = u * mask[..., None].to(u.dtype)
+        z, xBC, dt = torch.split(self.in_proj(u), [self.d_inner, self.d_inner + 4 * G * N, 2 * H], dim=-1)
+        xBC = F.silu(self.conv1d(xBC.transpose(1, 2)).transpose(1, 2))
+        dt2 = F.softplus(torch.cat([dt[..., :H], dt[..., H:].flip(1)], dim=0) + self.dt_bias)  # (2b, l, H)
+        if mask is not None:
+            m = mask[..., None].to(u.dtype)
+            xBC = xBC * m
+            dt2 = dt2 * torch.cat([m, m.flip(1)], dim=0)
+        x, BC = torch.split(xBC, [self.d_inner, 4 * G * N], dim=-1)
+        x_og = x
+        x2 = torch.cat([x, x.flip(1)], dim=0)
+        BC2 = torch.cat([BC[..., :2 * G * N], BC[..., 2 * G * N:].flip(1)], dim=0)
+        Bm, Cm = torch.split(BC2, [G * N, G * N], dim=-1)
+        A = -torch.exp(self.A_log.float())
+        Xh = x2.reshape(2 * b, l, H, P).float() * dt2.float()[..., None]
+        Ah = A[None, None, :] * dt2.float()
+        Bh = Bm.reshape(2 * b, l, G, N).float().repeat_interleave(H // G, dim=2)
+        Ch = Cm.reshape(2 * b, l, G, N).float().repeat_interleave(H // G, dim=2)
+        y = ssd_scan(Xh, Ah, Bh, Ch, self.chunk_size).reshape(2 * b, l, self.d_inner).to(u.dtype)
+        y = torch.cat([torch.zeros_like(y[:, :1]), y[:, :-1]], dim=1)  # shift by one, first position zero
+        y_fw, y_bw = y[:b], y[b:].flip(1)
+        d_term = F.linear(x_og, self.fc_D.weight, self.D).repeat_interleave(P, dim=-1)
+        y = y_fw + y_bw + x_og * d_term
+        y = self.norm(y, z)
+        return self.out_proj(y)
+
+
+class HydraBlock(nn.Module):
+    def __init__(self, d_model: int, mlp_mult: int = 2, **mixer_kwargs):
+        super().__init__()
+        self.norm1 = RMSNormGated(d_model)
+        self.mixer = HydraMixer(d_model, **mixer_kwargs)
+        self.norm2 = RMSNormGated(d_model)
+        hidden = mlp_mult * d_model
+        self.fc1 = nn.Linear(d_model, 2 * hidden)
+        self.fc2 = nn.Linear(hidden, d_model)
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
+        x = x + self.mixer(self.norm1(x), mask)
+        h, g = self.fc1(self.norm2(x)).chunk(2, dim=-1)
+        return x + self.fc2(F.silu(g) * h)

modeling_hev.py

@@ -1,163 +1,163 @@
-"""HevForDecisions: frozen DiffusionGemma backbone + Hev head, exposing Jev-shaped decide()."""
-from __future__ import annotations
-
-import math
-import random
-from pathlib import Path
-
-import torch
-from safetensors.torch import load_file
-from transformers import PreTrainedModel
-
-from .cache import collate
-from .configuration_hev import HevConfig
-from .features import NormCapture, _slot_token, encode_example, extra_layer_names
-from .head import HevEnsemble, HevHead, HevHeadConfig
-from .prompting import PROMPT_FORMAT_VERSION, alias_token_ids
-from .schema import Example, Question
-
-
-# Hub download options that from_pretrained forwards to snapshot_download and HevConfig.from_pretrained.
-HUB_KWARGS = ("revision", "token", "cache_dir", "local_files_only", "force_download")
-
-
-def _is_local_dir(name: str) -> bool:
-    try:
-        return Path(name).is_dir()
-    except (OSError, ValueError):
-        return False
-
-
-def _vocab_size(backbone) -> int:
-    """Backbone vocabulary size, read as hev.features.main reads it for random slot draws (default 262144)."""
-    return getattr(getattr(backbone.config, "text_config", None), "vocab_size", 262144)
-
-
-def _param_device(module) -> torch.device:
-    """Device of the module's first parameter, or CPU when it has none."""
-    params = module.parameters() if hasattr(module, "parameters") else iter(())
-    first = next(params, None)
-    return first.device if first is not None else torch.device("cpu")
-
-
-class HevForDecisions(PreTrainedModel):
-    config_class = HevConfig
-
-    def __init__(self, config: HevConfig, backbone=None, tokenizer=None):
-        super().__init__(config)
-        members = []
-        temps = list(getattr(config, "member_temperatures", None) or [config.temperatures])
-        for i in range(int(getattr(config, "n_heads", 1))):
-            head = HevHead(HevHeadConfig.from_dict(config.head_config))
-            with torch.no_grad():
-                head.log_temp.copy_(torch.log(torch.tensor(temps[i] if i < len(temps) else config.temperatures,
-                                                           dtype=torch.float32)))
-            members.append(head)
-        self.head = HevEnsemble(members)
-        self.backbone = backbone
-        self.tokenizer = tokenizer
-        self._cap = None
-        self._alias_ids = None
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, backbone_cls=None, tokenizer_cls=None,
-                        dtype=torch.bfloat16, device_map="auto", load_in_4bit: bool = False, config=None, **kwargs):
-        """Load the Hev config, head, frozen backbone and tokenizer from a local folder or a Hub repo.
-
-        `revision`, `token`, `cache_dir`, `local_files_only` and `force_download` are forwarded to the Hub
-        download and to HevConfig.from_pretrained. A HevConfig passed as `config` (AutoModel does this) is used as is.
-        """
-        hub = {k: kwargs[k] for k in HUB_KWARGS if kwargs.get(k) is not None}
-        name = str(pretrained_model_name_or_path)
-        path = Path(name)
-        from_hub = not path.exists()
-        if from_hub:
-            from huggingface_hub import snapshot_download
-            path = Path(snapshot_download(name, **hub))
-        if not isinstance(config, HevConfig):
-            config = HevConfig.from_pretrained(str(path), **hub)
-        if from_hub:
-            config._name_or_path = name  # the repo id, not the local snapshot folder
-        if backbone_cls is None or tokenizer_cls is None:
-            from transformers import AutoConfig, AutoTokenizer, DiffusionGemmaForBlockDiffusion
-            backbone_cls = backbone_cls or DiffusionGemmaForBlockDiffusion
-            tokenizer_cls = tokenizer_cls or AutoTokenizer
-            bb_cfg = AutoConfig.for_model(**config.backbone_config)
-        else:
-            bb_cfg = None
-        bb_kwargs = {"config": bb_cfg, "dtype": dtype, "device_map": device_map}
-        if load_in_4bit:
-            from transformers import BitsAndBytesConfig
-            bb_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=dtype)
-        backbone = backbone_cls.from_pretrained(str(path), **bb_kwargs).eval()
-        tokenizer = tokenizer_cls.from_pretrained(str(path))
-        model = cls(config, backbone=backbone, tokenizer=tokenizer)
-        if len(model.head.members) == 1 and (path / "hev_head.safetensors").exists():
-            model.head.members[0].load_state_dict(load_file(str(path / "hev_head.safetensors")))
-        else:
-            for i, member in enumerate(model.head.members):
-                member.load_state_dict(load_file(str(path / f"hev_head_{i}.safetensors")))
-        model.head.to(_param_device(backbone))
-        return model.eval()
-
-    @staticmethod
-    def confidence(probs) -> float:
-        k = len(probs)
-        if k < 2:
-            return 1.0
-        h = -sum(p * math.log(p) for p in probs if p > 0)
-        return max(0.0, min(1.0, 1.0 - h / math.log(k)))
-
-    def _extra_layers(self) -> dict:
-        extra = getattr(self.config, "extra_layers", None) or {}
-        return {"encoder": list(extra.get("encoder", [])), "decoder": list(extra.get("decoder", []))}
-
-    def _ensure_capture(self):
-        if self._cap is None:
-            self._cap = NormCapture(self.backbone, self.config.encoder_norm, self.config.decoder_norm,
-                                    extra=extra_layer_names({"extra_layers": self._extra_layers()}))
-            self._alias_ids = alias_token_ids(self.tokenizer, len(self.config.aliases))
-
-    @torch.no_grad()
-    def decide(self, state: str, questions: dict) -> dict:
-        """Answer Jev-shaped questions about `state` with the backbone features and the trained head.
-
-        The answer slots are seeded exactly as hev.features.extract seeded them in training, by
-        `config.seed_mode`: the pad token (0) for "pad", the first alias token for "alias", and an average
-        over RANDOM_DRAWS forward passes with random slot tokens (fixed rng seed 1234) for "random".
-        `config.slot_token_id` is not read here; seed_mode governs.
-        """
-        version = getattr(self.config, "prompt_format_version", 1)
-        if version != PROMPT_FORMAT_VERSION:
-            raise RuntimeError(
-                f"this head was trained with prompt format {version}, but the installed hev code builds prompt format "
-                f"{PROMPT_FORMAT_VERSION}; load the code revision the head was published with, or retrain the head")
-        self._ensure_capture()
-        qs = [Question.from_jev(qid, spec) for qid, spec in questions.items()]
-        ex = Example(id="request", source="api", state=state, questions=qs, gold={})
-        info = {"canvas_length": self.config.canvas_length, "encoder_norm": self.config.encoder_norm,
-                "decoder_norm": self.config.decoder_norm, "extra_layers": self._extra_layers()}
-        slot_tokens = _slot_token(self.config.seed_mode, self._alias_ids, random.Random(1234),
-                                  _vocab_size(self.backbone))
-        feats, rec = encode_example(self.backbone, self.tokenizer, ex, self._cap, info,
-                                    slot_tokens=slot_tokens, alias_ids=self._alias_ids,
-                                    max_tokens=self.config.max_prompt_tokens)
-        batch = collate([(rec, feats)]).to(next(self.head.parameters()).device)
-        logits = self.head(batch, use_temperature=True)
-        answers = {}
-        for q, lg in zip(qs, logits):
-            probs = torch.softmax(lg.float(), 0).tolist()
-            if q.type == "noul":
-                answers[q.qid] = {"type": "noul", "noul": probs[0]}
-            elif q.type == "choice":
-                best = max(range(q.k), key=lambda i: probs[i])
-                answers[q.qid] = {"type": "choice", "choice": q.options[best],
-                                  "probabilities": dict(zip(q.options, probs)), "confidence": self.confidence(probs)}
-            else:
-                answers[q.qid] = {"type": "score", "score": sum(i * p for i, p in enumerate(probs)),
-                                  "probabilities": {str(i): p for i, p in enumerate(probs)},
-                                  "confidence": self.confidence(probs), "legend": list(q.descriptions)}
-        name = getattr(self.config, "_name_or_path", "") or ""
-        # A local folder name would leak a filesystem path into every response, so only a Hub repo id is reported.
-        return {"model": name if name and not _is_local_dir(name) else "hev",
-                "answers": answers, "usage": {"input_tokens": rec.n_prompt, "output_tokens": 0}}
+"""HevForDecisions: frozen DiffusionGemma backbone + Hev head, exposing Jev-shaped decide()."""
+from __future__ import annotations
+
+import math
+import random
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+from transformers import PreTrainedModel
+
+from .cache import collate
+from .configuration_hev import HevConfig
+from .features import NormCapture, _slot_token, encode_example, extra_layer_names
+from .head import HevEnsemble, HevHead, HevHeadConfig
+from .prompting import PROMPT_FORMAT_VERSION, alias_token_ids
+from .schema import Example, Question
+
+
+# Hub download options that from_pretrained forwards to snapshot_download and HevConfig.from_pretrained.
+HUB_KWARGS = ("revision", "token", "cache_dir", "local_files_only", "force_download")
+
+
+def _is_local_dir(name: str) -> bool:
+    try:
+        return Path(name).is_dir()
+    except (OSError, ValueError):
+        return False
+
+
+def _vocab_size(backbone) -> int:
+    """Backbone vocabulary size, read as hev.features.main reads it for random slot draws (default 262144)."""
+    return getattr(getattr(backbone.config, "text_config", None), "vocab_size", 262144)
+
+
+def _param_device(module) -> torch.device:
+    """Device of the module's first parameter, or CPU when it has none."""
+    params = module.parameters() if hasattr(module, "parameters") else iter(())
+    first = next(params, None)
+    return first.device if first is not None else torch.device("cpu")
+
+
+class HevForDecisions(PreTrainedModel):
+    config_class = HevConfig
+
+    def __init__(self, config: HevConfig, backbone=None, tokenizer=None):
+        super().__init__(config)
+        members = []
+        temps = list(getattr(config, "member_temperatures", None) or [config.temperatures])
+        for i in range(int(getattr(config, "n_heads", 1))):
+            head = HevHead(HevHeadConfig.from_dict(config.head_config))
+            with torch.no_grad():
+                head.log_temp.copy_(torch.log(torch.tensor(temps[i] if i < len(temps) else config.temperatures,
+                                                           dtype=torch.float32)))
+            members.append(head)
+        self.head = HevEnsemble(members)
+        self.backbone = backbone
+        self.tokenizer = tokenizer
+        self._cap = None
+        self._alias_ids = None
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, backbone_cls=None, tokenizer_cls=None,
+                        dtype=torch.bfloat16, device_map="auto", load_in_4bit: bool = False, config=None, **kwargs):
+        """Load the Hev config, head, frozen backbone and tokenizer from a local folder or a Hub repo.
+
+        `revision`, `token`, `cache_dir`, `local_files_only` and `force_download` are forwarded to the Hub
+        download and to HevConfig.from_pretrained. A HevConfig passed as `config` (AutoModel does this) is used as is.
+        """
+        hub = {k: kwargs[k] for k in HUB_KWARGS if kwargs.get(k) is not None}
+        name = str(pretrained_model_name_or_path)
+        path = Path(name)
+        from_hub = not path.exists()
+        if from_hub:
+            from huggingface_hub import snapshot_download
+            path = Path(snapshot_download(name, **hub))
+        if not isinstance(config, HevConfig):
+            config = HevConfig.from_pretrained(str(path), **hub)
+        if from_hub:
+            config._name_or_path = name  # the repo id, not the local snapshot folder
+        if backbone_cls is None or tokenizer_cls is None:
+            from transformers import AutoConfig, AutoTokenizer, DiffusionGemmaForBlockDiffusion
+            backbone_cls = backbone_cls or DiffusionGemmaForBlockDiffusion
+            tokenizer_cls = tokenizer_cls or AutoTokenizer
+            bb_cfg = AutoConfig.for_model(**config.backbone_config)
+        else:
+            bb_cfg = None
+        bb_kwargs = {"config": bb_cfg, "dtype": dtype, "device_map": device_map}
+        if load_in_4bit:
+            from transformers import BitsAndBytesConfig
+            bb_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=dtype)
+        backbone = backbone_cls.from_pretrained(str(path), **bb_kwargs).eval()
+        tokenizer = tokenizer_cls.from_pretrained(str(path))
+        model = cls(config, backbone=backbone, tokenizer=tokenizer)
+        if len(model.head.members) == 1 and (path / "hev_head.safetensors").exists():
+            model.head.members[0].load_state_dict(load_file(str(path / "hev_head.safetensors")))
+        else:
+            for i, member in enumerate(model.head.members):
+                member.load_state_dict(load_file(str(path / f"hev_head_{i}.safetensors")))
+        model.head.to(_param_device(backbone))
+        return model.eval()
+
+    @staticmethod
+    def confidence(probs) -> float:
+        k = len(probs)
+        if k < 2:
+            return 1.0
+        h = -sum(p * math.log(p) for p in probs if p > 0)
+        return max(0.0, min(1.0, 1.0 - h / math.log(k)))
+
+    def _extra_layers(self) -> dict:
+        extra = getattr(self.config, "extra_layers", None) or {}
+        return {"encoder": list(extra.get("encoder", [])), "decoder": list(extra.get("decoder", []))}
+
+    def _ensure_capture(self):
+        if self._cap is None:
+            self._cap = NormCapture(self.backbone, self.config.encoder_norm, self.config.decoder_norm,
+                                    extra=extra_layer_names({"extra_layers": self._extra_layers()}))
+            self._alias_ids = alias_token_ids(self.tokenizer, len(self.config.aliases))
+
+    @torch.no_grad()
+    def decide(self, state: str, questions: dict) -> dict:
+        """Answer Jev-shaped questions about `state` with the backbone features and the trained head.
+
+        The answer slots are seeded exactly as hev.features.extract seeded them in training, by
+        `config.seed_mode`: the pad token (0) for "pad", the first alias token for "alias", and an average
+        over RANDOM_DRAWS forward passes with random slot tokens (fixed rng seed 1234) for "random".
+        `config.slot_token_id` is not read here; seed_mode governs.
+        """
+        version = getattr(self.config, "prompt_format_version", 1)
+        if version != PROMPT_FORMAT_VERSION:
+            raise RuntimeError(
+                f"this head was trained with prompt format {version}, but the installed hev code builds prompt format "
+                f"{PROMPT_FORMAT_VERSION}; load the code revision the head was published with, or retrain the head")
+        self._ensure_capture()
+        qs = [Question.from_jev(qid, spec) for qid, spec in questions.items()]
+        ex = Example(id="request", source="api", state=state, questions=qs, gold={})
+        info = {"canvas_length": self.config.canvas_length, "encoder_norm": self.config.encoder_norm,
+                "decoder_norm": self.config.decoder_norm, "extra_layers": self._extra_layers()}
+        slot_tokens = _slot_token(self.config.seed_mode, self._alias_ids, random.Random(1234),
+                                  _vocab_size(self.backbone))
+        feats, rec = encode_example(self.backbone, self.tokenizer, ex, self._cap, info,
+                                    slot_tokens=slot_tokens, alias_ids=self._alias_ids,
+                                    max_tokens=self.config.max_prompt_tokens)
+        batch = collate([(rec, feats)]).to(next(self.head.parameters()).device)
+        logits = self.head(batch, use_temperature=True)
+        answers = {}
+        for q, lg in zip(qs, logits):
+            probs = torch.softmax(lg.float(), 0).tolist()
+            if q.type == "noul":
+                answers[q.qid] = {"type": "noul", "noul": probs[0]}
+            elif q.type == "choice":
+                best = max(range(q.k), key=lambda i: probs[i])
+                answers[q.qid] = {"type": "choice", "choice": q.options[best],
+                                  "probabilities": dict(zip(q.options, probs)), "confidence": self.confidence(probs)}
+            else:
+                answers[q.qid] = {"type": "score", "score": sum(i * p for i, p in enumerate(probs)),
+                                  "probabilities": {str(i): p for i, p in enumerate(probs)},
+                                  "confidence": self.confidence(probs), "legend": list(q.descriptions)}
+        name = getattr(self.config, "_name_or_path", "") or ""
+        # A local folder name would leak a filesystem path into every response, so only a Hub repo id is reported.
+        return {"model": name if name and not _is_local_dir(name) else "hev",
+                "answers": answers, "usage": {"input_tokens": rec.n_prompt, "output_tokens": 0}}

prompting.py

@@ -1,118 +1,118 @@
-"""Prompt and canvas construction. Token spans are exact because every piece is tokenised separately."""
-from __future__ import annotations
-
-import random
-from dataclasses import dataclass
-
-from .schema import ALIASES, Example
-
-# Bump when anything below changes the token ids, option spans or canvas a trained head was fitted on; every
-# feature cache must then be re-extracted and the head retrained. tests/test_format_pin.py pins version 1.
-PROMPT_FORMAT_VERSION = 1
-SENTINEL = "<<HEV_USER_TEXT>>"
-HEADER = "State:\n"
-QUESTIONS_HEADER = "\nQuestions:\n"
-FOOTER = "\nAnswer as JSON with one option letter per question.\n"
-MIN_STATE_TOKENS = 16
-
-
-@dataclass
-class Piece:
-    text: str
-    kind: str  # header | state | question | option | footer
-    qidx: int = -1
-    oidx: int = -1
-    alias_idx: int = -1
-
-
-@dataclass
-class TokenizedPrompt:
-    ids: list[int]
-    option_spans: list[list[tuple[int, int]]]
-    option_alias: list[list[int]]
-    truncated: bool
-
-
-def render_pieces(example: Example, option_order: dict[int, list[int]] | None = None) -> list[Piece]:
-    pieces = [Piece(HEADER, "header"), Piece(example.state.strip() + "\n", "state"),
-              Piece(QUESTIONS_HEADER, "header")]
-    for qi, q in enumerate(example.questions):
-        pieces.append(Piece(f"{q.qid} ({q.type}): {q.instructions}\n", "question", qi))
-        order = (option_order or {}).get(qi) or list(range(q.k))
-        for pos, oi in enumerate(order):
-            pieces.append(Piece(f"  {ALIASES[pos]}: {q.descriptions[oi]}\n", "option", qi, oi, pos))
-    pieces.append(Piece(FOOTER, "footer"))
-    return pieces
-
-
-def chat_affixes(tok) -> tuple[list[int], list[int]]:
-    text = tok.apply_chat_template([{"role": "user", "content": SENTINEL}],
-                                   tokenize=False, add_generation_prompt=True)
-    pre, post = text.split(SENTINEL)
-    return tok.encode(pre, add_special_tokens=False), tok.encode(post, add_special_tokens=False)
-
-
-def tokenize_prompt(example: Example, tok, max_tokens: int = 1024,
-                    option_order: dict[int, list[int]] | None = None) -> TokenizedPrompt:
-    pre, post = chat_affixes(tok)
-    pieces = render_pieces(example, option_order)
-    enc = [tok.encode(p.text, add_special_tokens=False) for p in pieces]
-    budget = max_tokens - len(pre) - len(post)
-    fixed = sum(len(e) for e, p in zip(enc, pieces) if p.kind != "state")
-    truncated = False
-    for i, p in enumerate(pieces):
-        if p.kind == "state" and fixed + len(enc[i]) > budget:
-            keep = max(budget - fixed, MIN_STATE_TOKENS)
-            half = keep // 2
-            enc[i] = enc[i][:half] + enc[i][-(keep - half):]
-            truncated = True
-    ids = list(pre)
-    spans = [[(0, 0)] * q.k for q in example.questions]
-    alias = [[0] * q.k for q in example.questions]
-    for p, e in zip(pieces, enc):
-        start = len(ids)
-        ids.extend(e)
-        if p.kind == "option":
-            spans[p.qidx][p.oidx] = (start, len(ids))
-            alias[p.qidx][p.oidx] = p.alias_idx
-    ids.extend(post)
-    if len(ids) > max_tokens:
-        raise ValueError(f"{example.id}: prompt has {len(ids)} tokens even after truncation (max {max_tokens})")
-    return TokenizedPrompt(ids=ids, option_spans=spans, option_alias=alias, truncated=truncated)
-
-
-def canvas_ids(example: Example, tok, slot_token_id: int, canvas_length: int,
-               eos_id: int = 1) -> tuple[list[int], list[int]]:
-    ids: list[int] = []
-    slots: list[int] = []
-    for qi, q in enumerate(example.questions):
-        prefix = '{"' if qi == 0 else ', "'
-        ids.extend(tok.encode(prefix + q.qid + '": "', add_special_tokens=False))
-        slots.append(len(ids))
-        ids.append(slot_token_id)
-        ids.extend(tok.encode('"', add_special_tokens=False))
-    ids.extend(tok.encode("}", add_special_tokens=False))
-    if len(ids) > canvas_length:
-        raise ValueError(f"{example.id}: canvas needs {len(ids)} tokens, canvas_length is {canvas_length}")
-    ids.extend([eos_id] * (canvas_length - len(ids)))
-    return ids, slots
-
-
-def alias_token_ids(tok, n: int) -> list[int]:
-    out = []
-    for a in ALIASES[:n]:
-        ids = tok.encode(a, add_special_tokens=False)
-        if len(ids) != 1:
-            raise ValueError(f"alias {a!r} is not a single token: {ids}")
-        out.append(ids[0])
-    return out
-
-
-def shuffled_order(example: Example, rng: random.Random) -> dict[int, list[int]]:
-    order = {}
-    for qi, q in enumerate(example.questions):
-        if q.type == "choice":
-            perm = list(range(q.k))
-            rng.shuffle(perm)
-            order[qi] = perm
-    return order
+"""Prompt and canvas construction. Token spans are exact because every piece is tokenised separately."""
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass
+
+from .schema import ALIASES, Example
+
+# Bump when anything below changes the token ids, option spans or canvas a trained head was fitted on; every
+# feature cache must then be re-extracted and the head retrained. tests/test_format_pin.py pins version 1.
+PROMPT_FORMAT_VERSION = 1
+SENTINEL = "<<HEV_USER_TEXT>>"
+HEADER = "State:\n"
+QUESTIONS_HEADER = "\nQuestions:\n"
+FOOTER = "\nAnswer as JSON with one option letter per question.\n"
+MIN_STATE_TOKENS = 16
+
+
+@dataclass
+class Piece:
+    text: str
+    kind: str  # header | state | question | option | footer
+    qidx: int = -1
+    oidx: int = -1
+    alias_idx: int = -1
+
+
+@dataclass
+class TokenizedPrompt:
+    ids: list[int]
+    option_spans: list[list[tuple[int, int]]]
+    option_alias: list[list[int]]
+    truncated: bool
+
+
+def render_pieces(example: Example, option_order: dict[int, list[int]] | None = None) -> list[Piece]:
+    pieces = [Piece(HEADER, "header"), Piece(example.state.strip() + "\n", "state"),
+              Piece(QUESTIONS_HEADER, "header")]
+    for qi, q in enumerate(example.questions):
+        pieces.append(Piece(f"{q.qid} ({q.type}): {q.instructions}\n", "question", qi))
+        order = (option_order or {}).get(qi) or list(range(q.k))
+        for pos, oi in enumerate(order):
+            pieces.append(Piece(f"  {ALIASES[pos]}: {q.descriptions[oi]}\n", "option", qi, oi, pos))
+    pieces.append(Piece(FOOTER, "footer"))
+    return pieces
+
+
+def chat_affixes(tok) -> tuple[list[int], list[int]]:
+    text = tok.apply_chat_template([{"role": "user", "content": SENTINEL}],
+                                   tokenize=False, add_generation_prompt=True)
+    pre, post = text.split(SENTINEL)
+    return tok.encode(pre, add_special_tokens=False), tok.encode(post, add_special_tokens=False)
+
+
+def tokenize_prompt(example: Example, tok, max_tokens: int = 1024,
+                    option_order: dict[int, list[int]] | None = None) -> TokenizedPrompt:
+    pre, post = chat_affixes(tok)
+    pieces = render_pieces(example, option_order)
+    enc = [tok.encode(p.text, add_special_tokens=False) for p in pieces]
+    budget = max_tokens - len(pre) - len(post)
+    fixed = sum(len(e) for e, p in zip(enc, pieces) if p.kind != "state")
+    truncated = False
+    for i, p in enumerate(pieces):
+        if p.kind == "state" and fixed + len(enc[i]) > budget:
+            keep = max(budget - fixed, MIN_STATE_TOKENS)
+            half = keep // 2
+            enc[i] = enc[i][:half] + enc[i][-(keep - half):]
+            truncated = True
+    ids = list(pre)
+    spans = [[(0, 0)] * q.k for q in example.questions]
+    alias = [[0] * q.k for q in example.questions]
+    for p, e in zip(pieces, enc):
+        start = len(ids)
+        ids.extend(e)
+        if p.kind == "option":
+            spans[p.qidx][p.oidx] = (start, len(ids))
+            alias[p.qidx][p.oidx] = p.alias_idx
+    ids.extend(post)
+    if len(ids) > max_tokens:
+        raise ValueError(f"{example.id}: prompt has {len(ids)} tokens even after truncation (max {max_tokens})")
+    return TokenizedPrompt(ids=ids, option_spans=spans, option_alias=alias, truncated=truncated)
+
+
+def canvas_ids(example: Example, tok, slot_token_id: int, canvas_length: int,
+               eos_id: int = 1) -> tuple[list[int], list[int]]:
+    ids: list[int] = []
+    slots: list[int] = []
+    for qi, q in enumerate(example.questions):
+        prefix = '{"' if qi == 0 else ', "'
+        ids.extend(tok.encode(prefix + q.qid + '": "', add_special_tokens=False))
+        slots.append(len(ids))
+        ids.append(slot_token_id)
+        ids.extend(tok.encode('"', add_special_tokens=False))
+    ids.extend(tok.encode("}", add_special_tokens=False))
+    if len(ids) > canvas_length:
+        raise ValueError(f"{example.id}: canvas needs {len(ids)} tokens, canvas_length is {canvas_length}")
+    ids.extend([eos_id] * (canvas_length - len(ids)))
+    return ids, slots
+
+
+def alias_token_ids(tok, n: int) -> list[int]:
+    out = []
+    for a in ALIASES[:n]:
+        ids = tok.encode(a, add_special_tokens=False)
+        if len(ids) != 1:
+            raise ValueError(f"alias {a!r} is not a single token: {ids}")
+        out.append(ids[0])
+    return out
+
+
+def shuffled_order(example: Example, rng: random.Random) -> dict[int, list[int]]:
+    order = {}
+    for qi, q in enumerate(example.questions):
+        if q.type == "choice":
+            perm = list(range(q.k))
+            rng.shuffle(perm)
+            order[qi] = perm
+    return order

schema.py

@@ -1,104 +1,104 @@
-"""Request schema shared by data loading, prompting, the head and the client."""
-from __future__ import annotations
-
-import math
-from dataclasses import asdict, dataclass, field
-from typing import Optional
-
-QUESTION_TYPES = ("noul", "choice", "score")
-TYPE_INDEX = {t: i for i, t in enumerate(QUESTION_TYPES)}
-MAX_CHOICE_OPTIONS = 64
-MAX_SCORE_LEVELS = 10
-NOUL_OPTIONS = ["true", "false"]
-NOUL_DEFAULT_DESCRIPTIONS = ["yes, the statement is true", "no, the statement is false"]
-ALIASES = (
-    [chr(c) for c in range(ord("A"), ord("Z") + 1)]
-    + [chr(c) for c in range(ord("a"), ord("z") + 1)]
-    + ["A" + chr(c) for c in range(ord("A"), ord("L") + 1)]
-)
-
-
-@dataclass
-class Question:
-    qid: str
-    type: str
-    instructions: str
-    options: list[str]
-    descriptions: list[str]
-
-    @classmethod
-    def from_jev(cls, qid: str, spec: dict) -> "Question":
-        qtype = spec.get("type")
-        if qtype not in QUESTION_TYPES:
-            raise ValueError(f"{qid}: unknown question type {qtype!r}")
-        instructions = str(spec.get("instructions", "")).strip()
-        criteria = spec.get("criteria")
-        if qtype == "noul":
-            options = list(NOUL_OPTIONS)
-            descriptions = list(NOUL_DEFAULT_DESCRIPTIONS)
-            if criteria is not None:
-                if not isinstance(criteria, dict) or set(criteria) - set(NOUL_OPTIONS):
-                    raise ValueError(f"{qid}: noul criteria must be a mapping with 'true' and 'false' keys")
-                descriptions = [str(criteria.get("true") or descriptions[0]),
-                                str(criteria.get("false") or descriptions[1])]
-        elif qtype == "choice":
-            if not isinstance(criteria, dict):
-                raise ValueError(f"{qid}: choice needs a criteria mapping")
-            options = [str(k) for k in criteria.keys()]
-            descriptions = [str(v) if v else str(k) for k, v in criteria.items()]
-            if not (2 <= len(options) <= MAX_CHOICE_OPTIONS):
-                raise ValueError(f"{qid}: choice needs 2..{MAX_CHOICE_OPTIONS} options, got {len(options)}")
-        else:
-            if not isinstance(criteria, list):
-                raise ValueError(f"{qid}: score needs a criteria list")
-            if not (2 <= len(criteria) <= MAX_SCORE_LEVELS):
-                raise ValueError(f"{qid}: score needs 2..{MAX_SCORE_LEVELS} levels, got {len(criteria)}")
-            options = [str(i) for i in range(len(criteria))]
-            descriptions = [str(c) for c in criteria]
-        if len(set(options)) != len(options):
-            raise ValueError(f"{qid}: duplicate options")
-        return cls(qid=qid, type=qtype, instructions=instructions,
-                   options=options, descriptions=descriptions)
-
-    @property
-    def k(self) -> int:
-        return len(self.options)
-
-
-def argmax(values) -> int:
-    """Index of the first maximum, so ties go to the lowest index. Shared by evaluation and the zero-shot read."""
-    return max(range(len(values)), key=lambda i: values[i])
-
-
-def align_distribution(probs: dict, options: list[str]) -> list[float]:
-    vals = [float(probs.get(o, 0.0)) for o in options]
-    if any(v < 0 or not math.isfinite(v) for v in vals):
-        raise ValueError("distribution values must be finite and non-negative")
-    total = sum(vals)
-    if total <= 0:
-        raise ValueError("distribution has no mass on the listed options")
-    return [v / total for v in vals]
-
-
-@dataclass
-class Example:
-    id: str
-    source: str
-    state: str
-    questions: list[Question]
-    gold: dict[str, Optional[list[float]]]
-    meta: dict = field(default_factory=dict)
-
-    def to_dict(self) -> dict:
-        return {
-            "id": self.id, "source": self.source, "state": self.state,
-            "questions": [asdict(q) for q in self.questions],
-            "gold": self.gold, "meta": self.meta,
-        }
-
-    @classmethod
-    def from_dict(cls, d: dict) -> "Example":
-        return cls(id=d["id"], source=d["source"], state=d["state"],
-                   questions=[Question(**q) for q in d["questions"]],
-                   gold={k: (list(v) if v is not None else None) for k, v in d["gold"].items()},
-                   meta=dict(d.get("meta", {})))
+"""Request schema shared by data loading, prompting, the head and the client."""
+from __future__ import annotations
+
+import math
+from dataclasses import asdict, dataclass, field
+from typing import Optional
+
+QUESTION_TYPES = ("noul", "choice", "score")
+TYPE_INDEX = {t: i for i, t in enumerate(QUESTION_TYPES)}
+MAX_CHOICE_OPTIONS = 64
+MAX_SCORE_LEVELS = 10
+NOUL_OPTIONS = ["true", "false"]
+NOUL_DEFAULT_DESCRIPTIONS = ["yes, the statement is true", "no, the statement is false"]
+ALIASES = (
+    [chr(c) for c in range(ord("A"), ord("Z") + 1)]
+    + [chr(c) for c in range(ord("a"), ord("z") + 1)]
+    + ["A" + chr(c) for c in range(ord("A"), ord("L") + 1)]
+)
+
+
+@dataclass
+class Question:
+    qid: str
+    type: str
+    instructions: str
+    options: list[str]
+    descriptions: list[str]
+
+    @classmethod
+    def from_jev(cls, qid: str, spec: dict) -> "Question":
+        qtype = spec.get("type")
+        if qtype not in QUESTION_TYPES:
+            raise ValueError(f"{qid}: unknown question type {qtype!r}")
+        instructions = str(spec.get("instructions", "")).strip()
+        criteria = spec.get("criteria")
+        if qtype == "noul":
+            options = list(NOUL_OPTIONS)
+            descriptions = list(NOUL_DEFAULT_DESCRIPTIONS)
+            if criteria is not None:
+                if not isinstance(criteria, dict) or set(criteria) - set(NOUL_OPTIONS):
+                    raise ValueError(f"{qid}: noul criteria must be a mapping with 'true' and 'false' keys")
+                descriptions = [str(criteria.get("true") or descriptions[0]),
+                                str(criteria.get("false") or descriptions[1])]
+        elif qtype == "choice":
+            if not isinstance(criteria, dict):
+                raise ValueError(f"{qid}: choice needs a criteria mapping")
+            options = [str(k) for k in criteria.keys()]
+            descriptions = [str(v) if v else str(k) for k, v in criteria.items()]
+            if not (2 <= len(options) <= MAX_CHOICE_OPTIONS):
+                raise ValueError(f"{qid}: choice needs 2..{MAX_CHOICE_OPTIONS} options, got {len(options)}")
+        else:
+            if not isinstance(criteria, list):
+                raise ValueError(f"{qid}: score needs a criteria list")
+            if not (2 <= len(criteria) <= MAX_SCORE_LEVELS):
+                raise ValueError(f"{qid}: score needs 2..{MAX_SCORE_LEVELS} levels, got {len(criteria)}")
+            options = [str(i) for i in range(len(criteria))]
+            descriptions = [str(c) for c in criteria]
+        if len(set(options)) != len(options):
+            raise ValueError(f"{qid}: duplicate options")
+        return cls(qid=qid, type=qtype, instructions=instructions,
+                   options=options, descriptions=descriptions)
+
+    @property
+    def k(self) -> int:
+        return len(self.options)
+
+
+def argmax(values) -> int:
+    """Index of the first maximum, so ties go to the lowest index. Shared by evaluation and the zero-shot read."""
+    return max(range(len(values)), key=lambda i: values[i])
+
+
+def align_distribution(probs: dict, options: list[str]) -> list[float]:
+    vals = [float(probs.get(o, 0.0)) for o in options]
+    if any(v < 0 or not math.isfinite(v) for v in vals):
+        raise ValueError("distribution values must be finite and non-negative")
+    total = sum(vals)
+    if total <= 0:
+        raise ValueError("distribution has no mass on the listed options")
+    return [v / total for v in vals]
+
+
+@dataclass
+class Example:
+    id: str
+    source: str
+    state: str
+    questions: list[Question]
+    gold: dict[str, Optional[list[float]]]
+    meta: dict = field(default_factory=dict)
+
+    def to_dict(self) -> dict:
+        return {
+            "id": self.id, "source": self.source, "state": self.state,
+            "questions": [asdict(q) for q in self.questions],
+            "gold": self.gold, "meta": self.meta,
+        }
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "Example":
+        return cls(id=d["id"], source=d["source"], state=d["state"],
+                   questions=[Question(**q) for q in d["questions"]],
+                   gold={k: (list(v) if v is not None else None) for k, v in d["gold"].items()},
+                   meta=dict(d.get("meta", {})))