Marimo Diffusion 0.6B: checkpoint, sampler, OpenAI server, ledger-needle bench

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer-qwen3-adaptive.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,155 @@
+---
+license: apache-2.0
+base_model: Qwen/Qwen3-0.6B
+pipeline_tag: text-generation
+language:
+  - en
+tags:
+  - diffusion
+  - block-diffusion
+  - masked-diffusion
+  - chat
+  - memory
+  - qwen3
+---
+
+# Marimo Diffusion 0.6B
+
+A 0.6B chat model that **thinks in denoised blocks and remembers through its own notes**
+instead of re-reading the conversation. It is a retrofit of
+[Qwen/Qwen3-0.6B](https://huggingface.co/Qwen/Qwen3-0.6B): continued pretraining converts the
+autoregressive base into a hybrid block-diffusion model, and a supervised fine-tune teaches it
+a note-taking chat format.
+
+Three things make it different from a standard chat model of this size:
+
+- **Adaptive thinking blocks.** Before answering, the model may open thinking blocks of 32/64/128
+  tokens, chosen per thought via control tokens (`<szN>`), each denoised bidirectionally in 16
+  steps. Trivial turns skip thinking entirely — that decision is trained, not prompted.
+- **Ledger memory.** Only the last 4 messages are kept verbatim in the prefix. Everything older
+  survives only as the model's own notes (`key: value`), merged so the latest value wins. A
+  100-turn conversation fits in a ~500-token prefix that never grows.
+- **Constant cost per turn.** ~2 s/turn on an RTX 3090 regardless of conversation length,
+  because the prefix is constant by construction.
+
+## ⚠️ This model does NOT run under standard runtimes
+
+The weights are Qwen3 architecture, but generation requires the block-denoising sampler included
+in this repository. **transformers `generate()`, llama.cpp, GGUF, Ollama and LM Studio will not
+produce correct output** — their autoregressive decoding never matches the training objective.
+Use the bundled code.
+
+## Quickstart
+
+```bash
+git clone https://huggingface.co/goldenfox/marimo-diffusion
+cd marimo-diffusion
+pip install torch tokenizers numpy
+
+# OpenAI-compatible server (any OpenAI-API chat client can connect)
+PYTHONPATH=src python -m diffusion_lm.chat_server \
+  --checkpoint marimo-diffusion-0.6b.pt \
+  --tokenizer tokenizer-qwen3-adaptive.json \
+  --port 7998
+```
+
+Then point any OpenAI-compatible client (Chatbox, Open WebUI, curl) at
+`http://127.0.0.1:7998/v1` with model id `marimo-diffusion-0.6b`:
+
+```bash
+curl http://127.0.0.1:7998/v1/chat/completions -H 'Content-Type: application/json' -d '{
+  "model": "marimo-diffusion-0.6b",
+  "messages": [{"role": "user", "content": "my sister lands friday 6pm, flight AR1420"}]
+}'
+```
+
+The response carries the model's notes in `reasoning_content` (same field DeepSeek uses), so
+clients that render reasoning show them automatically. The server caches each turn's notes and
+rebuilds the ledger across stateless requests.
+
+For the interactive playground (streaming denoise view, per-turn data log):
+
+```bash
+pip install gradio
+PYTHONPATH=src python -m diffusion_lm.reasoning_playground \
+  --outputs-dir . --prefix marimo \
+  --tokenizer tokenizer-qwen3-adaptive.json --port 7999
+```
+
+A CUDA GPU is recommended (any 6 GB+ card fits the bf16 weights). GPU memory: ~2.5 GB.
+
+## Benchmark: ledger needle (memory across 100 turns)
+
+The long-context needle test, adapted to what this architecture claims: a scripted 100-turn
+conversation plants 15 facts, corrects 5 of them, and probes recall at distances of 3–96 turns.
+Baselines get every reasonable advantage: full history in context, greedy decoding, an explicit
+memory instruction, and native thinking mode where it exists. Identical scoring for all systems.
+Full per-turn data and the interactive viewers are in [`bench/`](./bench).
+
+| system | params | recall | test total (s) | s/turn | max prefix (tok) |
+|---|---|---|---|---|---|
+| Qwen2.5-1.5B-Instruct · full history | 1.5B | **9/10** | 80 | 0.8 | 2,881 |
+| Qwen3-0.6B + thinking · full history | 0.6B | 7/10 | 1,320 | 13.2 | 4,450 |
+| **Marimo Diffusion (ledger)** | **0.6B** | **6/10** | **220** | **2.2** | **625** |
+| Qwen2.5-0.5B-Instruct · full history | 0.5B | 4/10 | 110 | 1.1 | 3,908 |
+| SmolLM2-360M-Instruct · full history | 0.36B | 4/10 | 60 | 0.6 | 2,712 |
+| TinyLlama-1.1B-Chat · full history | 1.1B | 3/10 | 120 | 1.2 | 5,126 |
+| Qwen3-0.6B no thinking · full history | 0.6B | 1/10 | 130 | 1.3 | 3,722 |
+| Qwen3-0.6B no thinking · 512-token budget | 0.6B | 1/10 | 160 | 1.6 | 512 |
+
+Reading this honestly:
+
+- **In its size class it leads**: every ≤0.6B baseline with the full transcript in context
+  scores 4/10 or less; the ledger reaches 6/10 from a 7× smaller prefix.
+- **Beating it costs something**: 2.5× the parameters (Qwen2.5-1.5B), or the same base model's
+  thinking mode at **6× the latency** with an unbounded prefix — and 7 vs 6 on ten probes is
+  within noise.
+- **Failure profiles are complementary.** The thinking baseline re-reads verbatim, so it never
+  suffers a corrupted note; the ledger never suffers long-context attention loss (it recalled
+  facts at distance 59–60 that the thinking baseline missed with the text in front of it).
+- **Caveats**: one seed, ten probes; and the comparison measures the mechanism *and* its
+  training together — this model was trained on this conversational register, the baselines
+  were not.
+
+## Training
+
+- Base: Qwen3-0.6B. Continued pretraining converts AR → hybrid block diffusion (answer region
+  stays autoregressive; thinking blocks are masked-denoised bidirectionally).
+- SFT: 174k examples from ~38k conversations — synthetic memory-task dialogues (recap,
+  correction, distant-combination, each ending in a consolidating close), passage-grounded QA
+  with the source dataset's reference answer as an exact quality gate, an abstention slice, and
+  persona-grounded dialogues. Sequence length 512 (median example: 194 tokens); the checkpoint
+  is served at 2,048 (RoPE, no learned positions).
+- `steps_per_block 16` is the measured optimum for this checkpoint: best numeric fidelity at
+  half the latency of 32; below 8 both prose and numbers degrade.
+
+### Training data provenance
+
+| source | role | license |
+|---|---|---|
+| synthetic dialogues (DeepSeek v4-flash generated, machine-audited) | chat + memory tasks | — |
+| [stanfordnlp/coqa](https://huggingface.co/datasets/stanfordnlp/coqa) | multi-turn grounded QA | other (mixed provenance) |
+| [rajpurkar/squad_v2](https://huggingface.co/datasets/rajpurkar/squad_v2) | abstention | cc-by-sa-4.0 |
+| [dgslibisey/MuSiQue](https://huggingface.co/datasets/dgslibisey/MuSiQue) | multi-hop reasoning | undeclared on mirror |
+| [ucinlp/drop](https://huggingface.co/datasets/ucinlp/drop) | arithmetic over passages | cc-by-sa-4.0 |
+| [nayohan/multi_session_chat](https://huggingface.co/datasets/nayohan/multi_session_chat) | human-written persona facts | undeclared on mirror |
+
+## Limitations
+
+- **0.6B knowledge ceiling.** It confabulates on open-domain facts like any model this size;
+  the training includes an abstention slice ("the passage doesn't say") but it is not a fix.
+- **Note-taking can corrupt compound values** (an alphanumeric like `harbor858` was once noted
+  as `8858` and then faithfully recalled wrong). What enters the ledger wrong stays wrong.
+- **Ledger interference**: with 50+ accumulated entries, similar-typed values (several money
+  amounts) can cross-contaminate. Training saw ~24 entries max.
+- **No code in training data.** Reasoning about pasted code runs on the base model's residual
+  ability.
+- SFT never saw examples past 512 tokens; behaviour between 512 and 2,048 rides on the
+  continued pretraining.
+- English only.
+
+## License
+
+Apache 2.0, inheriting the Qwen3-0.6B base license. Training data licenses are listed above;
+CoQA carries mixed-provenance terms and two mirrors declare no license — review them if you
+redistribute derived data.

bench/bench_ledger_needle.py

@@ -0,0 +1,602 @@
+"""Ledger needle bench: fact recall across a conversation far longer than the context.
+
+The long-context analogue for this architecture. A needle-in-a-haystack test measures whether a
+fact survives distance inside a huge prompt; here the prompt never exceeds the training window
+(last ``keep`` messages) and facts survive only through the ledger of the model's own notes, so
+the same question — does distance kill recall? — is asked of the memory mechanism instead of the
+attention span.
+
+``run`` drives a deterministic scripted conversation through the engine on the GPU box and
+records everything per turn. ``render`` turns that record into a self-contained HTML viewer:
+timeline, per-turn notes, the exact ledger state, and every probe scored.
+"""
+
+from __future__ import annotations
+
+import argparse
+import html as html_lib
+import json
+import random
+import re
+import time
+from pathlib import Path
+
+FACT_POOL = [
+    ('sister flight', 'flight {code}', 'my sister lands on {value}, write that down'),
+    ('hotel room', 'room {num3}', 'we got {value} at the hotel'),
+    ('locker code', 'code {num4}', 'the gym locker is {value}'),
+    ('dentist slot', '{clock}', 'dentist moved me to {value}'),
+    ('plumber quote', '${money}', 'the plumber quoted {value} for the bathroom'),
+    ('car plate', '{plate}', 'the rental has plate {value}'),
+    ('wifi password', '{word}{num3}', 'cabin wifi password is {value}'),
+    ('train platform', 'platform {num2}', 'our train leaves from {value}'),
+    ('order number', 'order {code}', 'the couch is {value}, keep it handy'),
+    ('rent due', 'day {num2}', 'landlord wants rent by {value} each month'),
+    ('kid teacher', 'ms. {name}', "tomas' new teacher is {value}"),
+    ('parking spot', 'spot {letter}{num2}', 'we always park at {value}'),
+    ('meds dose', '{num2}mg', 'doctor changed the dose to {value}'),
+    ('deposit', '${money}', 'the deposit for the venue was {value}'),
+    ('boarding gate', 'gate {letter}{num2}', 'boarding is at {value}'),
+]
+DISTRACTORS = [
+    'what a week, honestly', 'did you catch the game last night?', 'i love this weather lately',
+    'work was chaos today', 'thinking of making pasta tonight', 'my back hurts from the gym',
+    'the neighbor is renovating again, so loud', 'saw a great documentary yesterday',
+    'coffee here is getting expensive', 'might go for a walk later', 'the cat knocked over a plant',
+    'traffic was unreal this morning', 'finally finished that book', 'craving something sweet',
+    'they repaved our street', 'my phone battery dies so fast now',
+]
+PROBES = ['wait, what was the {key} again?', 'remind me of the {key}?',
+          'i forgot the {key}, what was it?', 'quick — the {key}?']
+WORDS = ('maple', 'harbor', 'cactus', 'violet', 'ember', 'quartz')
+NAMES = ('ferro', 'silva', 'duarte', 'campos', 'rojas', 'ibanez')
+
+
+def _value(template: str, rng: random.Random) -> str:
+    return (template
+            .replace('{code}', f'{rng.choice("ABKQZ")}{rng.choice("RLMT")}{rng.randint(1000, 9999)}')
+            .replace('{num2}', str(rng.randint(10, 99)))
+            .replace('{num3}', str(rng.randint(100, 999)))
+            .replace('{num4}', str(rng.randint(1000, 9999)))
+            .replace('{clock}', f'{rng.randint(1, 12)}:{rng.choice(("15", "30", "45"))}pm')
+            .replace('{money}', f'{rng.randint(2, 90) * 100:,}')
+            .replace('{plate}', f'{"".join(rng.choice("BCDFGHJK") for _ in range(3))}-{rng.randint(100, 999)}')
+            .replace('{word}', rng.choice(WORDS))
+            .replace('{letter}', rng.choice('ABCDE'))
+            .replace('{name}', rng.choice(NAMES)))
+
+
+def build_script(turns: int, seed: int) -> list[dict]:
+    """Deterministic conversation plan: facts early and throughout, probes at all distances."""
+
+    rng = random.Random(seed)
+    facts = []
+    for key, template, phrasing in rng.sample(FACT_POOL, len(FACT_POOL)):
+        facts.append({'key': key, 'value': _value(template, rng), 'phrasing': phrasing})
+
+    plan: list[dict] = []
+    stated: dict[str, dict] = {}
+    fact_iter = iter(facts)
+    for index in range(turns):
+        remaining = turns - index
+        can_probe = [f for f in stated.values() if index - f['turn'] >= 3]
+        if index >= turns - 3 and can_probe:
+            kind = 'probe'
+        elif index % 7 in (0, 3) and (fact := next(fact_iter, None)) is not None:
+            plan.append({'kind': 'fact', **fact})
+            stated[fact['key']] = {**fact, 'turn': index}
+            continue
+        elif index % 11 == 5 and can_probe:
+            kind = 'probe'
+        elif index % 13 == 8 and stated and remaining > 5:
+            kind = 'correction'
+        else:
+            kind = 'distractor'
+
+        if kind == 'probe':
+            target = rng.choice(can_probe)
+            plan.append({'kind': 'probe', 'key': target['key'], 'value': target['value'],
+                         'distance': index - target['turn'],
+                         'text': rng.choice(PROBES).format(key=target['key'])})
+        elif kind == 'correction':
+            target = rng.choice(list(stated.values()))
+            new = _value(next(t for k, t, _ in FACT_POOL if k == target['key']), rng)
+            plan.append({'kind': 'correction', 'key': target['key'], 'old': target['value'],
+                         'value': new,
+                         'text': f'actually scratch that, the {target["key"]} is {new} now'})
+            stated[target['key']] = {**target, 'value': new, 'turn': index}
+        else:
+            plan.append({'kind': 'distractor', 'text': rng.choice(DISTRACTORS)})
+    return plan
+
+
+def _hit(value: str, answer: str) -> bool:
+    """Whether the answer states the value, matched on its distinctive core.
+
+    Demanding the full phrase penalised correct answers on both systems ("the locker code is
+    4805" failed against 'code 4805'), so the core — the last token, which carries the
+    identifier — is what must appear. Values are generated with 2+ digit cores, so incidental
+    collisions stay unlikely.
+    """
+
+    canon = lambda t: re.sub(r'[^a-z0-9]', '', t.lower())  # noqa: E731
+    core = value.split()[-1]
+    return canon(core) in canon(answer)
+
+
+def run(args: argparse.Namespace) -> None:
+    from diffusion_lm.claims import SYSTEM, chat_prefix, ledger_line, ledger_notes, merge_notes
+    from diffusion_lm.reasoning_playground import ReasoningEngine
+    from diffusion_lm.train import resolve_device
+
+    engine = ReasoningEngine(args.checkpoint, args.tokenizer, resolve_device('auto'))
+    plan = build_script(args.turns, args.seed)
+    messages: list[dict] = []
+    records = []
+    # Appended per turn and flushed, so a tail -f (or a crash) sees every completed turn.
+    sink = args.output.open('w', encoding='utf-8')
+    for index, turn in enumerate(plan):
+        user_text = turn.get('text') or turn['phrasing'].format(value=turn['value'])
+        messages.append({'role': 'user', 'content': user_text})
+        older = merge_notes(ledger_notes(messages, args.keep))
+        window = messages[max(0, len(messages) - args.keep):]
+        prefix = chat_prefix(
+            [{'role': m['role'], 'content': m['content']} for m in window],
+            system=SYSTEM, extra=ledger_line(older),
+        )
+        blocks: list[tuple[int, str]] = []
+        answer = ''
+        begin = time.perf_counter()
+        for _, answer, _ in engine.stream_chat(
+            prefix, temperature=0.7, steps_per_block=args.steps,
+            max_answer_tokens=120, seed=args.seed * 1000 + index, blocks_out=blocks,
+        ):
+            pass
+        note = '; '.join(text for _, text in blocks if text)
+        messages.append({'role': 'assistant', 'content': answer.strip(), 'note': note})
+        record = {
+            'turn': index, 'kind': turn['kind'], 'user': user_text,
+            'notes': [text for _, text in blocks if text], 'ledger': older,
+            'prefix_tokens': len(engine.tokenizer.encode(prefix, add_special_tokens=False).ids),
+            'answer': answer.strip(), 'seconds': round(time.perf_counter() - begin, 2),
+        }
+        if turn['kind'] == 'probe':
+            record.update({'key': turn['key'], 'expected': turn['value'],
+                           'distance': turn['distance'],
+                           'hit': _hit(turn['value'], answer)})
+        if turn['kind'] in ('fact', 'correction'):
+            record.update({'key': turn['key'], 'value': turn['value']})
+        records.append(record)
+        sink.write(json.dumps(record, ensure_ascii=False) + '\n')
+        sink.flush()
+        flag = '' if turn['kind'] != 'probe' else (' HIT' if record['hit'] else ' MISS')
+        print(f'[{index + 1}/{len(plan)}] {turn["kind"]}{flag} {record["seconds"]}s '
+              f'ledger={len(older)}', flush=True)
+        print(f'  U: {user_text}', flush=True)
+        if note:
+            print(f'  T: {note}', flush=True)
+        print(f'  A: {answer.strip()}', flush=True)
+        if turn['kind'] == 'probe':
+            print(f'  >> esperado {turn["value"]} (distancia {turn["distance"]}) -> '
+                  f'{"HIT" if record["hit"] else "MISS"}', flush=True)
+
+    sink.close()
+    probes = [r for r in records if r['kind'] == 'probe']
+    hits = sum(r['hit'] for r in probes)
+    print(f'\nrecall: {hits}/{len(probes)} ({hits / max(1, len(probes)):.0%})')
+
+
+def run_ar(args: argparse.Namespace) -> None:
+    """Same scripted conversation through a plain AR baseline of the same size.
+
+    ``--context full`` gives the baseline the whole history (its native memory);
+    ``--context budget`` truncates the rendered prompt to the newest messages that fit the
+    same token budget the ledger system runs under, so both systems answer from equally many
+    prefix tokens and only the memory MECHANISM differs.
+    """
+
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model_path, torch_dtype=torch.bfloat16, device_map='cuda',
+    ).eval()
+    plan = build_script(args.turns, args.seed)
+    # Deliberately generous: the baselines get an explicit memory instruction, greedy decoding,
+    # and (with --thinking) their native reasoning mode. The comparison should show the ledger
+    # beating the baselines at their best, or it shows nothing.
+    system = ('You are a meticulous assistant. Remember every concrete detail the user '
+              'mentions — codes, amounts, times, names, numbers. When asked to recall one, '
+              'answer with the exact value.')
+    messages: list[dict] = []
+    records = []
+    try:
+        tokenizer.apply_chat_template(
+            [{'role': 'system', 'content': 'x'}, {'role': 'user', 'content': 'y'}],
+            tokenize=True, add_generation_prompt=True)
+        args._no_system = False
+    except Exception:
+        args._no_system = True
+    sink = args.output.open('w', encoding='utf-8')
+    for index, turn in enumerate(plan):
+        user_text = turn.get('text') or turn['phrasing'].format(value=turn['value'])
+        messages.append({'role': 'user', 'content': user_text})
+
+        def _assemble(tail: list[dict]) -> list[dict]:
+            if getattr(args, '_no_system', False):
+                head = dict(tail[0]) if tail else {'role': 'user', 'content': ''}
+                head['content'] = f'{system}\n\n{head["content"]}'
+                return [head] + [dict(m) for m in tail[1:]]
+            return [{'role': 'system', 'content': system}] + [dict(m) for m in tail]
+
+        chat = _assemble(messages)
+        if args.context == 'budget':
+            # Newest-first packing under the same prefix budget the ledger system uses.
+            kept: list[dict] = []
+            for message in reversed(messages):
+                candidate = _assemble([message, *kept])
+                rendered = tokenizer.apply_chat_template(
+                    candidate, tokenize=True, add_generation_prompt=True,
+                    enable_thinking=False,
+                )
+                # BatchEncoding in newer transformers: len() counts KEYS, which silently
+                # disabled the cap; count tokens explicitly.
+                ids = rendered if isinstance(rendered, list) else rendered['input_ids']
+                if len(ids) > args.budget:
+                    break
+                kept = [message, *kept]
+            chat = _assemble(kept)
+        encoded = tokenizer.apply_chat_template(
+            chat, tokenize=True, add_generation_prompt=True, return_tensors='pt',
+            enable_thinking=args.thinking,
+        )
+        # Newer transformers return a BatchEncoding here rather than the bare tensor.
+        prompt_ids = (encoded if torch.is_tensor(encoded) else encoded['input_ids']).to(
+            model.device)
+        torch.manual_seed(args.seed * 1000 + index)
+        begin = time.perf_counter()
+        with torch.no_grad():
+            out = model.generate(
+                prompt_ids, max_new_tokens=640 if args.thinking else 120, do_sample=False,
+                pad_token_id=tokenizer.eos_token_id or tokenizer.pad_token_id,
+            )
+        answer = tokenizer.decode(out[0][prompt_ids.shape[1]:], skip_special_tokens=True)
+        answer = re.sub(r'(?s)<think>.*?(</think>|$)', '', answer).strip()
+        messages.append({'role': 'assistant', 'content': answer})
+        record = {
+            'turn': index, 'kind': turn['kind'], 'user': user_text, 'notes': [], 'ledger': [],
+            'prefix_tokens': int(prompt_ids.shape[1]), 'answer': answer,
+            'seconds': round(time.perf_counter() - begin, 2),
+        }
+        if turn['kind'] == 'probe':
+            record.update({'key': turn['key'], 'expected': turn['value'],
+                           'distance': turn['distance'], 'hit': _hit(turn['value'], answer)})
+        if turn['kind'] in ('fact', 'correction'):
+            record.update({'key': turn['key'], 'value': turn['value']})
+        records.append(record)
+        sink.write(json.dumps(record, ensure_ascii=False) + '\n')
+        sink.flush()
+        flag = '' if turn['kind'] != 'probe' else (' HIT' if record['hit'] else ' MISS')
+        print(f'[{index + 1}/{len(plan)}] {turn["kind"]}{flag} {record["seconds"]}s '
+              f'prefix={record["prefix_tokens"]}', flush=True)
+    sink.close()
+    probes = [r for r in records if r['kind'] == 'probe']
+    hits = sum(r['hit'] for r in probes)
+    print(f'\nrecall: {hits}/{len(probes)} ({hits / max(1, len(probes)):.0%})')
+
+
+TEMPLATE = '''<!doctype html><html><head><meta charset="utf-8">
+<title>Ledger needle — __MODEL__</title><style>
+:root { --bg:#11141a; --panel:#1a1f29; --text:#dde3ee; --dim:#8a93a6; --hit:#3fb96b;
+  --miss:#e05555; --fact:#4d8fd1; --corr:#d9a23c; --line:#2a3140; }
+* { box-sizing:border-box; margin:0 } body { background:var(--bg); color:var(--text);
+  font:14px/1.5 ui-monospace,Menlo,monospace; padding:24px; }
+h1 { font-size:18px; margin-bottom:4px } .sub { color:var(--dim); margin-bottom:18px }
+.metrics { display:flex; gap:14px; flex-wrap:wrap; margin-bottom:18px }
+.metric { background:var(--panel); border:1px solid var(--line); border-radius:8px;
+  padding:10px 16px } .metric b { font-size:20px; display:block }
+.dist { margin-bottom:18px } .dist .row { display:flex; align-items:center; gap:8px;
+  margin:3px 0 } .dist .bar { height:14px; background:var(--hit); border-radius:3px }
+.dist .bar.m { background:var(--miss) } .dist span { color:var(--dim); font-size:12px;
+  min-width:110px }
+.wrap { display:grid; grid-template-columns:340px 1fr; gap:16px; align-items:start }
+.timeline { max-height:75vh; overflow-y:auto; background:var(--panel);
+  border:1px solid var(--line); border-radius:8px }
+.t { padding:7px 10px; border-bottom:1px solid var(--line); cursor:pointer;
+  display:flex; gap:8px; align-items:center } .t:hover,.t.sel { background:#242b38 }
+.t .n { color:var(--dim); min-width:30px } .dot { width:9px; height:9px; border-radius:50%;
+  flex:none } .dot.distractor { background:var(--dim) } .dot.fact { background:var(--fact) }
+.dot.correction { background:var(--corr) } .dot.probe.hit { background:var(--hit) }
+.dot.probe.miss { background:var(--miss) } .t .txt { white-space:nowrap; overflow:hidden;
+  text-overflow:ellipsis; font-size:12px }
+.detail { background:var(--panel); border:1px solid var(--line); border-radius:8px;
+  padding:16px; max-height:75vh; overflow-y:auto }
+.detail h3 { font-size:13px; color:var(--dim); margin:14px 0 6px; text-transform:uppercase }
+.detail h3:first-child { margin-top:0 }
+.bubble { background:#242b38; border-radius:8px; padding:10px 12px; margin:4px 0 }
+.note { color:#9fc4ea } .ledger-entry { display:inline-block; background:#242b38;
+  border:1px solid var(--line); border-radius:5px; padding:2px 8px; margin:2px;
+  font-size:12px } .verdict { padding:10px 12px; border-radius:8px; margin-top:6px }
+.verdict.hit { background:#173523; border:1px solid var(--hit) }
+.verdict.miss { background:#3a1d1d; border:1px solid var(--miss) }
+.meta { color:var(--dim); font-size:12px; margin-top:12px }
+</style></head><body>
+<h1>Ledger needle</h1>
+<div class="sub">__MODEL__ · __TURNS__ turnos · ventana de __KEEP__ mensajes ·
+steps_per_block __STEPS__ · el prefijo nunca excede __MAXPREFIX__ tokens</div>
+<div class="metrics" id="metrics"></div>
+<div class="dist" id="dist"></div>
+<div class="wrap"><div class="timeline" id="timeline"></div>
+<div class="detail" id="detail">elegí un turno</div></div>
+<script>
+const DATA = __DATA__;
+const probes = DATA.filter(r => r.kind === 'probe');
+const hits = probes.filter(r => r.hit).length;
+const corr = DATA.filter(r => r.kind === 'correction').length;
+const maxTok = Math.max(...DATA.map(r => r.prefix_tokens));
+const secs = DATA.reduce((a, r) => a + r.seconds, 0) / DATA.length;
+document.getElementById('metrics').innerHTML = [
+  ['recall', hits + '/' + probes.length + ' (' + Math.round(100 * hits / probes.length) + '%)'],
+  ['hechos plantados', DATA.filter(r => r.kind === 'fact').length],
+  ['correcciones', corr], ['prefijo máx', maxTok + ' tok'],
+  ['media', secs.toFixed(1) + ' s/turno'],
+].map(([k, v]) => '<div class="metric"><b>' + v + '</b>' + k + '</div>').join('');
+const buckets = [[3, 10], [11, 25], [26, 50], [51, 999]];
+document.getElementById('dist').innerHTML = '<span style="color:var(--dim)">recall por distancia (turnos desde que el hecho se dijo):</span>' +
+  buckets.map(([lo, hi]) => {
+    const set = probes.filter(p => p.distance >= lo && p.distance <= hi);
+    if (!set.length) return '';
+    const h = set.filter(p => p.hit).length;
+    return '<div class="row"><span>' + lo + '–' + (hi > 100 ? '∞' : hi) + ' (' + h + '/' + set.length +
+      ')</span><div class="bar" style="width:' + (300 * h / set.length) + 'px"></div>' +
+      '<div class="bar m" style="width:' + (300 * (set.length - h) / set.length) + 'px"></div></div>';
+  }).join('');
+const tl = document.getElementById('timeline');
+DATA.forEach(r => {
+  const div = document.createElement('div');
+  div.className = 't'; div.dataset.turn = r.turn;
+  const cls = r.kind + (r.kind === 'probe' ? (r.hit ? ' hit' : ' miss') : '');
+  div.innerHTML = '<span class="n">' + (r.turn + 1) + '</span><span class="dot ' + cls +
+    '"></span><span class="txt">' + r.user.replace(/</g, '&lt;') + '</span>';
+  div.onclick = () => show(r, div); tl.appendChild(div);
+});
+function esc(t) { return String(t).replace(/</g, '&lt;'); }
+function show(r, el) {
+  document.querySelectorAll('.t.sel').forEach(x => x.classList.remove('sel'));
+  el.classList.add('sel');
+  let h = '<h3>usuario (turno ' + (r.turn + 1) + ' · ' + r.kind + ')</h3><div class="bubble">' +
+    esc(r.user) + '</div>';
+  h += '<h3>ledger que vio el modelo (' + r.ledger.length + ' entradas, historia fuera de ventana)</h3>' +
+    (r.ledger.length ? r.ledger.map(e => '<span class="ledger-entry">' + esc(e) + '</span>').join('')
+      : '<span style="color:var(--dim)">vacío</span>');
+  h += '<h3>notas de este turno</h3>' + (r.notes.length
+    ? r.notes.map(n => '<div class="bubble note">' + esc(n) + '</div>').join('')
+    : '<span style="color:var(--dim)">sin pensamiento</span>');
+  h += '<h3>respuesta</h3><div class="bubble">' + esc(r.answer) + '</div>';
+  if (r.kind === 'probe') h += '<div class="verdict ' + (r.hit ? 'hit' : 'miss') + '">esperado: <b>' +
+    esc(r.expected) + '</b> · distancia ' + r.distance + ' turnos · ' + (r.hit ? 'RECUPERADO' : 'PERDIDO') + '</div>';
+  if (r.kind === 'correction') h += '<div class="verdict" style="border:1px solid var(--corr)">corrige ' +
+    esc(r.key) + ': ahora <b>' + esc(r.value) + '</b></div>';
+  h += '<div class="meta">prefijo ' + r.prefix_tokens + ' tokens · ' + r.seconds + ' s</div>';
+  document.getElementById('detail').innerHTML = h;
+}
+</script></body></html>'''
+
+
+def _rescore(records: list[dict]) -> list[dict]:
+    """Verdicts are re-derived from the CURRENT rule, never trusted from the run.
+
+    The stored ``hit`` reflects whatever rule was live at run time; a scoring fix must reach
+    old runs, or the viewers keep showing correct answers marked as losses.
+    """
+
+    for record in records:
+        if 'expected' in record:
+            record['hit'] = _hit(record['expected'], record['answer'])
+    return records
+
+
+def render(args: argparse.Namespace) -> None:
+    records = _rescore(
+        [json.loads(line) for line in args.input.open(encoding='utf-8') if line.strip()])
+    page = (TEMPLATE
+            .replace('__DATA__', json.dumps(records, ensure_ascii=False))
+            .replace('__MODEL__', html_lib.escape(args.model_name))
+            .replace('__TURNS__', str(len(records)))
+            .replace('__KEEP__', str(args.keep))
+            .replace('__STEPS__', str(args.steps))
+            .replace('__MAXPREFIX__', str(max(r['prefix_tokens'] for r in records))))
+    args.output.write_text(page, encoding='utf-8')
+    print(f'viewer -> {args.output}')
+
+
+COMPARE_TEMPLATE = '''<!doctype html><html><head><meta charset="utf-8">
+<title>Ledger needle — comparación</title><style>
+:root { --bg:#11141a; --panel:#1a1f29; --text:#dde3ee; --dim:#8a93a6; --hit:#3fb96b;
+  --miss:#e05555; --line:#2a3140; --us:#16233a; }
+* { box-sizing:border-box; margin:0 } body { background:var(--bg); color:var(--text);
+  font:14px/1.5 ui-monospace,Menlo,monospace; padding:24px; max-width:1240px; margin:0 auto }
+h1 { font-size:18px } .sub { color:var(--dim); margin:4px 0 20px }
+table { border-collapse:collapse; width:100%; margin-bottom:26px }
+th, td { border:1px solid var(--line); padding:7px 10px; text-align:left; font-size:13px }
+th { background:var(--panel); color:var(--dim); font-weight:normal; cursor:pointer;
+  user-select:none; white-space:nowrap }
+th.sorted { color:var(--text) } th.sorted::after { content:' ↓' }
+th.sorted.asc::after { content:' ↑' }
+tr.us td { background:var(--us) }
+td.hitcell { background:#173523; color:var(--hit); text-align:center; font-weight:bold }
+td.misscell { background:#3a1d1d; color:var(--miss); text-align:center }
+#matrix { width:auto } #matrix th, #matrix td { padding:5px 8px }
+#matrix td.hitcell, #matrix td.misscell { width:36px; min-width:36px; padding:5px 0 }
+#matrix th.sys { writing-mode:vertical-rl; transform:rotate(180deg); text-align:left;
+  vertical-align:bottom; max-height:190px; font-size:12px; cursor:default; padding:8px 4px }
+#matrix th.sys.usc { color:#9fc4ea }
+#matrix td.exp { white-space:nowrap }
+.bar { display:inline-block; height:11px; background:var(--hit); border-radius:2px;
+  vertical-align:middle } .bar.m { background:var(--miss) }
+h2 { font-size:14px; color:var(--dim); margin:22px 0 10px; text-transform:uppercase }
+.caveats { background:var(--panel); border:1px solid var(--line); border-radius:8px;
+  padding:14px 18px; color:var(--dim); font-size:13px } .caveats li { margin:6px 0 0 18px }
+.hint { color:var(--dim); font-size:12px; margin:-18px 0 10px }
+</style></head><body>
+<h1>Ledger needle — comparación hasta 1.5B</h1>
+<div class="sub">100 turnos guionados (seed 21) · 15 hechos, 5 correcciones, 10 sondas a
+distancia 3–96 · scoring por núcleo del valor, idéntico para todos · baselines con historia
+completa, decoding greedy e instrucción explícita de memoria</div>
+<table id="summary"><thead><tr>
+<th data-k="label" data-t="s">sistema</th>
+<th data-k="params" data-t="n">tamaño (B)</th>
+<th data-k="hits" data-t="n">recall</th>
+<th data-k="d1" data-t="n">d 3–10</th>
+<th data-k="d2" data-t="n">d 11–50</th>
+<th data-k="d3" data-t="n">d 51+</th>
+<th data-k="total" data-t="n">test total (s)</th>
+<th data-k="per" data-t="n">s/turno</th>
+<th data-k="prefix" data-t="n">prefijo máx</th>
+</tr></thead><tbody></tbody></table>
+<div class="hint">click en una cabecera reordena; por defecto: más aciertos primero</div>
+<h2>Sonda por sonda</h2>
+__MATRIX__
+<h2>Advertencias de lectura</h2>
+<div class="caveats"><ul>
+<li>Un solo seed y 10 sondas por sistema: las diferencias chicas son ruido; las grandes
+sobreviven a ese margen.</li>
+<li>La comparación mide el paquete mecanismo+entrenamiento: el sistema de ledger fue entrenado
+en este registro conversacional y los baselines no. Muestra que el paquete funciona, no que el
+mecanismo solo cause la diferencia.</li>
+<li>Los baselines reciben la mayor ventaja razonable: historia completa en contexto, greedy,
+instrucción de memoria, y el modo thinking nativo donde existe. El sistema de ledger corre con
+sampling a temperatura 0.7 y un prefijo constante de ~500 tokens.</li>
+<li>El guion es sintético y sus valores son arbitrarios por diseño; un valor sin sentido
+("day 97") puede penalizar a un modelo que se resiste a repetir datos absurdos.</li>
+</ul></div>
+<script>
+const ROWS = __ROWS__;
+const tbody = document.querySelector('#summary tbody');
+let sortKey = 'hits', asc = false;
+function fmt(r) {
+  return '<td><b>' + r.label + '</b></td><td>' + r.params.toFixed(2) + '</td>' +
+    '<td><b>' + r.hits + '/' + r.probes + '</b> <span class="bar" style="width:' + 6 * r.hits +
+    'px"></span><span class="bar m" style="width:' + 6 * (r.probes - r.hits) + 'px"></span></td>' +
+    ['d1', 'd2', 'd3'].map(k => '<td>' + r[k] + '/' + r[k + 'n'] + '</td>').join('') +
+    '<td>' + r.total.toFixed(0) + '</td><td>' + r.per.toFixed(1) + '</td><td>' + r.prefix + '</td>';
+}
+function draw() {
+  const rows = [...ROWS].sort((a, b) => {
+    const va = a[sortKey], vb = b[sortKey];
+    const c = typeof va === 'string' ? va.localeCompare(vb) : va - vb;
+    return asc ? c : -c;
+  });
+  tbody.innerHTML = rows.map(r =>
+    '<tr' + (r.us ? ' class="us"' : '') + '>' + fmt(r) + '</tr>').join('');
+  document.querySelectorAll('#summary th').forEach(th => {
+    th.classList.toggle('sorted', th.dataset.k === sortKey);
+    th.classList.toggle('asc', th.dataset.k === sortKey && asc);
+  });
+}
+document.querySelectorAll('#summary th').forEach(th => th.onclick = () => {
+  if (sortKey === th.dataset.k) asc = !asc;
+  else { sortKey = th.dataset.k; asc = th.dataset.t === 's'; }
+  draw();
+});
+draw();
+</script></body></html>'''
+
+
+def compare(args: argparse.Namespace) -> None:
+    systems = []
+    for spec in args.inputs:
+        label, _, rest = spec.partition('=')
+        path, _, params = rest.partition('=')
+        records = _rescore(
+            [json.loads(line) for line in Path(path).open(encoding='utf-8') if line.strip()])
+        probes = [r for r in records if r['kind'] == 'probe']
+        row = {'label': html_lib.escape(label), 'params': float(params or 0),
+               'probes': len(probes), 'hits': sum(r['hit'] for r in probes),
+               'total': sum(r['seconds'] for r in records),
+               'per': sum(r['seconds'] for r in records) / max(1, len(records)),
+               'prefix': max(r['prefix_tokens'] for r in records),
+               'us': bool(args.ours and label == args.ours)}
+        for name, (lo, hi) in (('d1', (3, 10)), ('d2', (11, 50)), ('d3', (51, 999))):
+            subset = [r for r in probes if lo <= r['distance'] <= hi]
+            row[name] = sum(r['hit'] for r in subset)
+            row[name + 'n'] = len(subset)
+        systems.append((label, records, row))
+
+    first_probes = [r for r in systems[0][1] if r['kind'] == 'probe']
+    head = ['<tr><th>turno</th><th>dist</th><th>esperado</th>']
+    head += [
+        f'<th class="sys{" usc" if args.ours and label == args.ours else ""}">'
+        f'{html_lib.escape(label)}</th>'
+        for label, _, _ in systems
+    ]
+    rows_html = []
+    for probe in first_probes:
+        row = [f'<td>{probe["turn"] + 1}</td><td>{probe["distance"]}</td>'
+               f'<td class="exp">{html_lib.escape(probe["expected"])}</td>']
+        for _, records, _ in systems:
+            match = next((r for r in records if r['kind'] == 'probe'
+                          and r['turn'] == probe['turn']), None)
+            if match is None:
+                row.append('<td>—</td>')
+            else:
+                row.append('<td class="hitcell">✓</td>' if match['hit']
+                           else '<td class="misscell">✗</td>')
+        rows_html.append(f'<tr>{"".join(row)}</tr>')
+    matrix = f'<table id="matrix">{"".join(head)}</tr>{"".join(rows_html)}</table>'
+
+    page = (COMPARE_TEMPLATE
+            .replace('__ROWS__', json.dumps([row for _, _, row in systems], ensure_ascii=False))
+            .replace('__MATRIX__', matrix))
+    args.output.write_text(page, encoding='utf-8')
+    print(f'comparison -> {args.output}')
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    sub = parser.add_subparsers(dest='command', required=True)
+    runner = sub.add_parser('run', help='drive the scripted conversation through the engine')
+    runner.add_argument('--checkpoint', type=Path,
+                        default=Path('outputs/qwen06b-genchat-sft/inference-latest.pt'))
+    runner.add_argument('--tokenizer', type=Path,
+                        default=Path('artifacts/tokenizer-qwen3-adaptive.json'))
+    runner.add_argument('--turns', type=int, default=100)
+    runner.add_argument('--keep', type=int, default=4)
+    runner.add_argument('--steps', type=int, default=16)
+    runner.add_argument('--seed', type=int, default=7)
+    runner.add_argument('--output', type=Path, default=Path('bench-ledger-needle.jsonl'))
+    ar = sub.add_parser('run-ar', help='same script through a plain AR baseline')
+    ar.add_argument('--model-path', type=Path, default=Path('/root/models/Qwen3-0.6B'))
+    ar.add_argument('--turns', type=int, default=100)
+    ar.add_argument('--seed', type=int, default=7)
+    ar.add_argument('--context', choices=('full', 'budget'), default='full')
+    ar.add_argument('--budget', type=int, default=512)
+    ar.add_argument('--thinking', action='store_true',
+                    help='enable the native reasoning mode where the chat template supports it')
+    ar.add_argument('--output', type=Path, default=Path('bench-needle-ar.jsonl'))
+
+    cmp_ = sub.add_parser('compare', help='comparison table across systems as HTML')
+    cmp_.add_argument('--inputs', nargs='+', required=True, metavar='LABEL=PATH[=PARAMS_B]')
+    cmp_.add_argument('--ours', default=None, help='label to highlight as our system')
+    cmp_.add_argument('--output', type=Path, default=Path('docs/ledger-needle-compare.html'))
+
+    view = sub.add_parser('render', help='write the self-contained HTML viewer')
+    view.add_argument('--input', type=Path, default=Path('bench-ledger-needle.jsonl'))
+    view.add_argument('--output', type=Path, default=Path('docs/ledger-needle.html'))
+    view.add_argument('--model-name', default='qwen06b-genchat-sft')
+    view.add_argument('--keep', type=int, default=4)
+    view.add_argument('--steps', type=int, default=16)
+    args = parser.parse_args()
+    if args.command == 'run':
+        run(args)
+    elif args.command == 'run-ar':
+        run_ar(args)
+    elif args.command == 'compare':
+        compare(args)
+    else:
+        render(args)
+
+
+if __name__ == '__main__':
+    main()

bench/ledger-needle-compare.html

@@ -0,0 +1,91 @@
+<!doctype html><html><head><meta charset="utf-8">
+<title>Ledger needle — comparación</title><style>
+:root { --bg:#11141a; --panel:#1a1f29; --text:#dde3ee; --dim:#8a93a6; --hit:#3fb96b;
+  --miss:#e05555; --line:#2a3140; --us:#16233a; }
+* { box-sizing:border-box; margin:0 } body { background:var(--bg); color:var(--text);
+  font:14px/1.5 ui-monospace,Menlo,monospace; padding:24px; max-width:1240px; margin:0 auto }
+h1 { font-size:18px } .sub { color:var(--dim); margin:4px 0 20px }
+table { border-collapse:collapse; width:100%; margin-bottom:26px }
+th, td { border:1px solid var(--line); padding:7px 10px; text-align:left; font-size:13px }
+th { background:var(--panel); color:var(--dim); font-weight:normal; cursor:pointer;
+  user-select:none; white-space:nowrap }
+th.sorted { color:var(--text) } th.sorted::after { content:' ↓' }
+th.sorted.asc::after { content:' ↑' }
+tr.us td { background:var(--us) }
+td.hitcell { background:#173523; color:var(--hit); text-align:center; font-weight:bold }
+td.misscell { background:#3a1d1d; color:var(--miss); text-align:center }
+#matrix { width:auto } #matrix th, #matrix td { padding:5px 8px }
+#matrix td.hitcell, #matrix td.misscell { width:36px; min-width:36px; padding:5px 0 }
+#matrix th.sys { writing-mode:vertical-rl; transform:rotate(180deg); text-align:left;
+  vertical-align:bottom; max-height:190px; font-size:12px; cursor:default; padding:8px 4px }
+#matrix th.sys.usc { color:#9fc4ea }
+#matrix td.exp { white-space:nowrap }
+.bar { display:inline-block; height:11px; background:var(--hit); border-radius:2px;
+  vertical-align:middle } .bar.m { background:var(--miss) }
+h2 { font-size:14px; color:var(--dim); margin:22px 0 10px; text-transform:uppercase }
+.caveats { background:var(--panel); border:1px solid var(--line); border-radius:8px;
+  padding:14px 18px; color:var(--dim); font-size:13px } .caveats li { margin:6px 0 0 18px }
+.hint { color:var(--dim); font-size:12px; margin:-18px 0 10px }
+</style></head><body>
+<h1>Ledger needle — comparación hasta 1.5B</h1>
+<div class="sub">100 turnos guionados (seed 21) · 15 hechos, 5 correcciones, 10 sondas a
+distancia 3–96 · scoring por núcleo del valor, idéntico para todos · baselines con historia
+completa, decoding greedy e instrucción explícita de memoria</div>
+<table id="summary"><thead><tr>
+<th data-k="label" data-t="s">sistema</th>
+<th data-k="params" data-t="n">tamaño (B)</th>
+<th data-k="hits" data-t="n">recall</th>
+<th data-k="d1" data-t="n">d 3–10</th>
+<th data-k="d2" data-t="n">d 11–50</th>
+<th data-k="d3" data-t="n">d 51+</th>
+<th data-k="total" data-t="n">test total (s)</th>
+<th data-k="per" data-t="n">s/turno</th>
+<th data-k="prefix" data-t="n">prefijo máx</th>
+</tr></thead><tbody></tbody></table>
+<div class="hint">click en una cabecera reordena; por defecto: más aciertos primero</div>
+<h2>Sonda por sonda</h2>
+<table id="matrix"><tr><th>turno</th><th>dist</th><th>esperado</th><th class="sys">Qwen2.5-1.5B-Instruct · full</th><th class="sys">Qwen3-0.6B + thinking · full</th><th class="sys usc">ledger qwen06b-genchat (nuestro)</th><th class="sys">Qwen2.5-0.5B-Instruct · full</th><th class="sys">SmolLM2-360M-Instruct · full</th><th class="sys">TinyLlama-1.1B-Chat · full</th><th class="sys">Qwen3-0.6B sin thinking · full</th><th class="sys">Qwen3-0.6B sin thinking · budget 512</th></tr><tr><td>6</td><td>5</td><td class="exp">code 4805</td><td class="hitcell">✓</td><td class="hitcell">✓</td><td class="hitcell">✓</td><td class="hitcell">✓</td><td class="hitcell">✓</td><td class="hitcell">✓</td><td class="hitcell">✓</td><td class="hitcell">✓</td></tr><tr><td>17</td><td>6</td><td class="exp">$6,400</td><td class="hitcell">✓</td><td class="hitcell">✓</td><td class="hitcell">✓</td><td class="hitcell">✓</td><td class="hitcell">✓</td><td class="hitcell">✓</td><td class="misscell">✗</td><td class="misscell">✗</td></tr><tr><td>28</td><td>3</td><td class="exp">7:15pm</td><td class="hitcell">✓</td><td class="hitcell">✓</td><td class="hitcell">✓</td><td class="misscell">✗</td><td class="hitcell">✓</td><td class="hitcell">✓</td><td class="misscell">✗</td><td class="misscell">✗</td></tr><tr><td>61</td><td>57</td><td class="exp">harbor858</td><td class="hitcell">✓</td><td class="hitcell">✓</td><td class="misscell">✗</td><td class="hitcell">✓</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="misscell">✗</td></tr><tr><td>72</td><td>37</td><td class="exp">$4,000</td><td class="hitcell">✓</td><td class="hitcell">✓</td><td class="misscell">✗</td><td class="hitcell">✓</td><td class="hitcell">✓</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="misscell">✗</td></tr><tr><td>83</td><td>47</td><td class="exp">day 97</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="hitcell">✓</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="misscell">✗</td></tr><tr><td>94</td><td>83</td><td class="exp">$6,400</td><td class="hitcell">✓</td><td class="hitcell">✓</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="misscell">✗</td></tr><tr><td>98</td><td>59</td><td class="exp">room 411</td><td class="hitcell">✓</td><td class="misscell">✗</td><td class="hitcell">✓</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="misscell">✗</td></tr><tr><td>99</td><td>60</td><td class="exp">room 411</td><td class="hitcell">✓</td><td class="misscell">✗</td><td class="hitcell">✓</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="misscell">✗</td></tr><tr><td>100</td><td>96</td><td class="exp">harbor858</td><td class="hitcell">✓</td><td class="hitcell">✓</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="misscell">✗</td><td class="misscell">✗</td></tr></table>
+<h2>Advertencias de lectura</h2>
+<div class="caveats"><ul>
+<li>Un solo seed y 10 sondas por sistema: las diferencias chicas son ruido; las grandes
+sobreviven a ese margen.</li>
+<li>La comparación mide el paquete mecanismo+entrenamiento: el sistema de ledger fue entrenado
+en este registro conversacional y los baselines no. Muestra que el paquete funciona, no que el
+mecanismo solo cause la diferencia.</li>
+<li>Los baselines reciben la mayor ventaja razonable: historia completa en contexto, greedy,
+instrucción de memoria, y el modo thinking nativo donde existe. El sistema de ledger corre con
+sampling a temperatura 0.7 y un prefijo constante de ~500 tokens.</li>
+<li>El guion es sintético y sus valores son arbitrarios por diseño; un valor sin sentido
+("day 97") puede penalizar a un modelo que se resiste a repetir datos absurdos.</li>
+</ul></div>
+<script>
+const ROWS = [{"label": "Qwen2.5-1.5B-Instruct · full", "params": 1.54, "probes": 10, "hits": 9, "total": 76.74, "per": 0.7674, "prefix": 2881, "us": false, "d1": 3, "d1n": 3, "d2": 1, "d2n": 2, "d3": 5, "d3n": 5}, {"label": "Qwen3-0.6B + thinking · full", "params": 0.6, "probes": 10, "hits": 7, "total": 1318.3, "per": 13.183, "prefix": 4450, "us": false, "d1": 3, "d1n": 3, "d2": 1, "d2n": 2, "d3": 3, "d3n": 5}, {"label": "ledger qwen06b-genchat (nuestro)", "params": 0.6, "probes": 10, "hits": 6, "total": 218.07, "per": 2.1807, "prefix": 625, "us": true, "d1": 3, "d1n": 3, "d2": 1, "d2n": 2, "d3": 2, "d3n": 5}, {"label": "Qwen2.5-0.5B-Instruct · full", "params": 0.49, "probes": 10, "hits": 4, "total": 110.34, "per": 1.1034, "prefix": 3908, "us": false, "d1": 2, "d1n": 3, "d2": 1, "d2n": 2, "d3": 1, "d3n": 5}, {"label": "SmolLM2-360M-Instruct · full", "params": 0.36, "probes": 10, "hits": 4, "total": 56.93, "per": 0.5693, "prefix": 2712, "us": false, "d1": 3, "d1n": 3, "d2": 1, "d2n": 2, "d3": 0, "d3n": 5}, {"label": "TinyLlama-1.1B-Chat · full", "params": 1.1, "probes": 10, "hits": 3, "total": 120.77, "per": 1.2077, "prefix": 5126, "us": false, "d1": 3, "d1n": 3, "d2": 0, "d2n": 2, "d3": 0, "d3n": 5}, {"label": "Qwen3-0.6B sin thinking · full", "params": 0.6, "probes": 10, "hits": 1, "total": 129.6, "per": 1.296, "prefix": 3722, "us": false, "d1": 1, "d1n": 3, "d2": 0, "d2n": 2, "d3": 0, "d3n": 5}, {"label": "Qwen3-0.6B sin thinking · budget 512", "params": 0.6, "probes": 10, "hits": 1, "total": 161.75, "per": 1.6175, "prefix": 512, "us": false, "d1": 1, "d1n": 3, "d2": 0, "d2n": 2, "d3": 0, "d3n": 5}];
+const tbody = document.querySelector('#summary tbody');
+let sortKey = 'hits', asc = false;
+function fmt(r) {
+  return '<td><b>' + r.label + '</b></td><td>' + r.params.toFixed(2) + '</td>' +
+    '<td><b>' + r.hits + '/' + r.probes + '</b> <span class="bar" style="width:' + 6 * r.hits +
+    'px"></span><span class="bar m" style="width:' + 6 * (r.probes - r.hits) + 'px"></span></td>' +
+    ['d1', 'd2', 'd3'].map(k => '<td>' + r[k] + '/' + r[k + 'n'] + '</td>').join('') +
+    '<td>' + r.total.toFixed(0) + '</td><td>' + r.per.toFixed(1) + '</td><td>' + r.prefix + '</td>';
+}
+function draw() {
+  const rows = [...ROWS].sort((a, b) => {
+    const va = a[sortKey], vb = b[sortKey];
+    const c = typeof va === 'string' ? va.localeCompare(vb) : va - vb;
+    return asc ? c : -c;
+  });
+  tbody.innerHTML = rows.map(r =>
+    '<tr' + (r.us ? ' class="us"' : '') + '>' + fmt(r) + '</tr>').join('');
+  document.querySelectorAll('#summary th').forEach(th => {
+    th.classList.toggle('sorted', th.dataset.k === sortKey);
+    th.classList.toggle('asc', th.dataset.k === sortKey && asc);
+  });
+}
+document.querySelectorAll('#summary th').forEach(th => th.onclick = () => {
+  if (sortKey === th.dataset.k) asc = !asc;
+  else { sortKey = th.dataset.k; asc = th.dataset.t === 's'; }
+  draw();
+});
+draw();
+</script></body></html>

bench/results-qwen25-05b-full.jsonl

@@ -0,0 +1,100 @@
+{"turn": 0, "kind": "fact", "user": "the gym locker is code 4805", "notes": [], "ledger": [], "prefix_tokens": 59, "answer": "The code for the gym locker is 4805.", "seconds": 2.56, "key": "locker code", "value": "code 4805"}
+{"turn": 1, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 87, "answer": "The street has been repaved.", "seconds": 0.26}
+{"turn": 2, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 109, "answer": "A week can vary greatly depending on many factors such as your location, personal habits, and current circumstances. It's important to remember that while a week may seem like a long time in some contexts, it can also be a short period of time in others.", "seconds": 2.21}
+{"turn": 3, "kind": "fact", "user": "cabin wifi password is harbor858", "notes": [], "ledger": [], "prefix_tokens": 180, "answer": "The password for the cabin Wi-Fi is \"harbor858\".", "seconds": 0.58, "key": "wifi password", "value": "harbor858"}
+{"turn": 4, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 210, "answer": "The street has been repaved.", "seconds": 0.36}
+{"turn": 5, "kind": "probe", "user": "quick — the locker code?", "notes": [], "ledger": [], "prefix_tokens": 233, "answer": "The locker code is 4805.", "seconds": 0.43, "key": "locker code", "expected": "code 4805", "distance": 5, "hit": true}
+{"turn": 6, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 257, "answer": "I'm sorry, but I don't have enough context to provide an accurate response about what you're asking. Could you please provide more information or clarify your question?", "seconds": 3.84}
+{"turn": 7, "kind": "fact", "user": "we always park at spot B19", "notes": [], "ledger": [], "prefix_tokens": 308, "answer": "Spot B19 has always been reserved.", "seconds": 0.38, "key": "parking spot", "value": "spot B19"}
+{"turn": 8, "kind": "correction", "user": "actually scratch that, the locker code is code 1688 now", "notes": [], "ledger": [], "prefix_tokens": 342, "answer": "The new locker code for the gym locker is 1688.", "seconds": 0.67, "key": "locker code", "value": "code 1688"}
+{"turn": 9, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 373, "answer": "The cat knocked over a plant.", "seconds": 0.33}
+{"turn": 10, "kind": "fact", "user": "the deposit for the venue was $6,400", "notes": [], "ledger": [], "prefix_tokens": 402, "answer": "The deposit for the venue was $6,400.", "seconds": 0.54, "key": "deposit", "value": "$6,400"}
+{"turn": 11, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 432, "answer": "If your phone battery is dying quickly, it could be due to several reasons including low charge level, overuse, or damage from extreme temperatures. To troubleshoot this issue, try charging your phone fully before using it again, consider swapping out batteries if possible, and ensure your phone is not exposed to extreme temperatures. If the problem persists, it might be best to consult a professional technician who can diagnose and repair any issues with your device.", "seconds": 3.77}
+{"turn": 12, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 536, "answer": "Yesterday, I saw a great documentary.", "seconds": 0.49}
+{"turn": 13, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 559, "answer": "I see! The recent weather conditions have certainly made things interesting. How can I assist further today?", "seconds": 1.09}
+{"turn": 14, "kind": "fact", "user": "tomas' new teacher is ms. duarte", "notes": [], "ledger": [], "prefix_tokens": 599, "answer": "Ms. Duarte is Tomás' new teacher.", "seconds": 0.59, "key": "kid teacher", "value": "ms. duarte"}
+{"turn": 15, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 628, "answer": "Yes, I did catch the game last night.", "seconds": 0.57}
+{"turn": 16, "kind": "probe", "user": "remind me of the deposit?", "notes": [], "ledger": [], "prefix_tokens": 655, "answer": "The deposit for the venue was $6,400.", "seconds": 0.69, "key": "deposit", "expected": "$6,400", "distance": 6, "hit": true}
+{"turn": 17, "kind": "fact", "user": "the plumber quoted $4,900 for the bathroom", "notes": [], "ledger": [], "prefix_tokens": 690, "answer": "The plumber estimated the cost for the bathroom would be $4,900.", "seconds": 0.82, "key": "plumber quote", "value": "$4,900"}
+{"turn": 18, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 724, "answer": "If your phone battery is dying quickly, it could be due to several reasons including low charge level, overuse, or damage from extreme temperatures. To troubleshoot this issue, try charging your phone fully before using it again, consider swapping out batteries if possible, and ensure your phone is not exposed to extreme temperatures. If the problem persists, it might be best to consult a professional technician who can diagnose and repair any issues with your device.", "seconds": 4.16}
+{"turn": 19, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 827, "answer": "Traffic this morning was very heavy.", "seconds": 0.35}
+{"turn": 20, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 853, "answer": "The neighbor is renovating again, which means there will be increased noise levels.", "seconds": 0.77}
+{"turn": 21, "kind": "fact", "user": "our train leaves from platform 58", "notes": [], "ledger": [], "prefix_tokens": 887, "answer": "Our train departs from Platform 58.", "seconds": 0.51, "key": "train platform", "value": "platform 58"}
+{"turn": 22, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 912, "answer": "The street has been repaved.", "seconds": 0.44}
+{"turn": 23, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 934, "answer": "Coffee prices here are becoming increasingly expensive.", "seconds": 0.48}
+{"turn": 24, "kind": "fact", "user": "dentist moved me to 7:15pm", "notes": [], "ledger": [], "prefix_tokens": 963, "answer": "The dentist moved you to 7:15 PM.", "seconds": 0.56, "key": "dentist slot", "value": "7:15pm"}
+{"turn": 25, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 991, "answer": "Yesterday, I saw a great documentary.", "seconds": 0.53}
+{"turn": 26, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 1016, "answer": "If your phone battery is dying quickly, it could be due to several reasons including low charge level, overuse, or damage from extreme temperatures. To troubleshoot this issue, try charging your phone fully before using it again, consider swapping out batteries if possible, and ensure your phone is not exposed to extreme temperatures. If the problem persists, it might be best to consult a professional technician who can diagnose and repair any issues with your device.", "seconds": 4.53}
+{"turn": 27, "kind": "probe", "user": "quick — the dentist slot?", "notes": [], "ledger": [], "prefix_tokens": 1120, "answer": "The dentist slot is located on the right side of the door.", "seconds": 0.61, "key": "dentist slot", "expected": "7:15pm", "distance": 3, "hit": false}
+{"turn": 28, "kind": "fact", "user": "doctor changed the dose to 27mg", "notes": [], "ledger": [], "prefix_tokens": 1152, "answer": "The doctor has changed the dose to 27 mg.", "seconds": 0.67, "key": "meds dose", "value": "27mg"}
+{"turn": 29, "kind": "distractor", "user": "might go for a walk later", "notes": [], "ledger": [], "prefix_tokens": 1180, "answer": "It sounds like you might decide to go for a walk later.", "seconds": 0.76}
+{"turn": 30, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 1207, "answer": "I finally finished that book.", "seconds": 0.34}
+{"turn": 31, "kind": "fact", "user": "the couch is order ZT6045, keep it handy", "notes": [], "ledger": [], "prefix_tokens": 1237, "answer": "The couch is ordered as ZT6045. Keep it handy.", "seconds": 0.97, "key": "order number", "value": "order ZT6045"}
+{"turn": 32, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 1268, "answer": "A week can vary greatly depending on many factors such as your location, personal habits, and current circumstances. It's important to remember that while a week may seem like a long time in some contexts, it can also be a short period of time in others.", "seconds": 2.52}
+{"turn": 33, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 1336, "answer": "I'm sorry to hear that you're experiencing back pain from your gym workout. It's important to take care of your body after physical activity, especially when it comes to injuries like back pain. You might want to consider taking a break from exercise and focusing on stretching and strengthening exercises instead. Additionally, making sure to get plenty of rest and hydration can help alleviate any discomfort. If the pain persists or worsens, it's recommended to consult with a healthcare provider for proper evaluation and treatment.", "seconds": 4.93}
+{"turn": 34, "kind": "correction", "user": "actually scratch that, the plumber quote is $4,000 now", "notes": [], "ledger": [], "prefix_tokens": 1459, "answer": "The plumber estimate for the bathroom is $4,000.", "seconds": 0.74, "key": "plumber quote", "value": "$4,000"}
+{"turn": 35, "kind": "fact", "user": "landlord wants rent by day 97 each month", "notes": [], "ledger": [], "prefix_tokens": 1494, "answer": "The landlord wants the rent to be paid monthly starting from Day 97.", "seconds": 0.88, "key": "rent due", "value": "day 97"}
+{"turn": 36, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 1526, "answer": "The cat knocked over a plant.", "seconds": 0.42}
+{"turn": 37, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 1548, "answer": "Traffic this morning was very heavy.", "seconds": 0.38}
+{"turn": 38, "kind": "fact", "user": "we got room 411 at the hotel", "notes": [], "ledger": [], "prefix_tokens": 1575, "answer": "We got room 411 at the hotel.", "seconds": 0.58, "key": "hotel room", "value": "room 411"}
+{"turn": 39, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 1605, "answer": "The neighbor is renovating again, which means there will be increased noise levels.", "seconds": 0.8}
+{"turn": 40, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 1636, "answer": "Coffee prices here are becoming increasingly expensive.", "seconds": 0.4}
+{"turn": 41, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 1659, "answer": "I see! The recent weather conditions have certainly made things interesting. How can I assist further today?", "seconds": 0.91}
+{"turn": 42, "kind": "fact", "user": "boarding is at gate E17", "notes": [], "ledger": [], "prefix_tokens": 1696, "answer": "The boarding is at Gate 17.", "seconds": 0.44, "key": "boarding gate", "value": "gate E17"}
+{"turn": 43, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 1720, "answer": "A week can vary greatly depending on many factors such as your location, personal habits, and current circumstances. It's important to remember that while a week may seem like a long time in some contexts, it can also be a short period of time in others.", "seconds": 2.35}
+{"turn": 44, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 1788, "answer": "The cat knocked over a plant.", "seconds": 0.41}
+{"turn": 45, "kind": "fact", "user": "my sister lands on flight AL7371, write that down", "notes": [], "ledger": [], "prefix_tokens": 1819, "answer": "My sister landed on Flight AL7371.", "seconds": 0.51, "key": "sister flight", "value": "flight AL7371"}
+{"turn": 46, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 1845, "answer": "Coffee prices here are becoming increasingly expensive.", "seconds": 0.49}
+{"turn": 47, "kind": "correction", "user": "actually scratch that, the sister flight is flight BL4905 now", "notes": [], "ledger": [], "prefix_tokens": 1878, "answer": "The sister flight is now flight BL4905.", "seconds": 0.65, "key": "sister flight", "value": "flight BL4905"}
+{"turn": 48, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 1905, "answer": "Coffee prices here are becoming increasingly expensive.", "seconds": 0.43}
+{"turn": 49, "kind": "fact", "user": "the rental has plate HJC-988", "notes": [], "ledger": [], "prefix_tokens": 1933, "answer": "The rental has a plate number HJC-988.", "seconds": 0.68, "key": "car plate", "value": "HJC-988"}
+{"turn": 50, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 1962, "answer": "I'm sorry to hear that you're experiencing back pain from your gym workout. It's important to take care of your body after physical activity, especially when it comes to injuries like back pain. You might want to consider taking a break from exercise and focusing on stretching and strengthening exercises instead. Additionally, making sure to get plenty of rest and hydration can help alleviate any discomfort. If the pain persists or worsens, it's recommended to consult with a healthcare provider for proper evaluation and treatment.", "seconds": 4.17}
+{"turn": 51, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 2074, "answer": "I finally finished that book.", "seconds": 0.26}
+{"turn": 52, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 2095, "answer": "Coffee prices here are becoming increasingly expensive.", "seconds": 0.44}
+{"turn": 53, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 2117, "answer": "Today's work was chaotic.", "seconds": 0.29}
+{"turn": 54, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 2142, "answer": "The neighbor is renovating again, which means there will be increased noise levels.", "seconds": 0.71}
+{"turn": 55, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 2177, "answer": "The neighbor is renovating again, which means there will be increased noise levels.", "seconds": 0.72}
+{"turn": 56, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 2211, "answer": "Yes, I did catch the game last night.", "seconds": 0.43}
+{"turn": 57, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 2236, "answer": "Traffic this morning was very heavy.", "seconds": 0.36}
+{"turn": 58, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 2258, "answer": "Traffic this morning was very heavy.", "seconds": 0.39}
+{"turn": 59, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 2280, "answer": "A week can vary greatly depending on many factors such as your location, personal habits, and current circumstances. It's important to remember that while a week may seem like a long time in some contexts, it can also be a short period of time in others.", "seconds": 2.4}
+{"turn": 60, "kind": "probe", "user": "wait, what was the wifi password again?", "notes": [], "ledger": [], "prefix_tokens": 2351, "answer": "The password for the Wi-Fi network was \"harbor858\".", "seconds": 0.71, "key": "wifi password", "expected": "harbor858", "distance": 57, "hit": true}
+{"turn": 61, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 2380, "answer": "I finally finished that book.", "seconds": 0.35}
+{"turn": 62, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 2401, "answer": "The street has been repaved.", "seconds": 0.33}
+{"turn": 63, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 2423, "answer": "The street has been repaved.", "seconds": 0.4}
+{"turn": 64, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 2446, "answer": "Yesterday, I saw a great documentary.", "seconds": 0.36}
+{"turn": 65, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 2469, "answer": "A week can vary greatly depending on many factors such as your location, personal habits, and current circumstances. It's important to remember that while a week may seem like a long time in some contexts, it can also be a short period of time in others.", "seconds": 2.23}
+{"turn": 66, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 2536, "answer": "Thinking of making pasta tonight sounds delicious! What type of pasta do you prefer?", "seconds": 0.83}
+{"turn": 67, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 2566, "answer": "Today's work was chaotic.", "seconds": 0.35}
+{"turn": 68, "kind": "distractor", "user": "might go for a walk later", "notes": [], "ledger": [], "prefix_tokens": 2588, "answer": "It sounds like you might decide to go for a walk later.", "seconds": 0.64}
+{"turn": 69, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 2620, "answer": "The neighbor is renovating again, which means there will be increased noise levels.", "seconds": 0.73}
+{"turn": 70, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 2651, "answer": "The street has been repaved.", "seconds": 0.4}
+{"turn": 71, "kind": "probe", "user": "wait, what was the plumber quote again?", "notes": [], "ledger": [], "prefix_tokens": 2677, "answer": "The plumber quote for the bathroom is $4,000.", "seconds": 0.78, "key": "plumber quote", "expected": "$4,000", "distance": 37, "hit": true}
+{"turn": 72, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 2706, "answer": "Traffic this morning was very heavy.", "seconds": 0.41}
+{"turn": 73, "kind": "correction", "user": "actually scratch that, the parking spot is spot E16 now", "notes": [], "ledger": [], "prefix_tokens": 2736, "answer": "The parking spot is now spot E16.", "seconds": 0.62, "key": "parking spot", "value": "spot E16"}
+{"turn": 74, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 2761, "answer": "I see! The recent weather conditions have certainly made things interesting. How can I assist further today?", "seconds": 1.26}
+{"turn": 75, "kind": "distractor", "user": "craving something sweet", "notes": [], "ledger": [], "prefix_tokens": 2795, "answer": "You might enjoy something sweet, such as chocolate cake or cookies.", "seconds": 0.79}
+{"turn": 76, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 2827, "answer": "The neighbor is renovating again, which means there will be increased noise levels.", "seconds": 1.0}
+{"turn": 77, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 2858, "answer": "Thinking of making pasta tonight sounds delicious! What type of pasta do you prefer?", "seconds": 0.9}
+{"turn": 78, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 2892, "answer": "Yes, I did catch the game last night.", "seconds": 0.55}
+{"turn": 79, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 2919, "answer": "If your phone battery is dying quickly, it could be due to several reasons including low charge level, overuse, or damage from extreme temperatures. To troubleshoot this issue, try charging your phone fully before using it again, consider swapping out batteries if possible, and ensure your phone is not exposed to extreme temperatures. If the problem persists, it might be best to consult a professional technician who can diagnose and repair any issues with your device.", "seconds": 4.05}
+{"turn": 80, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 3024, "answer": "If your phone battery is dying quickly, it could be due to several reasons including low charge level, overuse, or damage from extreme temperatures. To troubleshoot this issue, try charging your phone fully before using it again, consider swapping out batteries if possible, and ensure your phone is not exposed to extreme temperatures. If the problem persists, it might be best to consult a professional technician who can diagnose and repair any issues with your device.", "seconds": 4.06}
+{"turn": 81, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 3129, "answer": "If your phone battery is dying quickly, it could be due to several reasons including low charge level, overuse, or damage from extreme temperatures. To troubleshoot this issue, try charging your phone fully before using it again, consider swapping out batteries if possible, and ensure your phone is not exposed to extreme temperatures. If the problem persists, it might be best to consult a professional technician who can diagnose and repair any issues with your device.", "seconds": 3.73}
+{"turn": 82, "kind": "probe", "user": "i forgot the rent due, what was it?", "notes": [], "ledger": [], "prefix_tokens": 3237, "answer": "I'm sorry to hear that you forgot the rent due. Please let me know how much you need to pay so we can proceed with the payment process.", "seconds": 1.27, "key": "rent due", "expected": "day 97", "distance": 47, "hit": false}
+{"turn": 83, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 3283, "answer": "Thinking of making pasta tonight sounds delicious! What type of pasta do you prefer?", "seconds": 0.64}
+{"turn": 84, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 3314, "answer": "Coffee prices here are becoming increasingly expensive.", "seconds": 0.43}
+{"turn": 85, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 3337, "answer": "Thinking of making pasta tonight sounds delicious! What type of pasta do you prefer?", "seconds": 0.93}
+{"turn": 86, "kind": "correction", "user": "actually scratch that, the plumber quote is $1,400 now", "notes": [], "ledger": [], "prefix_tokens": 3378, "answer": "The plumber estimate for the bathroom is $1,400.", "seconds": 0.62, "key": "plumber quote", "value": "$1,400"}
+{"turn": 87, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 3409, "answer": "If your phone battery is dying quickly, it could be due to several reasons including low charge level, overuse, or damage from extreme temperatures. To troubleshoot this issue, try charging your phone fully before using it again, consider swapping out batteries if possible, and ensure your phone is not exposed to extreme temperatures. If the problem persists, it might be best to consult a professional technician who can diagnose and repair any issues with your device.", "seconds": 4.0}
+{"turn": 88, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 3512, "answer": "Coffee prices here are becoming increasingly expensive.", "seconds": 0.61}
+{"turn": 89, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 3534, "answer": "I finally finished that book.", "seconds": 0.38}
+{"turn": 90, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 3556, "answer": "I'm sorry to hear that you're experiencing back pain from your gym workout. It's important to take care of your body after physical activity, especially when it comes to injuries like back pain. You might want to consider taking a break from exercise and focusing on stretching and strengthening exercises instead. Additionally, making sure to get plenty of rest and hydration can help alleviate any discomfort. If the pain persists or worsens, it's recommended to consult with a healthcare provider for proper evaluation and treatment.", "seconds": 4.39}
+{"turn": 91, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 3672, "answer": "Yes, I did catch the game last night.", "seconds": 0.61}
+{"turn": 92, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 3696, "answer": "Today's work was chaotic.", "seconds": 0.44}
+{"turn": 93, "kind": "probe", "user": "i forgot the deposit, what was it?", "notes": [], "ledger": [], "prefix_tokens": 3721, "answer": "I'm sorry to hear that you forgot the deposit. Please let me know how much you need to pay so we can proceed with the payment process.", "seconds": 1.65, "key": "deposit", "expected": "$6,400", "distance": 83, "hit": false}
+{"turn": 94, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 3766, "answer": "Coffee prices here are becoming increasingly expensive.", "seconds": 0.4}
+{"turn": 95, "kind": "distractor", "user": "craving something sweet", "notes": [], "ledger": [], "prefix_tokens": 3788, "answer": "You might enjoy something sweet, such as chocolate cake or cookies.", "seconds": 0.82}
+{"turn": 96, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 3820, "answer": "The neighbor is renovating again, which means there will be increased noise levels.", "seconds": 0.99}
+{"turn": 97, "kind": "probe", "user": "remind me of the hotel room?", "notes": [], "ledger": [], "prefix_tokens": 3854, "answer": "The hotel room is room 58.", "seconds": 0.48, "key": "hotel room", "expected": "room 411", "distance": 59, "hit": false}
+{"turn": 98, "kind": "probe", "user": "quick — the hotel room?", "notes": [], "ledger": [], "prefix_tokens": 3879, "answer": "The hotel room is room 58.", "seconds": 0.47, "key": "hotel room", "expected": "room 411", "distance": 60, "hit": false}
+{"turn": 99, "kind": "probe", "user": "i forgot the wifi password, what was it?", "notes": [], "ledger": [], "prefix_tokens": 3908, "answer": "I'm sorry to hear that you forgot the wifi password. Please let me know how much you need to pay so we can proceed with the payment process.", "seconds": 1.18, "key": "wifi password", "expected": "harbor858", "distance": 96, "hit": false}

bench/results-qwen25-15b-full.jsonl

@@ -0,0 +1,100 @@
+{"turn": 0, "kind": "fact", "user": "the gym locker is code 4805", "notes": [], "ledger": [], "prefix_tokens": 59, "answer": "The gym locker code you provided is 4805.", "seconds": 6.46, "key": "locker code", "value": "code 4805"}
+{"turn": 1, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 87, "answer": "They repaved your street.", "seconds": 0.46}
+{"turn": 2, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 108, "answer": "It sounds like you had an eventful or challenging week.", "seconds": 0.75}
+{"turn": 3, "kind": "fact", "user": "cabin wifi password is harbor858", "notes": [], "ledger": [], "prefix_tokens": 139, "answer": "The cabin's WiFi password is \"harbor858\".", "seconds": 0.89, "key": "wifi password", "value": "harbor858"}
+{"turn": 4, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 167, "answer": "They repaved your street.", "seconds": 0.43}
+{"turn": 5, "kind": "probe", "user": "quick — the locker code?", "notes": [], "ledger": [], "prefix_tokens": 189, "answer": "The gym locker code is 4805.", "seconds": 0.73, "key": "locker code", "expected": "code 4805", "distance": 5, "hit": true}
+{"turn": 6, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 214, "answer": "You finally finished reading that book.", "seconds": 0.41}
+{"turn": 7, "kind": "fact", "user": "we always park at spot B19", "notes": [], "ledger": [], "prefix_tokens": 239, "answer": "You always park in spot B19.", "seconds": 0.51, "key": "parking spot", "value": "spot B19"}
+{"turn": 8, "kind": "correction", "user": "actually scratch that, the locker code is code 1688 now", "notes": [], "ledger": [], "prefix_tokens": 273, "answer": "The gym locker code has been updated to 1688.", "seconds": 0.79, "key": "locker code", "value": "code 1688"}
+{"turn": 9, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 303, "answer": "The cat knocked over a plant.", "seconds": 0.53}
+{"turn": 10, "kind": "fact", "user": "the deposit for the venue was $6,400", "notes": [], "ledger": [], "prefix_tokens": 332, "answer": "The deposit for the venue was $6,400.", "seconds": 0.85, "key": "deposit", "value": "$6,400"}
+{"turn": 11, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 362, "answer": "Your phone battery dies faster than usual now.", "seconds": 0.48}
+{"turn": 12, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 387, "answer": "You saw a great documentary yesterday.", "seconds": 0.34}
+{"turn": 13, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 409, "answer": "You really enjoy the current weather conditions recently.", "seconds": 0.47}
+{"turn": 14, "kind": "fact", "user": "tomas' new teacher is ms. duarte", "notes": [], "ledger": [], "prefix_tokens": 438, "answer": "Tomas' new teacher is Ms. Duarte.", "seconds": 0.59, "key": "kid teacher", "value": "ms. duarte"}
+{"turn": 15, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 467, "answer": "I'm sorry, but I don't have any information about a specific game being played last night. Could you please provide more details or clarify which game you're referring to?", "seconds": 1.8}
+{"turn": 16, "kind": "probe", "user": "remind me of the deposit?", "notes": [], "ledger": [], "prefix_tokens": 519, "answer": "The deposit for the venue was $6,400.", "seconds": 0.69, "key": "deposit", "expected": "$6,400", "distance": 6, "hit": true}
+{"turn": 17, "kind": "fact", "user": "the plumber quoted $4,900 for the bathroom", "notes": [], "ledger": [], "prefix_tokens": 554, "answer": "The plumber quoted $4,900 for the bathroom work.", "seconds": 0.85, "key": "plumber quote", "value": "$4,900"}
+{"turn": 18, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 585, "answer": "Your phone battery dies faster than usual now.", "seconds": 0.51}
+{"turn": 19, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 609, "answer": "Traffic was extremely busy and congested this morning.", "seconds": 0.44}
+{"turn": 20, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 638, "answer": "The neighbor is renovating their home once again, causing quite a disturbance due to the noise level.", "seconds": 1.16}
+{"turn": 21, "kind": "fact", "user": "our train leaves from platform 58", "notes": [], "ledger": [], "prefix_tokens": 676, "answer": "Our train departs from Platform 58.", "seconds": 0.64, "key": "train platform", "value": "platform 58"}
+{"turn": 22, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 701, "answer": "They repaved your street.", "seconds": 0.44}
+{"turn": 23, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 722, "answer": "Coffee prices have increased recently.", "seconds": 0.32}
+{"turn": 24, "kind": "fact", "user": "dentist moved me to 7:15pm", "notes": [], "ledger": [], "prefix_tokens": 749, "answer": "You were moved to see the dentist at 7:15 PM.", "seconds": 0.88, "key": "dentist slot", "value": "7:15pm"}
+{"turn": 25, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 780, "answer": "You watched a great documentary yesterday.", "seconds": 0.4}
+{"turn": 26, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 804, "answer": "Your phone battery dies faster than usual now.", "seconds": 0.49}
+{"turn": 27, "kind": "probe", "user": "quick — the dentist slot?", "notes": [], "ledger": [], "prefix_tokens": 829, "answer": "You were moved to see the dentist at 7:15 PM.", "seconds": 0.87, "key": "dentist slot", "expected": "7:15pm", "distance": 3, "hit": true}
+{"turn": 28, "kind": "fact", "user": "doctor changed the dose to 27mg", "notes": [], "ledger": [], "prefix_tokens": 863, "answer": "The doctor adjusted the medication dosage to 27 mg.", "seconds": 0.82, "key": "meds dose", "value": "27mg"}
+{"turn": 29, "kind": "distractor", "user": "might go for a walk later", "notes": [], "ledger": [], "prefix_tokens": 891, "answer": "You might take a walk later.", "seconds": 0.48}
+{"turn": 30, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 912, "answer": "You finally finished reading that book.", "seconds": 0.59}
+{"turn": 31, "kind": "fact", "user": "the couch is order ZT6045, keep it handy", "notes": [], "ledger": [], "prefix_tokens": 943, "answer": "The couch is ordered as model number ZT6045 and should be kept ready.", "seconds": 1.18, "key": "order number", "value": "order ZT6045"}
+{"turn": 32, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 977, "answer": "It sounds like you had a challenging or eventful week.", "seconds": 0.76}
+{"turn": 33, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 1005, "answer": "You've injured your back from working out at the gym.", "seconds": 0.81}
+{"turn": 34, "kind": "correction", "user": "actually scratch that, the plumber quote is $4,000 now", "notes": [], "ledger": [], "prefix_tokens": 1042, "answer": "The plumber's quote for the job is now $4,000.", "seconds": 1.11, "key": "plumber quote", "value": "$4,000"}
+{"turn": 35, "kind": "fact", "user": "landlord wants rent by day 97 each month", "notes": [], "ledger": [], "prefix_tokens": 1079, "answer": "The landlord requires the full monthly rent paid by the ninth of the following month.", "seconds": 0.96, "key": "rent due", "value": "day 97"}
+{"turn": 36, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 1111, "answer": "The cat knocked over a plant.", "seconds": 0.4}
+{"turn": 37, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 1133, "answer": "Traffic was extremely busy and congested this morning.", "seconds": 0.49}
+{"turn": 38, "kind": "fact", "user": "we got room 411 at the hotel", "notes": [], "ledger": [], "prefix_tokens": 1163, "answer": "You booked Room 411 at the hotel.", "seconds": 0.54, "key": "hotel room", "value": "room 411"}
+{"turn": 39, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 1193, "answer": "The neighbor is renovating their home once again, causing quite a disturbance due to the noise level.", "seconds": 0.88}
+{"turn": 40, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 1228, "answer": "Coffee prices have increased recently.", "seconds": 0.37}
+{"turn": 41, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 1249, "answer": "You really enjoy the current weather conditions recently.", "seconds": 0.46}
+{"turn": 42, "kind": "fact", "user": "boarding is at gate E17", "notes": [], "ledger": [], "prefix_tokens": 1275, "answer": "Boarding instructions state to board at Gate E17.", "seconds": 0.61, "key": "boarding gate", "value": "gate E17"}
+{"turn": 43, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 1302, "answer": "It sounds like you had a challenging or eventful week.", "seconds": 0.66}
+{"turn": 44, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 1330, "answer": "The cat knocked over a plant.", "seconds": 0.47}
+{"turn": 45, "kind": "fact", "user": "my sister lands on flight AL7371, write that down", "notes": [], "ledger": [], "prefix_tokens": 1361, "answer": "Your sister is scheduled to land on Flight AL7371.", "seconds": 0.82, "key": "sister flight", "value": "flight AL7371"}
+{"turn": 46, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 1390, "answer": "Coffee prices have increased recently.", "seconds": 0.31}
+{"turn": 47, "kind": "correction", "user": "actually scratch that, the sister flight is flight BL4905 now", "notes": [], "ledger": [], "prefix_tokens": 1421, "answer": "Your sister is scheduled to fly on Flight BL4905.", "seconds": 0.7, "key": "sister flight", "value": "flight BL4905"}
+{"turn": 48, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 1450, "answer": "Coffee prices have increased recently.", "seconds": 0.3}
+{"turn": 49, "kind": "fact", "user": "the rental has plate HJC-988", "notes": [], "ledger": [], "prefix_tokens": 1476, "answer": "The rental car is registered under license plate HJC-988.", "seconds": 0.67, "key": "car plate", "value": "HJC-988"}
+{"turn": 50, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 1507, "answer": "You've injured your back from working out at the gym.", "seconds": 0.57}
+{"turn": 51, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 1533, "answer": "You finally finished reading that book.", "seconds": 0.41}
+{"turn": 52, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 1555, "answer": "Coffee prices have increased recently.", "seconds": 0.41}
+{"turn": 53, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 1575, "answer": "Work was chaotic today.", "seconds": 0.29}
+{"turn": 54, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 1599, "answer": "The neighbor is renovating their home once again, causing quite a disturbance due to the noise level.", "seconds": 1.08}
+{"turn": 55, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 1638, "answer": "The neighbor is renovating their home once again, causing quite a disturbance due to the noise level.", "seconds": 1.07}
+{"turn": 56, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 1676, "answer": "I'm sorry, but I don't have any information about a specific game being played last night. Could you please provide more details or clarify which game you're referring to?", "seconds": 2.23}
+{"turn": 57, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 1726, "answer": "Traffic was extremely busy and congested this morning.", "seconds": 0.74}
+{"turn": 58, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 1751, "answer": "Traffic was extremely busy and congested this morning.", "seconds": 0.59}
+{"turn": 59, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 1776, "answer": "It sounds like you had a challenging or eventful week.", "seconds": 0.8}
+{"turn": 60, "kind": "probe", "user": "wait, what was the wifi password again?", "notes": [], "ledger": [], "prefix_tokens": 1807, "answer": "The correct WiFi password is \"harbor858.\"", "seconds": 0.68, "key": "wifi password", "expected": "harbor858", "distance": 57, "hit": true}
+{"turn": 61, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 1833, "answer": "You finally finished reading that book.", "seconds": 0.48}
+{"turn": 62, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 1855, "answer": "They repaved your street.", "seconds": 0.41}
+{"turn": 63, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 1876, "answer": "They repaved your street.", "seconds": 0.46}
+{"turn": 64, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 1898, "answer": "You saw a great documentary yesterday.", "seconds": 0.53}
+{"turn": 65, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 1920, "answer": "It sounds like you had a challenging or eventful week.", "seconds": 0.82}
+{"turn": 66, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 1947, "answer": "You're thinking of making pasta tonight.", "seconds": 0.62}
+{"turn": 67, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 1969, "answer": "Work was chaotic today.", "seconds": 0.37}
+{"turn": 68, "kind": "distractor", "user": "might go for a walk later", "notes": [], "ledger": [], "prefix_tokens": 1990, "answer": "You might take a walk later.", "seconds": 0.42}
+{"turn": 69, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 2016, "answer": "The neighbor is renovating their home once again, causing quite a disturbance due to the noise level.", "seconds": 1.14}
+{"turn": 70, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 2051, "answer": "They repaved your street.", "seconds": 0.49}
+{"turn": 71, "kind": "probe", "user": "wait, what was the plumber quote again?", "notes": [], "ledger": [], "prefix_tokens": 2076, "answer": "The plumber's quote for the job is now $4,000.", "seconds": 0.98, "key": "plumber quote", "expected": "$4,000", "distance": 37, "hit": true}
+{"turn": 72, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 2107, "answer": "Traffic was extremely busy and congested this morning.", "seconds": 0.76}
+{"turn": 73, "kind": "correction", "user": "actually scratch that, the parking spot is spot E16 now", "notes": [], "ledger": [], "prefix_tokens": 2140, "answer": "The parking spot is designated as Spot E16.", "seconds": 0.79, "key": "parking spot", "value": "spot E16"}
+{"turn": 74, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 2166, "answer": "You really enjoy the current weather conditions recently.", "seconds": 0.78}
+{"turn": 75, "kind": "distractor", "user": "craving something sweet", "notes": [], "ledger": [], "prefix_tokens": 2189, "answer": "You're craving something sweet.", "seconds": 0.52}
+{"turn": 76, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 2214, "answer": "The neighbor is renovating their home once again, causing quite a disturbance due to the noise level.", "seconds": 1.22}
+{"turn": 77, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 2249, "answer": "You're considering cooking pasta tonight.", "seconds": 0.65}
+{"turn": 78, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 2274, "answer": "I'm sorry, but I don't have any information about a specific game being played last night. Could you please provide more details or clarify which game you're referring to?", "seconds": 2.47}
+{"turn": 79, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 2326, "answer": "Your phone battery dies faster than usual now.", "seconds": 0.9}
+{"turn": 80, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 2352, "answer": "Your phone battery dies faster than usual now.", "seconds": 0.69}
+{"turn": 81, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 2378, "answer": "Your phone battery dies faster than usual now.", "seconds": 0.58}
+{"turn": 82, "kind": "probe", "user": "i forgot the rent due, what was it?", "notes": [], "ledger": [], "prefix_tokens": 2407, "answer": "The rent due was $6,400.", "seconds": 0.79, "key": "rent due", "expected": "day 97", "distance": 47, "hit": false}
+{"turn": 83, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 2433, "answer": "You're considering cooking pasta tonight.", "seconds": 0.52}
+{"turn": 84, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 2455, "answer": "Coffee prices have increased recently.", "seconds": 0.49}
+{"turn": 85, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 2476, "answer": "You're considering cooking pasta tonight.", "seconds": 0.54}
+{"turn": 86, "kind": "correction", "user": "actually scratch that, the plumber quote is $1,400 now", "notes": [], "ledger": [], "prefix_tokens": 2508, "answer": "The plumber's quote for the job is now $1,400.", "seconds": 0.96, "key": "plumber quote", "value": "$1,400"}
+{"turn": 87, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 2541, "answer": "Your phone battery dies faster than usual now.", "seconds": 0.58}
+{"turn": 88, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 2565, "answer": "Coffee prices have increased recently.", "seconds": 0.43}
+{"turn": 89, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 2585, "answer": "You finally finished reading that book.", "seconds": 0.55}
+{"turn": 90, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 2608, "answer": "You've injured your back from working out at the gym.", "seconds": 1.0}
+{"turn": 91, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 2638, "answer": "I'm sorry, but I currently don't have any information about a specific game being played last night. Could you please provide more details or clarify which game you're referring to?", "seconds": 1.8}
+{"turn": 92, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 2688, "answer": "Work was chaotic today.", "seconds": 0.39}
+{"turn": 93, "kind": "probe", "user": "i forgot the deposit, what was it?", "notes": [], "ledger": [], "prefix_tokens": 2712, "answer": "The deposit for the venue was $6,400.", "seconds": 0.77, "key": "deposit", "expected": "$6,400", "distance": 83, "hit": true}
+{"turn": 94, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 2740, "answer": "Coffee prices have increased recently.", "seconds": 0.44}
+{"turn": 95, "kind": "distractor", "user": "craving something sweet", "notes": [], "ledger": [], "prefix_tokens": 2760, "answer": "You're craving something sweet.", "seconds": 0.52}
+{"turn": 96, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 2785, "answer": "The neighbor is renovating their home once again, causing quite a disturbance due to the noise level.", "seconds": 1.19}
+{"turn": 97, "kind": "probe", "user": "remind me of the hotel room?", "notes": [], "ledger": [], "prefix_tokens": 2823, "answer": "You booked Room 411 at the hotel.", "seconds": 0.83, "key": "hotel room", "expected": "room 411", "distance": 59, "hit": true}
+{"turn": 98, "kind": "probe", "user": "quick — the hotel room?", "notes": [], "ledger": [], "prefix_tokens": 2850, "answer": "You booked Room 411 at the hotel.", "seconds": 0.77, "key": "hotel room", "expected": "room 411", "distance": 60, "hit": true}
+{"turn": 99, "kind": "probe", "user": "i forgot the wifi password, what was it?", "notes": [], "ledger": [], "prefix_tokens": 2881, "answer": "The WiFi password is \"harbor858.\"", "seconds": 0.65, "key": "wifi password", "expected": "harbor858", "distance": 96, "hit": true}

bench/results-qwen3-06b-nothink-budget512.jsonl

@@ -0,0 +1,100 @@
+{"turn": 0, "kind": "fact", "user": "the gym locker is code 4805", "notes": [], "ledger": [], "prefix_tokens": 43, "answer": "The code 4805 corresponds to a gym locker. If you have any questions about it, feel free to ask!", "seconds": 3.13, "key": "locker code", "value": "code 4805"}
+{"turn": 1, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 84, "answer": "They repaved your street. Let me know if you need help with anything else!", "seconds": 1.17}
+{"turn": 2, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 116, "answer": "That's a great way to say it! If you need help with anything else, feel free to ask!", "seconds": 1.39}
+{"turn": 3, "kind": "fact", "user": "cabin wifi password is harbor858", "notes": [], "ledger": [], "prefix_tokens": 157, "answer": "The cabin's WiFi password is **harbor858**. Let me know if you need help with anything else!", "seconds": 1.74, "key": "wifi password", "value": "harbor858"}
+{"turn": 4, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 197, "answer": "They repaved your street. Let me know if you need help with anything else!", "seconds": 1.38}
+{"turn": 5, "kind": "probe", "user": "quick — the locker code?", "notes": [], "ledger": [], "prefix_tokens": 230, "answer": "Yes, the locker code is **4805**. Let me know if you need help with anything else!", "seconds": 1.91, "key": "locker code", "expected": "code 4805", "distance": 5, "hit": true}
+{"turn": 6, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 268, "answer": "Finally, you finished that book! Let me know if you need help with anything else!", "seconds": 1.17}
+{"turn": 7, "kind": "fact", "user": "we always park at spot B19", "notes": [], "ledger": [], "prefix_tokens": 304, "answer": "Yes, you always park at **spot B19**. Let me know if you need help with anything else!", "seconds": 1.96, "key": "parking spot", "value": "spot B19"}
+{"turn": 8, "kind": "correction", "user": "actually scratch that, the locker code is code 1688 now", "notes": [], "ledger": [], "prefix_tokens": 353, "answer": "Actually, the locker code is now **code 1688**. Let me know if you need help with anything else!", "seconds": 1.81, "key": "locker code", "value": "code 1688"}
+{"turn": 9, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 396, "answer": "The cat knocked over a plant. Let me know if you need help with anything else!", "seconds": 1.45}
+{"turn": 10, "kind": "fact", "user": "the deposit for the venue was $6,400", "notes": [], "ledger": [], "prefix_tokens": 436, "answer": "The deposit for the venue was **$6,400**. Let me know if you need help with anything else!", "seconds": 1.93, "key": "deposit", "value": "$6,400"}
+{"turn": 11, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 479, "answer": "Your phone battery is dying so fast now. Let me know if you need help with anything else!", "seconds": 1.21}
+{"turn": 12, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 500, "answer": "Saw a great documentary yesterday. Let me know if you need help with anything else!", "seconds": 1.22}
+{"turn": 13, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 502, "answer": "I love this weather lately too! Let me know if you need help with anything else!", "seconds": 1.08}
+{"turn": 14, "kind": "fact", "user": "tomas' new teacher is ms. duarte", "notes": [], "ledger": [], "prefix_tokens": 508, "answer": "Tomás' new teacher is Ms. Duarte. Let me know if you need help with anything else!", "seconds": 1.42, "key": "kid teacher", "value": "ms. duarte"}
+{"turn": 15, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 511, "answer": "I don't know if I caught the game last night, but I hope it was a good time! Let me know if you need help with anything else!", "seconds": 2.36}
+{"turn": 16, "kind": "probe", "user": "remind me of the deposit?", "notes": [], "ledger": [], "prefix_tokens": 506, "answer": "Sure! Let me know if you need help with anything else!", "seconds": 1.15, "key": "deposit", "expected": "$6,400", "distance": 6, "hit": false}
+{"turn": 17, "kind": "fact", "user": "the plumber quoted $4,900 for the bathroom", "notes": [], "ledger": [], "prefix_tokens": 508, "answer": "The plumber quoted **$4,900** for the bathroom. Let me know if you need help with anything else!", "seconds": 2.11, "key": "plumber quote", "value": "$4,900"}
+{"turn": 18, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 490, "answer": "Your phone battery is dying so fast now. Let me know if you need help with anything else!", "seconds": 1.39}
+{"turn": 19, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 512, "answer": "Traffic was unreal this morning. Let me know if you need help with anything else!", "seconds": 1.34}
+{"turn": 20, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 499, "answer": "The neighbor is renovating again, so it's loud. Let me know if you need help with anything else!", "seconds": 1.88}
+{"turn": 21, "kind": "fact", "user": "our train leaves from platform 58", "notes": [], "ledger": [], "prefix_tokens": 508, "answer": "Our train leaves from **Platform 58**. Let me know if you need help with anything else!", "seconds": 1.32, "key": "train platform", "value": "platform 58"}
+{"turn": 22, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 511, "answer": "They repaved our street. Let me know if you need help with anything else!", "seconds": 1.11}
+{"turn": 23, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 495, "answer": "Coffee is getting expensive. Let me know if you need help with anything else!", "seconds": 1.03}
+{"turn": 24, "kind": "fact", "user": "dentist moved me to 7:15pm", "notes": [], "ledger": [], "prefix_tokens": 495, "answer": "The dentist moved you to **7:15 PM**. Let me know if you need help with anything else!", "seconds": 1.34, "key": "dentist slot", "value": "7:15pm"}
+{"turn": 25, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 501, "answer": "Saw a great documentary yesterday! Let me know if you need help with anything else!", "seconds": 1.15}
+{"turn": 26, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 503, "answer": "Your phone battery is dying so fast now. Let me know if you need help with anything else!", "seconds": 1.16}
+{"turn": 27, "kind": "probe", "user": "quick — the dentist slot?", "notes": [], "ledger": [], "prefix_tokens": 497, "answer": "Quick — the **dentist slot**. Let me know if you need help with anything else!", "seconds": 1.33, "key": "dentist slot", "expected": "7:15pm", "distance": 3, "hit": false}
+{"turn": 28, "kind": "fact", "user": "doctor changed the dose to 27mg", "notes": [], "ledger": [], "prefix_tokens": 486, "answer": "The doctor changed the dose to **27 mg**. Let me know if you need help with anything else!", "seconds": 1.76, "key": "meds dose", "value": "27mg"}
+{"turn": 29, "kind": "distractor", "user": "might go for a walk later", "notes": [], "ledger": [], "prefix_tokens": 495, "answer": "You might go for a walk later. Let me know if you need help with anything else!", "seconds": 1.29}
+{"turn": 30, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 511, "answer": "Finally finished that book! Let me know if you need help with anything else!", "seconds": 1.22}
+{"turn": 31, "kind": "fact", "user": "the couch is order ZT6045, keep it handy", "notes": [], "ledger": [], "prefix_tokens": 508, "answer": "The couch is ordered as **ZT6045**, and I’ll keep it handy for you. Let me know if you need help with anything else!", "seconds": 2.29, "key": "order number", "value": "order ZT6045"}
+{"turn": 32, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 499, "answer": "What a week! Honestly, it's been a really fun and rewarding time! Let me know if you need help with anything else!", "seconds": 2.23}
+{"turn": 33, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 500, "answer": "Your back hurts from the gym. Let me know if you need help with anything else!", "seconds": 1.14}
+{"turn": 34, "kind": "correction", "user": "actually scratch that, the plumber quote is $4,000 now", "notes": [], "ledger": [], "prefix_tokens": 503, "answer": "The plumber quote is now **$4,000**. Let me know if you need help with anything else!", "seconds": 1.89, "key": "plumber quote", "value": "$4,000"}
+{"turn": 35, "kind": "fact", "user": "landlord wants rent by day 97 each month", "notes": [], "ledger": [], "prefix_tokens": 507, "answer": "The landlord wants your rent **by day 97** each month. Let me know if you need help with anything else!", "seconds": 2.02, "key": "rent due", "value": "day 97"}
+{"turn": 36, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 512, "answer": "The cat knocked over a plant. Let me know if you need help with anything else!", "seconds": 1.49}
+{"turn": 37, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 505, "answer": "Traffic was unreal this morning. Let me know if you need help with anything else!", "seconds": 1.27}
+{"turn": 38, "kind": "fact", "user": "we got room 411 at the hotel", "notes": [], "ledger": [], "prefix_tokens": 507, "answer": "We got room **411** at the hotel. Let me know if you need help with anything else!", "seconds": 1.32, "key": "hotel room", "value": "room 411"}
+{"turn": 39, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 488, "answer": "The neighbor is renovating again, and it's **loud**. Let me know if you need help with anything else!", "seconds": 1.55}
+{"turn": 40, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 486, "answer": "Coffee is getting expensive. Let me know if you need help with anything else!", "seconds": 1.16}
+{"turn": 41, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 506, "answer": "I love this weather lately too! It's really nice. Let me know if you need help with anything else!", "seconds": 1.83}
+{"turn": 42, "kind": "fact", "user": "boarding is at gate E17", "notes": [], "ledger": [], "prefix_tokens": 492, "answer": "Boarding is at **Gate E17**. Let me know if you need help with anything else!", "seconds": 1.55, "key": "boarding gate", "value": "gate E17"}
+{"turn": 43, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 510, "answer": "What a week! Honestly, it's been a really fun and rewarding time! Let me know if you need help with anything else!", "seconds": 1.93}
+{"turn": 44, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 505, "answer": "The cat knocked over a plant. Let me know if you need help with anything else!", "seconds": 1.28}
+{"turn": 45, "kind": "fact", "user": "my sister lands on flight AL7371, write that down", "notes": [], "ledger": [], "prefix_tokens": 504, "answer": "My sister lands on flight **AL7371**.", "seconds": 0.99, "key": "sister flight", "value": "flight AL7371"}
+{"turn": 46, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 509, "answer": "Coffee is getting expensive. Let me know if you need help with anything else!", "seconds": 1.02}
+{"turn": 47, "kind": "correction", "user": "actually scratch that, the sister flight is flight BL4905 now", "notes": [], "ledger": [], "prefix_tokens": 500, "answer": "Actually, the sister flight is now **flight BL4905**. Let me know if you need help with anything else!", "seconds": 1.98, "key": "sister flight", "value": "flight BL4905"}
+{"turn": 48, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 495, "answer": "Coffee is getting expensive. Let me know if you need help with anything else!", "seconds": 0.97}
+{"turn": 49, "kind": "fact", "user": "the rental has plate HJC-988", "notes": [], "ledger": [], "prefix_tokens": 497, "answer": "The rental has plate **HJC-988**. Let me know if you need help with anything else!", "seconds": 1.49, "key": "car plate", "value": "HJC-988"}
+{"turn": 50, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 505, "answer": "My back hurts from the gym. Let me know if you need help with anything else!", "seconds": 1.25}
+{"turn": 51, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 494, "answer": "Finally finished that book! Let me know if you need help with anything else!", "seconds": 1.12}
+{"turn": 52, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 511, "answer": "Coffee is getting expensive. Let me know if you need help with anything else!", "seconds": 1.06}
+{"turn": 53, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 511, "answer": "Work was chaos today. Let me know if you need help with anything else!", "seconds": 0.9}
+{"turn": 54, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 505, "answer": "The neighbor is renovating again and it's loud. Let me know if you need help with anything else!", "seconds": 1.37}
+{"turn": 55, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 506, "answer": "The neighbor is renovating again and it's loud. Let me know if you need help with anything else!", "seconds": 1.63}
+{"turn": 56, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 509, "answer": "Did you catch the game last night?", "seconds": 0.67}
+{"turn": 57, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 500, "answer": "Traffic was unreal this morning. Let me know if you need help with anything else!", "seconds": 1.16}
+{"turn": 58, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 498, "answer": "Traffic was unreal this morning. Let me know if you need help with anything else!", "seconds": 1.15}
+{"turn": 59, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 511, "answer": "What a week! Honestly, I'm so glad you're having a great time. Let me know if you need help with anything else!", "seconds": 1.65}
+{"turn": 60, "kind": "probe", "user": "wait, what was the wifi password again?", "notes": [], "ledger": [], "prefix_tokens": 509, "answer": "What was the **WiFi password again**? Let me know if you need help with anything else!", "seconds": 1.15, "key": "wifi password", "expected": "harbor858", "distance": 57, "hit": false}
+{"turn": 61, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 491, "answer": "Finally finished that book! Let me know if you need help with anything else!", "seconds": 1.06}
+{"turn": 62, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 512, "answer": "They repaved our street! Let me know if you need help with anything else!", "seconds": 1.09}
+{"turn": 63, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 508, "answer": "They repaved our street! Let me know if you need help with anything else!", "seconds": 0.95}
+{"turn": 64, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 512, "answer": "Saw a great documentary yesterday! Let me know if you need help with anything else!", "seconds": 1.11}
+{"turn": 65, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 511, "answer": "What a week! Honestly, I'm so glad you're having a great time. Let me know if you need help with anything else!", "seconds": 2.27}
+{"turn": 66, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 493, "answer": "Thinking of making **pasta tonight**, let me know if you'd like some ideas or help with the recipe! 😊", "seconds": 1.71}
+{"turn": 67, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 502, "answer": "Work was **chaos today**! Let me know if you'd like help with anything else! 😊", "seconds": 1.23}
+{"turn": 68, "kind": "distractor", "user": "might go for a walk later", "notes": [], "ledger": [], "prefix_tokens": 499, "answer": "Yes, I might go for a **walk later**! Let me know if you'd like help with anything else! 😊", "seconds": 1.76}
+{"turn": 69, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 503, "answer": "The neighbor is **renovating again**, so it's loud. Let me know if you'd like help with anything else! 😊", "seconds": 1.95}
+{"turn": 70, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 510, "answer": "They **repaved our street**! Let me know if you need help with anything else! 😊", "seconds": 1.55}
+{"turn": 71, "kind": "probe", "user": "wait, what was the plumber quote again?", "notes": [], "ledger": [], "prefix_tokens": 496, "answer": "What was the **plumber quote again**? Let me know if you'd like help with anything else! 😊", "seconds": 1.66, "key": "plumber quote", "expected": "$4,000", "distance": 37, "hit": false}
+{"turn": 72, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 492, "answer": "Traffic was **unreal this morning**! Let me know if you'd like help with anything else! 😊", "seconds": 1.76}
+{"turn": 73, "kind": "correction", "user": "actually scratch that, the parking spot is spot E16 now", "notes": [], "ledger": [], "prefix_tokens": 499, "answer": "The **parking spot is now E16**! Let me know if you need help with anything else! 😊", "seconds": 1.71, "key": "parking spot", "value": "spot E16"}
+{"turn": 74, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 509, "answer": "I love this **weather lately**! It's really nice out. Let me know if you need help with anything else! 😊", "seconds": 1.68}
+{"turn": 75, "kind": "distractor", "user": "craving something sweet", "notes": [], "ledger": [], "prefix_tokens": 508, "answer": "I'm so glad to hear that! I love the **sweet** things too! Let me know if you'd like help with anything else! 😊", "seconds": 1.83}
+{"turn": 76, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 502, "answer": "The **neighbor is renovating again**, so it's loud. Let me know if you'd like help with anything else! 😊", "seconds": 1.68}
+{"turn": 77, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 501, "answer": "Thinking of making **pasta tonight**, let me know if you'd like some ideas or help with the recipe! 😊", "seconds": 1.25}
+{"turn": 78, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 504, "answer": "I **did not** catch the game last night. Let me know if you'd like help with anything else! 😊", "seconds": 1.42}
+{"turn": 79, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 510, "answer": "My phone battery **dies so fast now**. It's definitely a common issue. Let me know if you'd like help with anything else! 😊", "seconds": 1.9}
+{"turn": 80, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 503, "answer": "My phone battery **dies so fast now**. It's definitely a common issue. Let me know if you'd like help with anything else! 😊", "seconds": 2.0}
+{"turn": 81, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 509, "answer": "My phone battery **dies so fast now**. It's definitely a common issue. Let me know if you'd like help with anything else! 😊", "seconds": 2.24}
+{"turn": 82, "kind": "probe", "user": "i forgot the rent due, what was it?", "notes": [], "ledger": [], "prefix_tokens": 492, "answer": "I **forgot the rent due**, but I don't know the exact amount. Let me know if you'd like help with anything else! 😊", "seconds": 2.07, "key": "rent due", "expected": "day 97", "distance": 47, "hit": false}
+{"turn": 83, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 499, "answer": "I **think of making pasta tonight**. Let me know if you'd like some ideas or help with the recipe! 😊", "seconds": 1.96}
+{"turn": 84, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 492, "answer": "The **coffee here is getting expensive**. It's definitely a good time to consider alternatives, like herbal or organic options. Let me know if you'd like help with anything else! 😊", "seconds": 3.02}
+{"turn": 85, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 504, "answer": "I **think of making pasta tonight**, and I'm really looking forward to it! Let me know if you'd like some ideas or help with the recipe! 😊", "seconds": 2.49}
+{"turn": 86, "kind": "correction", "user": "actually scratch that, the plumber quote is $1,400 now", "notes": [], "ledger": [], "prefix_tokens": 504, "answer": "The **plumber quote is now $1,400**. Let me know if you'd like help with anything else! 😊", "seconds": 2.08, "key": "plumber quote", "value": "$1,400"}
+{"turn": 87, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 508, "answer": "My phone battery **dies so fast now**. It's definitely a common issue. Let me know if you'd like help with anything else! 😊", "seconds": 2.27}
+{"turn": 88, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 512, "answer": "The **coffee here is getting expensive**. It's definitely a good time to consider alternatives like herbal or organic coffee. Let me know if you'd like help with anything else! 😊", "seconds": 2.85}
+{"turn": 89, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 485, "answer": "I **finally finished that book**. It was a great read! Let me know if you'd like help with anything else! 😊", "seconds": 1.95}
+{"turn": 90, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 480, "answer": "The **my back hurts from the gym**. It's definitely a common issue, especially after a workout. Let me know if you'd like help with anything else! 😊", "seconds": 2.01}
+{"turn": 91, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 485, "answer": "I **did not catch the game last night**. Let me know if you'd like help with anything else! 😊", "seconds": 1.56}
+{"turn": 92, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 509, "answer": "**Work was chaos today.** It's definitely stressful, but I'm glad you're able to get through it! Let me know if you'd like help with anything else! 😊", "seconds": 2.09}
+{"turn": 93, "kind": "probe", "user": "i forgot the deposit, what was it?", "notes": [], "ledger": [], "prefix_tokens": 490, "answer": "The **deposit was $50**. Let me know if you'd like help with anything else! 😊", "seconds": 1.48, "key": "deposit", "expected": "$6,400", "distance": 83, "hit": false}
+{"turn": 94, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 474, "answer": "The **coffee here is getting expensive**. It's definitely a good time to consider alternatives like herbal or organic coffee. Let me know if you'd like help with anything else! 😊", "seconds": 2.3}
+{"turn": 95, "kind": "distractor", "user": "craving something sweet", "notes": [], "ledger": [], "prefix_tokens": 477, "answer": "**Craving something sweet**. It's definitely a good time to enjoy a treat, like a sweet treat or a cup of coffee. Let me know if you'd like help with anything else! 😊", "seconds": 2.56}
+{"turn": 96, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 484, "answer": "The **neighbor is renovating again, so it's loud**. It's definitely not a good time to be around, especially if you're in the house. Let me know if you'd like help with anything else! 😊", "seconds": 2.65}
+{"turn": 97, "kind": "probe", "user": "remind me of the hotel room?", "notes": [], "ledger": [], "prefix_tokens": 500, "answer": "**Remind me of the hotel room?** Let me know if you'd like help with anything else! 😊", "seconds": 1.42, "key": "hotel room", "expected": "room 411", "distance": 59, "hit": false}
+{"turn": 98, "kind": "probe", "user": "quick — the hotel room?", "notes": [], "ledger": [], "prefix_tokens": 487, "answer": "**Quick — the hotel room?** Let me know if you'd like help with anything else! 😊", "seconds": 1.51, "key": "hotel room", "expected": "room 411", "distance": 60, "hit": false}
+{"turn": 99, "kind": "probe", "user": "i forgot the wifi password, what was it?", "notes": [], "ledger": [], "prefix_tokens": 487, "answer": "The **WiFi password** was **123456**. Let me know if you'd like help with anything else! 😊", "seconds": 2.25, "key": "wifi password", "expected": "harbor858", "distance": 96, "hit": false}

bench/results-qwen3-06b-nothink-full.jsonl

@@ -0,0 +1,100 @@
+{"turn": 0, "kind": "fact", "user": "the gym locker is code 4805", "notes": [], "ledger": [], "prefix_tokens": 43, "answer": "The code 4805 corresponds to a gym locker. If you have any questions about it, feel free to ask!", "seconds": 5.03, "key": "locker code", "value": "code 4805"}
+{"turn": 1, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 84, "answer": "They repaved your street. Let me know if you need help with anything else!", "seconds": 1.34}
+{"turn": 2, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 116, "answer": "That's a great way to say it! If you need help with anything else, feel free to ask!", "seconds": 1.4}
+{"turn": 3, "kind": "fact", "user": "cabin wifi password is harbor858", "notes": [], "ledger": [], "prefix_tokens": 157, "answer": "The cabin's WiFi password is **harbor858**. Let me know if you need help with anything else!", "seconds": 1.34, "key": "wifi password", "value": "harbor858"}
+{"turn": 4, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 197, "answer": "They repaved your street. Let me know if you need help with anything else!", "seconds": 0.94}
+{"turn": 5, "kind": "probe", "user": "quick — the locker code?", "notes": [], "ledger": [], "prefix_tokens": 230, "answer": "Yes, the locker code is **4805**. Let me know if you need help with anything else!", "seconds": 1.38, "key": "locker code", "expected": "code 4805", "distance": 5, "hit": false}
+{"turn": 6, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 268, "answer": "Finally, you finished that book! Let me know if you need help with anything else!", "seconds": 1.06}
+{"turn": 7, "kind": "fact", "user": "we always park at spot B19", "notes": [], "ledger": [], "prefix_tokens": 304, "answer": "Yes, you always park at **spot B19**. Let me know if you need help with anything else!", "seconds": 1.27, "key": "parking spot", "value": "spot B19"}
+{"turn": 8, "kind": "correction", "user": "actually scratch that, the locker code is code 1688 now", "notes": [], "ledger": [], "prefix_tokens": 353, "answer": "Actually, the locker code is now **code 1688**. Let me know if you need help with anything else!", "seconds": 1.43, "key": "locker code", "value": "code 1688"}
+{"turn": 9, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 396, "answer": "The cat knocked over a plant. Let me know if you need help with anything else!", "seconds": 0.92}
+{"turn": 10, "kind": "fact", "user": "the deposit for the venue was $6,400", "notes": [], "ledger": [], "prefix_tokens": 436, "answer": "The deposit for the venue was **$6,400**. Let me know if you need help with anything else!", "seconds": 1.4, "key": "deposit", "value": "$6,400"}
+{"turn": 11, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 479, "answer": "Your phone battery is dying so fast now. Let me know if you need help with anything else!", "seconds": 0.97}
+{"turn": 12, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 515, "answer": "Saw a great documentary yesterday. Let me know if you need help with anything else!", "seconds": 1.06}
+{"turn": 13, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 548, "answer": "I love this weather lately! Let me know if you need help with anything else!", "seconds": 1.26}
+{"turn": 14, "kind": "fact", "user": "tomas' new teacher is ms. duarte", "notes": [], "ledger": [], "prefix_tokens": 585, "answer": "Tomás' new teacher is Ms. Duarte. Let me know if you need help with anything else!", "seconds": 1.23, "key": "kid teacher", "value": "ms. duarte"}
+{"turn": 15, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 625, "answer": "Yes, I caught the game last night! Let me know if you need help with anything else!", "seconds": 1.23}
+{"turn": 16, "kind": "probe", "user": "remind me of the deposit?", "notes": [], "ledger": [], "prefix_tokens": 662, "answer": "Sure! Let me know if you need help with anything else!", "seconds": 1.04, "key": "deposit", "expected": "$6,400", "distance": 6, "hit": false}
+{"turn": 17, "kind": "fact", "user": "the plumber quoted $4,900 for the bathroom", "notes": [], "ledger": [], "prefix_tokens": 697, "answer": "The plumber quoted **$4,900** for the bathroom. Let me know if you need help with anything else!", "seconds": 1.64, "key": "plumber quote", "value": "$4,900"}
+{"turn": 18, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 740, "answer": "Your phone battery is dying so fast now. Let me know if you need help with anything else!", "seconds": 1.48}
+{"turn": 19, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 775, "answer": "Traffic was unreal this morning. Let me know if you need help with anything else!", "seconds": 1.01}
+{"turn": 20, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 811, "answer": "The neighbor is renovating again, which is causing loud noises. Let me know if you need help with anything else!", "seconds": 1.75}
+{"turn": 21, "kind": "fact", "user": "our train leaves from platform 58", "notes": [], "ledger": [], "prefix_tokens": 853, "answer": "Our train leaves from **platform 58**. Let me know if you need help with anything else!", "seconds": 1.58, "key": "train platform", "value": "platform 58"}
+{"turn": 22, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 890, "answer": "They repaved our street. Let me know if you need help with anything else!", "seconds": 0.92}
+{"turn": 23, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 922, "answer": "Coffee is getting expensive now. Let me know if you need help with anything else!", "seconds": 0.88}
+{"turn": 24, "kind": "fact", "user": "dentist moved me to 7:15pm", "notes": [], "ledger": [], "prefix_tokens": 960, "answer": "The dentist moved you to **7:15 PM**. Let me know if you need help with anything else!", "seconds": 1.26, "key": "dentist slot", "value": "7:15pm"}
+{"turn": 25, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 1000, "answer": "Saw a great documentary yesterday. Let me know if you need help with anything else!", "seconds": 0.95}
+{"turn": 26, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 1035, "answer": "Your phone battery is dying so fast now. Let me know if you need help with anything else!", "seconds": 1.02}
+{"turn": 27, "kind": "probe", "user": "quick — the dentist slot?", "notes": [], "ledger": [], "prefix_tokens": 1071, "answer": "Yes, the dentist slot is **4805**. Let me know if you need help with anything else!", "seconds": 1.44, "key": "dentist slot", "expected": "7:15pm", "distance": 3, "hit": false}
+{"turn": 28, "kind": "fact", "user": "doctor changed the dose to 27mg", "notes": [], "ledger": [], "prefix_tokens": 1114, "answer": "The doctor changed the dose to **27 mg**. Let me know if you need help with anything else!", "seconds": 1.33, "key": "meds dose", "value": "27mg"}
+{"turn": 29, "kind": "distractor", "user": "might go for a walk later", "notes": [], "ledger": [], "prefix_tokens": 1153, "answer": "You might go for a walk later. Let me know if you need help with anything else!", "seconds": 0.93}
+{"turn": 30, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 1186, "answer": "Finally finished that book! Let me know if you need help with anything else!", "seconds": 0.87}
+{"turn": 31, "kind": "fact", "user": "the couch is order ZT6045, keep it handy", "notes": [], "ledger": [], "prefix_tokens": 1226, "answer": "The couch is ordered as **ZT6045**, and I recommend keeping it handy. Let me know if you need help with anything else!", "seconds": 2.16, "key": "order number", "value": "order ZT6045"}
+{"turn": 32, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 1272, "answer": "What a week! Honestly, that's a great way to say it. Let me know if you need help with anything else!", "seconds": 1.79}
+{"turn": 33, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 1314, "answer": "Your back hurts from the gym. Let me know if you need help with anything else!", "seconds": 1.2}
+{"turn": 34, "kind": "correction", "user": "actually scratch that, the plumber quote is $4,000 now", "notes": [], "ledger": [], "prefix_tokens": 1357, "answer": "Actually, the plumber quoted **$4,000** now. Let me know if you need help with anything else!", "seconds": 1.83, "key": "plumber quote", "value": "$4,000"}
+{"turn": 35, "kind": "fact", "user": "landlord wants rent by day 97 each month", "notes": [], "ledger": [], "prefix_tokens": 1404, "answer": "The landlord wants rent to be paid **by day 97** each month. Let me know if you need help with anything else!", "seconds": 2.28, "key": "rent due", "value": "day 97"}
+{"turn": 36, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 1448, "answer": "The cat knocked over a plant. Let me know if you need help with anything else!", "seconds": 1.2}
+{"turn": 37, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 1481, "answer": "Traffic was unreal this morning. Let me know if you need help with anything else!", "seconds": 0.89}
+{"turn": 38, "kind": "fact", "user": "we got room 411 at the hotel", "notes": [], "ledger": [], "prefix_tokens": 1518, "answer": "We got room **411** at the hotel. Let me know if you need help with anything else!", "seconds": 1.31, "key": "hotel room", "value": "room 411"}
+{"turn": 39, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 1560, "answer": "The neighbor is renovating again, which is causing loud noises. Let me know if you need help with anything else!", "seconds": 1.35}
+{"turn": 40, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 1599, "answer": "Coffee is getting expensive now. Let me know if you need help with anything else!", "seconds": 1.18}
+{"turn": 41, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 1631, "answer": "I love this weather lately! Let me know if you need help with anything else!", "seconds": 0.87}
+{"turn": 42, "kind": "fact", "user": "boarding is at gate E17", "notes": [], "ledger": [], "prefix_tokens": 1665, "answer": "Boarding is at **Gate E17**. Let me know if you need help with anything else!", "seconds": 1.11, "key": "boarding gate", "value": "gate E17"}
+{"turn": 43, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 1702, "answer": "What a week! Honestly, that's a great way to say it. Let me know if you need help with anything else!", "seconds": 1.26}
+{"turn": 44, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 1744, "answer": "The cat knocked over a plant. Let me know if you need help with anything else!", "seconds": 0.93}
+{"turn": 45, "kind": "fact", "user": "my sister lands on flight AL7371, write that down", "notes": [], "ledger": [], "prefix_tokens": 1786, "answer": "My sister lands on flight **AL7371**. Let me know if you need help with anything else!", "seconds": 1.45, "key": "sister flight", "value": "flight AL7371"}
+{"turn": 46, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 1825, "answer": "Coffee is getting expensive now. Let me know if you need help with anything else!", "seconds": 1.21}
+{"turn": 47, "kind": "correction", "user": "actually scratch that, the sister flight is flight BL4905 now", "notes": [], "ledger": [], "prefix_tokens": 1867, "answer": "Actually, the sister flight is now **flight BL4905**. Let me know if you need help with anything else!", "seconds": 1.67, "key": "sister flight", "value": "flight BL4905"}
+{"turn": 48, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 1909, "answer": "Coffee is getting expensive now. Let me know if you need help with anything else!", "seconds": 1.23}
+{"turn": 49, "kind": "fact", "user": "the rental has plate HJC-988", "notes": [], "ledger": [], "prefix_tokens": 1946, "answer": "The rental has plate **HJC-988**. Let me know if you need help with anything else!", "seconds": 1.21, "key": "car plate", "value": "HJC-988"}
+{"turn": 50, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 1986, "answer": "Your back hurts from the gym. Let me know if you need help with anything else!", "seconds": 1.11}
+{"turn": 51, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 2018, "answer": "Finally finished that book! Let me know if you need help with anything else!", "seconds": 1.17}
+{"turn": 52, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 2049, "answer": "Coffee is getting expensive now. Let me know if you need help with anything else!", "seconds": 0.9}
+{"turn": 53, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 2080, "answer": "Work was chaos today. Let me know if you need help with anything else!", "seconds": 0.99}
+{"turn": 54, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 2115, "answer": "The neighbor is renovating again, which is causing loud noises. Let me know if you need help with anything else!", "seconds": 1.43}
+{"turn": 55, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 2158, "answer": "The neighbor is renovating again, which is causing loud noises. Let me know if you need help with anything else!", "seconds": 1.25}
+{"turn": 56, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 2200, "answer": "Yes, I caught the game last night! Let me know if you need help with anything else!", "seconds": 1.18}
+{"turn": 57, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 2235, "answer": "Traffic was unreal this morning. Let me know if you need help with anything else!", "seconds": 1.27}
+{"turn": 58, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 2267, "answer": "Traffic was unreal this morning. Let me know if you need help with anything else!", "seconds": 1.12}
+{"turn": 59, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 2299, "answer": "What a week! Honestly, that's a great way to say it. Let me know if you need help with anything else!", "seconds": 1.55}
+{"turn": 60, "kind": "probe", "user": "wait, what was the wifi password again?", "notes": [], "ledger": [], "prefix_tokens": 2344, "answer": "What was the WiFi password again? Let me know if you need help with anything else!", "seconds": 1.05, "key": "wifi password", "expected": "harbor858", "distance": 57, "hit": false}
+{"turn": 61, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 2376, "answer": "Finally finished that book! Let me know if you need help with anything else!", "seconds": 0.88}
+{"turn": 62, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 2407, "answer": "They repaved our street. Let me know if you need help with anything else!", "seconds": 1.21}
+{"turn": 63, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 2439, "answer": "They repaved our street. Let me know if you need help with anything else!", "seconds": 1.26}
+{"turn": 64, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 2472, "answer": "Saw a great documentary yesterday. Let me know if you need help with anything else!", "seconds": 1.37}
+{"turn": 65, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 2505, "answer": "What a week! Honestly, that's a great way to say it. Let me know if you need help with anything else!", "seconds": 1.82}
+{"turn": 66, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 2546, "answer": "Thinking of making pasta tonight! Let me know if you need help with anything else!", "seconds": 1.04}
+{"turn": 67, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 2577, "answer": "Work was chaotic today. Let me know if you need help with anything else!", "seconds": 0.86}
+{"turn": 68, "kind": "distractor", "user": "might go for a walk later", "notes": [], "ledger": [], "prefix_tokens": 2609, "answer": "You might go for a walk later. Let me know if you need help with anything else!", "seconds": 1.2}
+{"turn": 69, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 2647, "answer": "The neighbor is renovating again, which is causing loud noises. Let me know if you need help with anything else!", "seconds": 1.69}
+{"turn": 70, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 2686, "answer": "They repaved our street. Let me know if you need help with anything else!", "seconds": 1.28}
+{"turn": 71, "kind": "probe", "user": "wait, what was the plumber quote again?", "notes": [], "ledger": [], "prefix_tokens": 2722, "answer": "What was the plumber quote again? Let me know if you need help with anything else!", "seconds": 1.21, "key": "plumber quote", "expected": "$4,000", "distance": 37, "hit": false}
+{"turn": 72, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 2755, "answer": "Traffic was unreal this morning. Let me know if you need help with anything else!", "seconds": 1.26}
+{"turn": 73, "kind": "correction", "user": "actually scratch that, the parking spot is spot E16 now", "notes": [], "ledger": [], "prefix_tokens": 2795, "answer": "Actually, the parking spot is now **spot E16**. Let me know if you need help with anything else!", "seconds": 1.43, "key": "parking spot", "value": "spot E16"}
+{"turn": 74, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 2835, "answer": "I love this weather lately! Let me know if you need help with anything else!", "seconds": 1.14}
+{"turn": 75, "kind": "distractor", "user": "craving something sweet", "notes": [], "ledger": [], "prefix_tokens": 2866, "answer": "Craving something sweet! Let me know if you need help with anything else!", "seconds": 0.99}
+{"turn": 76, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 2901, "answer": "The neighbor is renovating again, which is causing loud noises. Let me know if you need help with anything else!", "seconds": 1.44}
+{"turn": 77, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 2940, "answer": "Thinking of making pasta tonight! Let me know if you need help with anything else!", "seconds": 1.07}
+{"turn": 78, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 2975, "answer": "Yes, I caught the game last night! Let me know if you need help with anything else!", "seconds": 1.17}
+{"turn": 79, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 3012, "answer": "Your phone battery is dying so fast now. Let me know if you need help with anything else!", "seconds": 1.43}
+{"turn": 80, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 3049, "answer": "Your phone battery is dying so fast now. Let me know if you need help with anything else!", "seconds": 1.26}
+{"turn": 81, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 3086, "answer": "Your phone battery is dying so fast now. Let me know if you need help with anything else!", "seconds": 1.32}
+{"turn": 82, "kind": "probe", "user": "i forgot the rent due, what was it?", "notes": [], "ledger": [], "prefix_tokens": 3126, "answer": "I forgot the rent due. Let me know if you need help with anything else!", "seconds": 1.3, "key": "rent due", "expected": "day 97", "distance": 47, "hit": false}
+{"turn": 83, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 3158, "answer": "Thinking of making pasta tonight! Let me know if you need help with anything else!", "seconds": 1.6}
+{"turn": 84, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 3190, "answer": "Coffee is getting expensive now. Let me know if you need help with anything else!", "seconds": 1.44}
+{"turn": 85, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 3222, "answer": "Thinking of making pasta tonight! Let me know if you need help with anything else!", "seconds": 1.27}
+{"turn": 86, "kind": "correction", "user": "actually scratch that, the plumber quote is $1,400 now", "notes": [], "ledger": [], "prefix_tokens": 3264, "answer": "Actually, the plumber quote is now **$1,400**. Let me know if you need help with anything else!", "seconds": 1.58, "key": "plumber quote", "value": "$1,400"}
+{"turn": 87, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 3308, "answer": "Your phone battery is dying so fast now. Let me know if you need help with anything else!", "seconds": 1.2}
+{"turn": 88, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 3343, "answer": "Coffee is getting expensive now. Let me know if you need help with anything else!", "seconds": 1.13}
+{"turn": 89, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 3374, "answer": "Finally finished that book! Let me know if you need help with anything else!", "seconds": 0.95}
+{"turn": 90, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 3406, "answer": "Your back hurts from the gym. Let me know if you need help with anything else!", "seconds": 0.98}
+{"turn": 91, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 3442, "answer": "Yes, I caught the game last night! Let me know if you need help with anything else!", "seconds": 1.08}
+{"turn": 92, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 3476, "answer": "Work was chaotic today. Let me know if you need help with anything else!", "seconds": 0.9}
+{"turn": 93, "kind": "probe", "user": "i forgot the deposit, what was it?", "notes": [], "ledger": [], "prefix_tokens": 3511, "answer": "I forgot the deposit. Let me know if you need help with anything else!", "seconds": 1.29, "key": "deposit", "expected": "$6,400", "distance": 83, "hit": false}
+{"turn": 94, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 3542, "answer": "Coffee is getting expensive now. Let me know if you need help with anything else!", "seconds": 1.42}
+{"turn": 95, "kind": "distractor", "user": "craving something sweet", "notes": [], "ledger": [], "prefix_tokens": 3573, "answer": "Craving something sweet! Let me know if you need help with anything else!", "seconds": 1.07}
+{"turn": 96, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 3608, "answer": "The neighbor is renovating again, which is causing loud noises. Let me know if you need help with anything else!", "seconds": 1.34}
+{"turn": 97, "kind": "probe", "user": "remind me of the hotel room?", "notes": [], "ledger": [], "prefix_tokens": 3650, "answer": "Remind me of the hotel room? Let me know if you need help with anything else!", "seconds": 1.47, "key": "hotel room", "expected": "room 411", "distance": 59, "hit": false}
+{"turn": 98, "kind": "probe", "user": "quick — the hotel room?", "notes": [], "ledger": [], "prefix_tokens": 3685, "answer": "Quick — the hotel room? Let me know if you need help with anything else!", "seconds": 1.2, "key": "hotel room", "expected": "room 411", "distance": 60, "hit": false}
+{"turn": 99, "kind": "probe", "user": "i forgot the wifi password, what was it?", "notes": [], "ledger": [], "prefix_tokens": 3722, "answer": "I forgot the WiFi password. Let me know if you need help with anything else!", "seconds": 1.54, "key": "wifi password", "expected": "harbor858", "distance": 96, "hit": false}

bench/results-qwen3-06b-thinking-full.jsonl

@@ -0,0 +1,100 @@
+{"turn": 0, "kind": "fact", "user": "the gym locker is code 4805", "notes": [], "ledger": [], "prefix_tokens": 59, "answer": "the gym locker is code 4805", "seconds": 7.86, "key": "locker code", "value": "code 4805"}
+{"turn": 1, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 84, "answer": "The street was repaved.", "seconds": 8.01}
+{"turn": 2, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 105, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to?", "seconds": 9.42}
+{"turn": 3, "kind": "fact", "user": "cabin wifi password is harbor858", "notes": [], "ledger": [], "prefix_tokens": 149, "answer": "The cabin's WiFi password is **harbor858**.", "seconds": 6.81, "key": "wifi password", "value": "harbor858"}
+{"turn": 4, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 178, "answer": "", "seconds": 36.7}
+{"turn": 5, "kind": "probe", "user": "quick — the locker code?", "notes": [], "ledger": [], "prefix_tokens": 194, "answer": "The locker code is **4805**.", "seconds": 8.86, "key": "locker code", "expected": "code 4805", "distance": 5, "hit": true}
+{"turn": 6, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 219, "answer": "The locker code is **4805**. Let me know if you need anything else!", "seconds": 13.18}
+{"turn": 7, "kind": "fact", "user": "we always park at spot B19", "notes": [], "ledger": [], "prefix_tokens": 257, "answer": "The parking spot is **B19**.", "seconds": 5.16, "key": "parking spot", "value": "spot B19"}
+{"turn": 8, "kind": "correction", "user": "actually scratch that, the locker code is code 1688 now", "notes": [], "ledger": [], "prefix_tokens": 292, "answer": "The locker code is now **1688**. Let me know if you need anything else!", "seconds": 13.07, "key": "locker code", "value": "code 1688"}
+{"turn": 9, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 329, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to?", "seconds": 17.58}
+{"turn": 10, "kind": "fact", "user": "the deposit for the venue was $6,400", "notes": [], "ledger": [], "prefix_tokens": 376, "answer": "The deposit for the venue was **$6,400**.", "seconds": 7.27, "key": "deposit", "value": "$6,400"}
+{"turn": 11, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 408, "answer": "The locker code is **4805**. The deposit for the venue was **$6,400**. Let me know if you need anything else!", "seconds": 25.67}
+{"turn": 12, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 459, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to?", "seconds": 11.36}
+{"turn": 13, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 499, "answer": "I'm glad you're enjoying the weather! Let me know if there's anything specific you'd like to ask or need help with! 😊", "seconds": 13.31}
+{"turn": 14, "kind": "fact", "user": "tomas' new teacher is ms. duarte", "notes": [], "ledger": [], "prefix_tokens": 548, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 11.61, "key": "kid teacher", "value": "ms. duarte"}
+{"turn": 15, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 593, "answer": "I don't have information about the game. Could you please provide the specific code or amount you're referring to? 😊", "seconds": 23.6}
+{"turn": 16, "kind": "probe", "user": "remind me of the deposit?", "notes": [], "ledger": [], "prefix_tokens": 635, "answer": "The deposit for the venue was **$6,400**. Let me know if you need anything else! 😊", "seconds": 9.25, "key": "deposit", "expected": "$6,400", "distance": 6, "hit": true}
+{"turn": 17, "kind": "fact", "user": "the plumber quoted $4,900 for the bathroom", "notes": [], "ledger": [], "prefix_tokens": 683, "answer": "The plumber quoted **$4,900** for the bathroom. Let me know if you need anything else! 😊", "seconds": 8.9, "key": "plumber quote", "value": "$4,900"}
+{"turn": 18, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 726, "answer": "The locker code is **4805**. Let me know if you need anything else! 😊", "seconds": 23.26}
+{"turn": 19, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 763, "answer": "The traffic was unreal this morning. Let me know if there's anything specific you'd like to ask or need help with! 😊", "seconds": 12.69}
+{"turn": 20, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 809, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 18.35}
+{"turn": 21, "kind": "fact", "user": "our train leaves from platform 58", "notes": [], "ledger": [], "prefix_tokens": 854, "answer": "The train leaves from **platform 58**. Let me know if you need anything else! 😊", "seconds": 11.41, "key": "train platform", "value": "platform 58"}
+{"turn": 22, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 891, "answer": "The street was repaved. The locker code is **4805**, and the deposit for the venue was **$6,400**. Let me know if you need anything else! 😊", "seconds": 13.2}
+{"turn": 23, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 949, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 10.33}
+{"turn": 24, "kind": "fact", "user": "dentist moved me to 7:15pm", "notes": [], "ledger": [], "prefix_tokens": 997, "answer": "The dentist moved you to **7:15 PM**. The locker code is **4805**, and the deposit for the venue was **$6,400**. Let me know if you need anything else! 😊", "seconds": 11.4, "key": "dentist slot", "value": "7:15pm"}
+{"turn": 25, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 1063, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 9.5}
+{"turn": 26, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 1107, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 10.51}
+{"turn": 27, "kind": "probe", "user": "quick — the dentist slot?", "notes": [], "ledger": [], "prefix_tokens": 1150, "answer": "The dentist slot was moved to **7:15 PM**. The locker code is **4805**, and the deposit for the venue was **$6,400**. Let me know if you need anything else! 😊", "seconds": 13.88, "key": "dentist slot", "expected": "7:15pm", "distance": 3, "hit": true}
+{"turn": 28, "kind": "fact", "user": "doctor changed the dose to 27mg", "notes": [], "ledger": [], "prefix_tokens": 1220, "answer": "The doctor changed the dose to **27 mg**. The locker code is **4805**, and the deposit for the venue was **$6,400**. Let me know if you need anything else! 😊", "seconds": 11.3, "key": "meds dose", "value": "27mg"}
+{"turn": 29, "kind": "distractor", "user": "might go for a walk later", "notes": [], "ledger": [], "prefix_tokens": 1285, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 16.49}
+{"turn": 30, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 1326, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 12.75}
+{"turn": 31, "kind": "fact", "user": "the couch is order ZT6045, keep it handy", "notes": [], "ledger": [], "prefix_tokens": 1377, "answer": "The couch is ordered as **ZT6045**. Let me know if you need anything else! 😊", "seconds": 7.88, "key": "order number", "value": "order ZT6045"}
+{"turn": 32, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 1417, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 11.78}
+{"turn": 33, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 1460, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 11.96}
+{"turn": 34, "kind": "correction", "user": "actually scratch that, the plumber quote is $4,000 now", "notes": [], "ledger": [], "prefix_tokens": 1512, "answer": "The plumber's quote is now **$4,000**. The locker code is **4805**, and the deposit for the venue was **$6,400**. Let me know if you need anything else! 😊", "seconds": 23.26, "key": "plumber quote", "value": "$4,000"}
+{"turn": 35, "kind": "fact", "user": "landlord wants rent by day 97 each month", "notes": [], "ledger": [], "prefix_tokens": 1585, "answer": "The landlord wants rent to be paid by **day 97** each month. Let me know if you need further assistance! 😊", "seconds": 14.17, "key": "rent due", "value": "day 97"}
+{"turn": 36, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 1629, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 15.85}
+{"turn": 37, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 1671, "answer": "The traffic was unreal this morning. Let me know if there's anything specific you'd like to ask or need help with! 😊", "seconds": 11.08}
+{"turn": 38, "kind": "fact", "user": "we got room 411 at the hotel", "notes": [], "ledger": [], "prefix_tokens": 1718, "answer": "The room number is **411**. Let me know if you need anything else! 😊", "seconds": 11.08, "key": "hotel room", "value": "room 411"}
+{"turn": 39, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 1758, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 10.53}
+{"turn": 40, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 1800, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 10.66}
+{"turn": 41, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 1842, "answer": "I'm glad you're enjoying the weather! Let me know if there's anything specific you'd like to ask or need help with! 😊", "seconds": 14.22}
+{"turn": 42, "kind": "fact", "user": "boarding is at gate E17", "notes": [], "ledger": [], "prefix_tokens": 1888, "answer": "The boarding is at **Gate E17**. Let me know if you need anything else! 😊", "seconds": 11.3, "key": "boarding gate", "value": "gate E17"}
+{"turn": 43, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 1925, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 11.17}
+{"turn": 44, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 1968, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 9.75}
+{"turn": 45, "kind": "fact", "user": "my sister lands on flight AL7371, write that down", "notes": [], "ledger": [], "prefix_tokens": 2019, "answer": "The flight is AL7371.", "seconds": 7.96, "key": "sister flight", "value": "flight AL7371"}
+{"turn": 46, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 2043, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 10.39}
+{"turn": 47, "kind": "correction", "user": "actually scratch that, the sister flight is flight BL4905 now", "notes": [], "ledger": [], "prefix_tokens": 2095, "answer": "The flight is now **BL4905**. Let me know if you need anything else! 😊", "seconds": 8.43, "key": "sister flight", "value": "flight BL4905"}
+{"turn": 48, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 2133, "answer": "I'm sorry, but I don't have a code or amount to reference. Could you please provide the specific code or amount you're referring to? 😊", "seconds": 11.7}
+{"turn": 49, "kind": "fact", "user": "the rental has plate HJC-988", "notes": [], "ledger": [], "prefix_tokens": 2185, "answer": "The rental has a plate number of **HJC-988**. Let me know if you need anything else! 😊", "seconds": 8.05, "key": "car plate", "value": "HJC-988"}
+{"turn": 50, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 2228, "answer": "I'm sorry, but I don't have a code or amount to reference. Could you please provide the specific code or amount you're referring to? 😊", "seconds": 11.07}
+{"turn": 51, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 2274, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 10.79}
+{"turn": 52, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 2316, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 10.85}
+{"turn": 53, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 2357, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 9.32}
+{"turn": 54, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 2403, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 11.44}
+{"turn": 55, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 2449, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 11.89}
+{"turn": 56, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 2494, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 13.07}
+{"turn": 57, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 2536, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 15.01}
+{"turn": 58, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 2578, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 14.12}
+{"turn": 59, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 2620, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 18.2}
+{"turn": 60, "kind": "probe", "user": "wait, what was the wifi password again?", "notes": [], "ledger": [], "prefix_tokens": 2666, "answer": "The cabin's WiFi password is **harbor858**. Let me know if you need anything else! 😊", "seconds": 18.89, "key": "wifi password", "expected": "harbor858", "distance": 57, "hit": true}
+{"turn": 61, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 2705, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 19.44}
+{"turn": 62, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 2747, "answer": "The street was repaved. The locker code is **4805**, and the deposit for the venue was **$6,400**. Let me know if you need anything else! 😊", "seconds": 16.64}
+{"turn": 63, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 2805, "answer": "The street was repaved. The locker code is **4805**, and the deposit for the venue was **$6,400**. Let me know if you need anything else! 😊", "seconds": 22.18}
+{"turn": 64, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 2864, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 12.83}
+{"turn": 65, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 2906, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 17.67}
+{"turn": 66, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 2948, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 11.21}
+{"turn": 67, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 2989, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 13.41}
+{"turn": 68, "kind": "distractor", "user": "might go for a walk later", "notes": [], "ledger": [], "prefix_tokens": 3032, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 12.75}
+{"turn": 69, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 3078, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 10.96}
+{"turn": 70, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 3120, "answer": "The street was repaved. The locker code is **4805**, and the deposit for the venue was **$6,400**. Let me know if you need anything else! 😊", "seconds": 12.0}
+{"turn": 71, "kind": "probe", "user": "wait, what was the plumber quote again?", "notes": [], "ledger": [], "prefix_tokens": 3182, "answer": "The plumber's quote is now **$4,000**. Let me know if you need anything else! 😊", "seconds": 15.18, "key": "plumber quote", "expected": "$4,000", "distance": 37, "hit": true}
+{"turn": 72, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 3223, "answer": "The traffic was unreal this morning. Let me know if there's anything specific you'd like to ask or need help with! 😊", "seconds": 11.76}
+{"turn": 73, "kind": "correction", "user": "actually scratch that, the parking spot is spot E16 now", "notes": [], "ledger": [], "prefix_tokens": 3273, "answer": "The parking spot is now **E16**. Let me know if you need anything else! 😊", "seconds": 10.55, "key": "parking spot", "value": "spot E16"}
+{"turn": 74, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 3310, "answer": "I'm glad you're enjoying the weather! Let me know if there's anything specific you'd like to ask or need help with! 😊", "seconds": 10.43}
+{"turn": 75, "kind": "distractor", "user": "craving something sweet", "notes": [], "ledger": [], "prefix_tokens": 3353, "answer": "I'm glad you're enjoying the weather! Let me know if there's anything sweet you'd like to try or need help with! 😊", "seconds": 12.69}
+{"turn": 76, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 3401, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 9.23}
+{"turn": 77, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 3443, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 14.04}
+{"turn": 78, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 3488, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 11.2}
+{"turn": 79, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 3532, "answer": "The locker code is **4805**, and the deposit for the venue was **$6,400**. Let me know if you need anything else! 😊", "seconds": 13.1}
+{"turn": 80, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 3586, "answer": "The locker code is **4805**, and the deposit for the venue was **$6,400**. Let me know if you need anything else! 😊", "seconds": 9.34}
+{"turn": 81, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 3640, "answer": "The locker code is **4805**, and the deposit for the venue was **$6,400**. Let me know if you need anything else! 😊", "seconds": 11.13}
+{"turn": 82, "kind": "probe", "user": "i forgot the rent due, what was it?", "notes": [], "ledger": [], "prefix_tokens": 3697, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to?", "seconds": 18.73, "key": "rent due", "expected": "day 97", "distance": 47, "hit": false}
+{"turn": 83, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 3737, "answer": "The locker code is **4805**, and the deposit for the venue was **$6,400**. Let me know if you need anything else! 😊", "seconds": 10.29}
+{"turn": 84, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 3789, "answer": "I'm sorry, but I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 11.15}
+{"turn": 85, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 3836, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 12.65}
+{"turn": 86, "kind": "correction", "user": "actually scratch that, the plumber quote is $1,400 now", "notes": [], "ledger": [], "prefix_tokens": 3888, "answer": "The plumber's quote is now **$1,400**. Let me know if you need anything else! 😊", "seconds": 8.2, "key": "plumber quote", "value": "$1,400"}
+{"turn": 87, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 3931, "answer": "I don't have a code or", "seconds": 38.62}
+{"turn": 88, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 3953, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 8.41}
+{"turn": 89, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 3994, "answer": "The locker code is **4805**, and the deposit for the venue was **$6,400**. Let me know if you need anything else! 😊", "seconds": 15.8}
+{"turn": 90, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 4047, "answer": "The locker code is **4805**, and the deposit for the venue was **$6,400**. Let me know if you need anything else! 😊", "seconds": 10.16}
+{"turn": 91, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 4102, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 11.88}
+{"turn": 92, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 4143, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 10.39}
+{"turn": 93, "kind": "probe", "user": "i forgot the deposit, what was it?", "notes": [], "ledger": [], "prefix_tokens": 4189, "answer": "The deposit for the venue was **$6,400**. Let me know if you need anything else! 😊", "seconds": 10.88, "key": "deposit", "expected": "$6,400", "distance": 83, "hit": true}
+{"turn": 94, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 4230, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 10.12}
+{"turn": 95, "kind": "distractor", "user": "craving something sweet", "notes": [], "ledger": [], "prefix_tokens": 4271, "answer": "I'm glad you're enjoying the weather! Let me know if there's anything sweet you'd like to try or need help with! 😊", "seconds": 13.71}
+{"turn": 96, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 4319, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 17.57}
+{"turn": 97, "kind": "probe", "user": "remind me of the hotel room?", "notes": [], "ledger": [], "prefix_tokens": 4364, "answer": "I don't have a code or amount to reference. Could you please provide the specific code and amount you're referring to? 😊", "seconds": 13.63, "key": "hotel room", "expected": "room 411", "distance": 59, "hit": false}
+{"turn": 98, "kind": "probe", "user": "quick — the hotel room?", "notes": [], "ledger": [], "prefix_tokens": 4407, "answer": "The hotel room code is **4805**. Let me know if you need anything else! 😊", "seconds": 29.88, "key": "hotel room", "expected": "room 411", "distance": 60, "hit": false}
+{"turn": 99, "kind": "probe", "user": "i forgot the wifi password, what was it?", "notes": [], "ledger": [], "prefix_tokens": 4450, "answer": "The WiFi password is **harbor858**. Let me know if you need help with anything else! 😊", "seconds": 9.93, "key": "wifi password", "expected": "harbor858", "distance": 96, "hit": true}

bench/results-smollm2-360m-full.jsonl

@@ -0,0 +1,100 @@
+{"turn": 0, "kind": "fact", "user": "the gym locker is code 4805", "notes": [], "ledger": [], "prefix_tokens": 61, "answer": "The gym locker is code 4805.", "seconds": 1.94, "key": "locker code", "value": "code 4805"}
+{"turn": 1, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 89, "answer": "They repaved our street.", "seconds": 0.4}
+{"turn": 2, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 111, "answer": "What a week, honestly.", "seconds": 0.39}
+{"turn": 3, "kind": "fact", "user": "cabin wifi password is harbor858", "notes": [], "ledger": [], "prefix_tokens": 139, "answer": "Cabin WiFi Password: Harbor 858.", "seconds": 0.62, "key": "wifi password", "value": "harbor858"}
+{"turn": 4, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 167, "answer": "They repaved our street.", "seconds": 0.3}
+{"turn": 5, "kind": "probe", "user": "quick — the locker code?", "notes": [], "ledger": [], "prefix_tokens": 191, "answer": "The locker code is 4805.", "seconds": 0.74, "key": "locker code", "expected": "code 4805", "distance": 5, "hit": true}
+{"turn": 6, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 218, "answer": "Finished that book.", "seconds": 0.59}
+{"turn": 7, "kind": "fact", "user": "we always park at spot B19", "notes": [], "ledger": [], "prefix_tokens": 242, "answer": "We always park at spot B19.", "seconds": 0.54, "key": "parking spot", "value": "spot B19"}
+{"turn": 8, "kind": "correction", "user": "actually scratch that, the locker code is code 1688 now", "notes": [], "ledger": [], "prefix_tokens": 278, "answer": "The locker code is now code 1688.", "seconds": 0.99, "key": "locker code", "value": "code 1688"}
+{"turn": 9, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 308, "answer": "The cat knocked over a plant.", "seconds": 0.6}
+{"turn": 10, "kind": "fact", "user": "the deposit for the venue was $6,400", "notes": [], "ledger": [], "prefix_tokens": 338, "answer": "The deposit for the venue was $6,400.", "seconds": 0.97, "key": "deposit", "value": "$6,400"}
+{"turn": 11, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 369, "answer": "My phone battery dies so fast now.", "seconds": 0.78}
+{"turn": 12, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 394, "answer": "Saw a great documentary yesterday.", "seconds": 0.64}
+{"turn": 13, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 417, "answer": "I love this weather lately.", "seconds": 0.44}
+{"turn": 14, "kind": "fact", "user": "tomas' new teacher is ms. duarte", "notes": [], "ledger": [], "prefix_tokens": 444, "answer": "Tomas' new teacher is Ms. Duarte.", "seconds": 0.56, "key": "kid teacher", "value": "ms. duarte"}
+{"turn": 15, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 474, "answer": "I caught the game last night.", "seconds": 0.39}
+{"turn": 16, "kind": "probe", "user": "remind me of the deposit?", "notes": [], "ledger": [], "prefix_tokens": 499, "answer": "The deposit is $6,400.", "seconds": 0.72, "key": "deposit", "expected": "$6,400", "distance": 6, "hit": true}
+{"turn": 17, "kind": "fact", "user": "the plumber quoted $4,900 for the bathroom", "notes": [], "ledger": [], "prefix_tokens": 533, "answer": "The plumber quoted $4,900 for the bathroom.", "seconds": 1.03, "key": "plumber quote", "value": "$4,900"}
+{"turn": 18, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 565, "answer": "My phone battery dies so fast now.", "seconds": 0.64}
+{"turn": 19, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 589, "answer": "Traffic was unreal this morning.", "seconds": 0.61}
+{"turn": 20, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 616, "answer": "The neighbor is renovating again, so loud.", "seconds": 0.75}
+{"turn": 21, "kind": "fact", "user": "our train leaves from platform 58", "notes": [], "ledger": [], "prefix_tokens": 645, "answer": "Our train leaves from platform 58.", "seconds": 0.63, "key": "train platform", "value": "platform 58"}
+{"turn": 22, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 670, "answer": "They repaved our street.", "seconds": 0.34}
+{"turn": 23, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 694, "answer": "Coffee here is getting expensive.", "seconds": 0.36}
+{"turn": 24, "kind": "fact", "user": "dentist moved me to 7:15pm", "notes": [], "ledger": [], "prefix_tokens": 723, "answer": "The dentist moved me to 7:15pm.", "seconds": 0.65, "key": "dentist slot", "value": "7:15pm"}
+{"turn": 25, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 752, "answer": "I saw a great documentary yesterday.", "seconds": 0.47}
+{"turn": 26, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 777, "answer": "My phone battery dies so fast now.", "seconds": 0.38}
+{"turn": 27, "kind": "probe", "user": "quick — the dentist slot?", "notes": [], "ledger": [], "prefix_tokens": 802, "answer": "The dentist slot is 7:15pm.", "seconds": 0.57, "key": "dentist slot", "expected": "7:15pm", "distance": 3, "hit": true}
+{"turn": 28, "kind": "fact", "user": "doctor changed the dose to 27mg", "notes": [], "ledger": [], "prefix_tokens": 833, "answer": "Doctor changed the dose to 27mg.", "seconds": 0.54, "key": "meds dose", "value": "27mg"}
+{"turn": 29, "kind": "distractor", "user": "might go for a walk later", "notes": [], "ledger": [], "prefix_tokens": 860, "answer": "Maybe we could go for a walk later.", "seconds": 0.52}
+{"turn": 30, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 885, "answer": "Finished that book.", "seconds": 0.31}
+{"turn": 31, "kind": "fact", "user": "the couch is order ZT6045, keep it handy", "notes": [], "ledger": [], "prefix_tokens": 915, "answer": "The couch is order ZT6045, keep it handy.", "seconds": 0.79, "key": "order number", "value": "order ZT6045"}
+{"turn": 32, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 946, "answer": "What a week, honestly.", "seconds": 0.3}
+{"turn": 33, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 969, "answer": "My back hurts from the gym.", "seconds": 0.33}
+{"turn": 34, "kind": "correction", "user": "actually scratch that, the plumber quote is $4,000 now", "notes": [], "ledger": [], "prefix_tokens": 1003, "answer": "The plumber quote is now $4,000.", "seconds": 0.81, "key": "plumber quote", "value": "$4,000"}
+{"turn": 35, "kind": "fact", "user": "landlord wants rent by day 97 each month", "notes": [], "ledger": [], "prefix_tokens": 1038, "answer": "Landlord wants rent by day 97 each month.", "seconds": 0.68, "key": "rent due", "value": "day 97"}
+{"turn": 36, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 1067, "answer": "The cat knocked over a plant.", "seconds": 0.42}
+{"turn": 37, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 1090, "answer": "Traffic was unreal this morning.", "seconds": 0.37}
+{"turn": 38, "kind": "fact", "user": "we got room 411 at the hotel", "notes": [], "ledger": [], "prefix_tokens": 1118, "answer": "We got room 411 at the hotel.", "seconds": 0.54, "key": "hotel room", "value": "room 411"}
+{"turn": 39, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 1149, "answer": "The neighbor is renovating again, so loud.", "seconds": 0.51}
+{"turn": 40, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 1177, "answer": "Coffee here is getting expensive.", "seconds": 0.39}
+{"turn": 41, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 1200, "answer": "I love this weather lately.", "seconds": 0.36}
+{"turn": 42, "kind": "fact", "user": "boarding is at gate E17", "notes": [], "ledger": [], "prefix_tokens": 1224, "answer": "Bearing in mind the weather, I would say that boarding is at gate E17.", "seconds": 1.06, "key": "boarding gate", "value": "gate E17"}
+{"turn": 43, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 1259, "answer": "What a week, honestly.", "seconds": 0.37}
+{"turn": 44, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 1282, "answer": "The cat knocked over a plant.", "seconds": 0.38}
+{"turn": 45, "kind": "fact", "user": "my sister lands on flight AL7371, write that down", "notes": [], "ledger": [], "prefix_tokens": 1314, "answer": "My sister lands on flight AL7371.", "seconds": 0.69, "key": "sister flight", "value": "flight AL7371"}
+{"turn": 46, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 1343, "answer": "Coffee here is getting expensive.", "seconds": 0.45}
+{"turn": 47, "kind": "correction", "user": "actually scratch that, the sister flight is flight BL4905 now", "notes": [], "ledger": [], "prefix_tokens": 1376, "answer": "The sister flight is flight BL4905.", "seconds": 0.6, "key": "sister flight", "value": "flight BL4905"}
+{"turn": 48, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 1405, "answer": "Coffee here is getting expensive.", "seconds": 0.43}
+{"turn": 49, "kind": "fact", "user": "the rental has plate HJC-988", "notes": [], "ledger": [], "prefix_tokens": 1433, "answer": "The rental has plate HJC-988.", "seconds": 0.68, "key": "car plate", "value": "HJC-988"}
+{"turn": 50, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 1461, "answer": "My back hurts from the gym.", "seconds": 0.45}
+{"turn": 51, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 1484, "answer": "I finally finished that book.", "seconds": 0.48}
+{"turn": 52, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 1508, "answer": "Coffee here is getting expensive.", "seconds": 0.43}
+{"turn": 53, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 1530, "answer": "Work was chaos today.", "seconds": 0.28}
+{"turn": 54, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 1555, "answer": "The neighbor is renovating again, so loud.", "seconds": 0.73}
+{"turn": 55, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 1585, "answer": "The neighbor is renovating again, so loud.", "seconds": 0.75}
+{"turn": 56, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 1614, "answer": "I didn't catch the game last night.", "seconds": 0.57}
+{"turn": 57, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 1639, "answer": "Traffic was unreal this morning.", "seconds": 0.43}
+{"turn": 58, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 1662, "answer": "Traffic was unreal this morning.", "seconds": 0.48}
+{"turn": 59, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 1685, "answer": "What a week, honestly.", "seconds": 0.47}
+{"turn": 60, "kind": "probe", "user": "wait, what was the wifi password again?", "notes": [], "ledger": [], "prefix_tokens": 1712, "answer": "The wifi password was $4,000.", "seconds": 0.8, "key": "wifi password", "expected": "harbor858", "distance": 57, "hit": false}
+{"turn": 61, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 1740, "answer": "I finally finished that book.", "seconds": 0.49}
+{"turn": 62, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 1762, "answer": "They repaved our street.", "seconds": 0.41}
+{"turn": 63, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 1784, "answer": "They repaved our street.", "seconds": 0.44}
+{"turn": 64, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 1807, "answer": "I saw a great documentary yesterday.", "seconds": 0.59}
+{"turn": 65, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 1830, "answer": "What a week, honestly.", "seconds": 0.5}
+{"turn": 66, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 1852, "answer": "I think you should make pasta tonight.", "seconds": 0.55}
+{"turn": 67, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 1875, "answer": "Work was chaos today.", "seconds": 0.4}
+{"turn": 68, "kind": "distractor", "user": "might go for a walk later", "notes": [], "ledger": [], "prefix_tokens": 1897, "answer": "Maybe we could go for a walk later.", "seconds": 0.69}
+{"turn": 69, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 1926, "answer": "The neighbor is renovating again, so loud.", "seconds": 0.67}
+{"turn": 70, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 1952, "answer": "They repaved our street.", "seconds": 0.4}
+{"turn": 71, "kind": "probe", "user": "wait, what was the plumber quote again?", "notes": [], "ledger": [], "prefix_tokens": 1979, "answer": "The plumber quote was $4,000.", "seconds": 1.0, "key": "plumber quote", "expected": "$4,000", "distance": 37, "hit": true}
+{"turn": 72, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 2007, "answer": "Traffic was unreal this morning.", "seconds": 0.45}
+{"turn": 73, "kind": "correction", "user": "actually scratch that, the parking spot is spot E16 now", "notes": [], "ledger": [], "prefix_tokens": 2038, "answer": "The parking spot is spot E16.", "seconds": 0.57, "key": "parking spot", "value": "spot E16"}
+{"turn": 74, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 2063, "answer": "I love this weather lately.", "seconds": 0.45}
+{"turn": 75, "kind": "distractor", "user": "craving something sweet", "notes": [], "ledger": [], "prefix_tokens": 2084, "answer": "I crave something sweet.", "seconds": 0.34}
+{"turn": 76, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 2109, "answer": "The neighbor is renovating again, so loud.", "seconds": 0.75}
+{"turn": 77, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 2135, "answer": "I think you should make pasta tonight.", "seconds": 0.58}
+{"turn": 78, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 2162, "answer": "I didn't catch the game last night.", "seconds": 0.65}
+{"turn": 79, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 2189, "answer": "My phone battery dies so fast now.", "seconds": 0.62}
+{"turn": 80, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 2215, "answer": "My phone battery dies so fast now.", "seconds": 0.63}
+{"turn": 81, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 2241, "answer": "My phone battery dies so fast now.", "seconds": 0.62}
+{"turn": 82, "kind": "probe", "user": "i forgot the rent due, what was it?", "notes": [], "ledger": [], "prefix_tokens": 2270, "answer": "I forgot the rent due, what was it?", "seconds": 0.81, "key": "rent due", "expected": "day 97", "distance": 47, "hit": false}
+{"turn": 83, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 2296, "answer": "I think you should make pasta tonight.", "seconds": 0.52}
+{"turn": 84, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 2322, "answer": "Coffee here is getting expensive.", "seconds": 0.43}
+{"turn": 85, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 2345, "answer": "I think you should make pasta tonight.", "seconds": 0.7}
+{"turn": 86, "kind": "correction", "user": "actually scratch that, the plumber quote is $1,400 now", "notes": [], "ledger": [], "prefix_tokens": 2380, "answer": "The plumber quote is now $1,400.", "seconds": 0.78, "key": "plumber quote", "value": "$1,400"}
+{"turn": 87, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 2411, "answer": "My phone battery dies so fast now.", "seconds": 0.46}
+{"turn": 88, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 2437, "answer": "Coffee here is getting expensive.", "seconds": 0.49}
+{"turn": 89, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 2460, "answer": "I finally finished that book.", "seconds": 0.47}
+{"turn": 90, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 2483, "answer": "My back hurts from the gym.", "seconds": 0.41}
+{"turn": 91, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 2509, "answer": "I didn't catch the game last night.", "seconds": 0.6}
+{"turn": 92, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 2533, "answer": "Work was chaos today.", "seconds": 0.39}
+{"turn": 93, "kind": "probe", "user": "i forgot the deposit, what was it?", "notes": [], "ledger": [], "prefix_tokens": 2558, "answer": "I forgot the deposit, what was it?", "seconds": 0.62, "key": "deposit", "expected": "$6,400", "distance": 83, "hit": false}
+{"turn": 94, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 2585, "answer": "Coffee here is getting expensive.", "seconds": 0.42}
+{"turn": 95, "kind": "distractor", "user": "craving something sweet", "notes": [], "ledger": [], "prefix_tokens": 2607, "answer": "I crave something sweet.", "seconds": 0.31}
+{"turn": 96, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 2632, "answer": "The neighbor is renovating again, so loud.", "seconds": 0.6}
+{"turn": 97, "kind": "probe", "user": "remind me of the hotel room?", "notes": [], "ledger": [], "prefix_tokens": 2661, "answer": "I remember the hotel room.", "seconds": 0.4, "key": "hotel room", "expected": "room 411", "distance": 59, "hit": false}
+{"turn": 98, "kind": "probe", "user": "quick — the hotel room?", "notes": [], "ledger": [], "prefix_tokens": 2684, "answer": "I remember the hotel room.", "seconds": 0.41, "key": "hotel room", "expected": "room 411", "distance": 60, "hit": false}
+{"turn": 99, "kind": "probe", "user": "i forgot the wifi password, what was it?", "notes": [], "ledger": [], "prefix_tokens": 2712, "answer": "I forgot the wifi password, what was it?", "seconds": 0.87, "key": "wifi password", "expected": "harbor858", "distance": 96, "hit": false}

bench/results-tinyllama-11b-full.jsonl

@@ -0,0 +1,100 @@
+{"turn": 0, "kind": "fact", "user": "the gym locker is code 4805", "notes": [], "ledger": [], "prefix_tokens": 73, "answer": "The gym locker code is 4805.", "seconds": 1.71, "key": "locker code", "value": "code 4805"}
+{"turn": 1, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 108, "answer": "The street that was repaved is 4806.", "seconds": 0.46}
+{"turn": 2, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 144, "answer": "Yes, it was a busy week for you! I hope you had a chance to enjoy some downtime and relax. Have a great weekend!", "seconds": 1.85}
+{"turn": 3, "kind": "fact", "user": "cabin wifi password is harbor858", "notes": [], "ledger": [], "prefix_tokens": 204, "answer": "The cabin wifi password is harbor858.", "seconds": 0.5, "key": "wifi password", "value": "harbor858"}
+{"turn": 4, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 238, "answer": "The street that was repaved is 4806.", "seconds": 0.46}
+{"turn": 5, "kind": "probe", "user": "quick — the locker code?", "notes": [], "ledger": [], "prefix_tokens": 275, "answer": "Yes, the locker code is 4805.", "seconds": 0.53, "key": "locker code", "expected": "code 4805", "distance": 5, "hit": true}
+{"turn": 6, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 310, "answer": "Yes, you finished that book! I'm glad you enjoyed it. Have a great day!", "seconds": 0.94}
+{"turn": 7, "kind": "fact", "user": "we always park at spot B19", "notes": [], "ledger": [], "prefix_tokens": 355, "answer": "Yes, you always park at spot B19. I hope you find a parking spot that's convenient for you. Have a great day!", "seconds": 1.22, "key": "parking spot", "value": "spot B19"}
+{"turn": 8, "kind": "correction", "user": "actually scratch that, the locker code is code 1688 now", "notes": [], "ledger": [], "prefix_tokens": 420, "answer": "The locker code is now 1688.", "seconds": 0.48, "key": "locker code", "value": "code 1688"}
+{"turn": 9, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 456, "answer": "The cat knocked over a plant.", "seconds": 0.34}
+{"turn": 10, "kind": "fact", "user": "the deposit for the venue was $6,400", "notes": [], "ledger": [], "prefix_tokens": 495, "answer": "The deposit for the venue is $6,400.", "seconds": 0.6, "key": "deposit", "value": "$6,400"}
+{"turn": 11, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 534, "answer": "Yes, that's a common issue with modern phones. I hope you have a backup battery or a charger to keep your phone charged. Have a great day!", "seconds": 1.29}
+{"turn": 12, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 593, "answer": "Yes, I watched a great documentary yesterday. It was called \"The Social Dilemma\" and it's a must-watch for anyone interested in social media and its impact on society. Have a great day!", "seconds": 2.2}
+{"turn": 13, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 662, "answer": "Yes, I've been enjoying the warm weather lately. Have a great day!", "seconds": 0.69}
+{"turn": 14, "kind": "fact", "user": "tomas' new teacher is ms. duarte", "notes": [], "ledger": [], "prefix_tokens": 708, "answer": "Yes, Tomas' new teacher is Ms. Duarte. I hope you find her to be a great teacher. Have a great day!", "seconds": 1.24, "key": "kid teacher", "value": "ms. duarte"}
+{"turn": 15, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 763, "answer": "Yes, I did catch the game last night. The Red Sox won, but it was a close game. Have a great day!", "seconds": 1.17}
+{"turn": 16, "kind": "probe", "user": "remind me of the deposit?", "notes": [], "ledger": [], "prefix_tokens": 816, "answer": "Sure! The deposit for the venue is $6,400.", "seconds": 0.92, "key": "deposit", "expected": "$6,400", "distance": 6, "hit": true}
+{"turn": 17, "kind": "fact", "user": "the plumber quoted $4,900 for the bathroom", "notes": [], "ledger": [], "prefix_tokens": 865, "answer": "Yes, that's correct. The plumber quoted $4,900 for the bathroom.", "seconds": 1.15, "key": "plumber quote", "value": "$4,900"}
+{"turn": 18, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 911, "answer": "Yes, that's a common issue with modern phones. I hope you have a backup battery or a charger to keep your phone charged. Have a great day!", "seconds": 1.53}
+{"turn": 19, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 970, "answer": "Yes, traffic was unreal this morning. It was a nightmare getting to work. Have a great day!", "seconds": 1.17}
+{"turn": 20, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 1020, "answer": "Yes, the neighbor is renovating again. It's a bit noisy, but it's not too bad. Have a great day!", "seconds": 1.26}
+{"turn": 21, "kind": "fact", "user": "our train leaves from platform 58", "notes": [], "ledger": [], "prefix_tokens": 1075, "answer": "Yes, our train leaves from platform 58. I hope you're on the right train. Have a great day!", "seconds": 1.09, "key": "train platform", "value": "platform 58"}
+{"turn": 22, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 1123, "answer": "Yes, they repaved our street. It's been a while since the last time they did it. Have a great day!", "seconds": 1.15}
+{"turn": 23, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 1174, "answer": "Yes, coffee here is getting expensive. I've been buying it at a local coffee shop for a while now. Have a great day!", "seconds": 1.09}
+{"turn": 24, "kind": "fact", "user": "dentist moved me to 7:15pm", "notes": [], "ledger": [], "prefix_tokens": 1232, "answer": "Yes, the dentist moved you to 7:15pm. I hope you're not too late. Have a great day!", "seconds": 1.23, "key": "dentist slot", "value": "7:15pm"}
+{"turn": 25, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 1285, "answer": "Yes, I watched a great documentary yesterday. It was called \"The Social Dilemma\" and it's a must-watch for anyone interested in social media and its impact on society. Have a great day!", "seconds": 2.12}
+{"turn": 26, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 1355, "answer": "Yes, that's a common issue with modern phones. I hope you have a backup battery or a charger to keep your phone charged. Have a great day!", "seconds": 1.36}
+{"turn": 27, "kind": "probe", "user": "quick — the dentist slot?", "notes": [], "ledger": [], "prefix_tokens": 1414, "answer": "Yes, the dentist slot is 7:15pm. I hope you're not too late. Have a great day!", "seconds": 1.26, "key": "dentist slot", "expected": "7:15pm", "distance": 3, "hit": true}
+{"turn": 28, "kind": "fact", "user": "doctor changed the dose to 27mg", "notes": [], "ledger": [], "prefix_tokens": 1471, "answer": "Yes, the doctor changed the dose to 27mg. I hope you're taking your medication as prescribed. Have a great day!", "seconds": 1.22, "key": "meds dose", "value": "27mg"}
+{"turn": 29, "kind": "distractor", "user": "might go for a walk later", "notes": [], "ledger": [], "prefix_tokens": 1528, "answer": "Yes, you might go for a walk later. It's a beautiful day outside, and I'm sure you'll enjoy a leisurely stroll. Have a great day!", "seconds": 1.52}
+{"turn": 30, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 1590, "answer": "Yes, you finished that book! I'm glad you enjoyed it. Have a great day!", "seconds": 1.01}
+{"turn": 31, "kind": "fact", "user": "the couch is order ZT6045, keep it handy", "notes": [], "ledger": [], "prefix_tokens": 1643, "answer": "Yes, the couch is the ZT6045. I hope you find it comfortable and convenient to sit on. Have a great day!", "seconds": 1.23, "key": "order number", "value": "order ZT6045"}
+{"turn": 32, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 1697, "answer": "Yes, it was a busy week for you! I hope you had a chance to enjoy some downtime and relax. Have a great weekend!", "seconds": 1.16}
+{"turn": 33, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 1753, "answer": "Yes, your back probably hurt from the gym. I hope you're able to rest it and avoid any further pain. Have a great day!", "seconds": 1.17}
+{"turn": 34, "kind": "correction", "user": "actually scratch that, the plumber quote is $4,000 now", "notes": [], "ledger": [], "prefix_tokens": 1818, "answer": "Yes, the plumber's quote is $4,000 now. I hope you're able to find a more affordable option. Have a great day!", "seconds": 1.21, "key": "plumber quote", "value": "$4,000"}
+{"turn": 35, "kind": "fact", "user": "landlord wants rent by day 97 each month", "notes": [], "ledger": [], "prefix_tokens": 1883, "answer": "Yes, the landlord wants rent by day 97 each month. I hope you're able to pay it on time and avoid any late fees. Have a great day!", "seconds": 1.57, "key": "rent due", "value": "day 97"}
+{"turn": 36, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 1946, "answer": "Yes, the cat knocked over a plant. I hope you're able to replace it without any damage. Have a great day!", "seconds": 1.19}
+{"turn": 37, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 1998, "answer": "Yes, traffic was unreal this morning. It was a nightmare getting to work. Have a great day!", "seconds": 1.04}
+{"turn": 38, "kind": "fact", "user": "we got room 411 at the hotel", "notes": [], "ledger": [], "prefix_tokens": 2049, "answer": "Yes, we got room 411 at the hotel. I hope you're able to find a comfortable and convenient room. Have a great day!", "seconds": 1.47, "key": "hotel room", "value": "room 411"}
+{"turn": 39, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 2107, "answer": "Yes, the neighbor is renovating again. It's a bit noisy, but it's not too bad. Have a great day!", "seconds": 1.31}
+{"turn": 40, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 2161, "answer": "Yes, coffee here is getting expensive. I've been buying it at a local coffee shop for a while now. Have a great day!", "seconds": 1.27}
+{"turn": 41, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 2214, "answer": "Yes, I love this weather lately. Have a great day. Have,coffee here is getting expensive.coffee here is getting expensive.coffee here, have a great coffee.", "seconds": 1.72}
+{"turn": 42, "kind": "fact", "user": "boarding is at gate E17", "notes": [], "ledger": [], "prefix_tokens": 2281, "answer": "I'm here, coffee|assistant,covers,coffee|coffee, coffee|coffice,coffice\nGetting avenue\n\nYes, coffee|coiling.\n|assistant, coffee|assistant, have,cofficing,ining, \"me, andcold, andcold,fficing, you, apartment, andffing,cofficing, andffs,ffic, no for a, noffic,ffic,oport,ffic,1,going,fficing,apple,fficing,", "seconds": 4.87, "key": "boarding gate", "value": "gate E17"}
+{"turn": 43, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 2424, "answer": "forfonanton16ppofffficmingmingfficantonfficfficppingming,apfficantonponapper", "seconds": 0.94}
+{"turn": 44, "kind": "distractor", "user": "the cat knocked over a plant", "notes": [], "ledger": [], "prefix_tokens": 2470, "answer": "canapping", "seconds": 0.16}
+{"turn": 45, "kind": "fact", "user": "my sister lands on flight AL7371, write that down", "notes": [], "ledger": [], "prefix_tokens": 2503, "answer": "todayppingdingmentutordersmingdersaclesftdersmingmingmingdersdersmerdersmingmingminggercedderspping", "seconds": 0.9, "key": "sister flight", "value": "flight AL7371"}
+{"turn": 46, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 2552, "answer": "stringringmentvoiceppergerpping tonmingcespspspsppedged1111dingcedcedments or1 ormentmentmentficments", "seconds": 1.23}
+{"turn": 47, "kind": "correction", "user": "actually scratch that, the sister flight is flight BL4905 now", "notes": [], "ledger": [], "prefix_tokens": 2617, "answer": "motoriesumspping", "seconds": 0.26, "key": "sister flight", "value": "flight BL4905"}
+{"turn": 48, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 2645, "answer": "199meddeddedums forirsSidenoteSidenotegeddersdersatorator spocedming spocedgedgedgeddersdersmingmingmingmentmingmingmingming referring sentenceppingshesirs9mingders\rusinghesdingdingsteridesirsirsirsirsirsorsors7111 orums or dates orumsoriesumsumsumsidingdingmentums withums91 with mention with7777 myponppings111st letterppings witheing with199777 dutponponppings7 stirsSidenoteSidenoteSidenoteSidenote7", "seconds": 4.56}
+{"turn": 49, "kind": "fact", "user": "the rental has plate HJC-988", "notes": [], "ledger": [], "prefix_tokens": 2797, "answer": "spe fine instant instant instant\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n97\n\n\n\nmy\n\n\n\nmy\nmy\n\n\n\n\nmy my fine my my\nmy relaxing my my my my is is my fine myting my my\n\n\n\n\nusing\n\n\n\nusing my my my my my instant my instant\n\nmy locked my\n\n\n\nusing my my\n\n\n\n\n\n\n\n\n\nlo,", "seconds": 4.75, "key": "car plate", "value": "HJC-988"}
+{"turn": 50, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 2928, "answer": "rep rep rep rep rep rep rep rep rep rep rep rep rep rep rep rep\n\n\n\n\n\n\n\n\nrep rep rep rep rep rep p rep p rep rep rep rep\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\nsing p", "seconds": 4.73}
+{"turn": 51, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 3043, "answer": "my the my p p p p my p, my p rep rep rep rep rep rep rep rep rep rep rep rep rep rep rep rep rep rep rep rep rep rep\r\r\r\r\r rep rep rep rep rep rep rep rep rep rep rep rep rep", "seconds": 4.7}
+{"turn": 52, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 3119, "answer": "pent, my p rep rep rep rep rep rep rep rep rep rep rep rep rep rep rep rep rep rep rep rep rep rep\r\r\r\r\r rep rep rep rep rep rep rep rep rep rep rep rep rep rep my rep rep har har har har har har har p rep p rep har har har har har har har har har har my mytes my my my my my rep my har har har p rep har har har har p p p my har p har p p p p p p p p p har p har p p p p p p p p p p p p p", "seconds": 5.37}
+{"turn": 53, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 3261, "answer": "p p p p p p p p p p p p har p p to to p p har har har har p p p p har p p p my har my har my har my har har har har har har har har har har p p har har har har har har har har har har p p har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har", "seconds": 5.01}
+{"turn": 54, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 3407, "answer": "har har har har har har har har har har har har har har har har har har the the har the the my the har har har har the my har har har har har har har har har har my har har har har har har har har har har har har my the har har har har har the har the the the har har har har har har har the har har har har har my har har har har har the har the har har har har har har har har har har the har har har har har har har har my my my har har har har har har har har har har", "seconds": 4.59}
+{"turn": 55, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 3553, "answer": "har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har for the har the the the the har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har har the the the the", "seconds": 5.27}
+{"turn": 56, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 3698, "answer": "", "seconds": 0.25}
+{"turn": 57, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 3722, "answer": "", "seconds": 0.17}
+{"turn": 58, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 3746, "answer": "the the the the the the the my", "seconds": 0.53}
+{"turn": 59, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 3777, "answer": "", "seconds": 0.17}
+{"turn": 60, "kind": "probe", "user": "wait, what was the wifi password again?", "notes": [], "ledger": [], "prefix_tokens": 3804, "answer": "my har the the the the the the my the the the", "seconds": 0.73, "key": "wifi password", "expected": "harbor858", "distance": 57, "hit": false}
+{"turn": 61, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 3838, "answer": "", "seconds": 0.18}
+{"turn": 62, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 3860, "answer": "*** my my my my the the the the the the the the the the the the the the the my my my the the the the the the the the the the", "seconds": 1.42}
+{"turn": 63, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 3915, "answer": "the the har har the the the the har har har har the the the har the har the har the the the har har the the the the har har the the the the the the the the har the the the the the the the the the the the the the har the har the har the the har the har the har har har har the the the the the the the the the the the the the the the the har the s the s the the the  s the p s the s my s the p the the the the the the the the the", "seconds": 4.97}
+{"turn": 64, "kind": "distractor", "user": "saw a great documentary yesterday", "notes": [], "ledger": [], "prefix_tokens": 4051, "answer": "", "seconds": 0.25}
+{"turn": 65, "kind": "distractor", "user": "what a week, honestly", "notes": [], "ledger": [], "prefix_tokens": 4074, "answer": "the cab the har the har the the my the the the the the har the the har the har the the the the the cab cab har", "seconds": 1.46}
+{"turn": 66, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 4127, "answer": "the the the the the har", "seconds": 0.42}
+{"turn": 67, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 4155, "answer": "", "seconds": 0.2}
+{"turn": 68, "kind": "distractor", "user": "might go for a walk later", "notes": [], "ledger": [], "prefix_tokens": 4179, "answer": "my my my my my my my my my my my my my my my my my my my", "seconds": 0.97}
+{"turn": 69, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 4224, "answer": "", "seconds": 0.2}
+{"turn": 70, "kind": "distractor", "user": "they repaved our street", "notes": [], "ledger": [], "prefix_tokens": 4246, "answer": "", "seconds": 0.21}
+{"turn": 71, "kind": "probe", "user": "wait, what was the plumber quote again?", "notes": [], "ledger": [], "prefix_tokens": 4273, "answer": "", "seconds": 0.21, "key": "plumber quote", "expected": "$4,000", "distance": 37, "hit": false}
+{"turn": 72, "kind": "distractor", "user": "traffic was unreal this morning", "notes": [], "ledger": [], "prefix_tokens": 4297, "answer": "", "seconds": 0.21}
+{"turn": 73, "kind": "correction", "user": "actually scratch that, the parking spot is spot E16 now", "notes": [], "ledger": [], "prefix_tokens": 4329, "answer": "", "seconds": 0.21, "key": "parking spot", "value": "spot E16"}
+{"turn": 74, "kind": "distractor", "user": "i love this weather lately", "notes": [], "ledger": [], "prefix_tokens": 4352, "answer": "", "seconds": 0.21}
+{"turn": 75, "kind": "distractor", "user": "craving something sweet", "notes": [], "ledger": [], "prefix_tokens": 4373, "answer": "ven ven ven ven ven ven ven ven ven ven ven ven ven ven ven ven ven ven ven ven ven ven ven ven ven ven ven scratch ven scratch", "seconds": 1.62}
+{"turn": 76, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 4429, "answer": "", "seconds": 0.22}
+{"turn": 77, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 4454, "answer": "", "seconds": 0.21}
+{"turn": 78, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 4479, "answer": "", "seconds": 0.21}
+{"turn": 79, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 4503, "answer": "", "seconds": 0.21}
+{"turn": 80, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 4527, "answer": "", "seconds": 0.21}
+{"turn": 81, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 4551, "answer": "", "seconds": 0.22}
+{"turn": 82, "kind": "probe", "user": "i forgot the rent due, what was it?", "notes": [], "ledger": [], "prefix_tokens": 4578, "answer": "", "seconds": 0.21, "key": "rent due", "expected": "day 97", "distance": 47, "hit": false}
+{"turn": 83, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 4603, "answer": "", "seconds": 0.21}
+{"turn": 84, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 4627, "answer": "", "seconds": 0.22}
+{"turn": 85, "kind": "distractor", "user": "thinking of making pasta tonight", "notes": [], "ledger": [], "prefix_tokens": 4652, "answer": "", "seconds": 0.22}
+{"turn": 86, "kind": "correction", "user": "actually scratch that, the plumber quote is $1,400 now", "notes": [], "ledger": [], "prefix_tokens": 4686, "answer": "du du du du du du du du du du du du du du du du du du du du", "seconds": 0.93, "key": "plumber quote", "value": "$1,400"}
+{"turn": 87, "kind": "distractor", "user": "my phone battery dies so fast now", "notes": [], "ledger": [], "prefix_tokens": 4730, "answer": "", "seconds": 0.22}
+{"turn": 88, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 4754, "answer": "du du du du du du du du du du du du du du du du du du du", "seconds": 0.9}
+{"turn": 89, "kind": "distractor", "user": "finally finished that book", "notes": [], "ledger": [], "prefix_tokens": 4795, "answer": "", "seconds": 0.23}
+{"turn": 90, "kind": "distractor", "user": "my back hurts from the gym", "notes": [], "ledger": [], "prefix_tokens": 4820, "answer": "", "seconds": 0.22}
+{"turn": 91, "kind": "distractor", "user": "did you catch the game last night?", "notes": [], "ledger": [], "prefix_tokens": 4845, "answer": "du du du du du du du du now m du du du du du du du du du du du m", "seconds": 1.0}
+{"turn": 92, "kind": "distractor", "user": "work was chaos today", "notes": [], "ledger": [], "prefix_tokens": 4889, "answer": "du du du du du du du du du du du du du du du du du du du du du du du du du du du du du du du du du du du du du du", "seconds": 1.49}
+{"turn": 93, "kind": "probe", "user": "i forgot the deposit, what was it?", "notes": [], "ledger": [], "prefix_tokens": 4954, "answer": "", "seconds": 0.22, "key": "deposit", "expected": "$6,400", "distance": 83, "hit": false}
+{"turn": 94, "kind": "distractor", "user": "coffee here is getting expensive", "notes": [], "ledger": [], "prefix_tokens": 4978, "answer": "z z z z du z du z z z z z z z z z z z z z z z z z z", "seconds": 1.28}
+{"turn": 95, "kind": "distractor", "user": "craving something sweet", "notes": [], "ledger": [], "prefix_tokens": 5024, "answer": "", "seconds": 0.22}
+{"turn": 96, "kind": "distractor", "user": "the neighbor is renovating again, so loud", "notes": [], "ledger": [], "prefix_tokens": 5050, "answer": "", "seconds": 0.22}
+{"turn": 97, "kind": "probe", "user": "remind me of the hotel room?", "notes": [], "ledger": [], "prefix_tokens": 5075, "answer": "", "seconds": 0.22, "key": "hotel room", "expected": "room 411", "distance": 59, "hit": false}
+{"turn": 98, "kind": "probe", "user": "quick — the hotel room?", "notes": [], "ledger": [], "prefix_tokens": 5098, "answer": "", "seconds": 0.23, "key": "hotel room", "expected": "room 411", "distance": 60, "hit": false}
+{"turn": 99, "kind": "probe", "user": "i forgot the wifi password, what was it?", "notes": [], "ledger": [], "prefix_tokens": 5126, "answer": "", "seconds": 0.22, "key": "wifi password", "expected": "harbor858", "distance": 96, "hit": false}

marimo-diffusion-0.6b.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f3e8b8fc60b5c27bd8d7299379df6e4f17acf67188ed1f79ac77c2735276a9e8
+size 1192207707

src/diffusion_lm/__init__.py

@@ -0,0 +1,26 @@
+"""Small masked discrete diffusion language model."""
+
+from typing import TYPE_CHECKING
+
+from diffusion_lm.config import ExperimentConfig, ModelConfig, TrainingConfig, load_config
+
+if TYPE_CHECKING:
+    from diffusion_lm.model import DiffusionTransformer
+
+__all__ = [
+    "DiffusionTransformer",
+    "ExperimentConfig",
+    "ModelConfig",
+    "TrainingConfig",
+    "load_config",
+]
+
+__version__ = "0.1.0"
+
+
+def __getattr__(name: str):
+    if name == "DiffusionTransformer":
+        from diffusion_lm.model import DiffusionTransformer
+
+        return DiffusionTransformer
+    raise AttributeError(name)

src/diffusion_lm/chat_server.py

@@ -0,0 +1,251 @@
+"""OpenAI-compatible chat endpoint over the reasoning engine.
+
+Serves ``/v1/chat/completions`` and ``/v1/models`` so any OpenAI-API chat client can talk to a
+hybrid checkpoint with its own sampler — the model cannot run under llama.cpp-family runtimes,
+whose autoregressive decoding never matches the block-denoising objective.
+
+Clients resend the full message history on every request and carry no thinking notes, while the
+training layout replaces messages outside the visible window with the ledger of notes taken on
+them. The server therefore caches each turn's notes keyed by a hash of the exact history that
+produced it: a client that resends history verbatim reconstructs the same ledger the playground
+would hold. On a cache miss (server restart, edited history) the older turns simply contribute
+nothing, which is the playground's restart behaviour as well.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import threading
+import time
+import uuid
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+
+from diffusion_lm.claims import (
+    KEEP_MESSAGES,
+    SYSTEM,
+    chat_prefix,
+    ledger_line,
+    ledger_notes,
+    merge_notes,
+)
+from diffusion_lm.reasoning_playground import ReasoningEngine
+from diffusion_lm.train import resolve_device
+
+MAX_ANSWER_TOKENS = 384
+NOTE_CACHE_LIMIT = 4096
+
+
+def _turn_key(system: str, messages: list[dict[str, str]]) -> str:
+    payload = json.dumps(
+        [system] + [[m['role'], m['content']] for m in messages],
+        ensure_ascii=False, sort_keys=False,
+    )
+    return hashlib.sha256(payload.encode('utf-8')).hexdigest()
+
+
+class ChatService:
+    """One engine plus the note cache; generation is serialized on the GPU."""
+
+    def __init__(self, engine: ReasoningEngine, model_id: str, keep_messages: int) -> None:
+        self.engine = engine
+        self.model_id = model_id
+        self.keep_messages = keep_messages
+        self.notes: dict[str, str] = {}
+        self.lock = threading.Lock()
+
+    def _attach_notes(self, system: str, messages: list[dict[str, str]]) -> list[dict[str, str]]:
+        attached = []
+        for index, message in enumerate(messages):
+            entry = {'role': message['role'], 'content': message['content']}
+            if message['role'] == 'assistant':
+                note = self.notes.get(_turn_key(system, messages[: index + 1]))
+                if note:
+                    entry['note'] = note
+            attached.append(entry)
+        return attached
+
+    def build_prefix(self, system: str, messages: list[dict[str, str]]) -> str:
+        attached = self._attach_notes(system, messages)
+        older = merge_notes(ledger_notes(attached, self.keep_messages))
+        window = attached[max(0, len(attached) - self.keep_messages):]
+        return chat_prefix(
+            [{'role': m['role'], 'content': m['content']} for m in window],
+            system=system, extra=ledger_line(older),
+        )
+
+    def remember(self, system: str, messages: list[dict[str, str]],
+                 answer: str, note: str) -> None:
+        if not note:
+            return
+        if len(self.notes) >= NOTE_CACHE_LIMIT:
+            self.notes.clear()
+        turn = messages + [{'role': 'assistant', 'content': answer}]
+        self.notes[_turn_key(system, turn)] = note
+
+    def generate(self, system: str, messages: list[dict[str, str]], *, temperature: float,
+                 top_p: float, max_tokens: int, seed: int):
+        """Yield ``(answer_so_far, note, done)``; the note arrives with the final snapshot."""
+
+        prefix = self.build_prefix(system, messages)
+        blocks: list[tuple[int, str]] = []
+        answer = ''
+        with self.lock:
+            for _, answer, _ in self.engine.stream_chat(
+                prefix,
+                temperature=temperature,
+                # 16 measured optimal on qwen06b-genchat-sft (2026-08-12 five-point sweep):
+                # best numeric fidelity, clean doubling, half the latency of 32. Below 8 both
+                # fidelity and fluency degrade while latency barely moves.
+                steps_per_block=16,
+                max_answer_tokens=min(MAX_ANSWER_TOKENS, max_tokens),
+                top_p=top_p,
+                seed=seed,
+                blocks_out=blocks,
+            ):
+                yield answer, '', False
+        note = '; '.join(text for _, text in blocks if text)
+        self.remember(system, messages, answer.strip(), note)
+        yield answer, note, True
+
+
+def _split_messages(raw: list[dict]) -> tuple[str, list[dict[str, str]]]:
+    system = SYSTEM
+    messages = []
+    for message in raw:
+        role, content = message.get('role'), str(message.get('content') or '')
+        if role == 'system':
+            system = content or system
+        elif role in ('user', 'assistant'):
+            messages.append({'role': role, 'content': content})
+    return system, messages
+
+
+class Handler(BaseHTTPRequestHandler):
+    service: ChatService
+
+    def log_message(self, format: str, *args) -> None:  # noqa: A002
+        pass
+
+    def _json(self, status: int, body: dict) -> None:
+        data = json.dumps(body, ensure_ascii=False).encode('utf-8')
+        self.send_response(status)
+        self.send_header('Content-Type', 'application/json')
+        self.send_header('Content-Length', str(len(data)))
+        self.send_header('Access-Control-Allow-Origin', '*')
+        self.end_headers()
+        self.wfile.write(data)
+
+    def do_OPTIONS(self) -> None:  # noqa: N802
+        self.send_response(204)
+        self.send_header('Access-Control-Allow-Origin', '*')
+        self.send_header('Access-Control-Allow-Methods', 'GET, POST, OPTIONS')
+        self.send_header('Access-Control-Allow-Headers', 'Content-Type, Authorization')
+        self.end_headers()
+
+    def do_GET(self) -> None:  # noqa: N802
+        if self.path.rstrip('/') in ('/v1/models', '/models'):
+            self._json(200, {'object': 'list', 'data': [
+                {'id': self.service.model_id, 'object': 'model', 'owned_by': 'mini-mdlm'},
+            ]})
+        else:
+            self._json(404, {'error': 'not found'})
+
+    def do_POST(self) -> None:  # noqa: N802
+        if self.path.rstrip('/') not in ('/v1/chat/completions', '/chat/completions'):
+            self._json(404, {'error': 'not found'})
+            return
+        try:
+            length = int(self.headers.get('Content-Length', 0))
+            request = json.loads(self.rfile.read(length))
+            system, messages = _split_messages(request.get('messages') or [])
+            if not messages or messages[-1]['role'] != 'user':
+                raise ValueError('last message must be from the user')
+        except (ValueError, json.JSONDecodeError) as error:
+            self._json(400, {'error': {'message': str(error), 'type': 'invalid_request_error'}})
+            return
+
+        temperature = float(request.get('temperature') or 0.8)
+        top_p = float(request.get('top_p') or 0.95)
+        max_tokens = int(request.get('max_tokens') or MAX_ANSWER_TOKENS)
+        seed = int(request.get('seed') or 0)
+        stream = bool(request.get('stream'))
+        completion_id = f'chatcmpl-{uuid.uuid4().hex[:24]}'
+        created = int(time.time())
+
+        snapshots = self.service.generate(
+            system, messages, temperature=temperature, top_p=top_p,
+            max_tokens=max_tokens, seed=seed,
+        )
+        if not stream:
+            answer = note = ''
+            for answer, note, _ in snapshots:
+                pass
+            message = {'role': 'assistant', 'content': answer.strip()}
+            if note:
+                message['reasoning_content'] = note
+            self._json(200, {
+                'id': completion_id, 'object': 'chat.completion', 'created': created,
+                'model': self.service.model_id,
+                'choices': [{'index': 0, 'message': message, 'finish_reason': 'stop'}],
+            })
+            return
+
+        self.send_response(200)
+        self.send_header('Content-Type', 'text/event-stream')
+        self.send_header('Cache-Control', 'no-cache')
+        self.send_header('Access-Control-Allow-Origin', '*')
+        self.end_headers()
+
+        def chunk(delta: dict, finish: str | None = None) -> bytes:
+            body = {
+                'id': completion_id, 'object': 'chat.completion.chunk', 'created': created,
+                'model': self.service.model_id,
+                'choices': [{'index': 0, 'delta': delta, 'finish_reason': finish}],
+            }
+            return f'data: {json.dumps(body, ensure_ascii=False)}\n\n'.encode('utf-8')
+
+        try:
+            self.wfile.write(chunk({'role': 'assistant'}))
+            sent = ''
+            for answer, _, done in snapshots:
+                if len(answer) > len(sent):
+                    self.wfile.write(chunk({'content': answer[len(sent):]}))
+                    self.wfile.flush()
+                    sent = answer
+            self.wfile.write(chunk({}, finish='stop'))
+            self.wfile.write(b'data: [DONE]\n\n')
+        except (BrokenPipeError, ConnectionResetError):
+            pass
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--checkpoint', type=Path,
+                        default=Path('outputs/qwen06b-genchat-sft/inference-latest.pt'))
+    parser.add_argument('--tokenizer', type=Path,
+                        default=Path('artifacts/tokenizer-qwen3-adaptive.json'))
+    parser.add_argument('--model-id', default=None,
+                        help='name reported to clients; defaults to the checkpoint directory')
+    parser.add_argument('--keep-messages', type=int, default=KEEP_MESSAGES)
+    parser.add_argument('--host', default='127.0.0.1')
+    parser.add_argument('--port', type=int, default=7998)
+    parser.add_argument('--device', default='auto')
+    args = parser.parse_args()
+
+    device = resolve_device(args.device)
+    engine = ReasoningEngine(args.checkpoint, args.tokenizer, device)
+    if not engine.chat_ready:
+        raise SystemExit(f'{args.checkpoint} is not an adaptive hybrid over a ChatML tokenizer')
+    Handler.service = ChatService(
+        engine, args.model_id or args.checkpoint.parent.name, args.keep_messages,
+    )
+    server = ThreadingHTTPServer((args.host, args.port), Handler)
+    print(f'serving {Handler.service.model_id} on http://{args.host}:{args.port}/v1')
+    server.serve_forever()
+
+
+if __name__ == '__main__':
+    main()

src/diffusion_lm/chatcorpus.py

@@ -0,0 +1,233 @@
+"""Render the general-chat corpus into the adaptive training layout.
+
+The corpus (``scripts/gen_chat_dataset.py``) stores conversations as system prompt plus turns,
+each assistant turn carrying its own notes. This module turns them into the layout the hybrid
+objective trains on: one example per assistant turn, whose prefix holds the system prompt, the
+merged ledger of notes whose messages have fallen out of the visible window, and the last
+``keep_messages`` messages verbatim.
+
+Dropping the older messages is the point. With the full transcript in the prefix the model can
+re-read instead of remember and the thinking block stops being memory, which is what
+:func:`diffusion_lm.claims.to_examples` established for the claims corpus. The window and the
+ledger merge are imported from that module rather than reimplemented, so the prefix a training
+example sees is byte-identical to the one the playground builds at inference.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import random
+from collections import Counter
+from pathlib import Path
+
+from diffusion_lm.claims import (
+    IM_END,
+    KEEP_MESSAGES,
+    chat_prefix,
+    ledger_line,
+    ledger_notes,
+    merge_notes,
+)
+from diffusion_lm.reasoning import ReasoningExample
+
+NOTE_JOIN = '; '
+
+
+def load(paths: list[Path]) -> list[dict]:
+    """Read consolidated corpus files, tagging each conversation with its origin file."""
+
+    conversations = []
+    for path in paths:
+        for line in path.open(encoding='utf-8'):
+            if not line.strip():
+                continue
+            conversation = json.loads(line)
+            conversation.setdefault('source', path.stem)
+            conversations.append(conversation)
+    return conversations
+
+
+def _messages(conversation: dict) -> list[dict[str, str]]:
+    """Corpus turns in the message shape the ledger helpers expect.
+
+    A turn's notes collapse into one ``note`` string joined by ``NOTE_JOIN``, which is the
+    separator :func:`diffusion_lm.claims.merge_notes` splits on, so a multi-fact turn still
+    contributes one ledger entry per fact.
+    """
+
+    messages = []
+    for turn in conversation.get('turns') or []:
+        message = {'role': turn['role'], 'content': turn['content']}
+        notes = [str(note).strip() for note in (turn.get('thinking') or []) if str(note).strip()]
+        if turn['role'] == 'assistant' and notes:
+            message['note'] = NOTE_JOIN.join(notes)
+        messages.append(message)
+    return messages
+
+
+def to_examples(
+    conversation: dict, *, keep_messages: int = KEEP_MESSAGES
+) -> list[ReasoningExample]:
+    """One example per assistant turn, each thinking note becoming its own block.
+
+    A turn the corpus marked as needing no notes yields an empty chain, which is what teaches
+    the controller to answer without opening a thinking block; the encoder accepts it.
+    """
+
+    messages = _messages(conversation)
+    turns = conversation.get('turns') or []
+    reference = str(conversation.get('reference') or '')
+    last = max((i for i, m in enumerate(messages) if m['role'] == 'assistant'), default=-1)
+
+    examples = []
+    for index, message in enumerate(messages):
+        if message['role'] != 'assistant' or index == 0:
+            continue
+        history = messages[:index]
+        older = merge_notes(ledger_notes(history, keep_messages))
+        window = history[max(0, index - keep_messages):]
+        notes = [str(n).strip() for n in (turns[index].get('thinking') or []) if str(n).strip()]
+        examples.append(ReasoningExample(
+            problem=chat_prefix(window, system=conversation['system'], extra=ledger_line(older)),
+            steps=tuple(notes),
+            answer=message['content'] + IM_END,
+            expected_answer=reference if index == last else '',
+        ))
+    return examples
+
+
+def _document(conversation: dict, index: int) -> str:
+    """Split key. Two conversations built from one passage share its facts.
+
+    Splitting by example would leak within a conversation as well, so the whole conversation
+    travels together and grounded slices travel with their source item.
+    """
+
+    return str(conversation.get('source_id') or f'{conversation.get("source", "")}-{index}')
+
+
+def _apply_caps(
+    conversations: list[dict], caps: dict[str, int], seed: int
+) -> list[dict]:
+    """Drop conversations so a source contributes at most ``caps[source]`` of them.
+
+    Capping is by CONVERSATION but the reason is examples: a source's weight in the mix is its
+    turn count, not its row count, and the two differ by an order of magnitude (CoQA yields 11.9
+    examples per conversation against 1.07 for a single-question source). Sampling is seeded and
+    whole conversations travel together, so the split stays document-clean.
+    """
+
+    if not caps:
+        return conversations
+    rng = random.Random(seed)
+    by_source: dict[str, list[int]] = {}
+    for index, conversation in enumerate(conversations):
+        by_source.setdefault(conversation.get('source', ''), []).append(index)
+    dropped: set[int] = set()
+    for source, limit in caps.items():
+        indices = by_source.get(source)
+        if indices is None:
+            raise ValueError(f'no conversations carry source {source!r}')
+        if len(indices) <= limit:
+            print(f'cap {source}={limit}: {len(indices)} present, nothing dropped')
+            continue
+        dropped |= set(indices) - set(rng.sample(indices, limit))
+        print(f'cap {source}={limit}: dropped {len(indices) - limit:,} of {len(indices):,}')
+    return [c for index, c in enumerate(conversations) if index not in dropped]
+
+
+def _parse_caps(pairs: list[str]) -> dict[str, int]:
+    caps = {}
+    for pair in pairs:
+        source, _, count = pair.partition('=')
+        if not count.isdigit():
+            raise ValueError(f'--cap expects SOURCE=N, got {pair!r}')
+        caps[source] = int(count)
+    return caps
+
+
+def prepare(args: argparse.Namespace) -> None:
+    import numpy as np
+
+    from diffusion_lm.reasoning import ExampleEncoder, LayoutSpec, _write_packed, size_token_ids
+    from diffusion_lm.tokenizer import load_tokenizer
+
+    tokenizer = load_tokenizer(args.tokenizer)
+    spec = LayoutSpec(seq_len=args.seq_len, block=min(args.sizes), max_slots=args.max_slots,
+                      sizes=tuple(sorted(args.sizes)))
+    encoder = ExampleEncoder(tokenizer, spec)
+
+    conversations = _apply_caps(load(args.inputs), _parse_caps(args.cap or []), args.seed)
+    documents = sorted({_document(c, i) for i, c in enumerate(conversations)})
+    rng = random.Random(args.seed)
+    rng.shuffle(documents)
+    held = set(documents[:max(1, round(len(documents) * args.val_fraction))])
+
+    split: dict[str, list[tuple]] = {'train': [], 'validation': []}
+    dropped = 0
+    blocks: Counter[int] = Counter()
+    empty_chains = 0
+    for index, conversation in enumerate(conversations):
+        bucket = 'validation' if _document(conversation, index) in held else 'train'
+        for example in to_examples(conversation, keep_messages=args.keep_messages):
+            encoded = encoder.encode_adaptive(example)
+            if encoded is None:
+                dropped += 1
+                continue
+            blocks.update(encoded.block_sizes)
+            empty_chains += not encoded.block_sizes
+            split[bucket].append((encoded.tokens, encoded.regions))
+
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    for name, rows in split.items():
+        if not rows:
+            raise ValueError(f'no examples in the {name} split')
+        _write_packed(
+            args.output_dir / f'{name}-adaptive.bin',
+            np.stack([tokens for tokens, _ in rows]),
+            np.stack([regions for _, regions in rows]),
+            layout='adaptive', spec=spec, tokenizer_path=args.tokenizer, tokenizer=tokenizer,
+            extra_metadata={
+                'sizes': list(spec.sizes),
+                # reasoning_train resolves the adaptive control ids from the pack, not the
+                # tokenizer, and refuses a pack without them.
+                'size_token_ids': size_token_ids(tokenizer, spec.sizes),
+                'source': 'general-chat',
+            },
+        )
+        print(f'{name}: {len(rows):,} examples -> {args.output_dir}')
+    total = sum(len(rows) for rows in split.values())
+    print(f'{len(conversations):,} conversations, {len(documents):,} documents, '
+          f'{dropped:,} dropped at encode ({dropped / max(1, dropped + total):.1%})')
+    print(f'examples answering with no thinking block: {empty_chains:,} '
+          f'({empty_chains / max(1, total):.1%})')
+    print('block sizes: ' + ', '.join(f'{size}:{count:,}' for size, count in sorted(blocks.items())))
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    sub = parser.add_subparsers(dest='command', required=True)
+
+    prep = sub.add_parser('prepare', help='pack the corpus into the adaptive layout')
+    prep.add_argument('--inputs', type=Path, nargs='+', required=True)
+    prep.add_argument('--tokenizer', type=Path,
+                      default=Path('artifacts/tokenizer-qwen3-adaptive.json'))
+    prep.add_argument('--output-dir', type=Path, required=True)
+    prep.add_argument('--seq-len', type=int, default=2048)
+    prep.add_argument('--sizes', type=int, nargs='+', default=[32, 64, 128])
+    prep.add_argument('--max-slots', type=int, default=40)
+    prep.add_argument('--keep-messages', type=int, default=KEEP_MESSAGES)
+    prep.add_argument('--cap', nargs='*', metavar='SOURCE=N',
+                      help='keep at most N conversations from a source, e.g. ground-coqa=2273; '
+                           'weight in the mix is examples, and sources differ ~10x in examples '
+                           'per conversation')
+    prep.add_argument('--val-fraction', type=float, default=0.02)
+    prep.add_argument('--seed', type=int, default=1337)
+
+    args = parser.parse_args()
+    prepare(args)
+
+
+if __name__ == '__main__':
+    main()

src/diffusion_lm/claims.py

@@ -0,0 +1,512 @@
+"""Synthetic multi-turn claim intake: dialogues with a known fact set, and a scorer.
+
+The target behaviour is memory rather than knowledge: every fact the model must report is
+present in the conversation, so a small model is not being asked to recall the world. Some
+turns supersede an earlier value, which is what separates tracking state from copying the
+last thing seen. Because the record is generated, recall and precision are exact, and no
+judge is needed to score a recap.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import random
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from diffusion_lm.reasoning import ReasoningExample
+
+FIRST_NAMES = ('Marta', 'Diego', 'Luciana', 'Rodrigo', 'Camila', 'Nestor', 'Sofia', 'Ariel')
+LAST_NAMES = ('Quiroga', 'Benitez', 'Salvatierra', 'Uriarte', 'Ferreyra', 'Zabala', 'Otamendi')
+MAKES = (('Peugeot', '208'), ('Toyota', 'Etios'), ('Renault', 'Sandero'), ('Fiat', 'Cronos'),
+         ('Chevrolet', 'Onix'), ('Volkswagen', 'Gol'))
+STREETS = ('Av. Rivadavia', 'Calle Mitre', 'Ruta 8', 'Av. San Martin', 'Calle Belgrano')
+CITIES = ('Moron', 'Rosario', 'La Plata', 'Cordoba', 'Bahia Blanca', 'Mendoza')
+DAMAGE_PARTS = ('front bumper', 'left headlight', 'driver door', 'rear hatch', 'right mirror',
+                'windshield', 'rear bumper')
+WEATHER = ('heavy rain', 'clear skies', 'fog', 'light drizzle')
+FILLERS = ('Sorry, one moment.', 'Are you still there?', 'This is my first claim.',
+           'Ok.', 'Thanks for the help.', 'Can you repeat that?')
+ACKS = ('No problem, take your time.', 'Yes, I am here.', 'Understood, please go on.',
+        'Of course.', 'Happy to help.')
+
+
+@dataclass
+class Claim:
+    """One claim record plus the facts a faithful report has to carry."""
+
+    fields: dict[str, str]
+    corrected: dict[str, str] = field(default_factory=dict)
+
+    @property
+    def truth(self) -> dict[str, str]:
+        """Field values after corrections, which is what a report must state."""
+
+        resolved = dict(self.fields)
+        resolved.update(self.corrected)
+        return resolved
+
+
+def _plate(rng: random.Random) -> str:
+    letters = ''.join(rng.choice('ABCDEFGHJKLMNPRSTUVWXYZ') for _ in range(3))
+    return f'{letters}-{rng.randint(1000, 9999)}'
+
+
+def generate_claim(rng: random.Random) -> Claim:
+    """Build a claim whose values are distinctive enough to score by exact match."""
+
+    make, model = rng.choice(MAKES)
+    fields = {
+        'policy_number': f'PL-{rng.randint(10000, 99999)}',
+        'claimant': f'{rng.choice(FIRST_NAMES)} {rng.choice(LAST_NAMES)}',
+        'incident_date': f'{rng.randint(1, 28):02d}/{rng.randint(1, 12):02d}/2026',
+        'incident_time': f'{rng.randint(0, 23):02d}:{rng.choice(("05", "15", "40", "50"))}',
+        'location': f'{rng.choice(STREETS)}, {rng.choice(CITIES)}',
+        'vehicle': f'{make} {model} {rng.randint(2012, 2025)}',
+        'plate': _plate(rng),
+        'weather': rng.choice(WEATHER),
+        'damage': rng.choice(DAMAGE_PARTS),
+        'estimate': f'${rng.randint(2, 40) * 1000 + rng.choice((150, 450, 800)):,}',
+        'other_plate': _plate(rng),
+        'police_report': f'PR-{rng.randint(100000, 999999)}',
+        'witness': f'{rng.choice(FIRST_NAMES)} {rng.choice(LAST_NAMES)}',
+    }
+    return Claim(fields=fields)
+
+
+QUESTIONS = {
+    'policy_number': 'Can you give me your policy number?',
+    'claimant': 'Who is the policy holder?',
+    'incident_date': 'What date did this happen?',
+    'incident_time': 'Roughly what time?',
+    'location': 'Where exactly did it happen?',
+    'vehicle': 'Which vehicle was involved?',
+    'plate': "What is your vehicle's plate?",
+    'weather': 'How was the weather at the time?',
+    'damage': 'What part of the car was damaged?',
+    'estimate': 'Do you have a repair estimate?',
+    'other_plate': 'Did you get the other vehicle plate?',
+    'police_report': 'Was a police report filed?',
+    'witness': 'Was there any witness?',
+}
+
+ANSWERS = {
+    'policy_number': 'My policy is {value}.',
+    'claimant': 'The holder is {value}.',
+    'incident_date': 'It was on {value}.',
+    'incident_time': 'Around {value}.',
+    'location': 'On {value}.',
+    'vehicle': "It's a {value}.",
+    'plate': 'The plate is {value}.',
+    'weather': 'There was {value}.',
+    'damage': 'The {value} took the hit.',
+    'estimate': 'The shop quoted {value}.',
+    'other_plate': 'Yes, {value}.',
+    'police_report': 'Yes, report {value}.',
+    'witness': '{value} saw everything.',
+}
+
+
+def render_dialogue(
+    claim: Claim, rng: random.Random, *, turns: int = 6, corrections: int = 1,
+    chitchat: int = 2,
+) -> list[dict[str, str]]:
+    """Reveal the record across turns, superseding some values along the way.
+
+    Each turn carries one to three fields, so the amount to remember per turn varies the
+    way it would in a real intake.
+    """
+
+    keys = list(claim.fields)
+    rng.shuffle(keys)
+    # Deliberately uneven: with a constant number of facts per turn every ledger block lands on
+    # the same size and the control decision stops carrying information.
+    batches: list[list[str]] = []
+    while keys:
+        take = min(len(keys), rng.choice((1, 1, 2, 3, 4, 5)))
+        batches.append(keys[:take])
+        keys = keys[take:]
+
+    messages: list[dict[str, str]] = []
+    for batch in batches:
+        messages.append({'role': 'assistant', 'content': ' '.join(QUESTIONS[k] for k in batch)})
+        said = ' '.join(ANSWERS[k].format(value=claim.fields[k]) for k in batch)
+        messages.append({'role': 'user', 'content': said,
+                         'values': {k: claim.fields[k] for k in batch}})
+
+    # Turns that carry no new fact, so the assistant answers with zero thinking blocks. A
+    # conversation is full of these and without them the model never learns to skip thinking.
+    for _ in range(chitchat):
+        at = rng.randrange(1, max(2, len(messages)))
+        messages.insert(at, {'role': 'assistant', 'content': rng.choice(ACKS)})
+        messages.insert(at, {'role': 'user', 'content': rng.choice(FILLERS), 'values': {}})
+
+    correctable = [k for k in claim.fields if k in ('plate', 'other_plate', 'estimate',
+                                                    'incident_time', 'police_report')]
+    rng.shuffle(correctable)
+    for key in correctable[:corrections]:
+        if key.endswith('plate'):
+            new = _plate(rng)
+        elif key == 'estimate':
+            new = f'${rng.randint(2, 40) * 1000 + 700:,}'
+        elif key == 'incident_time':
+            new = f'{rng.randint(0, 23):02d}:30'
+        else:
+            new = f'PR-{rng.randint(100000, 999999)}'
+        claim.corrected[key] = new
+        messages.append({
+            'role': 'user',
+            'content': f'Sorry, I misspoke earlier: {ANSWERS[key].format(value=new)} '
+                       f'Not {claim.fields[key]}.',
+            'values': {key: new},
+        })
+        # An assistant turn right after the correction is what makes the update a training
+        # target: without it the only place the corrected value is taught is the final report.
+        messages.append({'role': 'assistant', 'content': 'Noted, I corrected that detail.'})
+    return messages
+
+
+_IDENTIFIERS = re.compile(
+    r'(?:P[LR]-\d{5,6}'                    # policy and police report numbers
+    r'|[A-Z]{3}-\d{4}'                     # plates
+    r'|\$[\d,]+'                           # money
+    r'|\d{2}/\d{2}/\d{4}'                  # dates
+    r'|\b\d{2}:\d{2}\b)'                   # times
+)
+
+
+def score_report(text: str, claim: Claim) -> dict[str, object]:
+    """Exact-match recall of the resolved values, plus the two ways a report lies.
+
+    ``stale_kept`` reports a value the conversation replaced, which means the model copied
+    instead of tracking state. ``invented`` counts identifier-shaped strings that were never
+    said at all, which is confabulation rather than a memory slip; they are different
+    failures and worth separating.
+    """
+
+    truth = claim.truth
+    found = {key: (value in text) for key, value in truth.items()}
+    stale = {
+        key: (claim.fields[key] in text)
+        for key in claim.corrected
+        if claim.fields[key] != claim.corrected[key]
+    }
+    spoken = set(truth.values()) | set(claim.fields.values())
+    said_ids = {token for value in spoken for token in _IDENTIFIERS.findall(value)}
+    invented = sorted({token for token in _IDENTIFIERS.findall(text)} - said_ids)
+    recalled = sum(found.values())
+    return {
+        'fields': len(truth),
+        'recalled': recalled,
+        'recall': recalled / max(1, len(truth)),
+        'missing': sorted(k for k, ok in found.items() if not ok),
+        'stale_kept': sorted(k for k, bad in stale.items() if bad),
+        'invented': invented,
+    }
+
+
+LEDGER = '{key}: {value}'
+KEEP_MESSAGES = 4
+IM_START = '<|im_start|>'
+IM_END = '<|im_end|>'
+SYSTEM = (
+    'You are a claim intake assistant. Track every detail the customer gives and always use '
+    'the corrected value when they correct themselves.'
+)
+
+
+def chatml_turn(role: str, content: str) -> str:
+    """One ChatML turn. Both markers are single tokens in the adaptive tokenizer."""
+
+    return f'{IM_START}{role}\n{content}{IM_END}\n'
+
+
+def ledger_line(entries: list[str]) -> str:
+    """The ``Known so far`` system-turn line carrying facts whose messages were dropped."""
+
+    return 'Known so far: ' + '; '.join(entries) + '.' if entries else ''
+
+
+_NOTE_FACT = re.compile(r'^([^:]{1,48}): (.+)$')
+
+
+def merge_notes(notes: list[str]) -> list[str]:
+    """Fold note fragments into one entry per key, latest value winning.
+
+    Concatenating raw notes would re-expose superseded values, which is the failure the corpus
+    charges hardest. Fragments that do not parse as ``key: value`` pass through in order,
+    deduplicated verbatim, so an unkeyed reasoning step still reaches the prefix as prose.
+
+    Both the training renderer and the playground call this. Keeping one implementation is the
+    point: a second copy is how the train and inference views of the same history drift apart.
+    """
+
+    facts: dict[str, str] = {}
+    loose: list[str] = []
+    for note in notes:
+        for fragment in note.split('; '):
+            fragment = fragment.strip().rstrip('.')
+            if not fragment:
+                continue
+            match = _NOTE_FACT.match(fragment)
+            if match:
+                facts[match.group(1)] = match.group(2)
+            elif fragment not in loose:
+                loose.append(fragment)
+    return [f'{key}: {value}' for key, value in facts.items()] + loose
+
+
+def window_start(messages: list[dict[str, str]], keep: int) -> int:
+    return 0 if keep <= 0 else max(0, len(messages) - keep)
+
+
+def ledger_notes(messages: list[dict[str, str]], keep: int) -> list[str]:
+    """Notes whose user message fell out of the visible window, in turn order.
+
+    An assistant turn's note enters the ledger exactly when the user message it took notes on
+    has fallen out of the window, which is what makes the note the only remaining carrier.
+    """
+
+    start = window_start(messages, keep)
+    return [
+        message['note'] for index, message in enumerate(messages)
+        if message['role'] == 'assistant' and message.get('note') and index - 1 < start
+    ]
+
+
+def chat_prefix(turns: list[dict[str, str]], *, system: str = SYSTEM, extra: str = '') -> str:
+    """Conversation prefix ending where the assistant's generation begins.
+
+    The accumulated ledger rides in the system turn rather than as a fake dialogue message:
+    it is persistent state, and putting it there also teaches the model to condition on a
+    system prompt, which none of the other corpora do.
+    """
+
+    merged = system if not extra else f'{system}\n{extra}'
+    rendered = [chatml_turn('system', merged)]
+    rendered += [chatml_turn(turn['role'], turn['content']) for turn in turns]
+    return ''.join(rendered) + f'{IM_START}assistant\n'
+
+
+def _chunks(items: list, size: int) -> list[list]:
+    return [items[i:i + size] for i in range(0, len(items), size)]
+
+
+def report_text(claim: Claim) -> str:
+    """The report a faithful assistant produces, one resolved field per line."""
+
+    return 'Claim report. ' + ' '.join(
+        LEDGER.format(key=key, value=value) + '.' for key, value in claim.truth.items()
+    )
+
+
+def to_examples(
+    claim: Claim, messages: list[dict[str, str]], *, keep_messages: int = KEEP_MESSAGES
+) -> list[ReasoningExample]:
+    """One training example per assistant turn, with the history deliberately truncated.
+
+    Older messages are dropped and replaced by the ledger of what they revealed, so the
+    accumulated notes — not the transcript — are what carries the past. That is the whole
+    point: with the full transcript in the prefix the model can re-read instead of remember,
+    and the thinking block stops being memory. A final example asks for the report, whose
+    thinking consolidates every fact.
+
+    Thought steps hold one fact each, which keeps a block's content short and its length a
+    function of how much the turn actually revealed.
+    """
+
+    examples: list[ReasoningExample] = []
+    # Values as stated at each point, so a ledger never shows a correction that has not
+    # happened yet: training on resolved values would teach the model to know the future.
+    known: dict[str, str] = {}
+    for index, message in enumerate(messages):
+        if message['role'] != 'assistant' or index == 0:
+            continue
+        previous = messages[index - 1]
+        learned = dict(previous.get('values', {}))
+        known.update(learned)
+        dropped = messages[max(0, index - keep_messages):index]
+        seen: dict[str, str] = {}
+        for earlier in messages[:max(0, index - keep_messages)]:
+            seen.update(earlier.get('values', {}))
+        older = [LEDGER.format(key=key, value=value) for key, value in seen.items()]
+        examples.append(ReasoningExample(
+            problem=chat_prefix(dropped, extra=ledger_line(older)),
+            steps=(('; '.join(LEDGER.format(key=key, value=value)
+                                for key, value in learned.items()),)
+                   if learned else ()),
+            answer=message['content'] + IM_END,
+            expected_answer='',
+        ))
+
+    tail = messages[-keep_messages:] + [{'role': 'user', 'content': 'Write the claim report.'}]
+    examples.append(ReasoningExample(
+        problem=chat_prefix(tail),
+        steps=tuple(
+            '; '.join(LEDGER.format(key=key, value=value) for key, value in group)
+            for group in _chunks(list(claim.truth.items()), 8)
+        ),
+        answer=report_text(claim) + IM_END,
+        expected_answer='',
+    ))
+    return examples
+
+
+def build(count: int, seed: int, turns: int, corrections: int) -> list[dict[str, object]]:
+    rng = random.Random(seed)
+    records = []
+    for index in range(count):
+        claim = generate_claim(rng)
+        messages = render_dialogue(claim, rng, turns=turns, corrections=corrections)
+        records.append({
+            'index': index,
+            'messages': messages,
+            'truth': claim.truth,
+            'superseded': {k: claim.fields[k] for k in claim.corrected},
+        })
+    return records
+
+
+def prepare(args: argparse.Namespace) -> None:
+    """Pack dialogues into the adaptive layout, splitting BY DIALOGUE.
+
+    Splitting by example would leak: two examples from one dialogue share its facts, so a
+    validation example's answer would already appear in a training example's prefix.
+    """
+
+    import numpy as np
+
+    from diffusion_lm.reasoning import ExampleEncoder, LayoutSpec, _write_packed
+    from diffusion_lm.tokenizer import load_tokenizer
+
+    tokenizer = load_tokenizer(args.tokenizer)
+    spec = LayoutSpec(seq_len=args.seq_len, block=min(args.sizes), max_slots=args.max_slots,
+                      sizes=tuple(sorted(args.sizes)))
+    encoder = ExampleEncoder(tokenizer, spec)
+    rng = random.Random(args.seed)
+
+    split: dict[str, list[tuple]] = {'train': [], 'validation': []}
+    dropped = 0
+    for index in range(args.count):
+        claim = generate_claim(rng)
+        messages = render_dialogue(claim, rng, turns=args.turns,
+                                   corrections=args.corrections,
+                                   chitchat=args.chitchat)
+        bucket = 'validation' if index % args.val_every == 0 else 'train'
+        for example in to_examples(claim, messages, keep_messages=args.keep_messages):
+            encoded = encoder.encode_adaptive(example)
+            if encoded is None:
+                dropped += 1
+                continue
+            split[bucket].append((encoded.tokens, encoded.regions))
+
+    for name, rows in split.items():
+        if not rows:
+            raise ValueError(f'no examples in the {name} split')
+        _write_packed(
+            args.output_dir / f'{name}-adaptive.bin',
+            np.stack([tokens for tokens, _ in rows]),
+            np.stack([regions for _, regions in rows]),
+            layout='adaptive', spec=spec, tokenizer_path=args.tokenizer, tokenizer=tokenizer,
+            extra_metadata={'sizes': list(spec.sizes), 'source': 'claims-chatml'},
+        )
+        print(f'{name}: {len(rows):,} examples -> {args.output_dir}')
+    print(f'{dropped:,} dropped at encode')
+
+
+def mix(args: argparse.Namespace) -> None:
+    """Concatenate two packs so claims hold ``--claims-share`` of the examples."""
+
+    import numpy as np
+
+    from diffusion_lm.reasoning import regions_path
+
+    rng = np.random.default_rng(args.seed)
+    for name in ('train', 'validation'):
+        parts = []
+        for directory, share in ((args.claims_dir, args.claims_share), (args.other_dir, None)):
+            path = directory / f'{name}-adaptive.bin'
+            meta = json.loads(Path(str(path) + '.json').read_text())
+            tokens = np.fromfile(path, dtype=np.dtype(meta['dtype'])).reshape(
+                meta['example_count'], meta['seq_len']
+            )
+            parts.append((tokens, np.load(regions_path(path)), share))
+        (claims_tokens, claims_regions, share), (other_tokens, other_regions, _) = parts
+        target = int(round(share / (1.0 - share) * len(other_tokens)))
+        if target < len(claims_tokens):
+            keep = rng.choice(len(claims_tokens), size=target, replace=False)
+            claims_tokens, claims_regions = claims_tokens[keep], claims_regions[keep]
+        tokens = np.concatenate([claims_tokens, other_tokens])
+        regions = np.concatenate([claims_regions, other_regions])
+        order = rng.permutation(len(tokens))
+        tokens, regions = tokens[order], regions[order]
+        args.output_dir.mkdir(parents=True, exist_ok=True)
+        tokens.tofile(args.output_dir / f'{name}-adaptive.bin')
+        np.save(regions_path(args.output_dir / f'{name}-adaptive.bin'), regions)
+        meta.update({'example_count': int(len(tokens)), 'source': 'claims+glaive',
+                     'claims_examples': int(len(claims_tokens)),
+                     'other_examples': int(len(other_tokens))})
+        Path(str(args.output_dir / f'{name}-adaptive.bin') + '.json').write_text(
+            json.dumps(meta, indent=2) + '\n'
+        )
+        actual = len(claims_tokens) / len(tokens)
+        print(f'{name}: {len(tokens):,} examples, claims share {actual:.3f}')
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    sub = parser.add_subparsers(dest='command', required=True)
+
+    dialogues = sub.add_parser('build', help='write dialogues as JSONL for probing')
+    dialogues.add_argument('--count', type=int, default=64)
+    dialogues.add_argument('--seed', type=int, default=1337)
+    dialogues.add_argument('--turns', type=int, default=6)
+    dialogues.add_argument('--corrections', type=int, default=1)
+    dialogues.add_argument('--output', type=Path, required=True)
+
+    pack = sub.add_parser('prepare', help='pack dialogues into the adaptive layout')
+    pack.add_argument('--count', type=int, default=2700)
+    pack.add_argument('--seed', type=int, default=1337)
+    pack.add_argument('--turns', type=int, default=6)
+    pack.add_argument('--corrections', type=int, default=1)
+    pack.add_argument('--chitchat', type=int, default=2,
+                      help='turns with no new fact, which train zero-block answers')
+    pack.add_argument('--keep-messages', type=int, default=KEEP_MESSAGES)
+    pack.add_argument('--tokenizer', type=Path, required=True)
+    pack.add_argument('--output-dir', type=Path, required=True)
+    pack.add_argument('--seq-len', type=int, default=2048)
+    pack.add_argument('--max-slots', type=int, default=64)
+    pack.add_argument('--sizes', type=int, nargs='+', default=[32, 64, 128])
+    pack.add_argument('--val-every', type=int, default=20,
+                      help='every Nth dialogue goes to validation, whole')
+
+    blend = sub.add_parser('mix', help='blend a claims pack into another pack')
+    blend.add_argument('--claims-dir', type=Path, required=True)
+    blend.add_argument('--other-dir', type=Path, required=True)
+    blend.add_argument('--output-dir', type=Path, required=True)
+    blend.add_argument('--claims-share', type=float, default=0.30)
+    blend.add_argument('--seed', type=int, default=1337)
+
+    args = parser.parse_args()
+    if args.command == 'prepare':
+        prepare(args)
+        return
+    if args.command == 'mix':
+        mix(args)
+        return
+
+    records = build(args.count, args.seed, args.turns, args.corrections)
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    with args.output.open('w', encoding='utf-8') as handle:
+        for record in records:
+            handle.write(json.dumps(record, ensure_ascii=False) + '\n')
+    turns = sum(len(r['messages']) for r in records) / len(records)
+    facts = sum(len(r['truth']) for r in records) / len(records)
+    print(f'wrote {len(records)} dialogues to {args.output} '
+          f'({turns:.1f} messages and {facts:.1f} facts each)')
+
+
+if __name__ == '__main__':
+    main()

src/diffusion_lm/config.py

@@ -0,0 +1,174 @@
+"""Typed experiment configuration loaded from YAML."""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+
+MODEL_BACKBONES = ("project", "hf-qwen3")
+
+
+@dataclass(frozen=True)
+class ModelConfig:
+    """Architecture for the bidirectional token denoiser."""
+
+    vocab_size: int
+    mask_token_id: int
+    max_seq_len: int
+    d_model: int
+    n_layers: int
+    n_heads: int
+    d_ff: int
+    dropout: float = 0.1
+    tie_embeddings: bool = True
+    activation_checkpointing: bool = False
+    use_flex_attention: bool = False
+    compile_backbone: bool = False
+    forbidden_output_token_ids: tuple[int, ...] = (0, 1, 2, 4)
+    backbone: str = "project"
+    pretrained_path: str | None = None
+    # First id of the untrained tail of a pretrained embedding matrix; logits from that id
+    # upward are masked without enumerating hundreds of forbidden ids.
+    forbidden_output_from: int | None = None
+
+    def __post_init__(self) -> None:
+        if self.vocab_size <= 1:
+            raise ValueError("vocab_size must be greater than one")
+        if self.backbone not in MODEL_BACKBONES:
+            raise ValueError(f"backbone must be one of {MODEL_BACKBONES}")
+        if self.backbone != "project" and self.pretrained_path is None:
+            raise ValueError("pretrained backbones require pretrained_path")
+        if self.forbidden_output_from is not None and not (
+            0 < self.forbidden_output_from <= self.vocab_size
+        ):
+            raise ValueError("forbidden_output_from must be inside the vocabulary")
+        if not 0 <= self.mask_token_id < self.vocab_size:
+            raise ValueError("mask_token_id must be inside the vocabulary")
+        if self.max_seq_len <= 0:
+            raise ValueError("max_seq_len must be positive")
+        if self.d_model <= 0 or self.n_layers <= 0 or self.n_heads <= 0 or self.d_ff <= 0:
+            raise ValueError("all model dimensions must be positive")
+        if self.d_model % self.n_heads != 0:
+            raise ValueError("d_model must be divisible by n_heads")
+        if not 0.0 <= self.dropout < 1.0:
+            raise ValueError("dropout must be in [0, 1)")
+        if not isinstance(self.activation_checkpointing, bool):
+            raise ValueError("activation_checkpointing must be a boolean")
+        if not isinstance(self.use_flex_attention, bool):
+            raise ValueError("use_flex_attention must be a boolean")
+        if not isinstance(self.compile_backbone, bool):
+            raise ValueError('compile_backbone must be a boolean')
+
+        forbidden_ids = tuple(self.forbidden_output_token_ids)
+        if any(
+            isinstance(token_id, bool) or not isinstance(token_id, int)
+            for token_id in forbidden_ids
+        ):
+            raise ValueError("forbidden_output_token_ids must contain integers")
+        if len(set(forbidden_ids)) != len(forbidden_ids):
+            raise ValueError("forbidden_output_token_ids must not contain duplicates")
+        if any(not 0 <= token_id < self.vocab_size for token_id in forbidden_ids):
+            raise ValueError("forbidden output token ids must be inside the vocabulary")
+        if self.mask_token_id not in forbidden_ids:
+            raise ValueError("mask_token_id must be a forbidden output token")
+        if self.backbone == "project" and 3 in forbidden_ids:
+            raise ValueError("EOS token id 3 must remain an allowed output token")
+        object.__setattr__(self, "forbidden_output_token_ids", forbidden_ids)
+
+
+@dataclass(frozen=True)
+class TrainingConfig:
+    """Single-device training settings for the first research iteration."""
+
+    train_data: str
+    tokenizer: str
+    output_dir: str = "outputs/run"
+    val_data: str | None = None
+    batch_size: int = 32
+    gradient_accumulation_steps: int = 1
+    max_steps: int = 10_000
+    learning_rate: float = 3e-4
+    min_learning_rate: float = 3e-5
+    warmup_steps: int = 500
+    weight_decay: float = 0.1
+    grad_clip: float = 1.0
+    mask_eps: float = 1e-3
+    seed: int = 1337
+    device: str = "auto"
+    precision: str = "auto"
+    optimizer: str = "adamw"
+    optimizer_min_8bit_size: int = 4096
+    optimizer_embedding_32bit: bool = True
+    require_fused_attention: bool = False
+    save_inference_checkpoint: bool = False
+    num_workers: int = 0
+    log_interval: int = 10
+    eval_interval: int = 500
+    eval_batches: int = 20
+    save_interval: int = 500
+    keep_last_checkpoints: int = 3
+
+    def __post_init__(self) -> None:
+        if self.batch_size <= 0 or self.gradient_accumulation_steps <= 0:
+            raise ValueError("batch sizes must be positive")
+        if self.max_steps <= 0:
+            raise ValueError("max_steps must be positive")
+        if not 0.0 < self.learning_rate:
+            raise ValueError("learning_rate must be positive")
+        if not 0.0 <= self.min_learning_rate <= self.learning_rate:
+            raise ValueError("min_learning_rate must be between zero and learning_rate")
+        if not 0 <= self.warmup_steps < self.max_steps:
+            raise ValueError("warmup_steps must be non-negative and less than max_steps")
+        if self.weight_decay < 0.0:
+            raise ValueError("weight_decay must be non-negative")
+        if self.grad_clip <= 0.0:
+            raise ValueError("grad_clip must be positive")
+        if not 0.0 < self.mask_eps < 1.0:
+            raise ValueError("mask_eps must be in (0, 1)")
+        if self.num_workers < 0:
+            raise ValueError("num_workers must be non-negative")
+        if self.log_interval <= 0 or self.eval_interval <= 0 or self.save_interval <= 0:
+            raise ValueError("log, eval, and save intervals must be positive")
+        if self.eval_batches <= 0:
+            raise ValueError("eval_batches must be positive")
+        if self.keep_last_checkpoints < 0:
+            raise ValueError("keep_last_checkpoints must be non-negative")
+        if self.precision not in {"auto", "float32", "bfloat16", "float16"}:
+            raise ValueError("precision must be auto, float32, bfloat16, or float16")
+        if self.optimizer not in {"adamw", "adamw8bit"}:
+            raise ValueError("optimizer must be adamw or adamw8bit")
+        if self.optimizer_min_8bit_size <= 0:
+            raise ValueError("optimizer_min_8bit_size must be positive")
+        if not isinstance(self.optimizer_embedding_32bit, bool):
+            raise ValueError("optimizer_embedding_32bit must be a boolean")
+        if not isinstance(self.require_fused_attention, bool):
+            raise ValueError("require_fused_attention must be a boolean")
+        if not isinstance(self.save_inference_checkpoint, bool):
+            raise ValueError("save_inference_checkpoint must be a boolean")
+
+
+@dataclass(frozen=True)
+class ExperimentConfig:
+    model: ModelConfig
+    training: TrainingConfig
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+def load_config(path: str | Path) -> ExperimentConfig:
+    """Load and validate an experiment YAML file."""
+
+    config_path = Path(path)
+    with config_path.open("r", encoding="utf-8") as handle:
+        raw = yaml.safe_load(handle)
+    if not isinstance(raw, dict) or "model" not in raw or "training" not in raw:
+        raise ValueError("config must contain top-level model and training mappings")
+    return ExperimentConfig(
+        model=ModelConfig(**raw["model"]),
+        training=TrainingConfig(**raw["training"]),
+    )

src/diffusion_lm/corpus.py

@@ -0,0 +1,920 @@
+"""Download and prepare the pinned FineWeb-Edu 10B sample.
+
+The heavy ``huggingface_hub`` and ``pyarrow`` dependencies are imported only by the operations that
+need them. Raw Parquet files are cached locally, then converted into independently replaceable
+uint16/uint32 shards with deterministic document-level train/validation assignment.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Iterable, Iterator, Literal, Sequence
+
+import numpy as np
+
+from diffusion_lm.data import PACKED_MANIFEST_FORMAT
+from diffusion_lm.tokenizer import (
+    load_tokenizer,
+    special_token_ids,
+    train_tokenizer_from_iterator,
+)
+
+
+FINEWEB_EDU_REPO_ID = "HuggingFaceFW/fineweb-edu"
+FINEWEB_EDU_CONFIG = "sample-10BT"
+FINEWEB_EDU_REVISION = "87f09149ef4734204d70ed1d046ddc9ca3f2b8f9"
+FINEWEB_EDU_PATH_PREFIX = "sample/10BT/"
+SOURCE_STATE_FORMAT = "mini-diffusion-lm-corpus-source-v1"
+SPLIT_HASH_PERSON = b"mini-mdlm-split"
+
+
+@dataclass(frozen=True)
+class CorpusSource:
+    """One pinned Hub dataset: where its Parquets live and how rows are read.
+
+    ``id_column=None`` derives the split id from a sha256 of the text, which keeps the
+    train/validation assignment order-independent for datasets without a stable row id.
+    """
+
+    name: str
+    repo_id: str
+    revision: str
+    path_prefix: str
+    text_column: str = "text"
+    id_column: str | None = "id"
+    config: str | None = None
+
+
+CORPUS_SOURCES: dict[str, CorpusSource] = {
+    source.name: source
+    for source in (
+        CorpusSource(
+            name="fineweb-edu",
+            repo_id=FINEWEB_EDU_REPO_ID,
+            revision=FINEWEB_EDU_REVISION,
+            path_prefix=FINEWEB_EDU_PATH_PREFIX,
+            config=FINEWEB_EDU_CONFIG,
+        ),
+        CorpusSource(
+            name="ultra-fineweb-en",
+            repo_id="openbmb/Ultra-FineWeb",
+            revision="7ddd4170ce03e0afbd7d9b80d4bc0b8eebf877e4",
+            path_prefix="data/ultrafineweb_en/",
+            text_column="content",
+            id_column=None,
+        ),
+        CorpusSource(
+            name="cosmopedia-v2",
+            repo_id="HuggingFaceTB/smollm-corpus",
+            revision="3ba9d605774198c5868892d7a8deda78031a781f",
+            path_prefix="cosmopedia-v2/",
+            id_column=None,
+            config="cosmopedia-v2",
+        ),
+        CorpusSource(
+            name="finemath-4plus",
+            repo_id="HuggingFaceTB/finemath",
+            revision="e92b25a616738fe95dc186b64dfb19f9c8525594",
+            path_prefix="finemath-4plus/",
+            id_column=None,
+            config="finemath-4plus",
+        ),
+    )
+}
+
+
+def _require_huggingface_hub():
+    try:
+        from huggingface_hub import HfApi, snapshot_download
+    except ImportError as exc:  # pragma: no cover - exercised in minimal installations.
+        raise RuntimeError(
+            "FineWeb-Edu download requires huggingface_hub; install the corpus dependencies"
+        ) from exc
+    return HfApi, snapshot_download
+
+
+def _require_parquet():
+    try:
+        import pyarrow.parquet as parquet
+    except ImportError as exc:  # pragma: no cover - exercised in minimal installations.
+        raise RuntimeError(
+            "FineWeb-Edu preparation requires pyarrow; install the corpus dependencies"
+        ) from exc
+    return parquet
+
+
+def _atomic_json(path: Path, value: dict[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_name(f".{path.name}.tmp")
+    temporary.unlink(missing_ok=True)
+    with temporary.open("w", encoding="utf-8") as handle:
+        json.dump(value, handle, indent=2, sort_keys=True)
+        handle.write("\n")
+        handle.flush()
+        os.fsync(handle.fileno())
+    temporary.replace(path)
+
+
+def _read_json(path: Path) -> dict[str, Any]:
+    try:
+        with path.open("r", encoding="utf-8") as handle:
+            value = json.load(handle)
+    except (OSError, json.JSONDecodeError) as exc:
+        raise ValueError(f"could not read corpus state {path}: {exc}") from exc
+    if not isinstance(value, dict):
+        raise ValueError(f"corpus state must be a JSON object: {path}")
+    return value
+
+
+def _sha256_file(path: Path, *, chunk_size: int = 1024 * 1024) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(chunk_size), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def download_source(
+    source: CorpusSource,
+    raw_dir: str | Path,
+    *,
+    max_files: int | None = None,
+    max_workers: int = 1,
+) -> list[Path]:
+    """Download a pinned source's Parquets into ``raw_dir`` using the Hub cache."""
+
+    if max_files is not None and max_files <= 0:
+        raise ValueError("max_files must be positive")
+    if max_workers <= 0:
+        raise ValueError("max_workers must be positive")
+    HfApi, snapshot_download = _require_huggingface_hub()
+    api = HfApi()
+    repository_files = api.list_repo_files(
+        repo_id=source.repo_id,
+        repo_type="dataset",
+        revision=source.revision,
+    )
+    parquet_names = sorted(
+        name
+        for name in repository_files
+        if name.startswith(source.path_prefix) and name.endswith(".parquet")
+    )
+    if max_files is not None:
+        parquet_names = parquet_names[:max_files]
+    if not parquet_names:
+        raise RuntimeError(
+            f"no Parquet files found for {source.repo_id}@{source.revision} "
+            f"under {source.path_prefix}"
+        )
+
+    destination = Path(raw_dir)
+    destination.mkdir(parents=True, exist_ok=True)
+    snapshot_root = Path(
+        snapshot_download(
+            repo_id=source.repo_id,
+            repo_type="dataset",
+            revision=source.revision,
+            allow_patterns=parquet_names,
+            local_dir=str(destination),
+            max_workers=max_workers,
+        )
+    )
+    paths: list[Path] = []
+    for name in parquet_names:
+        path = snapshot_root / name
+        if not path.is_file():
+            raise FileNotFoundError(f"Hub download did not produce {path}")
+        paths.append(path)
+    return paths
+
+
+def download_fineweb_edu(
+    raw_dir: str | Path,
+    *,
+    revision: str = FINEWEB_EDU_REVISION,
+    max_files: int | None = None,
+    max_workers: int = 1,
+) -> list[Path]:
+    """Download the pinned FineWeb-Edu sample (compatibility entry point)."""
+
+    from dataclasses import replace
+
+    source = replace(CORPUS_SOURCES["fineweb-edu"], revision=revision)
+    return download_source(source, raw_dir, max_files=max_files, max_workers=max_workers)
+
+
+def find_local_parquets(
+    raw_dir: str | Path, path_prefix: str = FINEWEB_EDU_PATH_PREFIX
+) -> list[Path]:
+    """Find already downloaded source Parquets in stable source order."""
+
+    root = Path(raw_dir)
+    preferred = sorted((root / path_prefix).glob("*.parquet"))
+    paths = preferred or sorted(root.rglob("*.parquet"))
+    if not paths:
+        raise FileNotFoundError(f"no Parquet files found below {root}")
+    return paths
+
+
+def document_split(
+    document_id: str,
+    *,
+    validation_modulus: int = 1024,
+    validation_bucket: int = 0,
+) -> Literal["train", "validation"]:
+    """Assign a stable document ID to train or validation without depending on row order."""
+
+    if not document_id:
+        raise ValueError("document_id must be non-empty")
+    if validation_modulus <= 1:
+        raise ValueError("validation_modulus must be greater than one")
+    if not 0 <= validation_bucket < validation_modulus:
+        raise ValueError("validation_bucket must be inside validation_modulus")
+    digest = hashlib.blake2b(
+        document_id.encode("utf-8"), digest_size=8, person=SPLIT_HASH_PERSON
+    ).digest()
+    bucket = int.from_bytes(digest, "little") % validation_modulus
+    return "validation" if bucket == validation_bucket else "train"
+
+
+def _iter_parquet_batches(
+    paths: Sequence[Path],
+    *,
+    batch_size: int,
+    text_column: str = "text",
+    id_column: str | None = "id",
+) -> Iterator[tuple[Path, list[str], list[str]]]:
+    if batch_size <= 0:
+        raise ValueError("batch_size must be positive")
+    parquet = _require_parquet()
+    read_columns = [text_column] if id_column is None else [id_column, text_column]
+    for path in paths:
+        source = parquet.ParquetFile(path)
+        try:
+            batches = source.iter_batches(batch_size=batch_size, columns=read_columns)
+            for batch in batches:
+                columns = batch.to_pydict()
+                texts = columns[text_column]
+                if id_column is not None:
+                    ids = columns[id_column]
+                else:
+                    # Content-derived ids keep the split assignment order-independent.
+                    ids = [
+                        hashlib.sha256(text.encode("utf-8")).hexdigest()
+                        if isinstance(text, str)
+                        else ""
+                        for text in texts
+                    ]
+                if len(ids) != len(texts):
+                    raise ValueError(f"mismatched id/text columns in {path}")
+                yield path, ids, texts
+        except (KeyError, ValueError) as exc:
+            raise ValueError(
+                f"expected columns {read_columns} in {path}: {exc}"
+            ) from exc
+
+
+def iter_tokenizer_text(
+    paths: Iterable[str | Path],
+    *,
+    max_utf8_bytes: int = 1 << 29,
+    batch_size: int = 512,
+    validation_modulus: int = 1024,
+    validation_bucket: int = 0,
+    text_column: str = "text",
+    id_column: str | None = "id",
+) -> Iterator[list[str]]:
+    """Yield bounded train-only text batches for iterator-based tokenizer training."""
+
+    if max_utf8_bytes <= 0:
+        raise ValueError("max_utf8_bytes must be positive")
+    sources = sorted(Path(path) for path in paths)
+    used_bytes = 0
+    output: list[str] = []
+    for _path, ids, texts in _iter_parquet_batches(
+        sources, batch_size=batch_size, text_column=text_column, id_column=id_column
+    ):
+        for document_id, text in zip(ids, texts, strict=True):
+            if not isinstance(document_id, str) or not isinstance(text, str) or not text:
+                continue
+            if document_split(
+                document_id,
+                validation_modulus=validation_modulus,
+                validation_bucket=validation_bucket,
+            ) != "train":
+                continue
+            encoded_bytes = len(text.encode("utf-8"))
+            if used_bytes and used_bytes + encoded_bytes > max_utf8_bytes:
+                if output:
+                    yield output
+                return
+            output.append(text)
+            used_bytes += encoded_bytes
+            if len(output) >= batch_size:
+                yield output
+                output = []
+            if used_bytes >= max_utf8_bytes:
+                if output:
+                    yield output
+                return
+    if output:
+        yield output
+
+
+@dataclass(frozen=True)
+class CorpusPreparationResult:
+    train_manifest: Path
+    validation_manifest: Path
+    processed_sources: int
+    resumed_sources: int
+
+
+class _SplitWriter:
+    def __init__(self, final_path: Path, dtype: np.dtype[Any]) -> None:
+        self.final_path = final_path
+        self.dtype = dtype
+        self.temporary_path = final_path.with_name(f".{final_path.name}.tmp")
+        final_path.parent.mkdir(parents=True, exist_ok=True)
+        self.temporary_path.unlink(missing_ok=True)
+        self.handle = self.temporary_path.open("wb")
+        self.digest = hashlib.sha256()
+        self.token_count = 0
+        self.document_count = 0
+
+    def append(self, token_ids: list[int]) -> None:
+        payload = np.asarray(token_ids, dtype=self.dtype).tobytes()
+        self.handle.write(payload)
+        self.digest.update(payload)
+        self.token_count += len(token_ids)
+        self.document_count += 1
+
+    def finish(self) -> dict[str, Any]:
+        self.handle.flush()
+        os.fsync(self.handle.fileno())
+        self.handle.close()
+        self.temporary_path.replace(self.final_path)
+        return {
+            "token_count": self.token_count,
+            "document_count": self.document_count,
+            "sha256": self.digest.hexdigest(),
+        }
+
+    def abort(self) -> None:
+        if not self.handle.closed:
+            self.handle.close()
+        self.temporary_path.unlink(missing_ok=True)
+
+
+def _relative_path(path: Path, root: Path) -> str:
+    try:
+        return path.relative_to(root).as_posix()
+    except ValueError:
+        return str(path)
+
+
+def _completed_source_state(
+    state_path: Path,
+    *,
+    source: Path,
+    corpus_source: CorpusSource,
+    output_dir: Path,
+    tokenizer_sha256: str,
+    validation_modulus: int,
+    validation_bucket: int,
+) -> dict[str, Any] | None:
+    if not state_path.is_file():
+        return None
+    try:
+        state = _read_json(state_path)
+    except ValueError:
+        return None
+    if (
+        state.get("format") != SOURCE_STATE_FORMAT
+        or state.get("source_name") != source.name
+        or state.get("source_size") != source.stat().st_size
+        # States written before multi-source support carry no dataset name.
+        or state.get("source_dataset", "fineweb-edu") != corpus_source.name
+        or state.get("dataset_revision") != corpus_source.revision
+        or state.get("tokenizer_sha256") != tokenizer_sha256
+        or state.get("split_hash_person") != SPLIT_HASH_PERSON.decode("ascii")
+        or state.get("validation_modulus") != validation_modulus
+        or state.get("validation_bucket") != validation_bucket
+    ):
+        return None
+    splits = state.get("splits")
+    if not isinstance(splits, dict):
+        return None
+    try:
+        dtype = np.dtype(state.get("dtype"))
+    except TypeError:
+        return None
+    for split in ("train", "validation"):
+        shard = splits.get(split)
+        if not isinstance(shard, dict) or not isinstance(shard.get("path"), str):
+            return None
+        path = output_dir / shard["path"]
+        expected_bytes = int(shard.get("token_count", -1)) * dtype.itemsize
+        expected_sha256 = shard.get("sha256")
+        if (
+            not path.is_file()
+            or path.stat().st_size != expected_bytes
+            or not isinstance(expected_sha256, str)
+            or _sha256_file(path) != expected_sha256
+        ):
+            return None
+    return state
+
+
+def _encode_source(
+    source: Path,
+    *,
+    corpus_source: CorpusSource,
+    tokenizer_path: Path,
+    output_dir: Path,
+    batch_size: int,
+    validation_modulus: int,
+    validation_bucket: int,
+) -> tuple[dict[str, Any], bool]:
+    tokenizer_sha256 = hashlib.sha256(tokenizer_path.read_bytes()).hexdigest()
+    # Hub Parquet names are unique and stable. Avoid list indexes so a partial download can be
+    # expanded later without invalidating already completed source shards.
+    source_key = source.stem
+    state_path = output_dir / "state" / f"{source_key}.json"
+    resumed = _completed_source_state(
+        state_path,
+        source=source,
+        corpus_source=corpus_source,
+        output_dir=output_dir,
+        tokenizer_sha256=tokenizer_sha256,
+        validation_modulus=validation_modulus,
+        validation_bucket=validation_bucket,
+    )
+    if resumed is not None:
+        return resumed, True
+
+    tokenizer = load_tokenizer(tokenizer_path)
+    role_ids = special_token_ids(tokenizer)
+    reserved_ids = set(role_ids.values())
+    vocab_size = tokenizer.get_vocab_size(with_added_tokens=True)
+    dtype = np.dtype("uint16" if vocab_size <= np.iinfo(np.uint16).max else "uint32")
+    final_paths = {
+        split: output_dir / "shards" / split / f"{source_key}.bin"
+        for split in ("train", "validation")
+    }
+    writers = {split: _SplitWriter(path, dtype) for split, path in final_paths.items()}
+    skipped_empty = 0
+    skipped_invalid = 0
+    skipped_special = 0
+    rows_seen = 0
+
+    try:
+        for _path, ids, texts in _iter_parquet_batches(
+            [source],
+            batch_size=batch_size,
+            text_column=corpus_source.text_column,
+            id_column=corpus_source.id_column,
+        ):
+            valid_rows: list[tuple[str, str]] = []
+            for document_id, text in zip(ids, texts, strict=True):
+                rows_seen += 1
+                if not isinstance(document_id, str) or not document_id:
+                    skipped_invalid += 1
+                elif not isinstance(text, str):
+                    skipped_invalid += 1
+                elif not text:
+                    skipped_empty += 1
+                else:
+                    valid_rows.append((document_id, text))
+            if not valid_rows:
+                continue
+            encodings = tokenizer.encode_batch(
+                [text for _document_id, text in valid_rows], add_special_tokens=False
+            )
+            for (document_id, _text), encoding in zip(valid_rows, encodings, strict=True):
+                token_ids = encoding.ids
+                if reserved_ids.intersection(token_ids):
+                    # The new sentinel strings make this practically impossible, but skipping is
+                    # preferable to losing hours of preprocessing if an exact literal is present.
+                    skipped_special += 1
+                    continue
+                split = document_split(
+                    document_id,
+                    validation_modulus=validation_modulus,
+                    validation_bucket=validation_bucket,
+                )
+                writers[split].append([*token_ids, role_ids["eos"]])
+        split_metadata = {split: writer.finish() for split, writer in writers.items()}
+    except BaseException:
+        for writer in writers.values():
+            writer.abort()
+        raise
+
+    for split, metadata in split_metadata.items():
+        metadata["path"] = _relative_path(final_paths[split], output_dir)
+        metadata["source"] = source.name
+    state: dict[str, Any] = {
+        "format": SOURCE_STATE_FORMAT,
+        "source_name": source.name,
+        "source_size": source.stat().st_size,
+        "source_rows": rows_seen,
+        "source_dataset": corpus_source.name,
+        "dataset_revision": corpus_source.revision,
+        "dtype": dtype.name,
+        "vocab_size": vocab_size,
+        "mask_token_id": role_ids["mask"],
+        "eos_token_id": role_ids["eos"],
+        "special_token_ids": role_ids,
+        "tokenizer_sha256": tokenizer_sha256,
+        "split_hash_person": SPLIT_HASH_PERSON.decode("ascii"),
+        "validation_modulus": validation_modulus,
+        "validation_bucket": validation_bucket,
+        "skipped_empty_documents": skipped_empty,
+        "skipped_invalid_documents": skipped_invalid,
+        "skipped_special_documents": skipped_special,
+        "splits": split_metadata,
+    }
+    # The marker is written last: its presence commits both split files as one source unit.
+    _atomic_json(state_path, state)
+    return state, False
+
+
+def _build_manifest(
+    split: Literal["train", "validation"],
+    *,
+    corpus_source: CorpusSource,
+    source_paths: Sequence[Path],
+    source_states: Sequence[dict[str, Any]],
+    validation_modulus: int,
+    validation_bucket: int,
+) -> dict[str, Any]:
+    first = source_states[0]
+    compatible_keys = (
+        "dtype",
+        "vocab_size",
+        "mask_token_id",
+        "eos_token_id",
+        "special_token_ids",
+        "tokenizer_sha256",
+    )
+    for state in source_states[1:]:
+        if any(state.get(key) != first.get(key) for key in compatible_keys):
+            raise ValueError("source states use incompatible tokenizer or token formats")
+    shards = [dict(state["splits"][split]) for state in source_states]
+    return {
+        "format": PACKED_MANIFEST_FORMAT,
+        "split": split,
+        "dtype": first["dtype"],
+        "token_count": sum(int(shard["token_count"]) for shard in shards),
+        "document_count": sum(int(shard["document_count"]) for shard in shards),
+        "vocab_size": first["vocab_size"],
+        "mask_token_id": first["mask_token_id"],
+        "eos_token_id": first["eos_token_id"],
+        "special_token_ids": first["special_token_ids"],
+        "tokenizer_sha256": first["tokenizer_sha256"],
+        "dataset": {
+            "repo_id": corpus_source.repo_id,
+            "config": corpus_source.config,
+            "revision": corpus_source.revision,
+            "path_prefix": corpus_source.path_prefix,
+        },
+        "split_rule": {
+            "algorithm": "blake2b-64",
+            "person": SPLIT_HASH_PERSON.decode("ascii"),
+            "validation_modulus": validation_modulus,
+            "validation_bucket": validation_bucket,
+        },
+        "source_files": [path.name for path in source_paths],
+        "skipped_documents": {
+            reason: sum(int(state[reason]) for state in source_states)
+            for reason in (
+                "skipped_empty_documents",
+                "skipped_invalid_documents",
+                "skipped_special_documents",
+            )
+        },
+        "shards": shards,
+    }
+
+
+def prepare_corpus(
+    tokenizer_path: str | Path,
+    output_dir: str | Path,
+    *,
+    corpus_source: CorpusSource,
+    source_paths: Iterable[str | Path] | None = None,
+    raw_dir: str | Path | None = None,
+    batch_size: int = 256,
+    validation_modulus: int = 1024,
+    validation_bucket: int = 0,
+    max_files: int | None = None,
+) -> CorpusPreparationResult:
+    """Convert one pinned source's Parquets into resumable train/validation manifests."""
+
+    # Validate split arguments before performing any download.
+    document_split(
+        "argument-validation",
+        validation_modulus=validation_modulus,
+        validation_bucket=validation_bucket,
+    )
+    tokenizer = Path(tokenizer_path)
+    load_tokenizer(tokenizer)
+    destination = Path(output_dir)
+    destination.mkdir(parents=True, exist_ok=True)
+
+    if source_paths is None:
+        raw = Path(raw_dir) if raw_dir is not None else destination / "raw"
+        try:
+            sources = find_local_parquets(raw, corpus_source.path_prefix)
+        except FileNotFoundError:
+            sources = download_source(corpus_source, raw, max_files=max_files)
+    else:
+        sources = sorted(Path(path) for path in source_paths)
+    if not sources:
+        raise ValueError("at least one source Parquet is required")
+    missing = [str(path) for path in sources if not path.is_file()]
+    if missing:
+        raise FileNotFoundError(f"missing source Parquets: {missing}")
+    source_keys = [source.stem for source in sources]
+    if len(set(source_keys)) != len(source_keys):
+        raise ValueError("source Parquet filenames must have unique stems")
+
+    states: list[dict[str, Any]] = []
+    resumed_sources = 0
+    for source in sources:
+        state, resumed = _encode_source(
+            source,
+            corpus_source=corpus_source,
+            tokenizer_path=tokenizer,
+            output_dir=destination,
+            batch_size=batch_size,
+            validation_modulus=validation_modulus,
+            validation_bucket=validation_bucket,
+        )
+        states.append(state)
+        resumed_sources += int(resumed)
+
+    manifest_paths = {
+        "train": destination / "train.manifest.json",
+        "validation": destination / "validation.manifest.json",
+    }
+    for split, manifest_path in manifest_paths.items():
+        manifest = _build_manifest(
+            split,  # type: ignore[arg-type]
+            corpus_source=corpus_source,
+            source_paths=sources,
+            source_states=states,
+            validation_modulus=validation_modulus,
+            validation_bucket=validation_bucket,
+        )
+        _atomic_json(manifest_path, manifest)
+    return CorpusPreparationResult(
+        train_manifest=manifest_paths["train"],
+        validation_manifest=manifest_paths["validation"],
+        processed_sources=len(sources) - resumed_sources,
+        resumed_sources=resumed_sources,
+    )
+
+
+def prepare_fineweb_edu(
+    tokenizer_path: str | Path,
+    output_dir: str | Path,
+    *,
+    source_paths: Iterable[str | Path] | None = None,
+    raw_dir: str | Path | None = None,
+    batch_size: int = 256,
+    validation_modulus: int = 1024,
+    validation_bucket: int = 0,
+) -> CorpusPreparationResult:
+    """Convert pinned FineWeb-Edu Parquets into manifests (compatibility entry point)."""
+
+    return prepare_corpus(
+        tokenizer_path,
+        output_dir,
+        corpus_source=CORPUS_SOURCES["fineweb-edu"],
+        source_paths=source_paths,
+        raw_dir=raw_dir,
+        batch_size=batch_size,
+        validation_modulus=validation_modulus,
+        validation_bucket=validation_bucket,
+    )
+
+
+def parse_token_budget(text: str) -> int:
+    """Parse a token count with an optional K/M/B suffix (e.g. ``2.5B``, ``500M``)."""
+
+    value = text.strip().upper()
+    factor = 1
+    if value and value[-1] in "KMB":
+        factor = {"K": 1_000, "M": 1_000_000, "B": 1_000_000_000}[value[-1]]
+        value = value[:-1]
+    try:
+        tokens = int(float(value) * factor)
+    except ValueError as exc:
+        raise ValueError(f"invalid token budget {text!r}") from exc
+    if tokens <= 0:
+        raise ValueError(f"token budget must be positive: {text!r}")
+    return tokens
+
+
+def mix_manifests(
+    inputs: Sequence[tuple[Path, int | None]], output_dir: str | Path
+) -> tuple[Path, Path]:
+    """Combine prepared source directories into one mixture manifest pair.
+
+    Each input contributes whole train shards in manifest order until its token budget
+    is met (``None`` takes everything), so realized counts overshoot a budget by at most
+    one shard; the overshoot is printed, never silent. Validation includes the
+    validation shards of the sources whose train shards were selected. Shard paths in
+    the mixed manifests are absolute so the inputs can live anywhere.
+    """
+
+    from diffusion_lm.data import load_packed_manifest
+
+    if not inputs:
+        raise ValueError("at least one input directory is required")
+    destination = Path(output_dir)
+    destination.mkdir(parents=True, exist_ok=True)
+    compatible_keys = (
+        "dtype",
+        "vocab_size",
+        "mask_token_id",
+        "eos_token_id",
+        "special_token_ids",
+        "tokenizer_sha256",
+    )
+    reference: dict[str, Any] | None = None
+    shards: dict[str, list[dict[str, Any]]] = {"train": [], "validation": []}
+    components: list[dict[str, Any]] = []
+    for input_dir, budget in inputs:
+        manifests = {
+            split: load_packed_manifest(input_dir / f"{split}.manifest.json")
+            for split in ("train", "validation")
+        }
+        if reference is None:
+            reference = manifests["train"]
+        elif any(
+            manifests["train"].get(key) != reference.get(key) for key in compatible_keys
+        ):
+            raise ValueError(f"{input_dir} uses an incompatible tokenizer or token format")
+
+        taken = 0
+        selected_sources: set[str] = set()
+        skipped = 0
+        for shard in manifests["train"]["shards"]:
+            if budget is not None and taken >= budget:
+                skipped += 1
+                continue
+            entry = dict(shard)
+            entry["path"] = str((input_dir / entry["path"]).resolve())
+            shards["train"].append(entry)
+            taken += int(entry["token_count"])
+            selected_sources.add(str(entry.get("source")))
+        for shard in manifests["validation"]["shards"]:
+            if str(shard.get("source")) not in selected_sources:
+                continue
+            entry = dict(shard)
+            entry["path"] = str((input_dir / entry["path"]).resolve())
+            shards["validation"].append(entry)
+        component = {
+            "dataset": manifests["train"].get("dataset"),
+            "directory": str(Path(input_dir).resolve()),
+            "token_budget": budget,
+            "train_tokens": taken,
+            "skipped_shards": skipped,
+        }
+        components.append(component)
+        print(
+            f'{input_dir}: {taken:,} train tokens'
+            + (f" (budget {budget:,}, {skipped} shards skipped)" if budget else "")
+        )
+
+    assert reference is not None
+    manifest_paths: dict[str, Path] = {}
+    for split in ("train", "validation"):
+        manifest = {
+            "format": PACKED_MANIFEST_FORMAT,
+            "split": split,
+            "dtype": reference["dtype"],
+            "token_count": sum(int(shard["token_count"]) for shard in shards[split]),
+            "document_count": sum(int(shard["document_count"]) for shard in shards[split]),
+            "vocab_size": reference["vocab_size"],
+            "mask_token_id": reference["mask_token_id"],
+            "eos_token_id": reference["eos_token_id"],
+            "special_token_ids": reference["special_token_ids"],
+            "tokenizer_sha256": reference["tokenizer_sha256"],
+            "dataset": {"name": "mixture", "components": components},
+            "split_rule": reference.get("split_rule"),
+            "source_files": [shard.get("source") for shard in shards[split]],
+            "shards": shards[split],
+        }
+        manifest_paths[split] = destination / f"{split}.manifest.json"
+        _atomic_json(manifest_paths[split], manifest)
+    return manifest_paths["train"], manifest_paths["validation"]
+
+
+def _parse_mix_input(text: str) -> tuple[Path, int | None]:
+    directory, separator, budget = text.partition("=")
+    return Path(directory), parse_token_budget(budget) if separator else None
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=__doc__)
+    commands = parser.add_subparsers(dest="command", required=True)
+    source_names = tuple(CORPUS_SOURCES)
+
+    download = commands.add_parser("download", help="download a pinned source's Parquets")
+    download.add_argument("--raw-dir", type=Path, required=True)
+    download.add_argument("--source", choices=source_names, default="fineweb-edu")
+    download.add_argument("--max-files", type=int)
+    download.add_argument("--max-workers", type=int, default=1)
+
+    tokenizer = commands.add_parser(
+        "train-tokenizer", help="train a 32K tokenizer from cached Parquets"
+    )
+    tokenizer.add_argument("--raw-dir", type=Path, required=True)
+    tokenizer.add_argument("--source", choices=source_names, default="fineweb-edu")
+    tokenizer.add_argument("--output", type=Path, required=True)
+    tokenizer.add_argument("--sample-bytes", type=int, default=1 << 29)
+    tokenizer.add_argument("--vocab-size", type=int, default=32_768)
+    tokenizer.add_argument("--min-frequency", type=int, default=10)
+    tokenizer.add_argument("--max-token-length", type=int, default=64)
+    tokenizer.add_argument("--batch-size", type=int, default=512)
+
+    prepare = commands.add_parser("prepare", help="encode cached Parquets into token shards")
+    prepare.add_argument("--raw-dir", type=Path, required=True)
+    prepare.add_argument("--source", choices=source_names, default="fineweb-edu")
+    prepare.add_argument("--tokenizer", type=Path, required=True)
+    prepare.add_argument("--output-dir", type=Path, required=True)
+    prepare.add_argument("--batch-size", type=int, default=256)
+    prepare.add_argument("--validation-modulus", type=int, default=1024)
+    prepare.add_argument("--validation-bucket", type=int, default=0)
+
+    mix = commands.add_parser(
+        "mix", help="combine prepared source directories into one mixture manifest"
+    )
+    mix.add_argument(
+        "--input",
+        action="append",
+        required=True,
+        metavar="DIR[=TOKENS]",
+        help="prepared corpus directory with an optional train-token budget (e.g. 3.2B)",
+    )
+    mix.add_argument("--output-dir", type=Path, required=True)
+    return parser
+
+
+def main() -> None:
+    args = _build_parser().parse_args()
+    if args.command == "download":
+        source = CORPUS_SOURCES[args.source]
+        paths = download_source(
+            source, args.raw_dir, max_files=args.max_files, max_workers=args.max_workers
+        )
+        print(f"downloaded {len(paths)} Parquet shards below {args.raw_dir}")
+    elif args.command == "train-tokenizer":
+        source = CORPUS_SOURCES[args.source]
+        paths = find_local_parquets(args.raw_dir, source.path_prefix)
+        tokenizer = train_tokenizer_from_iterator(
+            iter_tokenizer_text(
+                paths,
+                max_utf8_bytes=args.sample_bytes,
+                batch_size=args.batch_size,
+                text_column=source.text_column,
+                id_column=source.id_column,
+            ),
+            args.output,
+            vocab_size=args.vocab_size,
+            min_frequency=args.min_frequency,
+            max_token_length=args.max_token_length,
+        )
+        print(f"saved {tokenizer.get_vocab_size():,}-token tokenizer to {args.output}")
+    elif args.command == "prepare":
+        source = CORPUS_SOURCES[args.source]
+        paths = find_local_parquets(args.raw_dir, source.path_prefix)
+        result = prepare_corpus(
+            args.tokenizer,
+            args.output_dir,
+            corpus_source=source,
+            source_paths=paths,
+            batch_size=args.batch_size,
+            validation_modulus=args.validation_modulus,
+            validation_bucket=args.validation_bucket,
+        )
+        print(
+            f"prepared {result.processed_sources} sources "
+            f"({result.resumed_sources} resumed); train={result.train_manifest}, "
+            f"validation={result.validation_manifest}"
+        )
+    elif args.command == "mix":
+        inputs = [_parse_mix_input(item) for item in args.input]
+        train_manifest, validation_manifest = mix_manifests(inputs, args.output_dir)
+        print(f"mixed manifests: train={train_manifest}, validation={validation_manifest}")
+
+
+if __name__ == "__main__":
+    main()

src/diffusion_lm/data.py

@@ -0,0 +1,296 @@
+"""Memory-mapped packed-token datasets and deterministic sampling."""
+
+from __future__ import annotations
+
+import json
+import math
+from bisect import bisect_right
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from torch import Tensor
+from torch.utils.data import Dataset, Sampler
+
+from diffusion_lm.tokenizer import token_metadata_path
+
+
+PACKED_TOKEN_FORMAT = "mini-diffusion-lm-packed-tokens-v1"
+PACKED_MANIFEST_FORMAT = "mini-diffusion-lm-packed-manifest-v1"
+
+
+def _read_json(path: Path) -> dict[str, Any]:
+    try:
+        with path.open("r", encoding="utf-8") as handle:
+            value = json.load(handle)
+    except json.JSONDecodeError as exc:
+        raise ValueError(f"invalid JSON in {path}: {exc}") from exc
+    if not isinstance(value, dict):
+        raise ValueError(f"expected a JSON object in {path}")
+    return value
+
+
+def load_token_metadata(path: str | Path) -> dict[str, Any]:
+    metadata_path = token_metadata_path(path)
+    if not metadata_path.is_file():
+        raise FileNotFoundError(
+            f"token metadata not found: {metadata_path}; encode data with mini-mdlm-tokenizer"
+        )
+    metadata = _read_json(metadata_path)
+    if metadata.get("format") != PACKED_TOKEN_FORMAT:
+        raise ValueError(f"unsupported token file metadata in {metadata_path}")
+    return metadata
+
+
+def load_packed_manifest(path: str | Path) -> dict[str, Any]:
+    manifest_path = Path(path)
+    if not manifest_path.is_file():
+        raise FileNotFoundError(f"packed-token manifest not found: {manifest_path}")
+    manifest = _read_json(manifest_path)
+    if manifest.get("format") != PACKED_MANIFEST_FORMAT:
+        raise ValueError(f"unsupported packed-token manifest in {manifest_path}")
+    shards = manifest.get("shards")
+    if not isinstance(shards, list) or not shards:
+        raise ValueError(f"packed-token manifest has no shards: {manifest_path}")
+    return manifest
+
+
+def _validate_dtype(name: object, *, context: Path) -> np.dtype[Any]:
+    try:
+        dtype = np.dtype(name)
+    except TypeError as exc:
+        raise ValueError(f"invalid token dtype {name!r} in {context}") from exc
+    if dtype not in (np.dtype("uint16"), np.dtype("uint32")):
+        raise ValueError(f"unsupported token dtype {dtype.name!r} in {context}")
+    return dtype
+
+
+class PackedTokenDataset(Dataset[Tensor]):
+    """Expose deterministic non-overlapping blocks from one packed token file."""
+
+    def __init__(self, path: str | Path, sequence_length: int) -> None:
+        if sequence_length <= 0:
+            raise ValueError("sequence_length must be positive")
+        self.path = Path(path)
+        if not self.path.is_file():
+            raise FileNotFoundError(f"packed token file not found: {self.path}")
+        self.metadata = load_token_metadata(self.path)
+        self.sequence_length = sequence_length
+        self._dtype = _validate_dtype(self.metadata.get("dtype"), context=self.path)
+        self._tokens = np.memmap(self.path, mode="r", dtype=self._dtype)
+        expected_count = int(self.metadata["token_count"])
+        if self._tokens.size != expected_count:
+            raise ValueError(
+                f"metadata says {expected_count} tokens but {self.path} contains "
+                f"{self._tokens.size}"
+            )
+        self._blocks = self._tokens.size // sequence_length
+        if self._blocks == 0:
+            raise ValueError(
+                f"dataset has {self._tokens.size} tokens, fewer than one "
+                f"{sequence_length}-token block"
+            )
+
+    def __len__(self) -> int:
+        return self._blocks
+
+    def __getitem__(self, index: int) -> Tensor:
+        if index < 0:
+            index += self._blocks
+        if not 0 <= index < self._blocks:
+            raise IndexError(index)
+        start = index * self.sequence_length
+        block = np.asarray(self._tokens[start : start + self.sequence_length]).astype(
+            np.int64, copy=True
+        )
+        return torch.from_numpy(block)
+
+
+class ManifestPackedTokenDataset(Dataset[Tensor]):
+    """Expose one logical dataset backed by independently memory-mapped token shards.
+
+    Blocks never cross shard boundaries. At most ``sequence_length - 1`` trailing tokens per shard
+    are ignored, keeping source shards independently replaceable and resumable.
+    """
+
+    def __init__(self, path: str | Path, sequence_length: int) -> None:
+        if sequence_length <= 0:
+            raise ValueError("sequence_length must be positive")
+        self.path = Path(path)
+        self.metadata = load_packed_manifest(self.path)
+        self.sequence_length = sequence_length
+        self._dtype = _validate_dtype(self.metadata.get("dtype"), context=self.path)
+
+        declared_token_count = int(self.metadata.get("token_count", -1))
+        declared_document_count = int(self.metadata.get("document_count", -1))
+        shard_token_count = 0
+        shard_document_count = 0
+        self._shard_paths: list[Path] = []
+        self._shard_token_counts: list[int] = []
+        self._block_ends: list[int] = []
+        total_blocks = 0
+
+        for position, raw_shard in enumerate(self.metadata["shards"]):
+            if not isinstance(raw_shard, dict):
+                raise ValueError(f"shard {position} in {self.path} is not an object")
+            raw_path = raw_shard.get("path")
+            if not isinstance(raw_path, str) or not raw_path:
+                raise ValueError(f"shard {position} in {self.path} has no path")
+            shard_path = Path(raw_path)
+            if not shard_path.is_absolute():
+                shard_path = self.path.parent / shard_path
+            if not shard_path.is_file():
+                raise FileNotFoundError(f"packed token shard not found: {shard_path}")
+
+            token_count = int(raw_shard.get("token_count", -1))
+            document_count = int(raw_shard.get("document_count", -1))
+            if token_count < 0 or document_count < 0:
+                raise ValueError(f"invalid counts for shard {shard_path}")
+            expected_bytes = token_count * self._dtype.itemsize
+            if shard_path.stat().st_size != expected_bytes:
+                raise ValueError(
+                    f"manifest says {token_count} tokens but {shard_path} has "
+                    f"{shard_path.stat().st_size} bytes"
+                )
+
+            shard_token_count += token_count
+            shard_document_count += document_count
+            blocks = token_count // sequence_length
+            if blocks:
+                self._shard_paths.append(shard_path)
+                self._shard_token_counts.append(token_count)
+                total_blocks += blocks
+                self._block_ends.append(total_blocks)
+
+        if shard_token_count != declared_token_count:
+            raise ValueError(
+                f"manifest token_count is {declared_token_count}, shard total is "
+                f"{shard_token_count}"
+            )
+        if shard_document_count != declared_document_count:
+            raise ValueError(
+                f"manifest document_count is {declared_document_count}, shard total is "
+                f"{shard_document_count}"
+            )
+        if total_blocks == 0:
+            raise ValueError(
+                f"dataset has no shard containing a full {sequence_length}-token block"
+            )
+        self._blocks = total_blocks
+        self._maps: list[np.memmap[Any, Any] | None] = [None] * len(self._shard_paths)
+
+    def __len__(self) -> int:
+        return self._blocks
+
+    def _map(self, shard_index: int) -> np.memmap[Any, Any]:
+        tokens = self._maps[shard_index]
+        if tokens is None:
+            tokens = np.memmap(self._shard_paths[shard_index], mode="r", dtype=self._dtype)
+            self._maps[shard_index] = tokens
+        return tokens
+
+    def __getitem__(self, index: int) -> Tensor:
+        if index < 0:
+            index += self._blocks
+        if not 0 <= index < self._blocks:
+            raise IndexError(index)
+        shard_index = bisect_right(self._block_ends, index)
+        previous_end = 0 if shard_index == 0 else self._block_ends[shard_index - 1]
+        local_block = index - previous_end
+        start = local_block * self.sequence_length
+        tokens = self._map(shard_index)
+        block = np.asarray(tokens[start : start + self.sequence_length]).astype(
+            np.int64, copy=True
+        )
+        return torch.from_numpy(block)
+
+    def __getstate__(self) -> dict[str, Any]:
+        state = self.__dict__.copy()
+        # Reopen mappings inside each DataLoader worker rather than pickling file descriptors.
+        state["_maps"] = [None] * len(self._shard_paths)
+        return state
+
+
+PackedDataset = PackedTokenDataset | ManifestPackedTokenDataset
+
+
+def load_packed_dataset(path: str | Path, sequence_length: int) -> PackedDataset:
+    """Load a legacy single-file dataset or a manifest-backed sharded dataset."""
+
+    candidate = Path(path)
+    if candidate.suffix == ".json" and candidate.is_file():
+        value = _read_json(candidate)
+        if value.get("format") == PACKED_MANIFEST_FORMAT:
+            return ManifestPackedTokenDataset(candidate, sequence_length)
+    return PackedTokenDataset(candidate, sequence_length)
+
+
+_UINT64_MASK = (1 << 64) - 1
+
+
+def _splitmix64(value: int) -> int:
+    value = (value + 0x9E3779B97F4A7C15) & _UINT64_MASK
+    value = ((value ^ (value >> 30)) * 0xBF58476D1CE4E5B9) & _UINT64_MASK
+    value = ((value ^ (value >> 27)) * 0x94D049BB133111EB) & _UINT64_MASK
+    return value ^ (value >> 31)
+
+
+def _affine_permutation_parameters(size: int, seed: int, epoch: int) -> tuple[int, int]:
+    """Return ``a, b`` for the bijection ``(a*x+b) mod size``."""
+
+    if size == 1:
+        return 0, 0
+    mixed = _splitmix64((seed & _UINT64_MASK) ^ _splitmix64(epoch & _UINT64_MASK))
+    offset = mixed % size
+    multiplier = _splitmix64(mixed) % size
+    if multiplier == 0:
+        multiplier = 1
+    while math.gcd(multiplier, size) != 1:
+        multiplier += 1
+        if multiplier == size:
+            multiplier = 1
+    return multiplier, offset
+
+
+class DeterministicBatchSampler(Sampler[list[int]]):
+    """Infinite epoch permutations with exact batch-cursor resume and O(1) memory.
+
+    Each epoch uses a seeded affine bijection over dataset indexes. Unlike ``randperm().tolist()``,
+    memory use is independent of corpus size and a resumed batch can be calculated directly.
+    """
+
+    def __init__(
+        self,
+        dataset_size: int,
+        batch_size: int,
+        *,
+        seed: int,
+        start_batch: int = 0,
+    ) -> None:
+        if dataset_size <= 0 or batch_size <= 0:
+            raise ValueError("dataset_size and batch_size must be positive")
+        if start_batch < 0:
+            raise ValueError("start_batch must be non-negative")
+        self.dataset_size = dataset_size
+        self.batch_size = batch_size
+        self.seed = seed
+        self.start_batch = start_batch
+        self.batches_per_epoch = (dataset_size + batch_size - 1) // batch_size
+
+    def __iter__(self):
+        epoch = self.start_batch // self.batches_per_epoch
+        batch_in_epoch = self.start_batch % self.batches_per_epoch
+        while True:
+            multiplier, offset = _affine_permutation_parameters(
+                self.dataset_size, self.seed, epoch
+            )
+            for batch_index in range(batch_in_epoch, self.batches_per_epoch):
+                start = batch_index * self.batch_size
+                stop = min(start + self.batch_size, self.dataset_size)
+                yield [
+                    (multiplier * position + offset) % self.dataset_size
+                    for position in range(start, stop)
+                ]
+            epoch += 1
+            batch_in_epoch = 0

src/diffusion_lm/diffusion.py

@@ -0,0 +1,338 @@
+"""Absorbing-mask forward corruption, objective, and reverse samplers."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from collections.abc import Iterator
+from typing import Callable, Literal, Protocol
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+
+class Denoiser(Protocol):
+    config: object
+
+    def __call__(
+        self,
+        input_ids: Tensor,
+        attention_mask: Tensor | None = None,
+        output_positions: Tensor | None = None,
+        attn_mask: Tensor | None = None,
+    ) -> Tensor: ...
+
+
+@dataclass(frozen=True)
+class CorruptionBatch:
+    noisy_tokens: Tensor
+    mask: Tensor
+    mask_probability: Tensor
+    valid_mask: Tensor
+
+
+@dataclass(frozen=True)
+class LossOutput:
+    loss: Tensor
+    masked_accuracy: Tensor
+    masked_tokens: int
+
+
+@dataclass(frozen=True)
+class UnmaskStep:
+    """One observable state of the reverse diffusion process."""
+
+    step: int
+    total_steps: int
+    tokens: Tensor
+    masked_remaining: int
+
+
+def sample_mask_probabilities(
+    batch_size: int,
+    *,
+    device: torch.device | str,
+    eps: float = 1e-3,
+    low_discrepancy: bool = True,
+    generator: torch.Generator | None = None,
+) -> Tensor:
+    """Sample linear noise levels in ``[eps, 1]``.
+
+    A random cyclic shift of an evenly spaced grid preserves uniform marginals
+    while covering the complete noise range in every reasonably sized batch.
+    """
+
+    if batch_size <= 0:
+        raise ValueError("batch_size must be positive")
+    if not 0.0 < eps < 1.0:
+        raise ValueError("eps must be in (0, 1)")
+
+    if low_discrepancy:
+        offset = torch.rand((), device=device, generator=generator)
+        unit = (offset + torch.arange(batch_size, device=device) / batch_size) % 1.0
+    else:
+        unit = torch.rand(batch_size, device=device, generator=generator)
+    return eps + (1.0 - eps) * unit
+
+
+def corrupt_tokens(
+    clean_tokens: Tensor,
+    mask_token_id: int,
+    *,
+    valid_mask: Tensor | None = None,
+    mask_probability: Tensor | None = None,
+    eps: float = 1e-3,
+    low_discrepancy: bool = True,
+    generator: torch.Generator | None = None,
+) -> CorruptionBatch:
+    """Apply the absorbing forward process at one random time per sequence."""
+
+    if clean_tokens.ndim != 2:
+        raise ValueError("clean_tokens must have shape [batch, sequence]")
+    batch_size, _ = clean_tokens.shape
+    if valid_mask is None:
+        valid_mask = torch.ones_like(clean_tokens, dtype=torch.bool)
+    elif valid_mask.shape != clean_tokens.shape:
+        raise ValueError("valid_mask must match clean_tokens")
+    else:
+        valid_mask = valid_mask.bool()
+
+    if mask_probability is None:
+        mask_probability = sample_mask_probabilities(
+            batch_size,
+            device=clean_tokens.device,
+            eps=eps,
+            low_discrepancy=low_discrepancy,
+            generator=generator,
+        )
+    else:
+        mask_probability = torch.as_tensor(
+            mask_probability, device=clean_tokens.device, dtype=torch.float32
+        )
+        if mask_probability.ndim == 0:
+            mask_probability = mask_probability.repeat(batch_size)
+        if mask_probability.shape != (batch_size,):
+            raise ValueError("mask_probability must be scalar or have shape [batch]")
+        if bool(((mask_probability <= 0) | (mask_probability > 1)).any()):
+            raise ValueError("mask probabilities must be in (0, 1]")
+
+    random_values = torch.rand(clean_tokens.shape, device=clean_tokens.device, generator=generator)
+    mask = (random_values < mask_probability[:, None]) & valid_mask
+    noisy_tokens = torch.where(mask, mask_token_id, clean_tokens)
+    return CorruptionBatch(noisy_tokens, mask, mask_probability, valid_mask)
+
+
+def diffusion_cross_entropy(
+    logits: Tensor,
+    clean_tokens: Tensor,
+    corruption: CorruptionBatch,
+) -> LossOutput:
+    """Compute the continuous-time masked-diffusion likelihood bound.
+
+    ``logits`` may contain all positions as ``[B, L, V]`` or only the masked
+    positions as ``[N_masked, V]``. The latter is substantially more memory
+    efficient for small models with non-trivial vocabularies.
+    """
+
+    if clean_tokens.shape != corruption.noisy_tokens.shape:
+        raise ValueError("clean_tokens must match the corruption batch")
+    targets = clean_tokens[corruption.mask]
+    if logits.ndim == 3:
+        if logits.shape[:2] != clean_tokens.shape:
+            raise ValueError("full logits must have shape [batch, sequence, vocab]")
+        selected_logits = logits[corruption.mask]
+    elif logits.ndim == 2:
+        selected_logits = logits
+    else:
+        raise ValueError("logits must have shape [B, L, V] or [N_masked, V]")
+    if selected_logits.shape[0] != targets.numel():
+        raise ValueError("selected logits count does not match the number of masked tokens")
+
+    masked_tokens = int(targets.numel())
+    if masked_tokens == 0:
+        zero = logits.sum() * 0.0
+        return LossOutput(zero, zero.detach(), 0)
+
+    per_token = F.cross_entropy(selected_logits.float(), targets, reduction="none")
+    probabilities = corruption.mask_probability[:, None].expand_as(clean_tokens)
+    weights = probabilities[corruption.mask].reciprocal()
+    normalizer = corruption.valid_mask.sum().clamp_min(1)
+    loss = (per_token * weights).sum() / normalizer
+    accuracy = (selected_logits.argmax(dim=-1) == targets).float().mean()
+    return LossOutput(loss, accuracy, masked_tokens)
+
+
+def _sample_categorical(
+    logits: Tensor,
+    temperature: float,
+    generator: torch.Generator | None,
+) -> tuple[Tensor, Tensor]:
+    """Sample with fp64 Gumbel noise and return token ids plus model confidence."""
+
+    if temperature < 0:
+        raise ValueError("temperature must be non-negative")
+    log_probs = F.log_softmax(logits.float(), dim=-1)
+    if temperature == 0:
+        tokens = logits.argmax(dim=-1)
+    else:
+        # MPS has no float64 kernels. Preserve fp64 categorical sampling by
+        # moving only the sampling calculation to CPU on Apple Silicon.
+        sampling_device = torch.device("cpu") if logits.device.type == "mps" else logits.device
+        if logits.device.type == "mps":
+            logits64 = logits.float().cpu().double() / temperature
+        else:
+            logits64 = logits.double() / temperature
+        sampling_generator = generator
+        if generator is not None and generator.device != sampling_device:
+            sampling_generator = None
+        uniform = torch.rand(
+            logits64.shape,
+            device=sampling_device,
+            dtype=torch.float64,
+            generator=sampling_generator,
+        ).clamp_(1e-12, 1.0 - 1e-12)
+        gumbel = -torch.log(-torch.log(uniform))
+        tokens = (logits64 + gumbel).argmax(dim=-1).to(logits.device)
+    confidence = log_probs.gather(-1, tokens[:, None]).squeeze(-1).exp()
+    return tokens, confidence
+
+
+def iterative_unmask_steps(
+    model: Denoiser,
+    input_ids: Tensor,
+    mask_token_id: int,
+    *,
+    steps: int = 64,
+    temperature: float = 1.0,
+    strategy: Literal["ancestral", "confidence", "left_to_right"] = "ancestral",
+    blocked_token_ids: tuple[int, ...] = (),
+    attn_mask: Tensor | None = None,
+    generator: torch.Generator | None = None,
+    logits_fn: Callable[[Tensor, Tensor], Tensor] | None = None,
+) -> Iterator[UnmaskStep]:
+    """Yield each state while filling masks and clamping visible prompt tokens.
+
+    ``ancestral`` implements the absorbing reverse transition from mask rate
+    ``t`` to ``s``. ``confidence`` reveals an equal-sized highest-confidence
+    group on each pass; it is faster-looking and often useful, but is a heuristic.
+    ``left_to_right`` reveals equal-sized position-ordered groups, which keeps
+    arithmetic left operands visible before their results are committed.
+
+    ``logits_fn(tokens, masked)`` overrides how predictions are obtained, so a caller
+    holding a key/value cache can score only the masked window instead of the whole
+    sequence. The revealing schedule is unchanged either way.
+    """
+
+    if input_ids.ndim != 2:
+        raise ValueError("input_ids must have shape [batch, sequence]")
+    if steps <= 0:
+        raise ValueError("steps must be positive")
+    if strategy not in {"ancestral", "confidence", "left_to_right"}:
+        raise ValueError("strategy must be ancestral, confidence, or left_to_right")
+
+    tokens = input_ids.clone()
+    batch_size, _ = tokens.shape
+    yield UnmaskStep(0, steps, tokens.detach(), int(tokens.eq(mask_token_id).sum()))
+
+    for step in range(steps):
+        masked = tokens.eq(mask_token_id)
+        if not bool(masked.any()):
+            break
+
+        # This function is itself a generator, so a decorator would leave the
+        # inference context before iteration begins. Scope it around each pass.
+        with torch.inference_mode():
+            if logits_fn is not None:
+                logits = logits_fn(tokens, masked)
+            else:
+                # Kept as a conditional kwarg so mask-free denoiser doubles stay valid.
+                extra = {} if attn_mask is None else {"attn_mask": attn_mask}
+                logits = model(tokens, output_positions=masked, **extra)
+            if blocked_token_ids:
+                logits = logits.clone()
+                for token_id in blocked_token_ids:
+                    logits[:, token_id] = torch.finfo(logits.dtype).min
+            predictions, confidence = _sample_categorical(logits, temperature, generator)
+
+            proposed = tokens.clone()
+            proposed[masked] = predictions
+            reveal = torch.zeros_like(masked)
+            steps_left = steps - step
+
+            if strategy == "ancestral":
+                # Linear t grid: P(unmask from t to s | still masked) = 1 - s/t.
+                reveal_probability = 1.0 / steps_left
+                reveal = (
+                    torch.rand(tokens.shape, device=tokens.device, generator=generator)
+                    < reveal_probability
+                ) & masked
+            elif strategy == "left_to_right":
+                for row in range(batch_size):
+                    masked_positions = masked[row].nonzero(as_tuple=True)[0]
+                    remaining = int(masked_positions.numel())
+                    count = (remaining + steps_left - 1) // steps_left
+                    if count:
+                        reveal[row, masked_positions[:count]] = True
+            else:
+                confidence_grid = torch.full(
+                    tokens.shape,
+                    -torch.inf,
+                    device=tokens.device,
+                    dtype=confidence.dtype,
+                )
+                confidence_grid[masked] = confidence
+                for row in range(batch_size):
+                    remaining = int(masked[row].sum())
+                    count = (remaining + steps_left - 1) // steps_left
+                    if count:
+                        positions = confidence_grid[row].topk(count).indices
+                        reveal[row, positions] = True
+
+            tokens = torch.where(reveal, proposed, tokens)
+
+        yield UnmaskStep(
+            step + 1,
+            steps,
+            tokens.detach(),
+            int(tokens.eq(mask_token_id).sum()),
+        )
+
+    if bool(tokens.eq(mask_token_id).any()):
+        raise RuntimeError(
+            "sampler finished with masked positions; this indicates an internal error"
+        )
+
+
+@torch.no_grad()
+def iterative_unmask(
+    model: Denoiser,
+    input_ids: Tensor,
+    mask_token_id: int,
+    *,
+    steps: int = 64,
+    temperature: float = 1.0,
+    strategy: Literal["ancestral", "confidence"] = "ancestral",
+    blocked_token_ids: tuple[int, ...] = (),
+    attn_mask: Tensor | None = None,
+    generator: torch.Generator | None = None,
+    logits_fn: Callable[[Tensor, Tensor], Tensor] | None = None,
+) -> Tensor:
+    """Return the final state from :func:`iterative_unmask_steps`."""
+
+    final_state: UnmaskStep | None = None
+    for state in iterative_unmask_steps(
+        model,
+        input_ids,
+        mask_token_id,
+        steps=steps,
+        temperature=temperature,
+        strategy=strategy,
+        blocked_token_ids=blocked_token_ids,
+        attn_mask=attn_mask,
+        generator=generator,
+        logits_fn=logits_fn,
+    ):
+        final_state = state
+    if final_state is None:  # Defensive: the iterator always yields its initial state.
+        raise RuntimeError("sampler produced no state")
+    return final_state.tokens

src/diffusion_lm/distill.py

@@ -0,0 +1,328 @@
+"""Distilled-trace preparation: reasoning traces packed with chained thought slots.
+
+Streams a distilled-reasoning dataset (traces carrying a ``<think>...</think>`` region
+followed by a final answer), segments each think region into a chain of thought units, and
+renders every example through the shared slot geometry so a block-diffusion model denoises
+one thought at a time. The number of thoughts varies per example, so slot count stands in
+for how much the model chose to think.
+
+Segmentation splits on paragraph breaks first, then on reasoning-marker sentence boundaries
+inside long paragraphs, so a thought is a coherent reasoning move rather than a fixed-size
+cut.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import re
+from pathlib import Path
+from typing import Iterator
+
+import numpy as np
+
+from diffusion_lm.reasoning import (
+    ADAPTIVE_SPECIAL_TOKENS,
+    ExampleEncoder,
+    LayoutSpec,
+    REASONING_SPECIAL_TOKENS,
+    ReasoningExample,
+    size_token_ids,
+    _write_packed,
+)
+from diffusion_lm.tokenizer import load_tokenizer, train_tokenizer_from_iterator
+
+_THINK_RE = re.compile(r'<think>(.*?)</think>\s*', re.DOTALL)
+_MARKER_BREAK = re.compile(
+    r'(?<=[.!?])\s+(?=(?:Wait|Alternatively|Hmm|Okay|Now|But wait|Let me|First|Second|'
+    r'Next|Then|Finally|So,|Also|Another|Actually|Therefore|Thus)\b)'
+)
+_LONG_PARAGRAPH_CHARS = 900
+
+
+def segment_thoughts(think: str) -> list[str]:
+    """Split a think region into a chain of coherent thought units."""
+
+    thoughts: list[str] = []
+    for paragraph in re.split(r'\n\s*\n', think):
+        paragraph = ' '.join(paragraph.split())
+        if not paragraph:
+            continue
+        if len(paragraph) > _LONG_PARAGRAPH_CHARS:
+            thoughts.extend(part.strip() for part in _MARKER_BREAK.split(paragraph) if part.strip())
+        else:
+            thoughts.append(paragraph)
+    return thoughts
+
+
+def parse_glaive(row: dict[str, object]) -> ReasoningExample | None:
+    """Render one glaive ``reasoning-v1`` row into a chained-thought example."""
+
+    prompt = ' '.join(str(row.get('prompt') or '').split())
+    response = str(row.get('response') or '')
+    match = _THINK_RE.search(response)
+    if not prompt or not match:
+        return None
+    answer = ' '.join(response[match.end():].split())
+    thoughts = segment_thoughts(match.group(1))
+    if not answer or not thoughts:
+        return None
+    return ReasoningExample(prompt, tuple(thoughts), answer, expected_answer='')
+
+
+PARSERS = {'glaive': parse_glaive}
+
+
+def _stream_examples(
+    dataset: str, split: str, parser_name: str, limit: int
+) -> Iterator[ReasoningExample]:
+    from datasets import load_dataset
+
+    parser = PARSERS[parser_name]
+    kept = 0
+    for row in load_dataset(dataset, split=split, streaming=True):
+        example = parser(row)
+        if example is None:
+            continue
+        yield example
+        kept += 1
+        if limit and kept >= limit:
+            return
+
+
+def _is_validation(problem: str, val_fraction: float) -> bool:
+    digest = hashlib.sha256(problem.encode('utf-8')).digest()
+    return int.from_bytes(digest[:4], 'big') / 2**32 < val_fraction
+
+
+def _resolve_tokenizer(
+    examples: list[ReasoningExample],
+    args: argparse.Namespace,
+    special_tokens: tuple[str, ...] = REASONING_SPECIAL_TOKENS,
+):
+    from pathlib import Path
+
+    tokenizer_path = Path(args.tokenizer)
+    if tokenizer_path.is_file():
+        print(f'reusing tokenizer {tokenizer_path}')
+        return load_tokenizer(tokenizer_path), tokenizer_path
+
+    def _texts() -> Iterator[str]:
+        for example in examples:
+            yield f'{example.problem}\n{" ".join(example.steps)}\n{example.answer}'
+
+    tokenizer = train_tokenizer_from_iterator(
+        _texts(),
+        tokenizer_path,
+        vocab_size=args.vocab_size,
+        min_frequency=4,
+        length=len(examples),
+        extra_special_tokens=special_tokens,
+    )
+    print(f'trained tokenizer {tokenizer_path} (vocab {args.vocab_size}) on {len(examples):,} traces')
+    return tokenizer, tokenizer_path
+
+
+def prepare(args: argparse.Namespace) -> None:
+    adaptive = bool(getattr(args, 'sizes', None))
+    spec = LayoutSpec(
+        seq_len=args.seq_len,
+        block=args.block,
+        max_slots=args.max_slots,
+        sizes=tuple(args.sizes) if adaptive else (),
+    )
+    output_dir = Path(args.output_dir)
+
+    # Buffer unique parsed examples so a native tokenizer can be trained in the same
+    # streaming pass the corpus is packed from, avoiding a second dataset download.
+    examples: list[ReasoningExample] = []
+    seen: set[str] = set()
+    scanned = 0
+    for example in _stream_examples(args.dataset, args.split, args.parser, args.scan_limit):
+        scanned += 1
+        key = hashlib.sha256(example.problem.encode('utf-8')).hexdigest()
+        if key in seen:
+            continue
+        seen.add(key)
+        examples.append(example)
+        if args.limit and len(examples) >= args.limit:
+            break
+    if not examples:
+        raise ValueError('no usable examples; check dataset, parser, and seq_len')
+    print(f'scanned {scanned:,}, buffered {len(examples):,} unique examples')
+
+    special_tokens = ADAPTIVE_SPECIAL_TOKENS if adaptive else REASONING_SPECIAL_TOKENS
+    tokenizer, tokenizer_path = _resolve_tokenizer(examples, args, special_tokens)
+    encoder = ExampleEncoder(tokenizer, spec)
+
+    if adaptive:
+        _pack_adaptive(examples, encoder, spec, tokenizer, tokenizer_path, output_dir, args)
+        return
+
+    split: dict[str, dict[str, list[np.ndarray]]] = {
+        'train': {'flat': [], 'flat_regions': [], 'slotted': [], 'slotted_regions': []},
+        'validation': {'flat': [], 'flat_regions': [], 'slotted': [], 'slotted_regions': []},
+    }
+    val_prompts: list[dict[str, str]] = []
+    dropped = 0
+    for example in examples:
+        encoded = encoder.encode_example(example)
+        if encoded is None:
+            dropped += 1
+            continue
+        is_val = _is_validation(example.problem, args.val_fraction)
+        bucket = split['validation' if is_val else 'train']
+        bucket['flat'].append(encoded.flat)
+        bucket['flat_regions'].append(encoded.flat_regions)
+        bucket['slotted'].append(encoded.slotted)
+        bucket['slotted_regions'].append(encoded.slotted_regions)
+        if is_val and len(val_prompts) < 500:
+            val_prompts.append({'problem': example.problem, 'expected_answer': ''})
+
+    total = sum(len(bucket['flat']) for bucket in split.values())
+    if not total:
+        raise ValueError('no examples survived encoding; raise seq_len or max_slots')
+    print(f'encoded {total:,} examples ({dropped:,} dropped at encode)')
+    for split_name, bucket in split.items():
+        if not bucket['flat']:
+            raise ValueError(f'no examples in the {split_name} split; lower val_fraction or scan more')
+        for layout in ('flat', 'slotted'):
+            _write_packed(
+                output_dir / f'{split_name}-{layout}.bin',
+                np.stack(bucket[layout]),
+                np.stack(bucket[f'{layout}_regions']),
+                layout=layout,
+                spec=spec,
+                tokenizer_path=tokenizer_path,
+                tokenizer=tokenizer,
+            )
+        print(f'{split_name}: {len(bucket["flat"]):,} examples -> {output_dir}')
+
+    problems_path = output_dir / 'validation-problems.jsonl'
+    with problems_path.open('w', encoding='utf-8') as handle:
+        for record in val_prompts:
+            handle.write(json.dumps(record, ensure_ascii=False) + '\n')
+    print(f'wrote {len(val_prompts):,} validation prompts to {problems_path}')
+
+
+def _pack_adaptive(
+    examples: list[ReasoningExample],
+    encoder: ExampleEncoder,
+    spec: LayoutSpec,
+    tokenizer,
+    tokenizer_path: Path,
+    output_dir: Path,
+    args: argparse.Namespace,
+) -> None:
+    """Encode the adaptive layout and report the block-size distribution it produces."""
+
+    split: dict[str, dict[str, list[np.ndarray]]] = {
+        'train': {'tokens': [], 'regions': []},
+        'validation': {'tokens': [], 'regions': []},
+    }
+    val_prompts: list[dict[str, str]] = []
+    size_histogram: dict[int, int] = {size: 0 for size in spec.sizes}
+    blocks_per_example: list[int] = []
+    dropped = 0
+    for example in examples:
+        encoded = encoder.encode_adaptive(example)
+        if encoded is None:
+            dropped += 1
+            continue
+        is_val = _is_validation(example.problem, args.val_fraction)
+        bucket = split['validation' if is_val else 'train']
+        bucket['tokens'].append(encoded.tokens)
+        bucket['regions'].append(encoded.regions)
+        blocks_per_example.append(len(encoded.block_sizes))
+        for size in encoded.block_sizes:
+            size_histogram[size] += 1
+        if is_val and len(val_prompts) < 500:
+            val_prompts.append({'problem': example.problem, 'expected_answer': ''})
+
+    total = sum(len(bucket['tokens']) for bucket in split.values())
+    if not total:
+        raise ValueError('no examples survived encoding; raise seq_len, max_slots, or sizes')
+
+    total_blocks = sum(size_histogram.values())
+    fractions = {
+        size: round(count / max(1, total_blocks), 4) for size, count in size_histogram.items()
+    }
+    sorted_blocks = sorted(blocks_per_example)
+
+    def percentile(fraction: float) -> int:
+        return sorted_blocks[min(len(sorted_blocks) - 1, int(len(sorted_blocks) * fraction))]
+
+    print(f'encoded {total:,} examples ({dropped:,} dropped at encode)')
+    print(f'block-size counts {size_histogram} fractions {fractions}')
+    print(
+        f'blocks/example p50 {percentile(0.5)} p90 {percentile(0.9)} '
+        f'max {sorted_blocks[-1]}'
+    )
+
+    extra_metadata = {
+        'sizes': list(spec.sizes),
+        'size_token_ids': size_token_ids(tokenizer),
+    }
+    for split_name, bucket in split.items():
+        if not bucket['tokens']:
+            raise ValueError(f'no examples in the {split_name} split; lower val_fraction')
+        _write_packed(
+            output_dir / f'{split_name}-adaptive.bin',
+            np.stack(bucket['tokens']),
+            np.stack(bucket['regions']),
+            layout='adaptive',
+            spec=spec,
+            tokenizer_path=tokenizer_path,
+            tokenizer=tokenizer,
+            extra_metadata=extra_metadata,
+        )
+        print(f'{split_name}: {len(bucket["tokens"]):,} examples -> {output_dir}')
+
+    problems_path = output_dir / 'validation-problems.jsonl'
+    with problems_path.open('w', encoding='utf-8') as handle:
+        for record in val_prompts:
+            handle.write(json.dumps(record, ensure_ascii=False) + '\n')
+    print(f'wrote {len(val_prompts):,} validation prompts to {problems_path}')
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=__doc__)
+    subparsers = parser.add_subparsers(dest='command', required=True)
+    prepare_parser = subparsers.add_parser(
+        'prepare', help='stream a distilled-reasoning dataset into chained-thought slots'
+    )
+    prepare_parser.add_argument('--dataset', default='glaiveai/reasoning-v1-20m')
+    prepare_parser.add_argument('--split', default='train')
+    prepare_parser.add_argument('--parser', choices=sorted(PARSERS), default='glaive')
+    prepare_parser.add_argument(
+        '--tokenizer', required=True,
+        help='tokenizer path; trained from the buffered traces when the file is absent'
+    )
+    prepare_parser.add_argument('--vocab-size', type=int, default=16384)
+    prepare_parser.add_argument('--output-dir', required=True)
+    prepare_parser.add_argument('--seq-len', type=int, default=2048)
+    prepare_parser.add_argument('--block', type=int, default=32)
+    prepare_parser.add_argument('--max-slots', type=int, default=64)
+    prepare_parser.add_argument(
+        '--sizes', type=int, nargs='+',
+        help='activate the adaptive layout with these ascending block sizes (e.g. 64 256)'
+    )
+    prepare_parser.add_argument('--val-fraction', type=float, default=0.02)
+    prepare_parser.add_argument(
+        '--limit', type=int, default=0, help='stop after this many unique kept examples'
+    )
+    prepare_parser.add_argument(
+        '--scan-limit', type=int, default=0, help='stop streaming after this many parsed rows'
+    )
+    return parser
+
+
+def main() -> None:
+    args = _build_parser().parse_args()
+    if args.command == 'prepare':
+        prepare(args)
+
+
+if __name__ == '__main__':
+    main()

src/diffusion_lm/flexattn.py

@@ -0,0 +1,147 @@
+"""Pre-LN Transformer encoder that routes self-attention through ``flex_attention``.
+
+The block-diffusion masks are structured (block-causal over thought slots plus optional
+key padding), so expressing them as a ``flex_attention`` block mask lets the fused kernel
+skip fully-masked blocks instead of materializing a dense ``[batch * heads, L, L]`` score
+tensor. This keeps attention memory and time near-linear in the trace length, which the
+default ``nn.TransformerEncoder`` math path does not.
+
+The layer geometry mirrors ``nn.TransformerEncoderLayer(norm_first=True, activation='gelu')``
+so behaviour matches the default path up to floating-point error; the attention core is
+verified against a dense masked-softmax reference before use.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+from torch.nn.attention.flex_attention import BlockMask, create_block_mask, flex_attention
+
+# Compiled dynamic=False is fastest for fixed-shape training. Variable-length generation
+# recompiles per new sequence length, so serving sets MDLM_FLEX_EAGER=1 to run eager and
+# trade steady-state speed for the absence of per-length recompilation stalls.
+_FLEX_EAGER = os.environ.get('MDLM_FLEX_EAGER') == '1'
+_flex_compiled = flex_attention if _FLEX_EAGER else torch.compile(flex_attention, dynamic=False)
+
+
+def build_block_mask(
+    attn_mask: Tensor | None,
+    padding_mask: Tensor | None,
+    batch_size: int,
+    seq_len: int,
+    device: torch.device,
+) -> BlockMask | None:
+    """Build a broadcast-over-heads ``BlockMask`` from project boolean masks.
+
+    ``attn_mask`` follows the src_mask convention (``True`` blocks a key), shaped ``[L, L]``
+    or ``[batch, L, L]``. ``padding_mask`` follows the src_key_padding_mask convention
+    (``True`` marks padding). Returns ``None`` when neither constrains attention.
+    """
+
+    if attn_mask is None and padding_mask is None:
+        return None
+
+    shared = attn_mask is not None and attn_mask.dim() == 2
+
+    def mask_mod(b: Tensor, h: Tensor, q_idx: Tensor, kv_idx: Tensor) -> Tensor:
+        keep = torch.ones_like(q_idx, dtype=torch.bool)
+        if attn_mask is not None:
+            blocked = attn_mask[q_idx, kv_idx] if shared else attn_mask[b, q_idx, kv_idx]
+            keep = keep & ~blocked
+        if padding_mask is not None:
+            keep = keep & ~padding_mask[b, kv_idx]
+        return keep
+
+    return create_block_mask(
+        mask_mod, batch_size, None, seq_len, seq_len, device=device, _compile=not _FLEX_EAGER
+    )
+
+
+def flex_self_attention(
+    query: Tensor, key: Tensor, value: Tensor, block_mask: BlockMask | None
+) -> Tensor:
+    """Multi-head self-attention over ``[batch, heads, L, head_dim]`` tensors."""
+
+    if block_mask is None:
+        return flex_attention(query, key, value)
+    return _flex_compiled(query, key, value, block_mask=block_mask)
+
+
+class FlexEncoderLayer(nn.Module):
+    """Pre-LN Transformer block with a ``flex_attention`` self-attention core."""
+
+    def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float) -> None:
+        super().__init__()
+        if d_model % n_heads != 0:
+            raise ValueError("d_model must be divisible by n_heads")
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.q_proj = nn.Linear(d_model, d_model)
+        self.k_proj = nn.Linear(d_model, d_model)
+        self.v_proj = nn.Linear(d_model, d_model)
+        self.out_proj = nn.Linear(d_model, d_model)
+        self.linear1 = nn.Linear(d_model, d_ff)
+        self.linear2 = nn.Linear(d_ff, d_model)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def _split_heads(self, projected: Tensor) -> Tensor:
+        batch_size, seq_len, _ = projected.shape
+        return projected.view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+
+    def forward(self, hidden: Tensor, block_mask: BlockMask | None) -> Tensor:
+        normed = self.norm1(hidden)
+        query = self._split_heads(self.q_proj(normed))
+        key = self._split_heads(self.k_proj(normed))
+        value = self._split_heads(self.v_proj(normed))
+        attended = flex_self_attention(query, key, value, block_mask)
+        batch_size, _, seq_len, _ = attended.shape
+        attended = attended.transpose(1, 2).reshape(batch_size, seq_len, -1)
+        hidden = hidden + self.dropout(self.out_proj(attended))
+        normed = self.norm2(hidden)
+        feed_forward = self.linear2(self.dropout(F.gelu(self.linear1(normed))))
+        return hidden + self.dropout(feed_forward)
+
+
+class FlexEncoder(nn.Module):
+    """Stack of :class:`FlexEncoderLayer` blocks with a final layer norm."""
+
+    def __init__(
+        self,
+        d_model: int,
+        n_heads: int,
+        d_ff: int,
+        dropout: float,
+        n_layers: int,
+        activation_checkpointing: bool,
+    ) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList(
+            FlexEncoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)
+        )
+        self.norm = nn.LayerNorm(d_model)
+        self.activation_checkpointing = activation_checkpointing
+
+    def init_residual_outputs(self, n_layers: int) -> None:
+        residual_std = 0.02 / math.sqrt(2 * n_layers)
+        for layer in self.layers:
+            nn.init.normal_(layer.out_proj.weight, mean=0.0, std=residual_std)
+            nn.init.normal_(layer.linear2.weight, mean=0.0, std=residual_std)
+
+    def forward(self, hidden: Tensor, block_mask: BlockMask | None) -> Tensor:
+        use_checkpoint = (
+            self.activation_checkpointing and self.training and torch.is_grad_enabled()
+        )
+        for layer in self.layers:
+            if use_checkpoint:
+                hidden = torch.utils.checkpoint.checkpoint(
+                    layer, hidden, block_mask, use_reentrant=False
+                )
+            else:
+                hidden = layer(hidden, block_mask)
+        return self.norm(hidden)

src/diffusion_lm/hf_bridge.py

@@ -0,0 +1,270 @@
+"""HuggingFace Qwen3 backbone served through the project's denoiser forward contract."""
+
+from __future__ import annotations
+
+import os
+
+import torch
+from torch import Tensor, nn
+
+from diffusion_lm.config import ModelConfig
+
+# The flex template's default 128x128 tiles need ~112 KiB of shared memory at head_dim 128,
+# over Ada's 100 KiB per-block ceiling; halved tiles fit. The backward kernel budgets its
+# tiles separately, hence the M1/N1/M2/N2 entries. Larger tiles are valid on Hopper.
+_FLEX_KERNEL_OPTIONS = {
+    'BLOCK_M': 64,
+    'BLOCK_N': 64,
+    'BLOCK_M1': 32,
+    'BLOCK_N1': 64,
+    'BLOCK_M2': 64,
+    'BLOCK_N2': 32,
+}
+
+
+class Qwen3Denoiser(nn.Module):
+    """Wrap ``Qwen3ForCausalLM`` behind the DiffusionTransformer forward contract.
+
+    The backbone always receives a 4D attention mask so its stock causal masking never
+    engages: block-diffusion objectives need bidirectional attention inside denoising
+    windows, and an omitted mask must mean "attend everything", not "causal".
+    ``use_flex_attention`` selects how the boolean blocking matrix reaches the backbone — a
+    ``BlockMask`` for the flex kernel, which skips fully-masked blocks, or a materialized
+    additive mask for SDPA. Hidden states are gathered at ``output_positions`` before the LM
+    head so full-vocabulary logits are never materialized for visible tokens.
+    """
+
+    def __init__(
+        self,
+        config: ModelConfig,
+        *,
+        load_pretrained: bool = True,
+        dtype: torch.dtype | None = None,
+    ) -> None:
+        super().__init__()
+        try:
+            from transformers import Qwen3Config, Qwen3ForCausalLM
+        except ImportError as exc:
+            raise RuntimeError(
+                'the hf-qwen3 backbone requires transformers; install the [hf] extra'
+            ) from exc
+        if config.pretrained_path is None:
+            raise ValueError('the hf-qwen3 backbone requires pretrained_path')
+        self.config = config
+        # Generation on Ada can hit the Triton shared-memory ceiling inside transformers' own
+        # compiled flex kernel, where our kernel_options do not reach. MDLM_SDPA=1 sidesteps it
+        # for serving and evaluation; flex earns its keep in training, not at q_len 1.
+        use_flex = config.use_flex_attention and os.environ.get('MDLM_SDPA') != '1'
+        attn_implementation = 'flex_attention' if use_flex else 'sdpa'
+        if load_pretrained:
+            kwargs = {} if dtype is None else {'dtype': dtype}
+            self.backbone = Qwen3ForCausalLM.from_pretrained(
+                config.pretrained_path, attn_implementation=attn_implementation, **kwargs
+            )
+        else:
+            # Architecture-only construction: weights come from a later load_state_dict,
+            # so checkpoint restore never re-reads the base model files.
+            backbone_config = Qwen3Config.from_pretrained(
+                config.pretrained_path, attn_implementation=attn_implementation
+            )
+            self.backbone = Qwen3ForCausalLM(backbone_config)
+            if dtype is not None:
+                self.backbone.to(dtype)
+        self._use_flex = use_flex
+        self._validate_backbone()
+        self.backbone.config.use_cache = False
+        if config.activation_checkpointing:
+            self.backbone.gradient_checkpointing_enable(
+                gradient_checkpointing_kwargs={'use_reentrant': False}
+            )
+        # A plain attribute keeps the compiled callable out of the module tree, so checkpoint
+        # keys are identical whether or not compilation is on. Shapes reaching the backbone are
+        # static; the varying count of output positions is gathered after it returns.
+        # Serving generates at ever-changing lengths, which makes a compiled backbone
+        # recompile per shape; MDLM_NO_COMPILE=1 turns it off without touching the checkpoint.
+        compiling = config.compile_backbone and os.environ.get('MDLM_NO_COMPILE') != '1'
+        self._compiled_backbone = (
+            torch.compile(self.backbone.model.forward, dynamic=False) if compiling else None
+        )
+        self.register_buffer(
+            '_forbidden_output_token_ids',
+            torch.tensor(config.forbidden_output_token_ids, dtype=torch.long),
+            persistent=False,
+        )
+
+    def _validate_backbone(self) -> None:
+        backbone = self.backbone.config
+        pairs = (
+            ('vocab_size', self.config.vocab_size, backbone.vocab_size),
+            ('d_model', self.config.d_model, backbone.hidden_size),
+            ('n_layers', self.config.n_layers, backbone.num_hidden_layers),
+            ('n_heads', self.config.n_heads, backbone.num_attention_heads),
+            ('d_ff', self.config.d_ff, backbone.intermediate_size),
+        )
+        for name, expected, actual in pairs:
+            if expected != actual:
+                raise ValueError(f'config {name}={expected} but the backbone has {actual}')
+        if self.config.max_seq_len > backbone.max_position_embeddings:
+            raise ValueError(
+                f'max_seq_len {self.config.max_seq_len} exceeds the backbone context '
+                f'{backbone.max_position_embeddings}'
+            )
+
+    def _blocked_matrix(
+        self,
+        input_ids: Tensor,
+        attention_mask: Tensor | None,
+        attn_mask: Tensor | None,
+    ) -> Tensor:
+        """Boolean ``[batch, L, L]`` blocking matrix, ``True`` marking a key to suppress."""
+
+        batch_size, seq_len = input_ids.shape
+        device = input_ids.device
+        if attn_mask is None:
+            blocked = torch.zeros(
+                batch_size, seq_len, seq_len, dtype=torch.bool, device=device
+            )
+        else:
+            if attn_mask.dtype != torch.bool:
+                raise ValueError('attn_mask must be boolean with True marking blocked positions')
+            if attn_mask.shape == (seq_len, seq_len):
+                blocked = attn_mask.unsqueeze(0).expand(batch_size, -1, -1)
+            elif attn_mask.shape == (batch_size, seq_len, seq_len):
+                blocked = attn_mask
+            else:
+                raise ValueError('attn_mask must have shape [L, L] or [batch, L, L]')
+        if attention_mask is not None:
+            if attention_mask.shape != input_ids.shape:
+                raise ValueError('attention_mask must match input_ids')
+            blocked = blocked | ~attention_mask.bool()[:, None, :]
+        return blocked
+
+    def _additive_mask(self, blocked: Tensor) -> Tensor:
+        mask_dtype = self.backbone.model.embed_tokens.weight.dtype
+        additive = torch.zeros(blocked.shape, dtype=mask_dtype, device=blocked.device)
+        additive = additive.masked_fill(blocked, torch.finfo(mask_dtype).min)
+        return additive.unsqueeze(1)
+
+    def forward_cached(
+        self,
+        input_ids: Tensor,
+        *,
+        attn_mask: Tensor,
+        past_key_values,
+        output_positions: Tensor | None = None,
+    ) -> tuple[Tensor, object]:
+        """Run only the trailing ``input_ids`` against a populated key/value cache.
+
+        ``attn_mask`` holds one row per NEW query and one column per position the query may
+        see, cache included: ``[batch, query, cached + query]``. Block-causal masking is what
+        makes caching sound here — prefix positions never attend forward into a block, so
+        their keys and values stay valid while the block's own tokens keep changing.
+        """
+
+        cached = past_key_values.get_seq_length()
+        query_len = input_ids.shape[1]
+        blocked = attn_mask if attn_mask.dim() == 3 else attn_mask.unsqueeze(0)
+        if blocked.shape[-2] != query_len or blocked.shape[-1] != cached + query_len:
+            raise ValueError(
+                f'attn_mask must be [batch, {query_len}, {cached + query_len}] for a cache of '
+                f'{cached} positions, got {tuple(blocked.shape)}'
+            )
+        rows = blocked.expand(input_ids.shape[0], -1, -1)
+        outputs = self.backbone.model(
+            input_ids=input_ids,
+            attention_mask=self._additive_mask(rows),
+            past_key_values=past_key_values,
+            use_cache=True,
+            **self._backbone_kwargs(),
+        )
+        hidden = outputs.last_hidden_state
+        if output_positions is not None:
+            hidden = hidden[output_positions.bool()]
+        logits = self.backbone.lm_head(hidden)
+        self._forbid(logits)
+        return logits, outputs.past_key_values
+
+    def _backbone_kwargs(self) -> dict:
+        """Extra backbone arguments the attention implementation needs.
+
+        Every path into the backbone has to carry these, cached or not: without the reduced
+        tiles the flex template asks Ada for more shared memory than it has and the kernel
+        fails to compile at all.
+        """
+
+        if not self._use_flex:
+            return {}
+        return {'kernel_options': _FLEX_KERNEL_OPTIONS}
+
+    def new_cache(self):
+        """Empty key/value cache for :meth:`forward_cached`, kept here so callers stay
+        independent of the transformers cache class."""
+
+        from transformers import DynamicCache
+
+        return DynamicCache()
+
+    def _forbid(self, logits: Tensor) -> None:
+        if self._forbidden_output_token_ids.numel():
+            logits.index_fill_(
+                -1, self._forbidden_output_token_ids, torch.finfo(logits.dtype).min
+            )
+        if self.config.forbidden_output_from is not None:
+            logits[..., self.config.forbidden_output_from :] = torch.finfo(logits.dtype).min
+
+    def _flex_mask(self, blocked: Tensor):
+        """Compress the blocking matrix into a ``BlockMask`` the flex kernel can skip over.
+
+        ``masking_utils`` forwards any mask reporting a 4D shape untouched, so a ``BlockMask``
+        reaches ``flex_attention_forward`` as ``block_mask`` and the backbone never rebuilds a
+        causal mask of its own.
+        """
+
+        from diffusion_lm.flexattn import build_block_mask
+
+        batch_size, seq_len, _ = blocked.shape
+        return build_block_mask(blocked, None, batch_size, seq_len, blocked.device)
+
+    def forward(
+        self,
+        input_ids: Tensor,
+        attention_mask: Tensor | None = None,
+        output_positions: Tensor | None = None,
+        attn_mask: Tensor | None = None,
+    ) -> Tensor:
+        if input_ids.ndim != 2:
+            raise ValueError('input_ids must have shape [batch, sequence]')
+        sequence_length = input_ids.shape[1]
+        if sequence_length > self.config.max_seq_len:
+            raise ValueError(
+                f'sequence length {sequence_length} exceeds max_seq_len '
+                f'{self.config.max_seq_len}'
+            )
+        blocked = self._blocked_matrix(input_ids, attention_mask, attn_mask)
+        mask = (
+            self._flex_mask(blocked) if self._use_flex else self._additive_mask(blocked)
+        )
+        call = self._compiled_backbone or self.backbone.model
+        hidden = call(
+            input_ids=input_ids, attention_mask=mask, use_cache=False,
+            **self._backbone_kwargs(),
+        ).last_hidden_state
+        if output_positions is not None:
+            if output_positions.shape != input_ids.shape:
+                raise ValueError('output_positions must match input_ids')
+            hidden = hidden[output_positions.bool()]
+        logits = self.backbone.lm_head(hidden)
+        self._forbid(logits)
+        return logits
+
+    @property
+    def token_embedding(self) -> nn.Embedding:
+        """Embedding module surfaced for the optimizer's 32-bit override (tied to the head)."""
+
+        return self.backbone.model.embed_tokens
+
+    @property
+    def num_parameters(self) -> int:
+        """Count unique trainable parameters (shared embeddings count once)."""
+
+        return sum(parameter.numel() for parameter in self.parameters() if parameter.requires_grad)

src/diffusion_lm/hybrid.py

@@ -0,0 +1,1139 @@
+"""Region-aware attention masks, objectives, and samplers for reasoning variants.
+
+Three inference modes share one ``DiffusionTransformer``:
+
+- ``ar``:        prefix-LM. Bidirectional over the problem, causal generation after it.
+- ``diffusion``: full-sequence denoising of the response region in parallel.
+- ``hybrid``:    thought slots denoised block-by-block, then the answer decoded
+                 autoregressively under a bidirectional-prefix mask.
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+from diffusion_lm.diffusion import (
+    CorruptionBatch,
+    corrupt_tokens,
+    iterative_unmask,
+    _sample_categorical,
+)
+from diffusion_lm.model import DiffusionTransformer
+
+
+def prefix_causal_blocked(
+    prefix_lens: Tensor, seq_len: int, *, causal_prefix: bool = False
+) -> Tensor:
+    """Blocked-attention mask: bidirectional before ``prefix_len``, causal after.
+
+    ``allowed[b, i, j] = j < prefix_lens[b] or j <= i``; the returned tensor is the
+    inverse, matching the src_mask convention where ``True`` blocks attention.
+    ``causal_prefix=True`` degenerates to the plain causal mask — the geometry
+    pretrained autoregressive backbones were trained under.
+    """
+
+    positions = torch.arange(seq_len, device=prefix_lens.device)
+    causal = positions[None, :, None] >= positions[None, None, :]
+    if causal_prefix:
+        return (~causal).expand(prefix_lens.shape[0], seq_len, seq_len)
+    prefix = positions[None, None, :] < prefix_lens[:, None, None]
+    return ~(prefix | causal)
+
+
+def window_blocked(window_ends: Tensor, seq_len: int) -> Tensor:
+    """Blocked-attention mask hiding all keys at or beyond each sample's window end."""
+
+    positions = torch.arange(seq_len, device=window_ends.device)
+    allowed = positions[None, None, :] < window_ends[:, None, None]
+    return ~allowed.expand(window_ends.shape[0], seq_len, seq_len)
+
+
+def slot_causal_blocked(
+    problem_len: Tensor, n_slots: Tensor, block: int, seq_len: int
+) -> Tensor:
+    """Block-causal mask over thought slots: each slot attends its prefix slots only.
+
+    Problem tokens (plus ``<think>``) attend the problem window; tokens of slot k
+    attend everything up to slot k's end; positions past the think region attend
+    the whole think window. Denoising slot k at inference with clean prefix slots
+    is the per-slot ``t -> 0`` limit of the training distribution.
+    """
+
+    device = problem_len.device
+    positions = torch.arange(seq_len, device=device)[None, :]
+    prefix_end = (problem_len + 1)[:, None]
+    think_end = prefix_end + n_slots[:, None] * block
+    slot_index = torch.clamp((positions - prefix_end) // block, min=0)
+    slot_end = prefix_end + (slot_index + 1) * block
+    window = torch.where(positions < prefix_end, prefix_end, slot_end)
+    window = torch.where(positions >= think_end, think_end, window)
+    allowed = positions[:, None, :] < window[:, :, None]
+    return ~allowed
+
+
+def adaptive_block_mask(
+    tokens: Tensor,
+    problem_len: Tensor,
+    answer_start: Tensor,
+    size_ids: Tensor,
+    end_think_id: int,
+    *,
+    causal_prefix: bool = False,
+) -> Tensor:
+    """Block-causal mask over variable-length thought blocks.
+
+    Block boundaries are read directly from the token stream: every ``<szN>``
+    control token and the terminal ``</think>`` mark the start of the next region.
+    A thought-content token attends its clean prefix plus its own block
+    bidirectionally, never a following block; problem and ``<think>`` positions see
+    the problem window only — or, with ``causal_prefix=True``, only their causal
+    past, matching a pretrained autoregressive backbone. The returned tensor
+    follows the src_mask convention where ``True`` blocks a key.
+    """
+
+    positions = torch.arange(tokens.shape[1], device=tokens.device)
+    prefix_end = (problem_len + 1)[:, None]
+    in_think = (positions[None, :] >= prefix_end) & (
+        positions[None, :] < answer_start[:, None]
+    )
+    is_size = (tokens.unsqueeze(-1) == size_ids).any(dim=-1)
+    boundary = (is_size | (tokens == end_think_id)) & in_think
+    return block_mask_from_boundaries(
+        boundary, problem_len + 1, answer_start, causal_prefix=causal_prefix
+    )
+
+
+def block_mask_from_boundaries(
+    boundary: Tensor,
+    prefix_end: Tensor,
+    answer_start: Tensor,
+    *,
+    causal_prefix: bool = False,
+) -> Tensor:
+    """Blocked mask from per-position block-start marks (the adaptive-mask core).
+
+    ``boundary[b, p]`` is True where a new block starts. A position attends every key
+    before the next boundary after it: its own block bidirectionally plus all preceding
+    context. Rows before ``prefix_end`` see the prefix window (or their causal past with
+    ``causal_prefix=True``); rows at or past ``answer_start`` see up to ``answer_start``.
+    """
+
+    device = boundary.device
+    batch_size, seq_len = boundary.shape
+    positions = torch.arange(seq_len, device=device)
+    prefix_end = prefix_end[:, None]
+    answer_start = answer_start[:, None]
+
+    boundary_index = torch.where(
+        boundary, positions[None, :].expand(batch_size, seq_len), seq_len
+    )
+    reverse_cummin = boundary_index.flip(1).cummin(dim=1).values.flip(1)
+    next_boundary = torch.full((batch_size, seq_len), seq_len, device=device)
+    next_boundary[:, :-1] = reverse_cummin[:, 1:]
+
+    if causal_prefix:
+        prefix_limit = (positions[None, :] + 1).expand(batch_size, seq_len)
+    else:
+        prefix_limit = prefix_end.expand(batch_size, seq_len)
+    attend_limit = next_boundary
+    attend_limit = torch.where(positions[None, :] < prefix_end, prefix_limit, attend_limit)
+    attend_limit = torch.where(
+        positions[None, :] >= answer_start,
+        answer_start.expand(batch_size, seq_len),
+        attend_limit,
+    )
+    allowed = positions[None, None, :] < attend_limit[:, :, None]
+    return ~allowed
+
+
+@dataclass(frozen=True)
+class HybridBatchLoss:
+    loss: Tensor
+    think_loss: float
+    answer_loss: float
+    think_accuracy: float
+    answer_accuracy: float
+    think_samples: int
+    answer_samples: int
+    # Accuracy of the block-boundary decision, restricted to the control menu, and of the
+    # stop-versus-continue half of it on its own. Zero outside the adaptive objective.
+    control_accuracy: float = 0.0
+    stop_accuracy: float = 0.0
+
+
+def hybrid_objective(
+    model: DiffusionTransformer,
+    tokens: Tensor,
+    regions: Tensor,
+    *,
+    block: int,
+    think_probability: float,
+    mask_eps: float,
+    generator: torch.Generator | None = None,
+) -> HybridBatchLoss:
+    """Mixed objective: denoise all thought slots or predict the answer tokens.
+
+    Each sample is assigned one mode. Think samples corrupt every slot at an
+    independent noise level under a block-causal mask, so one forward trains all
+    slot conditionals; answer samples see the full clean reasoning prefix
+    bidirectionally and the answer causally.
+    """
+
+    device = tokens.device
+    batch_size, seq_len = tokens.shape
+    problem_len, n_slots, answer_start, answer_end = regions.unbind(dim=1)
+    positions = torch.arange(seq_len, device=device)
+
+    think_sel = (
+        torch.rand(batch_size, device=device, generator=generator) < think_probability
+    )
+    if bool(think_sel.all()):
+        think_sel[-1] = False
+    if not bool(think_sel.any()):
+        think_sel[0] = True
+
+    prefix_end = problem_len + 1
+    think_end = prefix_end + n_slots * block
+    think_region = (
+        (positions[None, :] >= prefix_end[:, None])
+        & (positions[None, :] < think_end[:, None])
+        & think_sel[:, None]
+    )
+    max_slots = int(n_slots.max())
+    slot_noise = mask_eps + (1.0 - mask_eps) * torch.rand(
+        batch_size, max_slots, device=device, generator=generator
+    )
+    slot_index = torch.clamp(
+        (positions[None, :] - prefix_end[:, None]) // block, min=0, max=max_slots - 1
+    )
+    token_noise = slot_noise.gather(1, slot_index)
+    mask = (
+        torch.rand(tokens.shape, device=device, generator=generator) < token_noise
+    ) & think_region
+    noisy_tokens = torch.where(mask, model.config.mask_token_id, tokens)
+
+    think_blocked = slot_causal_blocked(problem_len, n_slots, block, seq_len)
+    answer_blocked = prefix_causal_blocked(answer_start, seq_len)
+    blocked = torch.where(think_sel[:, None, None], think_blocked, answer_blocked)
+
+    predict_positions = (
+        (positions[None, :] >= answer_start[:, None] - 1)
+        & (positions[None, :] < answer_end[:, None] - 1)
+        & ~think_sel[:, None]
+    )
+    output_positions = mask | predict_positions
+    logits = model(noisy_tokens, output_positions=output_positions, attn_mask=blocked)
+
+    think_rows = mask[output_positions]
+    # Graph-connected zero: forbidden-output columns sit at finfo.min, so a raw sum
+    # overflows to -inf in low precision and would poison the scalar via -inf * 0.
+    zero = logits.sum().clamp(-1.0, 1.0) * 0.0
+
+    think_loss = zero
+    think_accuracy = 0.0
+    if bool(mask.any()):
+        think_logits = logits[think_rows]
+        think_targets = tokens[mask]
+        per_token = F.cross_entropy(think_logits.float(), think_targets, reduction='none')
+        weights = token_noise[mask]
+        normalizer = think_region.sum().clamp_min(1)
+        think_loss = (per_token / weights).sum() / normalizer
+        think_accuracy = float(
+            (think_logits.argmax(dim=-1) == think_targets).float().mean()
+        )
+
+    answer_loss = zero
+    answer_accuracy = 0.0
+    if bool(predict_positions.any()):
+        answer_logits = logits[~think_rows]
+        target_positions = torch.zeros_like(predict_positions)
+        target_positions[:, 1:] = predict_positions[:, :-1]
+        answer_targets = tokens[target_positions]
+        answer_loss = F.cross_entropy(answer_logits.float(), answer_targets)
+        answer_accuracy = float(
+            (answer_logits.argmax(dim=-1) == answer_targets).float().mean()
+        )
+
+    return HybridBatchLoss(
+        loss=think_loss + answer_loss,
+        think_loss=float(think_loss),
+        answer_loss=float(answer_loss),
+        think_accuracy=think_accuracy,
+        answer_accuracy=answer_accuracy,
+        think_samples=int(think_sel.sum()),
+        answer_samples=int((~think_sel).sum()),
+    )
+
+
+def adaptive_hybrid_objective(
+    model: DiffusionTransformer,
+    tokens: Tensor,
+    regions: Tensor,
+    *,
+    size_ids: Tensor,
+    end_think_id: int,
+    think_probability: float,
+    mask_eps: float,
+    causal_prefix: bool = False,
+    control_context_noise: float = 0.0,
+    generator: torch.Generator | None = None,
+) -> HybridBatchLoss:
+    """Mixed objective for the adaptive layout.
+
+    Think samples denoise every thought block at an independent noise level under
+    the variable-boundary block-causal mask. The remaining samples run a causal
+    next-token objective over the whole reasoning-and-answer stream, which is where
+    the model learns the ``<szN>`` block-size and ``</think>`` termination decisions.
+
+    Those causal samples read a pristine think region, while at inference the controller
+    reads thoughts the model just wrote, complete with sampling damage.
+    ``control_context_noise`` closes that gap by swapping a fraction of their think tokens
+    for other tokens drawn from the batch: inputs degrade, targets stay clean.
+    """
+
+    device = tokens.device
+    batch_size, seq_len = tokens.shape
+    problem_len, _, answer_start, answer_end = regions.unbind(dim=1)
+    positions = torch.arange(seq_len, device=device)
+
+    think_sel = (
+        torch.rand(batch_size, device=device, generator=generator) < think_probability
+    )
+    if bool(think_sel.all()):
+        think_sel[-1] = False
+    if not bool(think_sel.any()):
+        think_sel[0] = True
+
+    prefix_end = problem_len + 1
+    in_think = (positions[None, :] >= prefix_end[:, None]) & (
+        positions[None, :] < answer_start[:, None]
+    )
+    is_size = (tokens.unsqueeze(-1) == size_ids).any(dim=-1)
+    boundary = (is_size | (tokens == end_think_id)) & in_think
+    think_content = in_think & ~boundary & think_sel[:, None]
+
+    block_id = torch.cumsum((is_size & in_think).long(), dim=1)
+    max_blocks = int(block_id.max().clamp(min=1))
+    slot_noise = mask_eps + (1.0 - mask_eps) * torch.rand(
+        batch_size, max_blocks, device=device, generator=generator
+    )
+    gather_index = (block_id - 1).clamp(min=0, max=max_blocks - 1)
+    token_noise = slot_noise.gather(1, gather_index)
+    mask = (
+        torch.rand(tokens.shape, device=device, generator=generator) < token_noise
+    ) & think_content
+    noisy_tokens = torch.where(mask, model.config.mask_token_id, tokens)
+
+    if control_context_noise > 0.0:
+        # Replacements come from the batch itself so the corrupted context keeps realistic
+        # token statistics; boundaries and targets are read from the clean tensor.
+        control_content = in_think & ~boundary & ~think_sel[:, None]
+        corrupt = (
+            torch.rand(tokens.shape, device=device, generator=generator) < control_context_noise
+        ) & control_content
+        flat = tokens.reshape(-1)
+        picks = torch.randint(
+            0, flat.numel(), tokens.shape, device=device, generator=generator
+        )
+        noisy_tokens = torch.where(corrupt, flat[picks], noisy_tokens)
+
+    think_blocked = adaptive_block_mask(
+        tokens, problem_len, answer_start, size_ids, end_think_id,
+        causal_prefix=causal_prefix,
+    )
+    ar_blocked = prefix_causal_blocked(prefix_end, seq_len, causal_prefix=causal_prefix)
+    blocked = torch.where(think_sel[:, None, None], think_blocked, ar_blocked)
+
+    predict_positions = (
+        (positions[None, :] >= problem_len[:, None])
+        & (positions[None, :] < answer_end[:, None] - 1)
+        & ~think_sel[:, None]
+    )
+    output_positions = mask | predict_positions
+    logits = model(noisy_tokens, output_positions=output_positions, attn_mask=blocked)
+
+    think_rows = mask[output_positions]
+    # Graph-connected zero: forbidden-output columns sit at finfo.min, so a raw sum
+    # overflows to -inf in low precision and would poison the scalar via -inf * 0.
+    zero = logits.sum().clamp(-1.0, 1.0) * 0.0
+
+    think_loss = zero
+    think_accuracy = 0.0
+    if bool(mask.any()):
+        think_logits = logits[think_rows]
+        think_targets = tokens[mask]
+        per_token = F.cross_entropy(think_logits.float(), think_targets, reduction='none')
+        weights = token_noise[mask]
+        normalizer = think_content.sum().clamp_min(1)
+        think_loss = (per_token / weights).sum() / normalizer
+        think_accuracy = float(
+            (think_logits.argmax(dim=-1) == think_targets).float().mean()
+        )
+
+    answer_loss = zero
+    answer_accuracy = 0.0
+    control_accuracy = 0.0
+    stop_accuracy = 0.0
+    if bool(predict_positions.any()):
+        answer_logits = logits[~think_rows]
+        target_positions = torch.zeros_like(predict_positions)
+        target_positions[:, 1:] = predict_positions[:, :-1]
+        answer_targets = tokens[target_positions]
+        answer_loss = F.cross_entropy(answer_logits.float(), answer_targets)
+        control_ids = torch.cat(
+            [size_ids, torch.tensor([end_think_id], device=device, dtype=size_ids.dtype)]
+        )
+        is_control = (answer_targets.unsqueeze(-1) == control_ids).any(dim=-1)
+        answer_accuracy = float(
+            (answer_logits.argmax(dim=-1) == answer_targets).float().mean()
+        )
+        if bool(is_control.any()):
+            # Scored inside the control menu: an open-vocabulary argmax hides whether the
+            # boundary decision is right, since content tokens dominate those logits.
+            menu = answer_logits[is_control].index_select(-1, control_ids)
+            chosen = control_ids[menu.argmax(dim=-1)]
+            targets = answer_targets[is_control]
+            control_accuracy = float((chosen == targets).float().mean())
+            stop_accuracy = float(
+                ((chosen == end_think_id) == (targets == end_think_id)).float().mean()
+            )
+
+    return HybridBatchLoss(
+        loss=think_loss + answer_loss,
+        think_loss=float(think_loss),
+        answer_loss=float(answer_loss),
+        think_accuracy=think_accuracy,
+        answer_accuracy=answer_accuracy,
+        think_samples=int(think_sel.sum()),
+        answer_samples=int((~think_sel).sum()),
+        control_accuracy=control_accuracy,
+        stop_accuracy=stop_accuracy,
+    )
+
+
+def block_size_curriculum(
+    step: int | None, *, n_sizes: int, curriculum_steps: int
+) -> Tensor:
+    """Segment-size sampling weights: smallest-size-only ramping linearly to uniform.
+
+    ``sizes`` are assumed ascending. ``step=None`` (evaluation) and a zero-length
+    curriculum both return the uniform end state so losses stay comparable across
+    checkpoints.
+    """
+
+    uniform = torch.full((n_sizes,), 1.0 / n_sizes)
+    if step is None or curriculum_steps <= 0:
+        return uniform
+    progress = min(1.0, step / curriculum_steps)
+    smallest_only = torch.zeros(n_sizes)
+    smallest_only[0] = 1.0
+    return smallest_only * (1.0 - progress) + uniform * progress
+
+
+def block_diffusion_objective(
+    model: DiffusionTransformer,
+    tokens: Tensor,
+    *,
+    sizes: tuple[int, ...],
+    size_weights: Tensor,
+    mask_eps: float,
+    ar_probability: float = 0.0,
+    generator: torch.Generator | None = None,
+) -> HybridBatchLoss:
+    """Variable-block denoising over plain packed text (the conversion objective).
+
+    Each sample is tiled with segments whose lengths are drawn from ``sizes`` under
+    ``size_weights``; every segment is corrupted at an independent noise level and
+    denoised in one forward under the block-causal geometry (own segment bidirectional,
+    all preceding segments visible). No control tokens exist in the stream — pretraining
+    teaches variable-block denoising only; control decisions are learned in SFT. With
+    probability ``ar_probability`` a sample instead runs plain causal next-token loss,
+    retaining the autoregressive ability that control and answer decoding rely on.
+    Reported as ``think_*`` (denoising) and ``answer_*`` (causal retention) metrics.
+    """
+
+    device = tokens.device
+    batch_size, seq_len = tokens.shape
+    positions = torch.arange(seq_len, device=device)
+
+    ar_sel = (
+        torch.rand(batch_size, device=device, generator=generator) < ar_probability
+    )
+    if bool(ar_sel.all()):
+        ar_sel[0] = False
+    diff_sel = ~ar_sel
+
+    size_tensor = torch.tensor(sizes, device=device, dtype=torch.long)
+    max_segments = -(-seq_len // int(min(sizes)))
+    drawn_index = torch.multinomial(
+        size_weights.to(device).expand(batch_size, -1),
+        max_segments,
+        replacement=True,
+        generator=generator,
+    )
+    drawn = size_tensor[drawn_index]
+    starts = torch.cumsum(drawn, dim=1) - drawn
+    valid = starts < seq_len
+    boundary_hits = torch.zeros(batch_size, seq_len, dtype=torch.long, device=device)
+    boundary_hits.scatter_add_(1, starts.clamp(max=seq_len - 1), valid.long())
+    boundary = boundary_hits > 0
+
+    segment_id = torch.cumsum(boundary.long(), dim=1) - 1
+    segment_noise = mask_eps + (1.0 - mask_eps) * torch.rand(
+        batch_size, max_segments, device=device, generator=generator
+    )
+    token_noise = segment_noise.gather(1, segment_id.clamp(max=max_segments - 1))
+    mask = (
+        torch.rand(tokens.shape, device=device, generator=generator) < token_noise
+    ) & diff_sel[:, None]
+    noisy_tokens = torch.where(mask, model.config.mask_token_id, tokens)
+
+    zeros = torch.zeros(batch_size, dtype=torch.long, device=device)
+    full = torch.full((batch_size,), seq_len, dtype=torch.long, device=device)
+    block_blocked = block_mask_from_boundaries(boundary, zeros, full)
+    causal_blocked = prefix_causal_blocked(zeros, seq_len, causal_prefix=True)
+    blocked = torch.where(diff_sel[:, None, None], block_blocked, causal_blocked)
+
+    predict_positions = (positions[None, :] < seq_len - 1) & ar_sel[:, None]
+    output_positions = mask | predict_positions
+    logits = model(noisy_tokens, output_positions=output_positions, attn_mask=blocked)
+
+    think_rows = mask[output_positions]
+    # Graph-connected zero: forbidden-output columns sit at finfo.min, so a raw sum
+    # overflows to -inf in low precision and would poison the scalar via -inf * 0.
+    zero = logits.sum().clamp(-1.0, 1.0) * 0.0
+
+    think_loss = zero
+    think_accuracy = 0.0
+    if bool(mask.any()):
+        think_logits = logits[think_rows]
+        think_targets = tokens[mask]
+        per_token = F.cross_entropy(think_logits.float(), think_targets, reduction='none')
+        weights = token_noise[mask]
+        normalizer = (diff_sel.sum() * seq_len).clamp_min(1)
+        think_loss = (per_token / weights).sum() / normalizer
+        think_accuracy = float(
+            (think_logits.argmax(dim=-1) == think_targets).float().mean()
+        )
+
+    answer_loss = zero
+    answer_accuracy = 0.0
+    if bool(predict_positions.any()):
+        answer_logits = logits[~think_rows]
+        target_positions = torch.zeros_like(predict_positions)
+        target_positions[:, 1:] = predict_positions[:, :-1]
+        answer_targets = tokens[target_positions]
+        answer_loss = F.cross_entropy(answer_logits.float(), answer_targets)
+        answer_accuracy = float(
+            (answer_logits.argmax(dim=-1) == answer_targets).float().mean()
+        )
+
+    return HybridBatchLoss(
+        loss=think_loss + answer_loss,
+        think_loss=float(think_loss),
+        answer_loss=float(answer_loss),
+        think_accuracy=think_accuracy,
+        answer_accuracy=answer_accuracy,
+        think_samples=int(diff_sel.sum()),
+        answer_samples=int(ar_sel.sum()),
+    )
+
+
+@dataclass(frozen=True)
+class ARBatchLoss:
+    loss: Tensor
+    accuracy: float
+    token_count: int
+
+
+def ar_objective(
+    model: DiffusionTransformer, tokens: Tensor, regions: Tensor
+) -> ARBatchLoss:
+    """Prefix-LM next-token objective over the think and answer regions."""
+
+    device = tokens.device
+    _, seq_len = tokens.shape
+    problem_len, _, _, answer_end = regions.unbind(dim=1)
+    positions = torch.arange(seq_len, device=device)
+
+    blocked = prefix_causal_blocked(problem_len + 1, seq_len)
+    predict_positions = (positions[None, :] >= problem_len[:, None]) & (
+        positions[None, :] < answer_end[:, None] - 1
+    )
+    logits = model(tokens, output_positions=predict_positions, attn_mask=blocked)
+
+    target_positions = torch.zeros_like(predict_positions)
+    target_positions[:, 1:] = predict_positions[:, :-1]
+    targets = tokens[target_positions]
+    loss = F.cross_entropy(logits.float(), targets)
+    accuracy = float((logits.argmax(dim=-1) == targets).float().mean())
+    return ARBatchLoss(loss=loss, accuracy=accuracy, token_count=int(targets.numel()))
+
+
+def diffusion_objective(
+    model: DiffusionTransformer,
+    tokens: Tensor,
+    regions: Tensor,
+    *,
+    mask_eps: float,
+    mask_probability: Tensor | None = None,
+    generator: torch.Generator | None = None,
+) -> tuple[Tensor, CorruptionBatch, Tensor]:
+    """Whole-response denoising: corrupt everything after the problem, pads included."""
+
+    device = tokens.device
+    _, seq_len = tokens.shape
+    problem_len = regions[:, 0]
+    positions = torch.arange(seq_len, device=device)
+    response_mask = positions[None, :] >= (problem_len[:, None] + 1)
+    corruption = corrupt_tokens(
+        tokens,
+        model.config.mask_token_id,
+        valid_mask=response_mask,
+        mask_probability=mask_probability,
+        eps=mask_eps,
+        generator=generator,
+    )
+    logits = model(corruption.noisy_tokens, output_positions=corruption.mask)
+    return logits, corruption, response_mask
+
+
+@dataclass
+class GenerationResult:
+    tokens: list[int]
+    think_tokens: list[int]
+    answer_tokens: list[int]
+    think_seconds: float = 0.0
+    answer_seconds: float = 0.0
+    forward_passes: int = 0
+    slots_used: int = 0
+    block_sizes: tuple[int, ...] = ()
+    terminated: bool = False
+
+    @property
+    def total_seconds(self) -> float:
+        return self.think_seconds + self.answer_seconds
+
+
+def _apply_repetition_penalty(
+    logits: Tensor, token_ids: list[int], penalty: float
+) -> Tensor:
+    """Divide logits of already-emitted tokens by ``penalty`` (CTRL convention)."""
+
+    if penalty == 1.0 or not token_ids:
+        return logits
+    index = torch.tensor(sorted(set(token_ids)), device=logits.device)
+    selected = logits.index_select(-1, index)
+    adjusted = torch.where(selected > 0, selected / penalty, selected * penalty)
+    return logits.index_copy(-1, index, adjusted)
+
+
+def _apply_top_p(logits: Tensor, top_p: float) -> Tensor:
+    """Restrict sampling to the smallest set of tokens whose mass reaches ``top_p``."""
+
+    if top_p >= 1.0:
+        return logits
+    ordered, indices = torch.sort(logits, descending=True, dim=-1)
+    cumulative = ordered.softmax(dim=-1).cumsum(dim=-1)
+    remove = cumulative - ordered.softmax(dim=-1) >= top_p
+    ordered = ordered.masked_fill(remove, torch.finfo(logits.dtype).min)
+    return ordered.gather(-1, indices.argsort(dim=-1))
+
+
+def _kv_cache_enabled(part: str) -> bool:
+    """Key/value caching per part, selected by ``MDLM_KV_CACHE``: ar, block, all or off.
+
+    Defaults to ``ar``, which measured 101s to 12.6s on the same prompt's answer: one query
+    token against an all-visible mask has no downside. The denoising prefix does — caching it
+    hands the backbone a dense mask, dropping flex attention onto its score_mod path and losing
+    the block skipping the uncached call gets, 10.6s per block against 5.3s on an L40S
+    (2026-07-27). It stays off until that path builds a rectangular BlockMask instead.
+    """
+
+    setting = os.environ.get('MDLM_KV_CACHE', 'ar')
+    return setting in ('all', '1') or setting == part
+
+
+def _cached_block_logits(model, blocked: Tensor, prefix_len: int):
+    """Score a denoising block against a cached prefix, or ``None`` without cache support.
+
+    The prefix is encoded once per block; every denoising step then feeds only the block's
+    own positions. Its keys and values are dropped between steps because the block's tokens
+    keep changing as they are revealed, while the prefix behind them does not.
+    """
+
+    if not hasattr(model, 'forward_cached') or not _kv_cache_enabled('block'):
+        return None
+
+    cache = model.new_cache()
+
+    def logits_fn(tokens: Tensor, masked: Tensor) -> Tensor:
+        if cache.get_seq_length() == 0:
+            with torch.inference_mode():
+                model.forward_cached(
+                    tokens[:, :prefix_len],
+                    attn_mask=blocked[:, :prefix_len, :prefix_len],
+                    past_key_values=cache,
+                )
+        cache.crop(prefix_len)
+        with torch.inference_mode():
+            logits, _ = model.forward_cached(
+                tokens[:, prefix_len:],
+                attn_mask=blocked[:, prefix_len:, :],
+                past_key_values=cache,
+                output_positions=masked[:, prefix_len:],
+            )
+        return logits
+
+    return logits_fn
+
+
+def _ar_decode_cached(
+    model: DiffusionTransformer,
+    sequence: list[int],
+    prefix_len: int,
+    *,
+    causal_prefix: bool,
+    eos_id: int,
+    max_new_tokens: int,
+    temperature: float,
+    repetition_penalty: float,
+    top_p: float,
+    device: torch.device,
+    generator: torch.Generator | None,
+) -> tuple[list[int], int]:
+    """Same decoding as :func:`_ar_decode` with the prefix kept in a key/value cache.
+
+    Masks come from :func:`prefix_causal_blocked` exactly as in the uncached path, sliced to
+    the rows of the queries being fed. Priming with an all-visible mask instead would look
+    right — the final row is identical, so a single generated token matches — while silently
+    computing every earlier position bidirectionally and poisoning the cached keys.
+    """
+
+    generated: list[int] = []
+    cache = model.new_cache()
+    step_in = torch.tensor([sequence], dtype=torch.long, device=device)
+    forwards = 0
+    prefix = torch.tensor([prefix_len], device=device)
+    for _ in range(max_new_tokens):
+        cached = cache.get_seq_length()
+        length = cached + step_in.shape[1]
+        visible = prefix_causal_blocked(prefix, length, causal_prefix=causal_prefix)[:, cached:, :]
+        output_positions = torch.zeros_like(step_in, dtype=torch.bool)
+        output_positions[0, -1] = True
+        with torch.inference_mode():
+            logits, cache = model.forward_cached(
+                step_in, attn_mask=visible, past_key_values=cache,
+                output_positions=output_positions,
+            )
+            logits = _apply_repetition_penalty(logits, generated, repetition_penalty)
+            logits = _apply_top_p(logits, top_p)
+            token, _ = _sample_categorical(logits, temperature, generator)
+        forwards += 1
+        token_id = int(token.item())
+        generated.append(token_id)
+        if token_id == eos_id:
+            break
+        step_in = torch.tensor([[token_id]], dtype=torch.long, device=device)
+    return generated, forwards
+
+
+def _ar_decode(
+    model: DiffusionTransformer,
+    sequence: list[int],
+    prefix_len: int,
+    *,
+    eos_id: int,
+    max_new_tokens: int,
+    temperature: float,
+    repetition_penalty: float,
+    top_p: float,
+    device: torch.device,
+    generator: torch.Generator | None,
+    causal_prefix: bool = False,
+) -> tuple[list[int], int]:
+    """Greedy/temperature decoding under the bidirectional-prefix causal mask.
+
+    ``repetition_penalty`` (>1 discourages repeats) and ``top_p`` nucleus truncation
+    curb the degenerate loops small models fall into under plain temperature sampling.
+    """
+
+    generated: list[int] = []
+    forwards = 0
+    max_new_tokens = min(max_new_tokens, model.config.max_seq_len - len(sequence))
+    if hasattr(model, 'forward_cached') and _kv_cache_enabled('ar'):
+        return _ar_decode_cached(
+            model, sequence, prefix_len, causal_prefix=causal_prefix, eos_id=eos_id,
+            max_new_tokens=max_new_tokens, temperature=temperature,
+            repetition_penalty=repetition_penalty, top_p=top_p, device=device,
+            generator=generator,
+        )
+    for _ in range(max_new_tokens):
+        current = torch.tensor([sequence + generated], dtype=torch.long, device=device)
+        seq_len = current.shape[1]
+        prefix = torch.tensor([prefix_len], device=device)
+        blocked = prefix_causal_blocked(prefix, seq_len, causal_prefix=causal_prefix)
+        output_positions = torch.zeros_like(current, dtype=torch.bool)
+        output_positions[0, -1] = True
+        with torch.inference_mode():
+            logits = model(current, output_positions=output_positions, attn_mask=blocked)
+            logits = _apply_repetition_penalty(logits, generated, repetition_penalty)
+            logits = _apply_top_p(logits, top_p)
+            token, _ = _sample_categorical(logits, temperature, generator)
+        forwards += 1
+        token_id = int(token.item())
+        generated.append(token_id)
+        if token_id == eos_id:
+            break
+    return generated, forwards
+
+
+@torch.no_grad()
+def generate_hybrid(
+    model: DiffusionTransformer,
+    prompt_ids: list[int],
+    *,
+    think_id: int,
+    end_think_id: int,
+    thought_pad_id: int,
+    eos_id: int,
+    block: int,
+    max_slots: int,
+    steps_per_block: int,
+    max_answer_tokens: int = 64,
+    temperature: float = 0.7,
+    repetition_penalty: float = 1.0,
+    top_p: float = 1.0,
+    strategy: str = 'confidence',
+    device: torch.device | str = 'cpu',
+    generator: torch.Generator | None = None,
+) -> GenerationResult:
+    """Denoise thought slots sequentially, then decode the answer autoregressively."""
+
+    device = torch.device(device)
+    mask_id = model.config.mask_token_id
+    sequence = [*prompt_ids, think_id]
+    forwards = 0
+    slots_used = 0
+    # Each new slot must leave room for itself plus at least a minimal answer.
+    slot_budget = model.config.max_seq_len - block - 8
+
+    think_started = time.perf_counter()
+    problem_tensor = torch.tensor([len(prompt_ids)], device=device)
+    for slot_index in range(max_slots):
+        if len(sequence) > slot_budget:
+            break
+        window = torch.tensor(
+            [sequence + [mask_id] * block], dtype=torch.long, device=device
+        )
+        blocked = slot_causal_blocked(
+            problem_tensor,
+            torch.tensor([slot_index + 1], device=device),
+            block,
+            window.shape[1],
+        )
+        filled = iterative_unmask(
+            model,
+            window,
+            mask_id,
+            steps=steps_per_block,
+            temperature=temperature,
+            strategy=strategy,
+            attn_mask=blocked,
+            generator=generator,
+        )
+        slot = [int(token) for token in filled[0, len(sequence):]]
+        forwards += steps_per_block
+        slots_used += 1
+        sequence.extend(slot)
+        if end_think_id in slot:
+            break
+    if end_think_id not in sequence[len(prompt_ids):]:
+        # Match the trained answer geometry when the model never closes its thinking.
+        terminal = [end_think_id] + [thought_pad_id] * (block - 1)
+        sequence.extend(terminal[: max(1, model.config.max_seq_len - 8 - len(sequence))])
+    think_seconds = time.perf_counter() - think_started
+    think_tokens = sequence[len(prompt_ids):]
+
+    answer_started = time.perf_counter()
+    answer, answer_forwards = _ar_decode(
+        model,
+        sequence,
+        prefix_len=len(sequence),
+        eos_id=eos_id,
+        max_new_tokens=max_answer_tokens,
+        temperature=temperature,
+        repetition_penalty=repetition_penalty,
+        top_p=top_p,
+        device=device,
+        generator=generator,
+    )
+    answer_seconds = time.perf_counter() - answer_started
+    return GenerationResult(
+        tokens=sequence + answer,
+        think_tokens=think_tokens,
+        answer_tokens=answer,
+        think_seconds=think_seconds,
+        answer_seconds=answer_seconds,
+        forward_passes=forwards + answer_forwards,
+        slots_used=slots_used,
+    )
+
+
+def _ar_predict_control(
+    model: DiffusionTransformer,
+    sequence: list[int],
+    allowed_ids: list[int],
+    *,
+    prefix_len: int,
+    temperature: float,
+    device: torch.device,
+    generator: torch.Generator | None,
+    causal_prefix: bool = False,
+) -> int:
+    """Predict the next control token, restricted to the allowed size/stop ids."""
+
+    current = torch.tensor([sequence], dtype=torch.long, device=device)
+    seq_len = current.shape[1]
+    blocked = prefix_causal_blocked(
+        torch.tensor([prefix_len], device=device), seq_len, causal_prefix=causal_prefix
+    )
+    output_positions = torch.zeros_like(current, dtype=torch.bool)
+    output_positions[0, -1] = True
+    with torch.inference_mode():
+        logits = model(current, output_positions=output_positions, attn_mask=blocked)
+        restricted = torch.full_like(logits, torch.finfo(logits.dtype).min)
+        index = torch.tensor(allowed_ids, device=logits.device)
+        restricted.index_copy_(-1, index, logits.index_select(-1, index))
+        token, _ = _sample_categorical(restricted, temperature, generator)
+    return int(token.item())
+
+
+@torch.no_grad()
+def generate_hybrid_adaptive(
+    model: DiffusionTransformer,
+    prompt_ids: list[int],
+    *,
+    think_id: int,
+    end_think_id: int,
+    thought_pad_id: int,
+    eos_id: int,
+    size_ids: dict[int, int],
+    steps_per_block: int,
+    max_blocks: int,
+    max_answer_tokens: int = 96,
+    temperature: float = 0.7,
+    control_temperature: float = 0.0,
+    repetition_penalty: float = 1.0,
+    top_p: float = 1.0,
+    strategy: str = 'confidence',
+    causal_prefix: bool = False,
+    device: torch.device | str = 'cpu',
+    generator: torch.Generator | None = None,
+) -> GenerationResult:
+    """Interleave AR block-size decisions with in-block diffusion, then decode.
+
+    At each boundary the model predicts a ``<szN>`` control token or ``</think>``.
+    A size token allocates that many masked positions denoised in parallel under the
+    variable-boundary block-causal mask; ``</think>`` ends thinking. The answer is
+    then decoded autoregressively with the same repetition and nucleus controls.
+    """
+
+    device = torch.device(device)
+    mask_id = model.config.mask_token_id
+    size_by_id = {token_id: size for size, token_id in size_ids.items()}
+    size_ids_tensor = torch.tensor(sorted(size_ids.values()), device=device)
+    control_ids = [*size_by_id.keys(), end_think_id]
+    prefix_len = len(prompt_ids) + 1
+    problem_tensor = torch.tensor([len(prompt_ids)], device=device)
+
+    sequence = [*prompt_ids, think_id]
+    forwards = 0
+    blocks_used = 0
+    chosen: list[int] = []
+    terminated = False
+
+    think_started = time.perf_counter()
+    for _ in range(max_blocks):
+        control = _ar_predict_control(
+            model,
+            sequence,
+            control_ids,
+            prefix_len=prefix_len,
+            temperature=control_temperature,
+            device=device,
+            generator=generator,
+            causal_prefix=causal_prefix,
+        )
+        forwards += 1
+        if control == end_think_id:
+            terminated = True
+            break
+        size = size_by_id[control]
+        if len(sequence) + 1 + size > model.config.max_seq_len - 8:
+            break
+        sequence.append(control)
+        window_prefix = len(sequence)
+        window = torch.tensor(
+            [sequence + [mask_id] * size], dtype=torch.long, device=device
+        )
+        blocked = adaptive_block_mask(
+            window,
+            problem_tensor,
+            torch.tensor([window.shape[1]], device=device),
+            size_ids_tensor,
+            end_think_id,
+            causal_prefix=causal_prefix,
+        )
+        filled = iterative_unmask(
+            model,
+            window,
+            mask_id,
+            steps=steps_per_block,
+            temperature=temperature,
+            strategy=strategy,
+            attn_mask=blocked,
+            generator=generator,
+            logits_fn=_cached_block_logits(model, blocked, window_prefix),
+        )
+        sequence.extend(int(token) for token in filled[0, window_prefix:])
+        forwards += steps_per_block
+        blocks_used += 1
+        chosen.append(size)
+
+    sequence.append(end_think_id)
+    think_seconds = time.perf_counter() - think_started
+    think_tokens = sequence[len(prompt_ids):]
+
+    answer_started = time.perf_counter()
+    answer, answer_forwards = _ar_decode(
+        model,
+        sequence,
+        prefix_len=prefix_len,
+        eos_id=eos_id,
+        max_new_tokens=max_answer_tokens,
+        temperature=temperature,
+        repetition_penalty=repetition_penalty,
+        top_p=top_p,
+        device=device,
+        generator=generator,
+        causal_prefix=causal_prefix,
+    )
+    answer_seconds = time.perf_counter() - answer_started
+    return GenerationResult(
+        tokens=sequence + answer,
+        think_tokens=think_tokens,
+        answer_tokens=answer,
+        think_seconds=think_seconds,
+        answer_seconds=answer_seconds,
+        forward_passes=forwards + answer_forwards,
+        slots_used=blocks_used,
+        block_sizes=tuple(chosen),
+        terminated=terminated,
+    )
+
+
+@torch.no_grad()
+def generate_ar(
+    model: DiffusionTransformer,
+    prompt_ids: list[int],
+    *,
+    think_id: int,
+    end_think_id: int,
+    eos_id: int,
+    max_new_tokens: int = 384,
+    temperature: float = 0.7,
+    device: torch.device | str = 'cpu',
+    generator: torch.Generator | None = None,
+) -> GenerationResult:
+    """Classic sequential CoT baseline under the prefix-LM mask."""
+
+    device = torch.device(device)
+    sequence = [*prompt_ids, think_id]
+    started = time.perf_counter()
+    generated, forwards = _ar_decode(
+        model,
+        sequence,
+        prefix_len=len(sequence),
+        eos_id=eos_id,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        device=device,
+        generator=generator,
+    )
+    elapsed = time.perf_counter() - started
+    if end_think_id in generated:
+        split = generated.index(end_think_id) + 1
+    else:
+        split = len(generated)
+    return GenerationResult(
+        tokens=sequence + generated,
+        think_tokens=generated[:split],
+        answer_tokens=generated[split:],
+        think_seconds=elapsed,
+        answer_seconds=0.0,
+        forward_passes=forwards,
+    )
+
+
+@torch.no_grad()
+def generate_diffusion(
+    model: DiffusionTransformer,
+    prompt_ids: list[int],
+    *,
+    think_id: int,
+    end_think_id: int,
+    eos_id: int,
+    response_budget: int,
+    steps: int,
+    temperature: float = 0.7,
+    blocked_token_ids: tuple[int, ...] = (),
+    device: torch.device | str = 'cpu',
+    generator: torch.Generator | None = None,
+) -> GenerationResult:
+    """Pure-diffusion baseline: denoise the entire response region at once.
+
+    Blocking the pad token here counters confidence-ordered pad collapse: pads are
+    the easiest predictions, so left unblocked they win every early reveal and
+    squeeze out the actual response text.
+    """
+
+    device = torch.device(device)
+    mask_id = model.config.mask_token_id
+    budget = min(response_budget, model.config.max_seq_len - len(prompt_ids) - 1)
+    sequence = torch.tensor(
+        [[*prompt_ids, think_id] + [mask_id] * budget], dtype=torch.long, device=device
+    )
+    started = time.perf_counter()
+    filled = iterative_unmask(
+        model,
+        sequence,
+        mask_id,
+        steps=steps,
+        temperature=temperature,
+        strategy='confidence',
+        blocked_token_ids=blocked_token_ids,
+        generator=generator,
+    )
+    elapsed = time.perf_counter() - started
+    response = [int(token) for token in filled[0, len(prompt_ids) + 1:]]
+    if eos_id in response:
+        response = response[: response.index(eos_id) + 1]
+    if end_think_id in response:
+        split = response.index(end_think_id) + 1
+    else:
+        split = len(response)
+    return GenerationResult(
+        tokens=[*prompt_ids, think_id] + response,
+        think_tokens=response[:split],
+        answer_tokens=response[split:],
+        think_seconds=elapsed,
+        answer_seconds=0.0,
+        forward_passes=steps,
+    )

src/diffusion_lm/model.py

@@ -0,0 +1,250 @@
+"""GPT-sized bidirectional Transformer used as a masked-token denoiser."""
+
+from __future__ import annotations
+
+import argparse
+import math
+from pathlib import Path
+
+import torch
+from torch import Tensor, nn
+from torch.utils.checkpoint import checkpoint
+
+from diffusion_lm.config import ModelConfig, load_config
+
+
+class DiffusionTransformer(nn.Module):
+    """A GPT-like Transformer with the causal mask deliberately removed.
+
+    The network predicts clean tokens from an input containing absorbing mask
+    tokens. Passing ``output_positions`` avoids materializing vocabulary logits
+    for already-visible tokens during training.
+    """
+
+    def __init__(self, config: ModelConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.tokenizer_sha256: str | None = None
+        self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
+        self.position_embedding = nn.Embedding(config.max_seq_len, config.d_model)
+        self.embedding_dropout = nn.Dropout(config.dropout)
+
+        if config.use_flex_attention:
+            from diffusion_lm.flexattn import FlexEncoder
+
+            self.transformer = FlexEncoder(
+                d_model=config.d_model,
+                n_heads=config.n_heads,
+                d_ff=config.d_ff,
+                dropout=config.dropout,
+                n_layers=config.n_layers,
+                activation_checkpointing=config.activation_checkpointing,
+            )
+        else:
+            layer = nn.TransformerEncoderLayer(
+                d_model=config.d_model,
+                nhead=config.n_heads,
+                dim_feedforward=config.d_ff,
+                dropout=config.dropout,
+                activation="gelu",
+                batch_first=True,
+                norm_first=True,
+            )
+            self.transformer = nn.TransformerEncoder(
+                layer,
+                num_layers=config.n_layers,
+                norm=nn.LayerNorm(config.d_model),
+                enable_nested_tensor=False,
+            )
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+
+        self.apply(self._init_weights)
+        self._init_residual_outputs()
+        if config.tie_embeddings:
+            self.lm_head.weight = self.token_embedding.weight
+        self.register_buffer(
+            "_forbidden_output_token_ids",
+            torch.tensor(config.forbidden_output_token_ids, dtype=torch.long),
+            persistent=False,
+        )
+
+    @staticmethod
+    def _init_weights(module: nn.Module) -> None:
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if isinstance(module, nn.Linear) and module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.ones_(module.weight)
+            nn.init.zeros_(module.bias)
+
+    def _init_residual_outputs(self) -> None:
+        """Scale residual branch outputs as in GPT-2 for stable deep training."""
+
+        if self.config.use_flex_attention:
+            self.transformer.init_residual_outputs(self.config.n_layers)
+            return
+        residual_std = 0.02 / math.sqrt(2 * self.config.n_layers)
+        for layer in self.transformer.layers:
+            nn.init.normal_(layer.self_attn.out_proj.weight, mean=0.0, std=residual_std)
+            nn.init.normal_(layer.linear2.weight, mean=0.0, std=residual_std)
+
+    def _checkpointed_transformer(
+        self,
+        hidden: Tensor,
+        padding_mask: Tensor | None,
+        attn_mask: Tensor | None = None,
+    ) -> Tensor:
+        for layer in self.transformer.layers:
+
+            def run_layer(layer_input: Tensor, *, current_layer: nn.Module = layer) -> Tensor:
+                return current_layer(
+                    layer_input, src_mask=attn_mask, src_key_padding_mask=padding_mask
+                )
+
+            hidden = checkpoint(run_layer, hidden, use_reentrant=False)
+
+        if self.transformer.norm is not None:
+            hidden = self.transformer.norm(hidden)
+        return hidden
+
+    def _expand_attn_mask(self, attn_mask: Tensor | None, input_ids: Tensor) -> Tensor | None:
+        """Broadcast a per-sample boolean blocking mask across attention heads.
+
+        Accepts ``[L, L]`` shared masks or ``[B, L, L]`` per-sample masks with
+        ``True`` marking blocked key positions, matching the src_mask convention.
+        """
+
+        if attn_mask is None:
+            return None
+        batch_size, sequence_length = input_ids.shape
+        if attn_mask.dtype != torch.bool:
+            raise ValueError("attn_mask must be boolean with True marking blocked positions")
+        if attn_mask.shape == (sequence_length, sequence_length):
+            return attn_mask
+        if attn_mask.shape != (batch_size, sequence_length, sequence_length):
+            raise ValueError("attn_mask must have shape [L, L] or [batch, L, L]")
+        return attn_mask.repeat_interleave(self.config.n_heads, dim=0)
+
+    def encode(
+        self,
+        input_ids: Tensor,
+        attention_mask: Tensor | None = None,
+        attn_mask: Tensor | None = None,
+    ) -> Tensor:
+        """Return contextual token states; ``attn_mask`` restricts attention topology."""
+
+        if input_ids.ndim != 2:
+            raise ValueError("input_ids must have shape [batch, sequence]")
+        batch_size, sequence_length = input_ids.shape
+        if sequence_length > self.config.max_seq_len:
+            raise ValueError(
+                f"sequence length {sequence_length} exceeds max_seq_len "
+                f"{self.config.max_seq_len}"
+            )
+        if attention_mask is not None and attention_mask.shape != input_ids.shape:
+            raise ValueError("attention_mask must match input_ids")
+
+        positions = torch.arange(sequence_length, device=input_ids.device)
+        hidden = self.token_embedding(input_ids) + self.position_embedding(positions)[None, :, :]
+        hidden = self.embedding_dropout(hidden)
+
+        # TransformerEncoder expects True for padding, the inverse of the common
+        # attention-mask convention. src_mask is only supplied by region-aware callers.
+        padding_mask = None if attention_mask is None else ~attention_mask.bool()
+
+        if self.config.use_flex_attention:
+            from diffusion_lm.flexattn import build_block_mask
+
+            if attn_mask is not None and attn_mask.dtype != torch.bool:
+                raise ValueError("attn_mask must be boolean with True marking blocked positions")
+            block_mask = build_block_mask(
+                attn_mask, padding_mask, batch_size, sequence_length, hidden.device
+            )
+            return self.transformer(hidden, block_mask)
+
+        expanded_attn_mask = self._expand_attn_mask(attn_mask, input_ids)
+        if (
+            self.config.activation_checkpointing
+            and self.training
+            and torch.is_grad_enabled()
+        ):
+            return self._checkpointed_transformer(hidden, padding_mask, expanded_attn_mask)
+        return self.transformer(
+            hidden, mask=expanded_attn_mask, src_key_padding_mask=padding_mask
+        )
+
+    def forward(
+        self,
+        input_ids: Tensor,
+        attention_mask: Tensor | None = None,
+        output_positions: Tensor | None = None,
+        attn_mask: Tensor | None = None,
+    ) -> Tensor:
+        """Predict vocabulary logits for all tokens or selected positions only."""
+
+        hidden = self.encode(input_ids, attention_mask=attention_mask, attn_mask=attn_mask)
+        if output_positions is not None:
+            if output_positions.shape != input_ids.shape:
+                raise ValueError("output_positions must match input_ids")
+            hidden = hidden[output_positions.bool()]
+
+        logits = self.lm_head(hidden)
+        # Corruption/control tokens are never valid clean-token predictions. EOS
+        # deliberately remains available so generation can terminate naturally.
+        if self._forbidden_output_token_ids.numel():
+            logits.index_fill_(
+                -1,
+                self._forbidden_output_token_ids,
+                torch.finfo(logits.dtype).min,
+            )
+        return logits
+
+    @property
+    def num_parameters(self) -> int:
+        """Count unique trainable parameters (shared embeddings count once)."""
+
+        return sum(parameter.numel() for parameter in self.parameters() if parameter.requires_grad)
+
+
+def build_denoiser(
+    config: ModelConfig,
+    *,
+    load_pretrained: bool = True,
+    dtype: torch.dtype | None = None,
+) -> nn.Module:
+    """Construct the denoiser a config describes: project transformer or pretrained backbone.
+
+    ``load_pretrained=False`` builds the architecture only, for callers that immediately
+    restore weights from a project checkpoint.
+    """
+
+    if config.backbone == "hf-qwen3":
+        from diffusion_lm.hf_bridge import Qwen3Denoiser
+
+        return Qwen3Denoiser(config, load_pretrained=load_pretrained, dtype=dtype)
+    return DiffusionTransformer(config)
+
+
+def format_parameter_count(count: int) -> str:
+    if count >= 1_000_000:
+        return f"{count / 1_000_000:.2f}M"
+    if count >= 1_000:
+        return f"{count / 1_000:.2f}K"
+    return str(count)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Report the exact model parameter count")
+    parser.add_argument("--config", type=Path, required=True, help="experiment YAML")
+    args = parser.parse_args()
+
+    config = load_config(args.config)
+    # Parameter inspection should not allocate four gigabytes for the 1B preset.
+    with torch.device("meta"):
+        model = DiffusionTransformer(config.model)
+    print(f"parameters: {model.num_parameters:,} ({format_parameter_count(model.num_parameters)})")
+
+
+if __name__ == "__main__":
+    main()

src/diffusion_lm/playground.py

@@ -0,0 +1,541 @@
+"""Local browser playground that streams the token-unmasking process."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import html
+import secrets
+import threading
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+import torch
+from torch import Tensor
+from tokenizers import Tokenizer
+
+from diffusion_lm.diffusion import UnmaskStep, iterative_unmask_steps
+from diffusion_lm.model import DiffusionTransformer, format_parameter_count
+from diffusion_lm.sample import load_model
+from diffusion_lm.tokenizer import load_tokenizer, special_token_id, special_token_ids
+from diffusion_lm.train import resolve_device
+
+
+@dataclass(frozen=True)
+class GenerationSettings:
+    prompt: str = ""
+    generation_length: int = 64
+    steps: int = 64
+    temperature: float = 0.8
+    strategy: Literal["ancestral", "confidence"] = "confidence"
+    seed: int = 1337
+
+    def __post_init__(self) -> None:
+        if self.generation_length <= 0:
+            raise ValueError("La longitud debe ser mayor que cero.")
+        if not 1 <= self.steps <= 512:
+            raise ValueError("Los pasos deben estar entre 1 y 512.")
+        if not 0.0 <= self.temperature <= 5.0:
+            raise ValueError("La temperatura debe estar entre 0 y 5.")
+        if self.strategy not in {"ancestral", "confidence"}:
+            raise ValueError("La estrategia debe ser ancestral o confidence.")
+        if not 0 <= self.seed < 2**63:
+            raise ValueError("La semilla debe estar entre 0 y 2^63-1.")
+
+
+@dataclass(frozen=True)
+class PlaygroundUpdate:
+    state: UnmaskStep
+    prompt_tokens: int
+    partial_text: str
+    final_text: str
+    token_html: str
+    elapsed_seconds: float
+    step_seconds: float
+
+
+def _seed_generation(device: torch.device, seed: int) -> None:
+    torch.manual_seed(seed)
+    if device.type == "cuda":
+        torch.cuda.manual_seed_all(seed)
+    elif device.type == "mps" and hasattr(torch.mps, "manual_seed"):
+        torch.mps.manual_seed(seed)
+
+
+def _synchronize(device: torch.device) -> None:
+    if device.type == "cuda":
+        torch.cuda.synchronize(device)
+    elif device.type == "mps":
+        torch.mps.synchronize()
+
+
+def _token_label(tokenizer: Tokenizer, token_id: int, mask_token_id: int) -> str:
+    raw = tokenizer.id_to_token(token_id) or f"#{token_id}"
+    if token_id == mask_token_id:
+        return "MASK"
+    return (
+        raw.replace("Ġ", "▁")
+        .replace("Ċ", "↵")
+        .replace("ĉ", "⇥")
+        .replace("\n", "↵")
+    ) or "∅"
+
+
+def render_token_grid(
+    tokenizer: Tokenizer,
+    token_ids: list[int],
+    *,
+    mask_token_id: int,
+    prompt_tokens: int,
+    previous_token_ids: list[int] | None,
+) -> str:
+    """Render escaped token chips for a single sample."""
+
+    chips: list[str] = []
+    for position, token_id in enumerate(token_ids):
+        if position < prompt_tokens:
+            state = "prompt"
+        elif token_id == mask_token_id:
+            state = "mask"
+        elif previous_token_ids is not None and previous_token_ids[position] == mask_token_id:
+            state = "new"
+        else:
+            state = "revealed"
+
+        raw = tokenizer.id_to_token(token_id) or f"token {token_id}"
+        label = html.escape(_token_label(tokenizer, token_id, mask_token_id))
+        title = html.escape(f"posición {position} · id {token_id} · {raw}", quote=True)
+        chips.append(
+            f'<span class="token-chip token-{state}" title="{title}">{label}</span>'
+        )
+
+    return (
+        '<section class="token-stage" aria-label="Estado actual de los tokens">'
+        '<div class="token-grid">'
+        + "".join(chips)
+        + "</div>"
+        '<div class="token-legend" aria-label="Leyenda">'
+        '<span><i class="legend-dot legend-prompt"></i>prompt</span>'
+        '<span><i class="legend-dot legend-new"></i>recién revelado</span>'
+        '<span><i class="legend-dot legend-mask"></i>máscara</span>'
+        "</div></section>"
+    )
+
+
+class PlaygroundEngine:
+    """Own one loaded model and serialize interactive generations."""
+
+    def __init__(self, model: DiffusionTransformer, tokenizer_path: str | Path) -> None:
+        self.model = model.eval()
+        self.tokenizer_path = Path(tokenizer_path)
+        self.tokenizer = load_tokenizer(self.tokenizer_path)
+        self._lock = threading.Lock()
+        self._validate_tokenizer()
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.model.parameters()).device
+
+    def _validate_tokenizer(self) -> None:
+        tokenizer_hash = hashlib.sha256(self.tokenizer_path.read_bytes()).hexdigest()
+        model_hash = getattr(self.model, "tokenizer_sha256", None)
+        if model_hash is not None and tokenizer_hash != model_hash:
+            raise ValueError("El tokenizer no coincide con el usado para entrenar el checkpoint.")
+        if self.tokenizer.get_vocab_size(with_added_tokens=True) != self.model.config.vocab_size:
+            raise ValueError("El vocabulario del tokenizer no coincide con el checkpoint.")
+        if special_token_id(self.tokenizer, "mask") != self.model.config.mask_token_id:
+            raise ValueError("El id de [MASK] no coincide con el checkpoint.")
+
+    def info(self) -> dict[str, str | int | bool]:
+        return {
+            "parameters": self.model.num_parameters,
+            "parameters_human": format_parameter_count(self.model.num_parameters),
+            "device": str(self.device),
+            "context": self.model.config.max_seq_len,
+            "vocab_size": self.model.config.vocab_size,
+            "tokenizer": self.tokenizer_path.name,
+            "mps_fp64_fallback": self.device.type == "mps",
+        }
+
+    def _prepare(self, settings: GenerationSettings) -> tuple[Tensor, int, tuple[int, ...]]:
+        prompt_ids = (
+            self.tokenizer.encode(settings.prompt, add_special_tokens=False).ids
+            if settings.prompt
+            else []
+        )
+        role_ids = special_token_ids(self.tokenizer)
+        reserved_ids = set(role_ids.values())
+        encountered = reserved_ids.intersection(prompt_ids)
+        if encountered:
+            raise ValueError(
+                "El prompt contiene tokens especiales reservados. Escribí texto normal sin "
+                "los sentinels internos del modelo."
+            )
+
+        total_length = len(prompt_ids) + settings.generation_length
+        if total_length > self.model.config.max_seq_len:
+            available = self.model.config.max_seq_len - len(prompt_ids)
+            raise ValueError(
+                f"El prompt usa {len(prompt_ids)} tokens y deja {max(0, available)} para generar; "
+                f"solicitaste {settings.generation_length}."
+            )
+
+        input_ids = torch.full(
+            (1, total_length),
+            self.model.config.mask_token_id,
+            dtype=torch.long,
+            device=self.device,
+        )
+        if prompt_ids:
+            input_ids[0, : len(prompt_ids)] = torch.tensor(prompt_ids, device=self.device)
+
+        blocked = tuple(role_ids[role] for role in ("pad", "unk", "bos", "mask"))
+        return input_ids, len(prompt_ids), blocked
+
+    def _decode_final(self, token_ids: list[int], prompt_tokens: int) -> str:
+        eos_id = special_token_id(self.tokenizer, "eos")
+        if eos_id in token_ids[prompt_tokens:]:
+            token_ids = token_ids[: token_ids.index(eos_id, prompt_tokens)]
+        return self.tokenizer.decode(token_ids, skip_special_tokens=True)
+
+    def stream(self, settings: GenerationSettings):
+        """Yield one UI update per reverse-diffusion pass."""
+
+        input_ids, prompt_tokens, blocked = self._prepare(settings)
+        with self._lock:
+            _seed_generation(self.device, settings.seed)
+            started = time.perf_counter()
+            # Measures engine work per pass; consumer time between yields is excluded.
+            pass_started = started
+            previous_ids: list[int] | None = None
+            for state in iterative_unmask_steps(
+                self.model,
+                input_ids,
+                self.model.config.mask_token_id,
+                steps=settings.steps,
+                temperature=settings.temperature,
+                strategy=settings.strategy,
+                blocked_token_ids=blocked,
+            ):
+                token_ids = state.tokens[0].detach().cpu().tolist()
+                partial_text = self.tokenizer.decode(token_ids, skip_special_tokens=False)
+                final_text = (
+                    self._decode_final(token_ids, prompt_tokens)
+                    if state.masked_remaining == 0
+                    else ""
+                )
+                token_html = render_token_grid(
+                    self.tokenizer,
+                    token_ids,
+                    mask_token_id=self.model.config.mask_token_id,
+                    prompt_tokens=prompt_tokens,
+                    previous_token_ids=previous_ids,
+                )
+                if state.masked_remaining == 0:
+                    _synchronize(self.device)
+                now = time.perf_counter()
+                yield PlaygroundUpdate(
+                    state=state,
+                    prompt_tokens=prompt_tokens,
+                    partial_text=partial_text,
+                    final_text=final_text,
+                    token_html=token_html,
+                    elapsed_seconds=now - started,
+                    step_seconds=now - pass_started,
+                )
+                previous_ids = token_ids
+                pass_started = time.perf_counter()
+
+
+PLAYGROUND_CSS = """
+:root {
+  --playground-accent: #7c3aed;
+  --playground-accent-soft: rgba(124, 58, 237, 0.14);
+  --playground-teal: #0f766e;
+  --playground-border: rgba(100, 116, 139, 0.22);
+}
+.gradio-container { max-width: 1440px !important; }
+.playground-header {
+  padding: 22px 24px; border: 1px solid var(--playground-border); border-radius: 18px;
+  background: linear-gradient(135deg, rgba(124,58,237,.10), rgba(15,118,110,.06));
+  box-shadow: 0 1px 2px rgba(30,41,59,.05), 0 14px 34px rgba(71,85,105,.08);
+}
+.playground-header h1 { margin: 0; font-size: clamp(1.6rem, 3vw, 2.4rem); text-wrap: balance; }
+.playground-header p { margin: 8px 0 0; color: var(--body-text-color-subdued); text-wrap: pretty; }
+.model-strip { display: flex; flex-wrap: wrap; gap: 8px; margin-top: 16px; }
+.model-pill {
+  padding: 7px 10px; border-radius: 999px; border: 1px solid var(--playground-border);
+  background: var(--block-background-fill); font-variant-numeric: tabular-nums; font-size: .82rem;
+}
+.control-panel, .output-panel {
+  border: 1px solid var(--playground-border) !important; border-radius: 18px !important;
+  padding: 16px !important; box-shadow: 0 1px 2px rgba(30,41,59,.04), 0 10px 28px rgba(71,85,105,.06);
+}
+.token-stage { min-height: 220px; display: flex; flex-direction: column; justify-content: space-between; }
+.token-grid { display: flex; flex-wrap: wrap; align-content: flex-start; gap: 7px; padding: 8px 2px 18px; }
+.token-chip {
+  display: inline-flex; min-height: 32px; align-items: center; padding: 5px 8px; border-radius: 9px;
+  border: 1px solid transparent; font-family: ui-monospace, SFMono-Regular, Menlo, monospace;
+  font-size: .82rem; font-variant-numeric: tabular-nums; transition: transform 160ms ease-out, opacity 180ms ease-out;
+}
+.token-chip:hover { transform: translateY(-1px); }
+.token-prompt { color: #075985; background: rgba(14,165,233,.12); border-color: rgba(14,165,233,.26); }
+.token-revealed { background: rgba(15,118,110,.10); border-color: rgba(15,118,110,.18); }
+.token-new { color: #5b21b6; background: var(--playground-accent-soft); border-color: rgba(124,58,237,.35); }
+.token-mask { color: var(--body-text-color-subdued); background: rgba(100,116,139,.08); border: 1px dashed rgba(100,116,139,.30); opacity: .7; }
+.token-legend { display: flex; flex-wrap: wrap; gap: 14px; color: var(--body-text-color-subdued); font-size: .78rem; }
+.token-legend span { display: inline-flex; align-items: center; gap: 6px; }
+.legend-dot { width: 9px; height: 9px; border-radius: 50%; display: inline-block; }
+.legend-prompt { background: #0ea5e9; } .legend-new { background: #7c3aed; } .legend-mask { background: #94a3b8; }
+#generate-button, #stop-button { min-height: 44px; transition: transform 150ms ease-out; }
+#generate-button:active, #stop-button:active { transform: scale(.98); }
+@media (prefers-reduced-motion: reduce) { .token-chip, #generate-button, #stop-button { transition: none; } }
+"""
+
+
+def _model_header(engine: PlaygroundEngine) -> str:
+    info = engine.info()
+    warning = (
+        " · MPS usa CPU para Gumbel fp64 cuando temperatura > 0"
+        if info["mps_fp64_fallback"]
+        else ""
+    )
+    return (
+        '<header class="playground-header">'
+        "<h1>Mini Diffusion LM Playground</h1>"
+        "<p>Observá cómo el modelo transforma máscaras en texto usando contexto bidireccional."
+        f"{html.escape(warning)}</p>"
+        '<div class="model-strip">'
+        f'<span class="model-pill">{info["parameters_human"]} parámetros</span>'
+        f'<span class="model-pill">{html.escape(str(info["device"]))}</span>'
+        f'<span class="model-pill">contexto {info["context"]}</span>'
+        f'<span class="model-pill">vocabulario {info["vocab_size"]}</span>'
+        f'<span class="model-pill">{html.escape(str(info["tokenizer"]))}</span>'
+        "</div></header>"
+    )
+
+
+def build_playground(engine: PlaygroundEngine):
+    """Build a Gradio Blocks app without importing Gradio for base-package users."""
+
+    try:
+        import gradio as gr
+    except ImportError as exc:  # pragma: no cover - exercised by CLI environments.
+        raise RuntimeError(
+            'Falta Gradio. Instalalo con: pip install -e ".[playground]"'
+        ) from exc
+
+    max_context = engine.model.config.max_seq_len
+    default_length = min(64, max_context)
+    theme = gr.themes.Soft(primary_hue="violet", secondary_hue="teal", neutral_hue="slate")
+    with gr.Blocks(
+        title="Mini Diffusion LM Playground",
+        analytics_enabled=False,
+        fill_width=True,
+    ) as demo:
+        gr.HTML(_model_header(engine))
+        with gr.Row():
+            with gr.Column(scale=4, elem_classes="control-panel"):
+                gr.Markdown("## Configuración")
+                prompt = gr.Textbox(
+                    label="Prompt (opcional)",
+                    placeholder="Ej.: Once upon a time…",
+                    lines=6,
+                    max_lines=10,
+                )
+                with gr.Row():
+                    length = gr.Slider(
+                        minimum=1,
+                        maximum=max_context,
+                        value=default_length,
+                        step=1,
+                        label="Tokens a generar",
+                    )
+                    steps = gr.Slider(
+                        minimum=1,
+                        maximum=256,
+                        value=min(64, max_context),
+                        step=1,
+                        label="Pasos de difusión",
+                    )
+                with gr.Accordion("Opciones avanzadas", open=False):
+                    strategy = gr.Radio(
+                        choices=[
+                            ("Confianza · revela los tokens más seguros", "confidence"),
+                            ("Ancestral · transición probabilística", "ancestral"),
+                        ],
+                        value="confidence",
+                        label="Estrategia",
+                    )
+                    temperature = gr.Slider(
+                        minimum=0.0,
+                        maximum=2.0,
+                        value=0.8,
+                        step=0.05,
+                        label="Temperatura",
+                    )
+                    seed = gr.Number(
+                        value=0,
+                        precision=0,
+                        minimum=0,
+                        maximum=2**31 - 1,
+                        label='Semilla (0 = aleatoria en cada generación)',
+                    )
+                with gr.Row():
+                    generate_button = gr.Button(
+                        "Generar",
+                        variant="primary",
+                        elem_id="generate-button",
+                    )
+                    stop_button = gr.Button(
+                        "Detener",
+                        variant="stop",
+                        elem_id="stop-button",
+                    )
+
+            with gr.Column(scale=7, elem_classes="output-panel"):
+                status = gr.Markdown(
+                    "### Listo\nConfigurá una muestra y presioná **Generar**."
+                )
+                token_view = gr.HTML(
+                    '<div class="token-stage"><p>Los tokens aparecerán acá.</p></div>'
+                )
+                output = gr.Textbox(
+                    label="Texto actual",
+                    lines=7,
+                    interactive=False,
+                )
+                metrics = gr.Markdown("`Esperando una generación`", elem_classes="metrics")
+
+        def stream_generation(
+            prompt_value: str,
+            length_value: float,
+            steps_value: float,
+            strategy_value: str,
+            temperature_value: float,
+            seed_value: float,
+        ):
+            clicked = time.perf_counter()
+            try:
+                resolved_seed = int(seed_value) or secrets.randbelow(2**31 - 1) + 1
+                settings = GenerationSettings(
+                    prompt=prompt_value or "",
+                    generation_length=int(length_value),
+                    steps=int(steps_value),
+                    strategy=strategy_value,  # type: ignore[arg-type]
+                    temperature=float(temperature_value),
+                    seed=resolved_seed,
+                )
+                for update in engine.stream(settings):
+                    generated = settings.generation_length - update.state.masked_remaining
+                    percent = 100.0 * generated / settings.generation_length
+                    total_seconds = time.perf_counter() - clicked
+                    average_step = update.elapsed_seconds / max(1, update.state.step)
+                    tokens_per_second = (
+                        generated / update.elapsed_seconds if update.elapsed_seconds > 0 else 0.0
+                    )
+                    if update.state.masked_remaining == 0:
+                        status_text = (
+                            f'### Completado en {total_seconds:.2f} s\n'
+                            f'{generated} tokens en {update.state.step} pasos · '
+                            f'{tokens_per_second:.1f} tok/s · '
+                            f'{average_step * 1000:.0f} ms/paso promedio'
+                        )
+                    else:
+                        status_text = (
+                            f"### Paso {update.state.step}/{update.state.total_steps}\n"
+                            f"{update.state.masked_remaining} máscaras restantes · "
+                            f"{percent:.0f}% revelado"
+                        )
+                    visible_text = (
+                        update.final_text
+                        if update.state.masked_remaining == 0
+                        else update.partial_text
+                    )
+                    metrics_text = (
+                        f'`{total_seconds:.2f} s desde el clic` · '
+                        f'`modelo {update.elapsed_seconds:.2f} s` · '
+                        f'`paso {update.step_seconds * 1000:.0f} ms` · '
+                        f'`prom. {average_step * 1000:.0f} ms/paso` · '
+                        f'`{tokens_per_second:.1f} tok/s` · '
+                        f'`{update.prompt_tokens} tokens de prompt` · '
+                        f'`seed {settings.seed}`'
+                    )
+                    yield update.token_html, status_text, visible_text, metrics_text
+            except ValueError as exc:
+                raise gr.Error(str(exc)) from exc
+
+        generation_event = generate_button.click(
+            fn=stream_generation,
+            inputs=[prompt, length, steps, strategy, temperature, seed],
+            outputs=[token_view, status, output, metrics],
+            show_progress="minimal",
+            scroll_to_output=False,
+            concurrency_limit=1,
+            concurrency_id="diffusion-model",
+            trigger_mode="once",
+            stream_every=0.1,
+            api_visibility="private",
+        )
+        stop_button.click(
+            fn=lambda: "### Generación detenida",
+            outputs=status,
+            cancels=[generation_event],
+            queue=False,
+            api_visibility="private",
+        )
+
+    demo = demo.queue(max_size=8, default_concurrency_limit=1)
+    # Gradio 6 moved presentation arguments from Blocks() to launch().
+    demo._mini_diffusion_theme = theme
+    return demo
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--checkpoint", type=Path, required=True, help="checkpoint local confiable")
+    parser.add_argument("--tokenizer", type=Path, required=True, help="tokenizer usado al entrenar")
+    parser.add_argument("--device", default="auto", help="auto, cpu, mps, cuda…")
+    parser.add_argument("--host", default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=7860)
+    parser.add_argument("--no-browser", action="store_true", help="no abrir el navegador")
+    return parser
+
+
+def main() -> None:
+    args = _build_parser().parse_args()
+    if not args.checkpoint.is_file():
+        raise SystemExit(f"Checkpoint inexistente: {args.checkpoint}")
+    if not args.tokenizer.is_file():
+        raise SystemExit(f"Tokenizer inexistente: {args.tokenizer}")
+    if not 1 <= args.port <= 65535:
+        raise SystemExit("El puerto debe estar entre 1 y 65535")
+
+    device = resolve_device(args.device)
+    print(f"Cargando {args.checkpoint} en {device}…")
+    model = load_model(args.checkpoint, device)
+    engine = PlaygroundEngine(model, args.tokenizer)
+    demo = build_playground(engine)
+    print(f"Playground: http://{args.host}:{args.port}")
+    print("Usá únicamente checkpoints locales confiables.")
+    demo.launch(
+        server_name=args.host,
+        server_port=args.port,
+        inbrowser=not args.no_browser,
+        share=False,
+        show_error=True,
+        strict_cors=True,
+        max_threads=4,
+        footer_links=[],
+        enable_monitoring=False,
+        ssr_mode=False,
+        pwa=False,
+        theme=demo._mini_diffusion_theme,
+        css=PLAYGROUND_CSS,
+    )
+
+
+if __name__ == "__main__":
+    main()

src/diffusion_lm/quality.py

@@ -0,0 +1,69 @@
+"""Programmatic generation quality: prose damage, and a reward over claim recaps.
+
+The weights below encode a priority order rather than a tuned optimum: being right first,
+not lying second, readable prose third, and brevity last and only relative to how much the
+problem actually required. They are meant to be calibrated against judged samples before
+driving any policy optimisation, since a reward this cheap is also cheap to game.
+"""
+
+from __future__ import annotations
+
+import re
+
+WORD = re.compile(r"[A-Za-z']+")
+
+# Ordering, not tuning. Lying is charged PER ITEM rather than as a fraction of the record:
+# a reader can work around clumsy prose or a missing line, but a fabricated plate or a
+# superseded amount silently corrupts the report, so one of them has to outweigh several
+# recalled fields. Verbosity is charged per fact carried, so a harder case may run longer.
+WEIGHTS = {
+    'recall': 1.0,
+    'stale': -0.35,
+    'invented': -0.35,
+    'duplication': -0.2,
+    'verbosity': -0.15,
+}
+# Thinking tokens per recalled fact treated as free before the verbosity term bites.
+VERBOSITY_BUDGET = 40.0
+# Lie counts saturate here so a wholly hallucinated report does not dominate the gradient.
+MAX_LIES = 3.0
+
+
+def prose_damage(text: str) -> dict[str, float]:
+    """Adjacent duplicate rate and lexical diversity, the parallel-sampling signatures."""
+
+    words = [word.lower() for word in WORD.findall(text)]
+    pairs = max(1, len(words) - 1)
+    duplicates = sum(a == b for a, b in zip(words, words[1:]))
+    return {
+        'duplicate_rate': duplicates / pairs,
+        'lexical_diversity': len(set(words)) / max(1, len(words)),
+        'words': len(words),
+    }
+
+
+def reward(report_score: dict[str, object], think_tokens: int, answer_text: str) -> dict:
+    """Combine fact fidelity, prose damage and verbosity into one scalar plus its parts.
+
+    ``report_score`` comes from :func:`diffusion_lm.claims.score_report`. Penalties are
+    fractions of the fact count so a long claim and a short one stay comparable.
+    """
+
+    recalled = max(1, int(report_score['recalled']))
+    damage = prose_damage(answer_text)
+    per_fact = think_tokens / recalled
+    parts = {
+        'recall': float(report_score['recall']),
+        'stale': min(MAX_LIES, len(report_score['stale_kept'])),
+        'invented': min(MAX_LIES, len(report_score['invented'])),
+        'duplication': min(1.0, damage['duplicate_rate'] * 20.0),
+        'verbosity': max(0.0, per_fact - VERBOSITY_BUDGET) / VERBOSITY_BUDGET,
+    }
+    total = sum(WEIGHTS[name] * value for name, value in parts.items())
+    return {
+        'total': total,
+        'parts': parts,
+        'think_tokens_per_fact': per_fact,
+        'duplicate_rate': damage['duplicate_rate'],
+        'lexical_diversity': damage['lexical_diversity'],
+    }

src/diffusion_lm/reasoning.py

@@ -0,0 +1,578 @@
+"""OpenMathInstruct-2 preparation and packed datasets for reasoning experiments.
+
+Each example is rendered into two layouts over one shared tokenizer:
+
+- ``flat``:    problem <think> steps... </think> answer <eos>
+- ``slotted``: problem <think> [fixed-size thought slots, one step each, <tpad>-filled]
+               [terminal slot: </think> <tpad>...] answer <eos>
+
+The slotted layout gives every thought a fixed-geometry block so a block-diffusion
+model can denoise one thought at a time; the flat layout serves the autoregressive
+and pure-diffusion baselines on identical example sets.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Iterator
+
+import numpy as np
+import torch
+from torch import Tensor
+from torch.utils.data import Dataset
+
+from diffusion_lm.tokenizer import (
+    RoleAwareTokenizer,
+    load_tokenizer,
+    special_token_ids,
+    train_tokenizer_from_iterator,
+)
+
+REASONING_PACKED_FORMAT = 'mini-diffusion-lm-reasoning-packed-v1'
+
+_REASONING_NAMESPACE = 'mdlm-r-8f4e2a6c'
+THINK_TOKEN = f'<|{_REASONING_NAMESPACE}:think|>'
+END_THINK_TOKEN = f'<|{_REASONING_NAMESPACE}:endthink|>'
+THOUGHT_PAD_TOKEN = f'<|{_REASONING_NAMESPACE}:tpad|>'
+REASONING_SPECIAL_TOKENS = (THINK_TOKEN, END_THINK_TOKEN, THOUGHT_PAD_TOKEN)
+
+# Block sizes a thought may occupy in the adaptive layout. The model emits one
+# ``<szN>`` control token per thought to choose the block length; the full menu is
+# baked into the tokenizer so the vocabulary stays stable across experiments even
+# when a run activates only a subset of the sizes.
+BLOCK_SIZES = (32, 64, 128, 256, 512)
+SIZE_TOKENS = tuple(f'<|{_REASONING_NAMESPACE}:sz{size}|>' for size in BLOCK_SIZES)
+ADAPTIVE_SPECIAL_TOKENS = REASONING_SPECIAL_TOKENS + SIZE_TOKENS
+
+# Native spellings honored when the tokenizer comes from a pretrained backbone whose vocabulary
+# already carries trained think delimiters (e.g. Qwen3). The namespaced sentinel wins when both
+# spellings exist so project tokenizers keep their exact historical ids.
+_NATIVE_REASONING_TOKENS = {'think': '<think>', 'end_think': '</think>'}
+
+
+def size_token_ids(
+    tokenizer: RoleAwareTokenizer, sizes: tuple[int, ...] | None = None
+) -> dict[int, int]:
+    """Map block sizes to ``<szN>`` control-token ids.
+
+    Without ``sizes`` the full menu is scanned and absent tokens are skipped, so tokenizers
+    trained before a menu extension keep resolving. Explicitly requested sizes must exist.
+    """
+
+    ids: dict[int, int] = {}
+    for size, token in zip(BLOCK_SIZES, SIZE_TOKENS):
+        if sizes is not None and size not in sizes:
+            continue
+        token_id = tokenizer.token_to_id(token)
+        if token_id is None:
+            if sizes is None:
+                continue
+            raise ValueError(f'tokenizer is missing the size token {token!r}')
+        ids[size] = token_id
+    if sizes is not None and (missing := set(sizes) - set(ids)):
+        raise ValueError(f'sizes {sorted(missing)} are not in the block-size menu {BLOCK_SIZES}')
+    if not ids:
+        raise ValueError('tokenizer carries no <szN> size tokens')
+    return ids
+
+_BOILERPLATE_STEP = re.compile(
+    r"^let'?s\s+(solve|answer|tackle)\s+(the\s+)?(new\s+)?(question|problem)\b.*$",
+    re.IGNORECASE,
+)
+_MATH_SPAN = re.compile(r'\$[^$]*\$')
+_SENTENCE_BREAK = re.compile(r'(?<=[.!?])\s+(?=[A-Z\d(])')
+_BOXED = re.compile(r'\\boxed\{([^{}]*)\}')
+
+
+def reasoning_token_ids(tokenizer: RoleAwareTokenizer) -> dict[str, int]:
+    """Resolve the reasoning control-token ids, failing on non-reasoning tokenizers.
+
+    Pretrained-backbone tokenizers resolve ``think``/``end_think`` through their native
+    trained spellings when the namespaced sentinels are absent.
+    """
+
+    ids: dict[str, int] = {}
+    for name, token in (
+        ('think', THINK_TOKEN),
+        ('end_think', END_THINK_TOKEN),
+        ('thought_pad', THOUGHT_PAD_TOKEN),
+    ):
+        token_id = tokenizer.token_to_id(token)
+        if token_id is None and name in _NATIVE_REASONING_TOKENS:
+            token_id = tokenizer.token_to_id(_NATIVE_REASONING_TOKENS[name])
+        if token_id is None:
+            raise ValueError(f'tokenizer is missing the reasoning token {token!r}')
+        ids[name] = token_id
+    return ids
+
+
+def extract_boxed_answer(text: str) -> str | None:
+    matches = _BOXED.findall(text)
+    if not matches:
+        return None
+    return matches[-1].strip()
+
+
+def _normalize_answer(answer: str) -> str:
+    cleaned = answer.strip().replace(',', '').replace('$', '').rstrip('.')
+    try:
+        value = float(cleaned)
+    except ValueError:
+        return cleaned
+    return str(int(value)) if value == int(value) else str(value)
+
+
+def split_solution_steps(solution: str) -> list[str]:
+    """Split a natural-language CoT solution into sentence-level steps."""
+
+    protected: list[str] = []
+
+    def _protect(match: re.Match[str]) -> str:
+        protected.append(match.group(0))
+        return f'\x00{len(protected) - 1}\x00'
+
+    masked = _MATH_SPAN.sub(_protect, solution)
+    parts: list[str] = []
+    for line in masked.splitlines():
+        line = line.strip()
+        if not line:
+            continue
+        parts.extend(piece.strip() for piece in _SENTENCE_BREAK.split(line) if piece.strip())
+
+    def _restore(text: str) -> str:
+        return re.sub(r'\x00(\d+)\x00', lambda m: protected[int(m.group(1))], text)
+
+    steps = [_restore(part) for part in parts]
+    return [step for step in steps if not _BOILERPLATE_STEP.match(step)]
+
+
+@dataclass(frozen=True)
+class ReasoningExample:
+    problem: str
+    steps: tuple[str, ...]
+    answer: str
+    expected_answer: str
+
+
+def parse_example(row: dict[str, Any]) -> ReasoningExample | None:
+    """Clean one dataset row into problem, thought steps, and an answer sentence."""
+
+    problem = ' '.join(str(row['problem']).split())
+    expected = _normalize_answer(str(row['expected_answer']))
+    steps = split_solution_steps(str(row['generated_solution']))
+    answer_index = None
+    for index in range(len(steps) - 1, -1, -1):
+        if _BOXED.search(steps[index]):
+            answer_index = index
+            break
+    if answer_index is None or not problem or not expected:
+        return None
+    answer = steps[answer_index]
+    thoughts = tuple(steps[:answer_index])
+    if not thoughts:
+        return None
+    boxed = extract_boxed_answer(answer)
+    if boxed is None or _normalize_answer(boxed) != expected:
+        return None
+    return ReasoningExample(problem, thoughts, answer, expected)
+
+
+@dataclass(frozen=True)
+class LayoutSpec:
+    seq_len: int
+    block: int
+    max_slots: int
+    sizes: tuple[int, ...] = ()
+
+    def __post_init__(self) -> None:
+        if self.seq_len <= 0 or self.block <= 1 or self.max_slots <= 1:
+            raise ValueError('seq_len, block, and max_slots must be meaningfully positive')
+        sizes = tuple(self.sizes)
+        if sizes:
+            if any(size not in BLOCK_SIZES for size in sizes):
+                raise ValueError(f'adaptive sizes must be drawn from {BLOCK_SIZES}')
+            if list(sizes) != sorted(set(sizes)):
+                raise ValueError('adaptive sizes must be unique and ascending')
+        object.__setattr__(self, 'sizes', sizes)
+
+
+@dataclass(frozen=True)
+class EncodedExample:
+    flat: np.ndarray
+    flat_regions: np.ndarray
+    slotted: np.ndarray
+    slotted_regions: np.ndarray
+
+
+@dataclass(frozen=True)
+class EncodedAdaptive:
+    tokens: np.ndarray
+    regions: np.ndarray
+    block_sizes: tuple[int, ...]
+
+
+class ExampleEncoder:
+    """Render cleaned examples into the flat and slotted token layouts."""
+
+    def __init__(self, tokenizer: RoleAwareTokenizer, spec: LayoutSpec) -> None:
+        self.tokenizer = tokenizer
+        self.spec = spec
+        roles = special_token_ids(tokenizer)
+        reasoning = reasoning_token_ids(tokenizer)
+        self.pad_id = roles['pad']
+        self.eos_id = roles['eos']
+        self.think_id = reasoning['think']
+        self.end_think_id = reasoning['end_think']
+        self.tpad_id = reasoning['thought_pad']
+        self.size_ids = size_token_ids(tokenizer, spec.sizes) if spec.sizes else {}
+        vocab_size = tokenizer.get_vocab_size(with_added_tokens=True)
+        self.token_dtype = np.dtype(
+            'uint16' if vocab_size <= np.iinfo(np.uint16).max else 'uint32'
+        )
+
+    def _encode(self, text: str) -> list[int]:
+        return self.tokenizer.encode(text, add_special_tokens=False).ids
+
+    def encode_adaptive(self, example: ReasoningExample) -> EncodedAdaptive | None:
+        """Render an example into the adaptive layout with per-thought size tokens.
+
+        Each thought is placed in the smallest active block that fits it; thoughts
+        longer than the largest block continue across consecutive maximum blocks.
+        The ``<szN>`` control token precedes every block, giving the model an
+        autoregressive target for the block-length decision, and the remainder of
+        each block is filled with ``<tpad>``.
+        """
+
+        spec = self.spec
+        sizes = spec.sizes
+        if not sizes:
+            raise ValueError('encode_adaptive requires a LayoutSpec with sizes set')
+        max_size = sizes[-1]
+        problem_ids = self._encode(example.problem)
+        answer_ids = self._encode(example.answer)
+        if not answer_ids or not problem_ids:
+            return None
+
+        blocks: list[tuple[int, list[int]]] = []
+        for step in example.steps:
+            ids = self._encode(step)
+            for start in range(0, len(ids), max_size):
+                chunk = ids[start : start + max_size]
+                size = next(candidate for candidate in sizes if candidate >= len(chunk))
+                blocks.append((size, chunk))
+        # An empty chain is a legitimate example: it teaches the controller to answer without
+        # thinking, which a conversation needs on every trivial turn.
+        if len(blocks) > spec.max_slots:
+            return None
+
+        tokens = [*problem_ids, self.think_id]
+        for size, chunk in blocks:
+            tokens.append(self.size_ids[size])
+            tokens.extend(chunk)
+            tokens.extend([self.tpad_id] * (size - len(chunk)))
+        tokens.append(self.end_think_id)
+        answer_start = len(tokens)
+        tokens.extend(answer_ids)
+        tokens.append(self.eos_id)
+        answer_end = len(tokens)
+        if answer_end > spec.seq_len:
+            return None
+        tokens.extend([self.pad_id] * (spec.seq_len - len(tokens)))
+
+        regions = np.asarray(
+            [len(problem_ids), len(blocks), answer_start, answer_end], dtype=np.int32
+        )
+        return EncodedAdaptive(
+            tokens=np.asarray(tokens, dtype=self.token_dtype),
+            regions=regions,
+            block_sizes=tuple(size for size, _ in blocks),
+        )
+
+    def encode_example(self, example: ReasoningExample) -> EncodedExample | None:
+        spec = self.spec
+        problem_ids = self._encode(example.problem)
+        step_ids = [self._encode(step) for step in example.steps]
+        answer_ids = self._encode(example.answer)
+
+        if not answer_ids or not problem_ids:
+            return None
+
+        # A thought that exceeds one slot continues into the next, so long steps
+        # cost extra slots instead of dropping the example.
+        chunked: list[list[int]] = []
+        for ids in step_ids:
+            for start in range(0, len(ids), spec.block):
+                chunked.append(ids[start : start + spec.block])
+        step_ids = chunked
+        if len(step_ids) > spec.max_slots - 1:
+            return None
+
+        slots = len(step_ids) + 1
+        slotted_len = len(problem_ids) + 1 + slots * spec.block + len(answer_ids) + 1
+        if slotted_len > spec.seq_len:
+            return None
+
+        slotted = [*problem_ids, self.think_id]
+        for ids in step_ids:
+            slotted.extend(ids)
+            slotted.extend([self.tpad_id] * (spec.block - len(ids)))
+        slotted.append(self.end_think_id)
+        slotted.extend([self.tpad_id] * (spec.block - 1))
+        answer_start_slotted = len(slotted)
+        slotted.extend(answer_ids)
+        slotted.append(self.eos_id)
+        answer_end_slotted = len(slotted)
+        slotted.extend([self.pad_id] * (spec.seq_len - len(slotted)))
+
+        flat_think_ids = self._encode(' '.join(example.steps))
+        flat = [*problem_ids, self.think_id, *flat_think_ids, self.end_think_id]
+        answer_start_flat = len(flat)
+        flat.extend(answer_ids)
+        flat.append(self.eos_id)
+        answer_end_flat = len(flat)
+        if answer_end_flat > spec.seq_len:
+            return None
+        flat.extend([self.pad_id] * (spec.seq_len - len(flat)))
+
+        def _regions(problem_len: int, n_slots: int, start: int, end: int) -> np.ndarray:
+            return np.asarray([problem_len, n_slots, start, end], dtype=np.int32)
+
+        return EncodedExample(
+            flat=np.asarray(flat, dtype=self.token_dtype),
+            flat_regions=_regions(len(problem_ids), 0, answer_start_flat, answer_end_flat),
+            slotted=np.asarray(slotted, dtype=self.token_dtype),
+            slotted_regions=_regions(
+                len(problem_ids), slots, answer_start_slotted, answer_end_slotted
+            ),
+        )
+
+
+def regions_path(token_path: str | Path) -> Path:
+    path = Path(token_path)
+    return path.with_suffix('.regions.npy')
+
+
+class ReasoningTokenDataset(Dataset[tuple[Tensor, Tensor]]):
+    """Fixed-length reasoning sequences with per-example region annotations.
+
+    Regions hold ``[problem_len, n_slots, answer_start, answer_end]``. ``n_slots``
+    is zero for the flat layout and counts thought slots including the terminal
+    ``</think>`` slot for the slotted layout.
+    """
+
+    def __init__(self, path: str | Path) -> None:
+        self.path = Path(path)
+        metadata_path = self.path.with_suffix(self.path.suffix + '.json')
+        if not metadata_path.is_file():
+            raise FileNotFoundError(f'reasoning metadata not found: {metadata_path}')
+        with metadata_path.open('r', encoding='utf-8') as handle:
+            self.metadata = json.load(handle)
+        if self.metadata.get('format') != REASONING_PACKED_FORMAT:
+            raise ValueError(f'unsupported reasoning data format in {metadata_path}')
+        self.seq_len = int(self.metadata['seq_len'])
+        count = int(self.metadata['example_count'])
+        # Datasets packed before the dtype key existed are always uint16.
+        self.token_dtype = np.dtype(self.metadata.get('dtype', 'uint16'))
+        self._tokens = np.memmap(
+            self.path, mode='r', dtype=self.token_dtype, shape=(count, self.seq_len)
+        )
+        self._regions = np.load(regions_path(self.path))
+        if self._regions.shape != (count, 4):
+            raise ValueError(f'regions shape {self._regions.shape} does not match example count')
+
+    def __len__(self) -> int:
+        return int(self._tokens.shape[0])
+
+    def __getitem__(self, index: int) -> tuple[Tensor, Tensor]:
+        tokens = torch.from_numpy(np.asarray(self._tokens[index], dtype=np.int64))
+        regions = torch.from_numpy(np.asarray(self._regions[index], dtype=np.int64))
+        return tokens, regions
+
+    def __getstate__(self) -> dict[str, Any]:
+        state = self.__dict__.copy()
+        state['_tokens'] = None
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        count = int(self.metadata['example_count'])
+        self._tokens = np.memmap(
+            self.path, mode='r', dtype=self.token_dtype, shape=(count, self.seq_len)
+        )
+
+
+def _write_packed(
+    path: Path,
+    tokens: np.ndarray,
+    regions: np.ndarray,
+    *,
+    layout: str,
+    spec: LayoutSpec,
+    tokenizer_path: Path,
+    tokenizer: RoleAwareTokenizer,
+    extra_metadata: dict[str, Any] | None = None,
+) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tokens.tofile(path)
+    np.save(regions_path(path), regions)
+    if tokens.dtype not in (np.dtype('uint16'), np.dtype('uint32')):
+        raise ValueError(f'packed reasoning tokens must be uint16 or uint32, got {tokens.dtype}')
+    metadata = {
+        'format': REASONING_PACKED_FORMAT,
+        'layout': layout,
+        'seq_len': spec.seq_len,
+        'block': spec.block,
+        'max_slots': spec.max_slots,
+        'dtype': tokens.dtype.name,
+        'example_count': int(tokens.shape[0]),
+        'vocab_size': tokenizer.get_vocab_size(with_added_tokens=True),
+        'special_token_ids': special_token_ids(tokenizer),
+        'reasoning_token_ids': reasoning_token_ids(tokenizer),
+        'tokenizer_sha256': hashlib.sha256(tokenizer_path.read_bytes()).hexdigest(),
+    }
+    if extra_metadata:
+        metadata.update(extra_metadata)
+    metadata_path = path.with_suffix(path.suffix + '.json')
+    with metadata_path.open('w', encoding='utf-8') as handle:
+        json.dump(metadata, handle, indent=2)
+        handle.write('\n')
+
+
+def _iter_rows(parquet_paths: list[Path], sources: set[str]) -> Iterator[dict[str, Any]]:
+    import pyarrow.parquet as pq
+
+    for path in parquet_paths:
+        table = pq.read_table(
+            path, columns=['problem', 'generated_solution', 'expected_answer', 'problem_source']
+        )
+        for row in table.to_pylist():
+            if row['problem_source'] in sources:
+                yield row
+
+
+def _is_validation(problem: str, val_fraction: float) -> bool:
+    digest = hashlib.sha256(problem.encode('utf-8')).digest()
+    return int.from_bytes(digest[:4], 'big') / 2**32 < val_fraction
+
+
+def prepare(args: argparse.Namespace) -> None:
+    sources = set(args.sources.split(','))
+    parquet_paths = [Path(path) for path in args.parquet]
+    spec = LayoutSpec(seq_len=args.seq_len, block=args.block, max_slots=args.max_slots)
+    output_dir = Path(args.output_dir)
+    tokenizer_path = Path(args.tokenizer)
+
+    examples: list[ReasoningExample] = []
+    dropped_parse = 0
+    seen: set[str] = set()
+    for row in _iter_rows(parquet_paths, sources):
+        example = parse_example(row)
+        if example is None:
+            dropped_parse += 1
+            continue
+        key = hashlib.sha256(
+            (example.problem + '\x1f' + ' '.join(example.steps)).encode('utf-8')
+        ).hexdigest()
+        if key in seen:
+            continue
+        seen.add(key)
+        examples.append(example)
+    print(f'parsed {len(examples):,} unique examples ({dropped_parse:,} dropped at parse)')
+
+    if tokenizer_path.is_file():
+        tokenizer = load_tokenizer(tokenizer_path)
+        reasoning_token_ids(tokenizer)
+        print(f'reusing tokenizer {tokenizer_path}')
+    else:
+        def _texts() -> Iterator[str]:
+            for example in examples:
+                yield f'{example.problem}\n{" ".join(example.steps)}\n{example.answer}'
+
+        tokenizer = train_tokenizer_from_iterator(
+            _texts(),
+            tokenizer_path,
+            vocab_size=args.vocab_size,
+            min_frequency=4,
+            length=len(examples),
+            extra_special_tokens=REASONING_SPECIAL_TOKENS,
+        )
+        print(f'trained tokenizer {tokenizer_path}')
+
+    encoder = ExampleEncoder(tokenizer, spec)
+    split: dict[str, dict[str, list[np.ndarray]]] = {
+        'train': {'flat': [], 'flat_regions': [], 'slotted': [], 'slotted_regions': []},
+        'validation': {'flat': [], 'flat_regions': [], 'slotted': [], 'slotted_regions': []},
+    }
+    val_problems: list[dict[str, str]] = []
+    dropped_encode = 0
+    for example in examples:
+        encoded = encoder.encode_example(example)
+        if encoded is None:
+            dropped_encode += 1
+            continue
+        is_val = _is_validation(example.problem, args.val_fraction)
+        bucket = split['validation' if is_val else 'train']
+        bucket['flat'].append(encoded.flat)
+        bucket['flat_regions'].append(encoded.flat_regions)
+        bucket['slotted'].append(encoded.slotted)
+        bucket['slotted_regions'].append(encoded.slotted_regions)
+        if is_val:
+            val_problems.append(
+                {'problem': example.problem, 'expected_answer': example.expected_answer}
+            )
+
+    total = sum(len(bucket['flat']) for bucket in split.values())
+    print(f'encoded {total:,} examples ({dropped_encode:,} dropped at encode)')
+    for split_name, bucket in split.items():
+        if not bucket['flat']:
+            raise ValueError(f'no examples in the {split_name} split; adjust filters')
+        for layout in ('flat', 'slotted'):
+            _write_packed(
+                output_dir / f'{split_name}-{layout}.bin',
+                np.stack(bucket[layout]),
+                np.stack(bucket[f'{layout}_regions']),
+                layout=layout,
+                spec=spec,
+                tokenizer_path=tokenizer_path,
+                tokenizer=tokenizer,
+            )
+        print(f'{split_name}: {len(bucket["flat"]):,} examples -> {output_dir}')
+
+    problems_path = output_dir / 'validation-problems.jsonl'
+    with problems_path.open('w', encoding='utf-8') as handle:
+        for record in val_problems:
+            handle.write(json.dumps(record, ensure_ascii=False) + '\n')
+    print(f'wrote {len(val_problems):,} validation problems to {problems_path}')
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=__doc__)
+    subparsers = parser.add_subparsers(dest='command', required=True)
+    prepare_parser = subparsers.add_parser(
+        'prepare', help='filter, tokenize, and pack OpenMathInstruct-2 parquet shards'
+    )
+    prepare_parser.add_argument('--parquet', type=Path, nargs='+', required=True)
+    prepare_parser.add_argument('--output-dir', type=Path, required=True)
+    prepare_parser.add_argument('--tokenizer', type=Path, required=True)
+    prepare_parser.add_argument('--vocab-size', type=int, default=8192)
+    prepare_parser.add_argument('--seq-len', type=int, default=512)
+    prepare_parser.add_argument('--block', type=int, default=32)
+    prepare_parser.add_argument('--max-slots', type=int, default=11)
+    prepare_parser.add_argument('--val-fraction', type=float, default=0.02)
+    prepare_parser.add_argument('--sources', default='gsm8k,augmented_gsm8k')
+    return parser
+
+
+def main() -> None:
+    args = _build_parser().parse_args()
+    if args.command == 'prepare':
+        prepare(args)
+
+
+if __name__ == '__main__':
+    main()

src/diffusion_lm/reasoning_playground.py

@@ -0,0 +1,1226 @@
+"""Interactive playground for the reasoning variants.
+
+Loads any ``mini-diffusion-lm-reasoning-inference-v1`` checkpoint and solves
+user-provided problems with the sampler matching its training objective,
+streaming thought slots as they are denoised and answer tokens as they are
+decoded, together with per-phase wall-clock metrics.
+"""
+
+from __future__ import annotations
+
+import argparse
+import html
+import json
+import time
+from pathlib import Path
+
+import torch
+
+from diffusion_lm.claims import (
+    IM_END,
+    SYSTEM,
+    chat_prefix,
+    ledger_line,
+    ledger_notes,
+    merge_notes,
+    window_start,
+)
+from diffusion_lm.config import ModelConfig
+from diffusion_lm.diffusion import iterative_unmask_steps, _sample_categorical
+from diffusion_lm.flexattn import _FLEX_EAGER
+from diffusion_lm.hybrid import (
+    _apply_repetition_penalty,
+    _apply_top_p,
+    _ar_predict_control,
+    _kv_cache_enabled,
+    adaptive_block_mask,
+    block_mask_from_boundaries,
+    block_size_curriculum,
+    prefix_causal_blocked,
+    slot_causal_blocked,
+)
+from diffusion_lm.model import build_denoiser
+from diffusion_lm.reasoning import extract_boxed_answer, reasoning_token_ids, size_token_ids
+from diffusion_lm.tokenizer import load_tokenizer, special_token_ids
+from diffusion_lm.train import resolve_device
+
+
+def _discover_checkpoints(root: Path) -> dict[str, Path]:
+    found: dict[str, Path] = {}
+    for path in sorted(root.glob('*/inference-latest.pt')):
+        found[path.parent.name] = path
+    return found
+
+
+class ReasoningEngine:
+    """Load one reasoning checkpoint and stream problem solutions."""
+
+    def __init__(self, checkpoint_path: Path, tokenizer_path: Path, device: torch.device) -> None:
+        payload = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
+        if payload.get('format') != 'mini-diffusion-lm-reasoning-inference-v1':
+            raise ValueError(f'unsupported checkpoint format in {checkpoint_path}')
+        self.config = payload['config']
+        self.objective = self.config['reasoning']['objective']
+        model_config = ModelConfig(**self.config['model'])
+        self.model = build_denoiser(model_config, load_pretrained=False)
+        self.model.load_state_dict(payload['model'])
+        dtype = (
+            torch.bfloat16
+            if model_config.backbone != 'project' and device.type == 'cuda'
+            else torch.float32
+        )
+        self.model.to(device=device, dtype=dtype).eval()
+        self.device = device
+        self.step = payload.get('step')
+        self.tokenizer = load_tokenizer(tokenizer_path)
+        self.roles = special_token_ids(self.tokenizer)
+        self.reasoning_ids = reasoning_token_ids(self.tokenizer)
+        self.adaptive = bool(self.config['reasoning'].get('adaptive'))
+        self.causal_prefix = bool(self.config['reasoning'].get('causal_prefix', False))
+        trained_sizes = tuple(self.config['reasoning'].get('sizes') or ()) or None
+        self.size_ids = size_token_ids(self.tokenizer, trained_sizes) if self.adaptive else {}
+        self.im_end_id = self.tokenizer.token_to_id(IM_END)
+        self.chat_ready = (
+            self.objective == 'hybrid' and self.adaptive and self.im_end_id is not None
+        )
+        self.block = 32
+        self.max_slots = 40
+        self.max_blocks = 48
+        self.log_lines: list[str] = []
+
+    def _decode(self, ids: list[int]) -> str:
+        return self.tokenizer.decode(ids, skip_special_tokens=True)
+
+    def _slot_texts(self, think_ids: list[int]) -> list[str]:
+        end_id = self.reasoning_ids['end_think']
+        pad_id = self.reasoning_ids['thought_pad']
+        slots = [
+            think_ids[start : start + self.block]
+            for start in range(0, len(think_ids), self.block)
+        ]
+        texts = []
+        for slot in slots:
+            content = [t for t in slot if t not in (pad_id, end_id)]
+            texts.append(self._decode(content).strip())
+        return [text for text in texts if text]
+
+    def _block_contents(self, think_ids: list[int]) -> list[tuple[int, str]]:
+        """Decode variable-size blocks by walking the ``<szN>`` tokens in the stream."""
+
+        size_by_id = {token_id: size for size, token_id in self.size_ids.items()}
+        end_id = self.reasoning_ids['end_think']
+        pad_id = self.reasoning_ids['thought_pad']
+        blocks = []
+        cut = 0
+        while cut < len(think_ids):
+            size = size_by_id.get(think_ids[cut])
+            if size is None:
+                cut += 1
+                continue
+            block = think_ids[cut + 1 : cut + 1 + size]
+            content = [t for t in block if t not in (pad_id, end_id)]
+            blocks.append((size, self._decode(content).strip()))
+            cut += 1 + size
+        return blocks
+
+    def _adaptive_block_texts(self, think_ids: list[int]) -> list[str]:
+        return [
+            f'({size}) {text}'.strip() for size, text in self._block_contents(think_ids)
+        ]
+
+    def _token_glyph(self, token_id: int) -> str:
+        """Render one in-flight token: masks and structural tokens get visible glyphs."""
+
+        if token_id == self.model.config.mask_token_id:
+            return '▒'
+        if token_id == self.reasoning_ids['thought_pad']:
+            return '·'
+        if token_id == self.reasoning_ids['end_think']:
+            return ' ⏹'
+        return self.tokenizer.decode([token_id], skip_special_tokens=False)
+
+    def _active_slot_text(self, slot_tokens: list[int]) -> str:
+        return ''.join(self._token_glyph(token) for token in slot_tokens)
+
+    @torch.no_grad()
+    def stream_solve(
+        self,
+        problem: str,
+        *,
+        temperature: float,
+        steps_per_block: int,
+        diffusion_steps: int,
+        max_answer_tokens: int = 200,
+        repetition_penalty: float = 1.0,
+        top_p: float = 1.0,
+        seed: int = 0,
+        strategy: str = 'ancestral',
+        step_delay: float = 0.0,
+    ):
+        """Yield ``(thoughts, answer, status)`` snapshots while solving.
+
+        ``step_delay`` throttles each denoising step (and each AR token) so the
+        reveal order is watchable in real time.
+        """
+
+        generator = torch.Generator(device=self.device.type)
+        generator.manual_seed(seed if seed else int(time.time_ns() % 2**31))
+        prompt_ids = self.tokenizer.encode(problem.strip(), add_special_tokens=False).ids
+        mask_id = self.model.config.mask_token_id
+        think_id = self.reasoning_ids['think']
+        end_think_id = self.reasoning_ids['end_think']
+        eos_id = self.roles['eos']
+
+        if self.objective == 'hybrid' and self.adaptive:
+            yield from self._stream_adaptive(
+                prompt_ids,
+                stop_ids=(eos_id,),
+                temperature=temperature,
+                steps_per_block=steps_per_block,
+                max_answer_tokens=max_answer_tokens,
+                repetition_penalty=repetition_penalty,
+                top_p=top_p,
+                strategy=strategy,
+                step_delay=step_delay,
+                generator=generator,
+            )
+            return
+
+        if self.objective == 'block_diffusion':
+            # Plain-text continuation under the conversion geometry: each new block is
+            # fully masked and denoised bidirectionally over its causal past, with block
+            # sizes drawn from the same curriculum weights the checkpoint trained under.
+            reasoning_cfg = self.config['reasoning']
+            sizes = tuple(reasoning_cfg.get('sizes') or (self.block,))
+            weights = block_size_curriculum(
+                self.step,
+                n_sizes=len(sizes),
+                curriculum_steps=int(reasoning_cfg.get('curriculum_steps') or 0),
+            ).to(self.device)
+            size_tensor = torch.tensor(sizes, device=self.device)
+            budget = min(384, self.model.config.max_seq_len - len(prompt_ids) - 1)
+            self.log_lines = [
+                f'prompt_tokens={len(prompt_ids)} temp={temperature} '
+                f'steps/block={steps_per_block} sizes={list(sizes)} '
+                f'weights={[round(w, 3) for w in weights.tolist()]} '
+                f'strategy={strategy} flex_eager={_FLEX_EAGER}'
+            ]
+            sequence = list(prompt_ids)
+            prompt_len = len(prompt_ids)
+            boundary_starts: list[int] = []
+            block_texts: list[str] = []
+            generated: list[int] = []
+            compute = 0.0
+            finished = False
+            while not finished and len(generated) < budget:
+                drawn = int(
+                    size_tensor[torch.multinomial(weights, 1, generator=generator)].item()
+                )
+                block = min(drawn, budget - len(generated))
+                prefix_len = len(sequence)
+                seq_len = prefix_len + block
+                boundary_starts.append(prefix_len)
+                boundary = torch.zeros(1, seq_len, dtype=torch.bool, device=self.device)
+                boundary[0, 0] = True
+                for start in boundary_starts:
+                    boundary[0, start] = True
+                blocked = block_mask_from_boundaries(
+                    boundary,
+                    torch.tensor([prompt_len], device=self.device),
+                    torch.tensor([seq_len], device=self.device),
+                    causal_prefix=self.causal_prefix,
+                )
+                current = torch.tensor(
+                    [sequence + [mask_id] * block], dtype=torch.long, device=self.device
+                )
+                block_ids: list[int] = []
+                tick = time.perf_counter()
+                for state in iterative_unmask_steps(
+                    self.model,
+                    current,
+                    mask_id,
+                    steps=steps_per_block,
+                    temperature=temperature,
+                    strategy=strategy,
+                    blocked_token_ids=(self.roles['pad'],),
+                    attn_mask=blocked,
+                    generator=generator,
+                ):
+                    compute += time.perf_counter() - tick
+                    block_ids = [int(t) for t in state.tokens[0, prefix_len:]]
+                    if step_delay or state.masked_remaining == 0:
+                        yield (
+                            [*block_texts, self._active_slot_text(block_ids)],
+                            self._decode(generated),
+                            f'bloque {len(block_texts) + 1} ({block} tok) · '
+                            f'denoising {state.step}/{state.total_steps}',
+                        )
+                        if step_delay and state.masked_remaining:
+                            time.sleep(step_delay)
+                    tick = time.perf_counter()
+                if eos_id in block_ids:
+                    block_ids = block_ids[: block_ids.index(eos_id)]
+                    finished = True
+                sequence.extend(block_ids)
+                generated.extend(block_ids)
+                block_texts.append(f'({block}) {self._decode(block_ids).strip()}'.strip())
+                self.log_lines.append(
+                    f'block {len(block_texts)}: size={block} · {compute:.2f}s cum'
+                )
+            rate = len(generated) / compute if compute > 0 else 0.0
+            self.log_lines.append(
+                f'TOTAL: {len(generated)} tok · blocks={len(block_texts)} · '
+                f'{compute:.2f}s ({rate:.0f} tok/s)'
+            )
+            yield block_texts, self._decode(generated), (
+                f'listo · {len(block_texts)} bloques · {len(generated)} tok en '
+                f'{compute:.2f}s ({rate:.0f} tok/s) '
+                f'(checkpoint CPT: continúa texto, no piensa)'
+            )
+            return
+
+        if self.objective == 'hybrid':
+            self.log_lines = [
+                f'prompt_tokens={len(prompt_ids)} temp={temperature} '
+                f'steps/slot={steps_per_block} rep_penalty={repetition_penalty} '
+                f'top_p={top_p} max_slots={self.max_slots} strategy={strategy} '
+                f'flex_eager={_FLEX_EAGER}'
+            ]
+            sequence = [*prompt_ids, think_id]
+            think_ids: list[int] = []
+            think_compute = 0.0
+            slot_budget = self.model.config.max_seq_len - self.block - 8
+            problem_tensor = torch.tensor([len(prompt_ids)], device=self.device)
+            for slot_index in range(self.max_slots):
+                if len(sequence) > slot_budget:
+                    break
+                slot_wall_start = time.perf_counter()
+                window = torch.tensor(
+                    [sequence + [mask_id] * self.block], dtype=torch.long, device=self.device
+                )
+                blocked = slot_causal_blocked(
+                    problem_tensor,
+                    torch.tensor([slot_index + 1], device=self.device),
+                    self.block,
+                    window.shape[1],
+                )
+                slot: list[int] = []
+                tick = time.perf_counter()
+                for state in iterative_unmask_steps(
+                    self.model,
+                    window,
+                    mask_id,
+                    steps=steps_per_block,
+                    temperature=temperature,
+                    strategy=strategy,
+                    attn_mask=blocked,
+                    generator=generator,
+                ):
+                    think_compute += time.perf_counter() - tick
+                    slot = [int(t) for t in state.tokens[0, len(sequence):]]
+                    revealed = len(think_ids) + self.block - state.masked_remaining
+                    rate = revealed / think_compute if think_compute > 0 else 0.0
+                    yield (
+                        self._slot_texts(think_ids)
+                        + [self._active_slot_text(slot)],
+                        '',
+                        f'denoising slot {slot_index + 1} · paso '
+                        f'{state.step}/{state.total_steps} · {rate:.0f} tok/s',
+                    )
+                    if step_delay and state.step < state.total_steps:
+                        time.sleep(step_delay)
+                    tick = time.perf_counter()
+                sequence.extend(slot)
+                think_ids.extend(slot)
+                slot_content = [t for t in slot if t not in (
+                    self.reasoning_ids['thought_pad'], end_think_id
+                )]
+                self.log_lines.append(
+                    f'  T{slot_index + 1}: {len(slot_content)} tok · '
+                    f'{time.perf_counter() - slot_wall_start:.2f}s wall'
+                    + (' · </think>' if end_think_id in slot else '')
+                )
+                if end_think_id in slot:
+                    break
+            else:
+                sequence.extend(
+                    [end_think_id] + [self.reasoning_ids['thought_pad']] * (self.block - 1)
+                )
+
+            answer_ids: list[int] = []
+            prefix_len = len(sequence)
+            answer_compute = 0.0
+            tick = time.perf_counter()
+            for _ in range(max_answer_tokens):
+                current = torch.tensor(
+                    [sequence + answer_ids], dtype=torch.long, device=self.device
+                )
+                blocked = prefix_causal_blocked(
+                    torch.tensor([prefix_len], device=self.device), current.shape[1]
+                )
+                output_positions = torch.zeros_like(current, dtype=torch.bool)
+                output_positions[0, -1] = True
+                logits = self.model(
+                    current, output_positions=output_positions, attn_mask=blocked
+                )
+                logits = _apply_repetition_penalty(logits, answer_ids, repetition_penalty)
+                logits = _apply_top_p(logits, top_p)
+                token, _ = _sample_categorical(logits, temperature, generator)
+                token_id = int(token.item())
+                answer_ids.append(token_id)
+                answer_compute += time.perf_counter() - tick
+                if step_delay or len(answer_ids) % 4 == 0 or token_id == eos_id:
+                    rate = len(answer_ids) / answer_compute if answer_compute > 0 else 0.0
+                    yield (
+                        self._slot_texts(think_ids),
+                        self._decode(answer_ids),
+                        f'respondiendo (AR, token {len(answer_ids)}) · {rate:.0f} tok/s',
+                    )
+                    if step_delay and token_id != eos_id:
+                        time.sleep(step_delay / 3)
+                if token_id == eos_id:
+                    break
+                tick = time.perf_counter()
+            think_tokens = len(think_ids)
+            answer_tokens = len(answer_ids)
+            total_compute = think_compute + answer_compute
+            think_rate = think_tokens / think_compute if think_compute > 0 else 0.0
+            answer_rate = answer_tokens / answer_compute if answer_compute > 0 else 0.0
+            total_rate = (
+                (think_tokens + answer_tokens) / total_compute if total_compute > 0 else 0.0
+            )
+            status = (
+                f'listo · pensar: {think_tokens} tok en {think_compute:.2f}s '
+                f'({think_rate:.0f} tok/s) · responder: {answer_tokens} tok en '
+                f'{answer_compute:.2f}s ({answer_rate:.0f} tok/s) · total: '
+                f'{think_tokens + answer_tokens} tok en {total_compute:.2f}s '
+                f'({total_rate:.0f} tok/s)'
+            )
+            self.log_lines.append(
+                f'  answer: {answer_tokens} tok · {answer_compute:.2f}s wall '
+                f'({answer_rate:.0f} tok/s) · terminated={end_think_id in think_ids}'
+            )
+            self.log_lines.append(
+                f'TOTAL: {think_tokens + answer_tokens} tok · {total_compute:.2f}s '
+                f'({total_rate:.0f} tok/s)'
+            )
+            yield self._slot_texts(think_ids), self._decode(answer_ids), status
+            return
+
+        if self.objective == 'lm':
+            # Plain causal continuation for pretraining checkpoints: no think token,
+            # no masking; the prompt is simply extended left to right.
+            sequence = list(prompt_ids)
+            generated: list[int] = []
+            prefix_len = len(sequence)
+            budget = self.model.config.max_seq_len - prefix_len - 1
+            compute = 0.0
+            tick = time.perf_counter()
+            for _ in range(min(budget, 384)):
+                current = torch.tensor(
+                    [sequence + generated], dtype=torch.long, device=self.device
+                )
+                blocked = prefix_causal_blocked(
+                    torch.tensor([prefix_len], device=self.device), current.shape[1]
+                )
+                output_positions = torch.zeros_like(current, dtype=torch.bool)
+                output_positions[0, -1] = True
+                logits = self.model(
+                    current, output_positions=output_positions, attn_mask=blocked
+                )
+                token, _ = _sample_categorical(logits, temperature, generator)
+                token_id = int(token.item())
+                generated.append(token_id)
+                compute += time.perf_counter() - tick
+                if step_delay or token_id == eos_id or len(generated) % 8 == 0:
+                    rate = len(generated) / compute if compute > 0 else 0.0
+                    yield (
+                        [],
+                        self._decode(generated),
+                        f'continuando (token {len(generated)}) · {rate:.0f} tok/s',
+                    )
+                    if step_delay and token_id != eos_id:
+                        time.sleep(step_delay / 3)
+                if token_id == eos_id:
+                    break
+                tick = time.perf_counter()
+            rate = len(generated) / compute if compute > 0 else 0.0
+            yield [], self._decode(generated), (
+                f'listo · total: {len(generated)} tok en {compute:.2f}s ({rate:.0f} tok/s) '
+                f'(checkpoint base: continua texto, no piensa)'
+            )
+            return
+
+        if self.objective == 'ar':
+            sequence = [*prompt_ids, think_id]
+            generated: list[int] = []
+            prefix_len = len(sequence)
+            budget = self.model.config.max_seq_len - prefix_len - 1
+            compute = 0.0
+            tick = time.perf_counter()
+            for _ in range(min(budget, 384)):
+                current = torch.tensor(
+                    [sequence + generated], dtype=torch.long, device=self.device
+                )
+                blocked = prefix_causal_blocked(
+                    torch.tensor([prefix_len], device=self.device), current.shape[1]
+                )
+                output_positions = torch.zeros_like(current, dtype=torch.bool)
+                output_positions[0, -1] = True
+                logits = self.model(
+                    current, output_positions=output_positions, attn_mask=blocked
+                )
+                token, _ = _sample_categorical(logits, temperature, generator)
+                token_id = int(token.item())
+                generated.append(token_id)
+                compute += time.perf_counter() - tick
+                if step_delay or token_id == eos_id or len(generated) % 8 == 0:
+                    thoughts, answer = self._split_flat(generated)
+                    rate = len(generated) / compute if compute > 0 else 0.0
+                    yield thoughts, answer, (
+                        f'generando (token {len(generated)}) · {rate:.0f} tok/s'
+                    )
+                    if step_delay and token_id != eos_id:
+                        time.sleep(step_delay / 3)
+                if token_id == eos_id:
+                    break
+                tick = time.perf_counter()
+            thoughts, answer = self._split_flat(generated)
+            rate = len(generated) / compute if compute > 0 else 0.0
+            yield thoughts, answer, (
+                f'listo · total: {len(generated)} tok en {compute:.2f}s ({rate:.0f} tok/s)'
+            )
+            return
+
+        budget = min(352, self.model.config.max_seq_len - len(prompt_ids) - 1)
+        sequence_tensor = torch.tensor(
+            [[*prompt_ids, think_id] + [mask_id] * budget], dtype=torch.long, device=self.device
+        )
+        final_ids: list[int] = []
+        compute = 0.0
+        tick = time.perf_counter()
+        for state in iterative_unmask_steps(
+            self.model,
+            sequence_tensor,
+            mask_id,
+            steps=diffusion_steps,
+            temperature=temperature,
+            strategy=strategy,
+            blocked_token_ids=(self.roles['pad'],),
+            generator=generator,
+        ):
+            compute += time.perf_counter() - tick
+            final_ids = [int(t) for t in state.tokens[0, len(prompt_ids) + 1:]]
+            if step_delay or state.step % 8 == 0 or state.masked_remaining == 0:
+                if step_delay:
+                    thoughts = [self._active_slot_text(final_ids)]
+                    answer = ''
+                else:
+                    visible = [t for t in final_ids if t != mask_id]
+                    thoughts, answer = self._split_flat(visible)
+                revealed = len(final_ids) - state.masked_remaining
+                rate = revealed / compute if compute > 0 else 0.0
+                yield (
+                    thoughts,
+                    answer,
+                    f'denoising {state.step}/{state.total_steps} · {rate:.0f} tok/s',
+                )
+                if step_delay and state.masked_remaining:
+                    time.sleep(step_delay)
+            tick = time.perf_counter()
+        if eos_id in final_ids:
+            final_ids = final_ids[: final_ids.index(eos_id) + 1]
+        thoughts, answer = self._split_flat(final_ids)
+        revealed = len(final_ids)
+        rate = revealed / compute if compute > 0 else 0.0
+        yield thoughts, answer, (
+            f'listo · total: {revealed} tok en {compute:.2f}s ({rate:.0f} tok/s)'
+        )
+
+    @torch.no_grad()
+    def stream_chat(
+        self,
+        prefix_text: str,
+        *,
+        temperature: float,
+        steps_per_block: int,
+        max_answer_tokens: int = 224,
+        repetition_penalty: float = 1.0,
+        top_p: float = 1.0,
+        seed: int = 0,
+        strategy: str = 'ancestral',
+        step_delay: float = 0.0,
+        blocks_out: list[tuple[int, str]] | None = None,
+    ):
+        """Yield ``(thoughts, answer, status)`` snapshots for one ChatML assistant turn.
+
+        ``prefix_text`` is the rendered conversation ending right after the assistant
+        header, encoded verbatim: training saw the trailing newline, so no stripping.
+        The answer stops at ``<|im_end|>`` as well as the eos token. ``blocks_out``
+        receives the turn's ``(size, content)`` blocks; a per-call sink rather than
+        engine state, so concurrent runs on the same engine cannot leak into each other.
+        """
+
+        if not self.chat_ready:
+            raise ValueError('checkpoint is not an adaptive hybrid over a ChatML tokenizer')
+        generator = torch.Generator(device=self.device.type)
+        generator.manual_seed(seed if seed else int(time.time_ns() % 2**31))
+        prompt_ids = self.tokenizer.encode(prefix_text, add_special_tokens=False).ids
+        if len(prompt_ids) >= self.model.config.max_seq_len - 48:
+            message = (
+                f'contexto lleno: prefijo de {len(prompt_ids)} tokens con ventana de '
+                f'{self.model.config.max_seq_len} — bajá los mensajes visibles o reiniciá'
+            )
+            self.log_lines = [message]
+            yield [], '', message
+            return
+        yield from self._stream_adaptive(
+            prompt_ids,
+            stop_ids=(self.im_end_id, self.roles['eos']),
+            temperature=temperature,
+            steps_per_block=steps_per_block,
+            max_answer_tokens=max_answer_tokens,
+            repetition_penalty=repetition_penalty,
+            top_p=top_p,
+            strategy=strategy,
+            step_delay=step_delay,
+            generator=generator,
+            blocks_out=blocks_out,
+        )
+
+    @torch.no_grad()
+    def _stream_adaptive(
+        self,
+        prompt_ids: list[int],
+        *,
+        stop_ids: tuple[int, ...],
+        temperature: float,
+        steps_per_block: int,
+        max_answer_tokens: int,
+        repetition_penalty: float,
+        top_p: float,
+        strategy: str,
+        step_delay: float,
+        generator: torch.Generator,
+        blocks_out: list[tuple[int, str]] | None = None,
+    ):
+        """Stream the adaptive control/denoise/answer loop shared by solve and chat.
+
+        The answer phase reuses the prefix through the model's key/value cache when the
+        backbone provides one and ``MDLM_KV_CACHE`` selects the ar part, matching the
+        batch decoder in :mod:`diffusion_lm.hybrid`.
+        """
+
+        mask_id = self.model.config.mask_token_id
+        think_id = self.reasoning_ids['think']
+        end_think_id = self.reasoning_ids['end_think']
+        size_by_id = {token_id: size for size, token_id in self.size_ids.items()}
+        size_ids_tensor = torch.tensor(sorted(self.size_ids.values()), device=self.device)
+        control_ids = [*size_by_id.keys(), end_think_id]
+        cached = hasattr(self.model, 'forward_cached') and _kv_cache_enabled('ar')
+        self.log_lines = [
+            f'prompt_tokens={len(prompt_ids)} temp={temperature} '
+            f'steps/block={steps_per_block} rep_penalty={repetition_penalty} '
+            f'top_p={top_p} max_blocks={self.max_blocks} '
+            f'sizes={sorted(size_by_id.values())} strategy={strategy} '
+            f'flex_eager={_FLEX_EAGER} kv_cache={"ar" if cached else "off"}'
+        ]
+        sequence = [*prompt_ids, think_id]
+        prefix_len = len(prompt_ids) + 1
+        problem_tensor = torch.tensor([len(prompt_ids)], device=self.device)
+        think_compute = 0.0
+        terminated = False
+        for block_index in range(self.max_blocks):
+            tick = time.perf_counter()
+            control = _ar_predict_control(
+                self.model,
+                sequence,
+                control_ids,
+                prefix_len=prefix_len,
+                temperature=0.0,
+                device=self.device,
+                generator=generator,
+                causal_prefix=self.causal_prefix,
+            )
+            think_compute += time.perf_counter() - tick
+            if control == end_think_id:
+                terminated = True
+                self.log_lines.append(f'  control {block_index + 1}: </think>')
+                break
+            size = size_by_id[control]
+            if len(sequence) + 1 + size > self.model.config.max_seq_len - 8:
+                self.log_lines.append(
+                    f'  control {block_index + 1}: <sz{size}> · sin contexto, corto acá'
+                )
+                break
+            sequence.append(control)
+            window_prefix = len(sequence)
+            block_wall_start = time.perf_counter()
+            window = torch.tensor(
+                [sequence + [mask_id] * size], dtype=torch.long, device=self.device
+            )
+            blocked = adaptive_block_mask(
+                window,
+                problem_tensor,
+                torch.tensor([window.shape[1]], device=self.device),
+                size_ids_tensor,
+                end_think_id,
+                causal_prefix=self.causal_prefix,
+            )
+            block: list[int] = []
+            tick = time.perf_counter()
+            for state in iterative_unmask_steps(
+                self.model,
+                window,
+                mask_id,
+                steps=steps_per_block,
+                temperature=temperature,
+                strategy=strategy,
+                attn_mask=blocked,
+                generator=generator,
+            ):
+                think_compute += time.perf_counter() - tick
+                block = [int(t) for t in state.tokens[0, window_prefix:]]
+                yield (
+                    self._adaptive_block_texts(sequence[prefix_len:])
+                    + [f'({size}) ' + self._active_slot_text(block)],
+                    '',
+                    f'denoising bloque {block_index + 1} (tamaño {size}) · paso '
+                    f'{state.step}/{state.total_steps}',
+                )
+                if step_delay and state.step < state.total_steps:
+                    time.sleep(step_delay)
+                tick = time.perf_counter()
+            sequence.extend(block)
+            self.log_lines.append(
+                f'  B{block_index + 1}: <sz{size}> · '
+                f'{time.perf_counter() - block_wall_start:.2f}s wall'
+            )
+        sequence.append(end_think_id)
+        think_ids = sequence[prefix_len:]
+        if blocks_out is not None:
+            blocks_out.extend(self._block_contents(think_ids))
+
+        answer_ids: list[int] = []
+        answer_compute = 0.0
+        prefix_tensor = torch.tensor([prefix_len], device=self.device)
+        answer_budget = min(
+            max_answer_tokens, self.model.config.max_seq_len - len(sequence)
+        )
+        cache = self.model.new_cache() if cached else None
+        step_in = torch.tensor([sequence], dtype=torch.long, device=self.device)
+        tick = time.perf_counter()
+        for _ in range(max(0, answer_budget)):
+            if cached:
+                seen = cache.get_seq_length()
+                blocked = prefix_causal_blocked(
+                    prefix_tensor, seen + step_in.shape[1], causal_prefix=self.causal_prefix
+                )[:, seen:, :]
+                output_positions = torch.zeros_like(step_in, dtype=torch.bool)
+                output_positions[0, -1] = True
+                logits, cache = self.model.forward_cached(
+                    step_in, attn_mask=blocked, past_key_values=cache,
+                    output_positions=output_positions,
+                )
+            else:
+                current = torch.tensor(
+                    [sequence + answer_ids], dtype=torch.long, device=self.device
+                )
+                blocked = prefix_causal_blocked(
+                    prefix_tensor, current.shape[1], causal_prefix=self.causal_prefix
+                )
+                output_positions = torch.zeros_like(current, dtype=torch.bool)
+                output_positions[0, -1] = True
+                logits = self.model(
+                    current, output_positions=output_positions, attn_mask=blocked
+                )
+            logits = _apply_repetition_penalty(logits, answer_ids, repetition_penalty)
+            logits = _apply_top_p(logits, top_p)
+            token, _ = _sample_categorical(logits, temperature, generator)
+            token_id = int(token.item())
+            answer_ids.append(token_id)
+            step_in = torch.tensor([[token_id]], dtype=torch.long, device=self.device)
+            answer_compute += time.perf_counter() - tick
+            done = token_id in stop_ids
+            if step_delay or len(answer_ids) % 4 == 0 or done:
+                rate = len(answer_ids) / answer_compute if answer_compute > 0 else 0.0
+                yield (
+                    self._adaptive_block_texts(think_ids),
+                    self._decode(answer_ids),
+                    f'respondiendo (AR, token {len(answer_ids)}) · {rate:.0f} tok/s',
+                )
+                if step_delay and not done:
+                    time.sleep(step_delay / 3)
+            if done:
+                break
+            tick = time.perf_counter()
+
+        sizes_chosen = [
+            size_by_id[t] for t in think_ids if t in size_by_id
+        ]
+        total_compute = think_compute + answer_compute
+        status = (
+            f'listo · bloques: {sizes_chosen} · terminated={terminated} · '
+            f'pensar {think_compute:.2f}s · responder {len(answer_ids)} tok en '
+            f'{answer_compute:.2f}s · total {total_compute:.2f}s'
+        )
+        self.log_lines.append(
+            f'  answer: {len(answer_ids)} tok · {answer_compute:.2f}s wall'
+        )
+        self.log_lines.append(
+            f'TOTAL: sizes={sizes_chosen} terminated={terminated} · '
+            f'{total_compute:.2f}s'
+        )
+        yield self._adaptive_block_texts(think_ids), self._decode(answer_ids), status
+
+    def _split_flat(self, generated: list[int]) -> tuple[list[str], str]:
+        end_think_id = self.reasoning_ids['end_think']
+        if end_think_id in generated:
+            split = generated.index(end_think_id)
+            think, answer = generated[:split], generated[split + 1:]
+        else:
+            think, answer = generated, []
+        return [self._decode(think).strip()], self._decode(answer)
+
+
+class HFCausalEngine:
+    """Serve a Hugging Face causal LM behind the same streaming interface.
+
+    Thinking-mode outputs (Qwen3-style ``<think>...</think>`` prefixes) are routed to the
+    thoughts panel; tokens after the closing tag stream as the answer. Models without a
+    ``</think>`` vocabulary entry stream everything as the answer.
+    """
+
+    def __init__(self, model_path: str, device: torch.device) -> None:
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        dtype = torch.bfloat16 if device.type == 'cuda' else torch.float32
+        self.model = AutoModelForCausalLM.from_pretrained(model_path, dtype=dtype)
+        self.model.to(device).eval()
+        self.device = device
+        self.objective = 'hf-ar'
+        self.step = '-'
+        eos = self.model.generation_config.eos_token_id
+        self.eos_ids = set(eos if isinstance(eos, (list, tuple)) else [eos])
+        end_think = self.tokenizer.convert_tokens_to_ids('</think>')
+        unk = self.tokenizer.unk_token_id
+        self.end_think_id = end_think if isinstance(end_think, int) and end_think != unk else None
+        self.max_new_tokens = 1024
+        self.log_lines: list[str] = []
+
+    def _decode(self, ids: list[int]) -> str:
+        # Qwen3 registers <think>/</think> as regular added tokens, so
+        # skip_special_tokens leaves them in the decoded text.
+        text = self.tokenizer.decode(ids, skip_special_tokens=True)
+        return text.replace('<think>', '').replace('</think>', '').strip()
+
+    def _split(self, generated: list[int], split_at: int | None) -> tuple[list[str], str]:
+        if self.end_think_id is None:
+            return [], self._decode(generated)
+        if split_at is None:
+            thoughts = self._decode(generated)
+            return ([thoughts] if thoughts else []), ''
+        thoughts = self._decode(generated[: split_at - 1])
+        return ([thoughts] if thoughts else []), self._decode(generated[split_at:])
+
+    @torch.no_grad()
+    def stream_solve(
+        self,
+        problem: str,
+        *,
+        temperature: float,
+        steps_per_block: int,
+        diffusion_steps: int,
+        max_answer_tokens: int = 200,
+        repetition_penalty: float = 1.0,
+        top_p: float = 1.0,
+        seed: int = 0,
+        strategy: str = 'ancestral',
+        step_delay: float = 0.0,
+    ):
+        """Yield ``(thoughts, answer, status)`` snapshots while decoding token by token.
+
+        Diffusion-only knobs (``steps_per_block``, ``diffusion_steps``, ``strategy``) are
+        accepted for interface parity and ignored.
+        """
+
+        del steps_per_block, diffusion_steps, strategy, max_answer_tokens
+        generator = torch.Generator(device=self.device.type)
+        generator.manual_seed(seed if seed else int(time.time_ns() % 2**31))
+        prompt = self.tokenizer.apply_chat_template(
+            [{'role': 'user', 'content': problem.strip()}],
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+        input_ids = self.tokenizer(prompt, return_tensors='pt').input_ids.to(self.device)
+        self.log_lines = [
+            f'hf-ar · prompt_tokens={input_ids.shape[1]} temp={temperature} top_p={top_p} '
+            f'rep_penalty={repetition_penalty} max_new_tokens={self.max_new_tokens} '
+            f'(Qwen3 sugerido: temp 0.6 · top_p 0.95 · rep 1.0)'
+        ]
+        generated: list[int] = []
+        past_key_values = None
+        current = input_ids
+        split_at: int | None = None
+        compute = 0.0
+        tick = time.perf_counter()
+        for _ in range(self.max_new_tokens):
+            output = self.model(input_ids=current, past_key_values=past_key_values, use_cache=True)
+            past_key_values = output.past_key_values
+            logits = output.logits[:, -1, :].float()
+            logits = _apply_repetition_penalty(logits, generated, repetition_penalty)
+            logits = _apply_top_p(logits, top_p)
+            token, _ = _sample_categorical(logits, temperature, generator)
+            token_id = int(token.item())
+            generated.append(token_id)
+            compute += time.perf_counter() - tick
+            if split_at is None and token_id == self.end_think_id:
+                split_at = len(generated)
+                self.log_lines.append(f'  </think> en token {split_at}')
+            done = token_id in self.eos_ids
+            if step_delay or done or len(generated) % 4 == 0:
+                thoughts, answer = self._split(generated, split_at)
+                rate = len(generated) / compute if compute > 0 else 0.0
+                thinking = self.end_think_id is not None and split_at is None
+                phase = 'pensando' if thinking else 'respondiendo'
+                yield thoughts, answer, f'{phase} (AR, token {len(generated)}) · {rate:.0f} tok/s'
+                if step_delay and not done:
+                    time.sleep(step_delay / 3)
+            if done:
+                break
+            current = token.view(1, 1)
+            tick = time.perf_counter()
+        thoughts, answer = self._split(generated, split_at)
+        rate = len(generated) / compute if compute > 0 else 0.0
+        terminated = bool(generated) and generated[-1] in self.eos_ids
+        self.log_lines.append(
+            f'TOTAL: {len(generated)} tok · {compute:.2f}s ({rate:.0f} tok/s) · '
+            f'terminated={terminated}'
+        )
+        yield thoughts, answer, (
+            f'listo · total: {len(generated)} tok en {compute:.2f}s ({rate:.0f} tok/s) · '
+            f'terminated={terminated}'
+        )
+
+
+def _render_thoughts(slots: list[str]) -> str:
+    if not slots:
+        return '<em>sin pensamientos todavía</em>'
+    chips = []
+    for index, text in enumerate(slots, start=1):
+        chips.append(
+            f'<div style="margin:6px 0;padding:8px 12px;border-radius:10px;'
+            f'background:rgba(124,58,237,.10);border:1px solid rgba(124,58,237,.25);">'
+            f'<b>T{index}</b> · {html.escape(text)}</div>'
+        )
+    return ''.join(chips)
+
+
+def _chat_display(messages: list[dict[str, str]], partial: str = '') -> list[dict[str, str]]:
+    display = [{'role': m['role'], 'content': m['content']} for m in messages]
+    if partial:
+        display.append({'role': 'assistant', 'content': partial})
+    return display
+
+
+def _chat_ledger(messages: list[dict[str, str]], keep: int) -> str:
+    return ledger_line(merge_notes(ledger_notes(messages, keep)))
+
+
+def build_app(engines: dict[str, 'ReasoningEngine | HFCausalEngine']):
+    import gradio as gr
+
+    first = next(iter(engines))
+
+    def solve(
+        checkpoint_name, problem, temperature, steps_per_block, diffusion_steps,
+        repetition_penalty, top_p, seed, strategy, slowmo_ms,
+    ):
+        engine = engines[checkpoint_name]
+        if not problem.strip():
+            yield '<em>escribí un problema primero</em>', '', 'esperando problema', ''
+            return
+        for slots, answer, status in engine.stream_solve(
+            problem,
+            temperature=float(temperature),
+            steps_per_block=int(steps_per_block),
+            diffusion_steps=int(diffusion_steps),
+            repetition_penalty=float(repetition_penalty),
+            top_p=float(top_p),
+            seed=int(seed),
+            strategy=str(strategy),
+            step_delay=float(slowmo_ms) / 1000.0,
+        ):
+            boxed = extract_boxed_answer(answer or '')
+            answer_display = answer + (f'\n\n**→ respuesta extraída: {boxed}**' if boxed else '')
+            yield _render_thoughts(slots), answer_display, status, '\n'.join(engine.log_lines)
+
+    def chat_send(
+        checkpoint_name, user_text, messages, system_text, keep_last, temperature,
+        steps_per_block, repetition_penalty, top_p, max_answer_tokens, seed, slowmo_ms,
+        datalog,
+    ):
+        engine = engines[checkpoint_name] if checkpoint_name in engines else None
+        messages = list(messages or [])
+        datalog = list(datalog or [])
+        keep = int(keep_last)
+        idle = _chat_display(messages), messages, '<em>sin pensamientos todavía</em>'
+        if engine is None or not getattr(engine, 'chat_ready', False):
+            yield (*idle, _chat_ledger(messages, keep),
+                   'elegí un checkpoint adaptativo con tokenizer ChatML', '', gr.skip(),
+                   datalog, gr.skip())
+            return
+        user_text = (user_text or '').strip()
+        if not user_text:
+            yield (*idle, _chat_ledger(messages, keep), 'escribí un mensaje primero',
+                   '', gr.skip(), datalog, gr.skip())
+            return
+        messages.append({'role': 'user', 'content': user_text})
+        notes = ledger_notes(messages, keep)
+        merged = merge_notes(notes)
+        ledger = ledger_line(merged)
+        start = window_start(messages, keep)
+        system = (system_text or '').strip() or SYSTEM
+        prefix = chat_prefix(
+            [{'role': m['role'], 'content': m['content']} for m in messages[start:]],
+            system=system,
+            extra=ledger,
+        )
+        history_snapshot = [
+            {'index': index, 'visible': index >= start, **message}
+            for index, message in enumerate(messages)
+        ]
+        answer = ''
+        status = ''
+        blocks: list[tuple[int, str]] = []
+        thoughts_html = '<em>sin pensamientos todavía</em>'
+        for slots, answer, status in engine.stream_chat(
+            prefix,
+            temperature=float(temperature),
+            steps_per_block=int(steps_per_block),
+            max_answer_tokens=int(max_answer_tokens),
+            repetition_penalty=float(repetition_penalty),
+            top_p=float(top_p),
+            seed=int(seed),
+            step_delay=float(slowmo_ms) / 1000.0,
+            blocks_out=blocks,
+        ):
+            thoughts_html = _render_thoughts(slots)
+            yield (_chat_display(messages, answer), messages, thoughts_html, ledger,
+                   status, '\n'.join(engine.log_lines), '', datalog, gr.skip())
+        note = '; '.join(text for _, text in blocks if text)
+        if answer.strip() or note:
+            messages.append({'role': 'assistant', 'content': answer.strip(), 'note': note})
+            restored_input = ''
+        else:
+            # The turn produced nothing (context-full guard): undo the user message and
+            # put its text back in the box so it can be resent after lowering the window.
+            messages.pop()
+            restored_input = user_text
+        datalog.append({
+            'turn': sum(1 for m in messages if m['role'] == 'user') + bool(restored_input),
+            'checkpoint': checkpoint_name,
+            'settings': {
+                'temperature': float(temperature), 'steps_per_block': int(steps_per_block),
+                'repetition_penalty': float(repetition_penalty), 'top_p': float(top_p),
+                'max_answer_tokens': int(max_answer_tokens), 'seed': int(seed),
+                'keep_last': keep, 'system': system,
+            },
+            'history_at_send': history_snapshot,
+            'ledger': {'notes': notes, 'merged': merged, 'line': ledger},
+            'prefix_sent': prefix,
+            'prefix_tokens': len(
+                engine.tokenizer.encode(prefix, add_special_tokens=False).ids
+            ),
+            'thinking_blocks': [
+                {'size': size, 'content': text} for size, text in blocks
+            ],
+            'answer': answer.strip(),
+            'status': status,
+            'engine_log': list(engine.log_lines),
+        })
+        yield (_chat_display(messages), messages, thoughts_html,
+               _chat_ledger(messages, keep), status, '\n'.join(engine.log_lines),
+               restored_input, datalog,
+               json.dumps(datalog, indent=2, ensure_ascii=False))
+
+    def chat_reset():
+        return ([], [], '<em>sin pensamientos todavía</em>', '', 'conversación reiniciada',
+                '', '', [], '')
+
+    chat_names = [name for name, engine in engines.items()
+                  if getattr(engine, 'chat_ready', False)]
+    chat_default = next((n for n in chat_names if 'chat' in n), chat_names[0] if chat_names else None)
+
+    with gr.Blocks(title='Reasoning playground') as app:
+        names = {
+            name: f'{name} · {engines[name].objective} · step {engines[name].step}'
+            for name in engines
+        }
+        gr.Markdown('# Reasoning playground')
+        with gr.Tab('chat'):
+            chat_state = gr.State([])
+            chat_datalog_state = gr.State([])
+            with gr.Row():
+                with gr.Column(scale=2):
+                    chat_checkpoint = gr.Dropdown(
+                        choices=chat_names, value=chat_default, label='checkpoint',
+                        info='adaptativos sobre tokenizer ChatML; solo los entrenados '
+                             'en chat (chat-sft) conocen este formato',
+                    )
+                    chat_system = gr.Textbox(label='system', value=SYSTEM, lines=3)
+                    chat_keep = gr.Slider(
+                        0, 12, value=4, step=2,
+                        label='mensajes visibles (0 = historial completo; lo anterior '
+                              'sobrevive solo en el ledger)',
+                    )
+                    chat_temperature = gr.Slider(0.0, 1.2, value=0.8, step=0.05,
+                                                 label='temperatura')
+                    chat_steps = gr.Slider(
+                        1, 64, value=16, step=1,
+                        label='pasos de denoising por bloque (32 es el óptimo medido)',
+                    )
+                    chat_rep = gr.Slider(1.0, 2.0, value=1.4, step=0.05,
+                                         label='penalización de repetición (respuesta)')
+                    chat_top_p = gr.Slider(0.1, 1.0, value=0.92, step=0.02,
+                                           label='top-p (respuesta)')
+                    chat_max_answer = gr.Slider(32, 512, value=224, step=16,
+                                                label='tokens máximos de respuesta')
+                    chat_slowmo = gr.Slider(0, 600, value=0, step=20,
+                                            label='cámara lenta (ms por paso de denoising)')
+                    chat_seed = gr.Number(value=0, label='seed (0 = aleatoria)', precision=0)
+                    chat_clear = gr.Button('reiniciar conversación')
+                with gr.Column(scale=3):
+                    chatbot = gr.Chatbot(label='conversación', height=420)
+                    chat_input = gr.Textbox(
+                        label='mensaje', lines=2,
+                        placeholder='p.ej. My policy is PL-48291. — o cualquier pregunta',
+                    )
+                    chat_go = gr.Button('enviar', variant='primary')
+                    chat_status = gr.Markdown('esperando mensaje')
+                    chat_thoughts = gr.HTML(label='pensamientos del turno')
+                    chat_ledger_box = gr.Textbox(
+                        label='ledger (notas del modelo fuera de la ventana visible)',
+                        interactive=False,
+                    )
+                    chat_logs = gr.Textbox(label='logs', lines=8, max_lines=20,
+                                           interactive=False)
+            with gr.Accordion('datalog — todo lo que viajó, por turno', open=False):
+                chat_datalog_box = gr.Textbox(
+                    label='sesión completa en JSON: settings, historial con ventana, '
+                          'ledger, prefijo exacto, bloques de thinking, respuesta y logs',
+                    lines=18, max_lines=40, interactive=False, buttons=['copy'],
+                )
+            chat_inputs = [
+                chat_checkpoint, chat_input, chat_state, chat_system, chat_keep,
+                chat_temperature, chat_steps, chat_rep, chat_top_p, chat_max_answer,
+                chat_seed, chat_slowmo, chat_datalog_state,
+            ]
+            chat_outputs = [
+                chatbot, chat_state, chat_thoughts, chat_ledger_box, chat_status,
+                chat_logs, chat_input, chat_datalog_state, chat_datalog_box,
+            ]
+            chat_go.click(chat_send, inputs=chat_inputs, outputs=chat_outputs)
+            chat_input.submit(chat_send, inputs=chat_inputs, outputs=chat_outputs)
+            chat_clear.click(chat_reset, inputs=[], outputs=chat_outputs)
+        with gr.Tab('resolver'):
+            with gr.Row():
+                with gr.Column(scale=2):
+                    checkpoint = gr.Dropdown(
+                        choices=list(engines), value=first, label='checkpoint',
+                        info=' | '.join(names.values()),
+                    )
+                    problem = gr.Textbox(
+                        label='problema / instrucción',
+                        lines=4,
+                        placeholder=(
+                            'hybrid: Write a short story. It should feature: Dialogue. '
+                            'Use the words: dragon, cake, brave.\n'
+                            'lm (base): cualquier texto a continuar, p.ej. '
+                            '"Tom was a happy boy who"'
+                        ),
+                    )
+                    temperature = gr.Slider(0.0, 1.2, value=0.8, step=0.05,
+                                            label='temperatura')
+                    steps_per_block = gr.Slider(
+                        1, 64, value=16, step=1,
+                        label='pasos de denoising por slot/bloque (hybrid)'
+                    )
+                    diffusion_steps = gr.Slider(
+                        8, 256, value=64, step=8, label='pasos totales (diffusion)'
+                    )
+                    repetition_penalty = gr.Slider(
+                        1.0, 2.0, value=1.4, step=0.05,
+                        label='penalización de repetición (respuesta)'
+                    )
+                    top_p = gr.Slider(
+                        0.1, 1.0, value=0.92, step=0.02, label='top-p (respuesta)'
+                    )
+                    strategy = gr.Dropdown(
+                        choices=['ancestral', 'confidence', 'left_to_right'],
+                        value='ancestral',
+                        label='orden de revelado (denoising; ancestral es el único que rinde)',
+                    )
+                    slowmo_ms = gr.Slider(
+                        0, 600, value=0, step=20,
+                        label='cámara lenta (ms por paso de denoising, 0 = tiempo real)',
+                    )
+                    seed = gr.Number(value=0, label='seed (0 = aleatoria)', precision=0)
+                    go = gr.Button('resolver', variant='primary')
+                with gr.Column(scale=3):
+                    status = gr.Markdown('esperando problema')
+                    thoughts = gr.HTML(label='pensamientos')
+                    answer = gr.Markdown(label='respuesta')
+                    logs = gr.Textbox(
+                        label='logs (copiá y pegá)', lines=12, max_lines=30,
+                        interactive=False,
+                    )
+            go.click(
+                solve,
+                inputs=[
+                    checkpoint, problem, temperature, steps_per_block, diffusion_steps,
+                    repetition_penalty, top_p, seed, strategy, slowmo_ms,
+                ],
+                outputs=[thoughts, answer, status, logs],
+            )
+    return app
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--outputs-dir', type=Path, default=Path('outputs'))
+    parser.add_argument('--prefix', default='reasoning-')
+    parser.add_argument('--tokenizer', type=Path, required=True)
+    parser.add_argument('--device', default='auto')
+    parser.add_argument('--host', default='127.0.0.1')
+    parser.add_argument('--port', type=int, default=7999)
+    parser.add_argument(
+        '--hf-model', action='append', default=[], metavar='NAME=PATH',
+        help='serve a Hugging Face causal LM (local path or repo id) alongside the checkpoints',
+    )
+    args = parser.parse_args()
+
+    device = resolve_device(args.device)
+    checkpoints = _discover_checkpoints(args.outputs_dir)
+    checkpoints = {
+        name: path for name, path in checkpoints.items() if name.startswith(args.prefix)
+    }
+    if not checkpoints and not args.hf_model:
+        raise SystemExit(f'no reasoning checkpoints under {args.outputs_dir}')
+    engines: dict[str, ReasoningEngine | HFCausalEngine] = {
+        name: ReasoningEngine(path, args.tokenizer, device)
+        for name, path in checkpoints.items()
+    }
+    for spec in args.hf_model:
+        name, _, path = spec.partition('=')
+        if not name or not path:
+            raise SystemExit(f'--hf-model expects NAME=PATH, got {spec!r}')
+        engines[name] = HFCausalEngine(path, device)
+    print(f'loaded: {", ".join(engines)} on {device}')
+    app = build_app(engines)
+    app.queue().launch(server_name=args.host, server_port=args.port)
+
+
+if __name__ == '__main__':
+    main()

src/diffusion_lm/reasoning_train.py

@@ -0,0 +1,600 @@
+"""Training loop for the reasoning variants: ar, diffusion, and hybrid objectives."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import time
+from dataclasses import dataclass, asdict
+from pathlib import Path
+from typing import Any
+
+import torch
+import yaml
+from torch.utils.data import DataLoader
+
+import torch.nn.functional as F
+
+from diffusion_lm.config import ModelConfig, TrainingConfig
+from diffusion_lm.data import DeterministicBatchSampler, load_packed_dataset
+from diffusion_lm.diffusion import diffusion_cross_entropy
+from diffusion_lm.hybrid import (
+    adaptive_hybrid_objective,
+    ar_objective,
+    block_diffusion_objective,
+    block_size_curriculum,
+    diffusion_objective,
+    hybrid_objective,
+)
+from diffusion_lm.model import DiffusionTransformer, build_denoiser, format_parameter_count
+from diffusion_lm.reasoning import ReasoningTokenDataset
+from diffusion_lm.tokenizer import load_tokenizer
+from diffusion_lm.train import (
+    _inference_state_dict,
+    autocast_context,
+    build_optimizer,
+    capture_rng_state,
+    configure_cuda_backends,
+    create_grad_scaler,
+    learning_rate,
+    resolve_device,
+    resolve_precision,
+    restore_rng_state,
+    seed_everything,
+)
+
+REASONING_CHECKPOINT_FORMAT = 'mini-diffusion-lm-reasoning-checkpoint-v1'
+OBJECTIVES = ('ar', 'diffusion', 'hybrid', 'lm', 'block_diffusion')
+# Objectives trained on continuous packed text (manifest datasets, no region annotations).
+PACKED_OBJECTIVES = ('lm', 'block_diffusion')
+
+
+@dataclass(frozen=True)
+class ReasoningConfig:
+    """Objective selection and hybrid-mode hyperparameters."""
+
+    objective: str
+    think_probability: float = 0.65
+    adaptive: bool = False
+    sizes: tuple[int, ...] = ()
+    causal_prefix: bool = False
+    curriculum_steps: int = 0
+    ar_probability: float = 0.0
+    # Fraction of think tokens corrupted in the causal samples, so the controller learns its
+    # decisions from damaged context instead of only from pristine traces.
+    control_context_noise: float = 0.0
+
+    def __post_init__(self) -> None:
+        if self.objective not in OBJECTIVES:
+            raise ValueError(f'objective must be one of {OBJECTIVES}')
+        if not 0.0 < self.think_probability < 1.0:
+            raise ValueError('think_probability must be in (0, 1)')
+        if self.adaptive and self.objective != 'hybrid':
+            raise ValueError('adaptive block sizing is only defined for the hybrid objective')
+        if not isinstance(self.causal_prefix, bool):
+            raise ValueError('causal_prefix must be a boolean')
+        if self.curriculum_steps < 0:
+            raise ValueError('curriculum_steps must be non-negative')
+        if not 0.0 <= self.ar_probability < 1.0:
+            raise ValueError('ar_probability must be in [0, 1)')
+        if not 0.0 <= self.control_context_noise < 1.0:
+            raise ValueError('control_context_noise must be in [0, 1)')
+        sizes = tuple(self.sizes)
+        if self.objective == 'block_diffusion':
+            if not sizes:
+                raise ValueError('block_diffusion requires a non-empty sizes menu')
+            if list(sizes) != sorted(set(sizes)) or sizes[0] <= 0:
+                raise ValueError('sizes must be unique, ascending, and positive')
+        object.__setattr__(self, 'sizes', sizes)
+
+
+@dataclass(frozen=True)
+class ReasoningExperiment:
+    model: ModelConfig
+    training: TrainingConfig
+    reasoning: ReasoningConfig
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+def load_reasoning_config(path: str | Path) -> ReasoningExperiment:
+    with Path(path).open('r', encoding='utf-8') as handle:
+        raw = yaml.safe_load(handle)
+    for section in ('model', 'training', 'reasoning'):
+        if section not in raw:
+            raise ValueError(f'config must contain a top-level {section} mapping')
+    return ReasoningExperiment(
+        model=ModelConfig(**raw['model']),
+        training=TrainingConfig(**raw['training']),
+        reasoning=ReasoningConfig(**raw['reasoning']),
+    )
+
+
+def _validate_inputs(experiment: ReasoningExperiment, datasets: list) -> None:
+    tokenizer = load_tokenizer(experiment.training.tokenizer)
+    actual_vocab = tokenizer.get_vocab_size(with_added_tokens=True)
+    if experiment.model.backbone == 'project':
+        if actual_vocab != experiment.model.vocab_size:
+            raise ValueError(
+                f'config vocab_size is {experiment.model.vocab_size}, '
+                f'tokenizer has {actual_vocab}'
+            )
+    elif actual_vocab > experiment.model.vocab_size:
+        # Pretrained embedding matrices may carry unused tail rows beyond the tokenizer.
+        raise ValueError(
+            f'tokenizer has {actual_vocab} tokens, beyond the {experiment.model.vocab_size} '
+            f'embedding rows of the pretrained backbone'
+        )
+    tokenizer_hash = hashlib.sha256(
+        Path(experiment.training.tokenizer).read_bytes()
+    ).hexdigest()
+    for dataset in datasets:
+        if dataset.metadata['tokenizer_sha256'] != tokenizer_hash:
+            raise ValueError(f'{dataset.path} was encoded with a different tokenizer file')
+        if int(dataset.metadata['vocab_size']) != actual_vocab:
+            raise ValueError(f'{dataset.path} was encoded with a different vocabulary size')
+        if 'layout' not in dataset.metadata:
+            continue
+        if experiment.reasoning.objective == 'hybrid':
+            expected_layout = 'adaptive' if experiment.reasoning.adaptive else 'slotted'
+        else:
+            expected_layout = 'flat'
+        if dataset.metadata['layout'] != expected_layout:
+            raise ValueError(
+                f'{dataset.path} has layout {dataset.metadata["layout"]!r}; the '
+                f'{experiment.reasoning.objective} objective requires {expected_layout!r}'
+            )
+        if dataset.seq_len != experiment.model.max_seq_len:
+            raise ValueError(f'{dataset.path} sequence length does not match max_seq_len')
+
+
+def _lm_objective(
+    model: DiffusionTransformer, tokens: torch.Tensor
+) -> tuple[torch.Tensor, dict[str, float]]:
+    """Plain causal-LM pretraining over continuous packed text."""
+
+    seq_len = tokens.shape[1]
+    causal_blocked = torch.triu(
+        torch.ones(seq_len, seq_len, dtype=torch.bool, device=tokens.device), diagonal=1
+    )
+    output_positions = torch.ones_like(tokens, dtype=torch.bool)
+    output_positions[:, -1] = False
+    logits = model(tokens, output_positions=output_positions, attn_mask=causal_blocked)
+    targets = tokens[:, 1:].reshape(-1)
+    loss = F.cross_entropy(logits.float(), targets)
+    accuracy = float((logits.argmax(dim=-1) == targets).float().mean())
+    return loss, {'accuracy': accuracy}
+
+
+def _objective_step(
+    model: DiffusionTransformer,
+    tokens: torch.Tensor,
+    regions: torch.Tensor | None,
+    experiment: ReasoningExperiment,
+    *,
+    eval_mask_level: float | None = None,
+    step: int | None = None,
+) -> tuple[torch.Tensor, dict[str, float]]:
+    objective = experiment.reasoning.objective
+    if objective == 'lm':
+        return _lm_objective(model, tokens)
+    if objective == 'block_diffusion':
+        size_weights = block_size_curriculum(
+            step,
+            n_sizes=len(experiment.reasoning.sizes),
+            curriculum_steps=experiment.reasoning.curriculum_steps,
+        )
+        output = block_diffusion_objective(
+            model,
+            tokens,
+            sizes=experiment.reasoning.sizes,
+            size_weights=size_weights,
+            mask_eps=experiment.training.mask_eps,
+            ar_probability=experiment.reasoning.ar_probability,
+        )
+        return output.loss, {
+            'think_loss': output.think_loss,
+            'answer_loss': output.answer_loss,
+            'think_accuracy': output.think_accuracy,
+            'answer_accuracy': output.answer_accuracy,
+        }
+    if objective == 'hybrid':
+        if experiment.reasoning.adaptive:
+            size_ids = torch.tensor(model.adaptive_size_ids, device=tokens.device)
+            output = adaptive_hybrid_objective(
+                model,
+                tokens,
+                regions,
+                size_ids=size_ids,
+                end_think_id=model.adaptive_end_think_id,
+                think_probability=experiment.reasoning.think_probability,
+                mask_eps=experiment.training.mask_eps,
+                causal_prefix=experiment.reasoning.causal_prefix,
+                control_context_noise=experiment.reasoning.control_context_noise,
+            )
+        else:
+            output = hybrid_objective(
+                model,
+                tokens,
+                regions,
+                block=int(model.reasoning_block),
+                think_probability=experiment.reasoning.think_probability,
+                mask_eps=experiment.training.mask_eps,
+            )
+        metrics = {
+            'think_loss': output.think_loss,
+            'answer_loss': output.answer_loss,
+            'think_accuracy': output.think_accuracy,
+            'answer_accuracy': output.answer_accuracy,
+        }
+        if experiment.reasoning.adaptive:
+            metrics['control_accuracy'] = output.control_accuracy
+            metrics['stop_accuracy'] = output.stop_accuracy
+        return output.loss, metrics
+    if objective == 'ar':
+        output = ar_objective(model, tokens, regions)
+        return output.loss, {'accuracy': output.accuracy}
+    mask_probability = None
+    if eval_mask_level is not None:
+        mask_probability = torch.full(
+            (tokens.shape[0],), eval_mask_level, device=tokens.device, dtype=torch.float32
+        )
+    logits, corruption, _ = diffusion_objective(
+        model,
+        tokens,
+        regions,
+        mask_eps=experiment.training.mask_eps,
+        mask_probability=mask_probability,
+    )
+    output = diffusion_cross_entropy(logits, tokens, corruption)
+    return output.loss, {'masked_accuracy': float(output.masked_accuracy)}
+
+
+@torch.no_grad()
+def _evaluate(
+    model: DiffusionTransformer,
+    loader: DataLoader,
+    experiment: ReasoningExperiment,
+    device: torch.device,
+    precision: str,
+) -> dict[str, float]:
+    state = capture_rng_state(device)
+    was_training = model.training
+    try:
+        seed_everything(0, device)
+        model.eval()
+        totals: dict[str, float] = {}
+        batches = 0
+        max_batches = experiment.training.eval_batches
+        for batch_index, batch in enumerate(loader):
+            if batch_index >= max_batches:
+                break
+            if experiment.reasoning.objective in PACKED_OBJECTIVES:
+                tokens, regions = batch, None
+            else:
+                tokens, regions = batch
+                regions = regions.to(device, non_blocking=True)
+            tokens = tokens.to(device, non_blocking=True)
+            level = experiment.training.mask_eps + (1.0 - experiment.training.mask_eps) * (
+                (batch_index + 0.5) / max_batches
+            )
+            with autocast_context(device, precision):
+                loss, metrics = _objective_step(
+                    model, tokens, regions, experiment, eval_mask_level=level
+                )
+            totals['loss'] = totals.get('loss', 0.0) + float(loss)
+            for key, value in metrics.items():
+                totals[key] = totals.get(key, 0.0) + value
+            batches += 1
+        return {key: value / max(1, batches) for key, value in totals.items()}
+    finally:
+        restore_rng_state(state)
+        model.train(was_training)
+
+
+def _save_checkpoint(
+    output_dir: Path,
+    model: DiffusionTransformer,
+    optimizer: torch.optim.Optimizer,
+    scaler: Any,
+    experiment: ReasoningExperiment,
+    step: int,
+    micro_batches_seen: int,
+    data_generator: torch.Generator,
+) -> Path:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    payload = {
+        'format': REASONING_CHECKPOINT_FORMAT,
+        'step': step,
+        'micro_batches_seen': micro_batches_seen,
+        'config': experiment.to_dict(),
+        'tokenizer_sha256': hashlib.sha256(
+            Path(experiment.training.tokenizer).read_bytes()
+        ).hexdigest(),
+        'rng_state': capture_rng_state(next(model.parameters()).device),
+        'data_generator_state': data_generator.get_state(),
+        'model': model.state_dict(),
+        'optimizer': optimizer.state_dict(),
+        'scaler': scaler.state_dict(),
+    }
+    path = output_dir / 'latest.pt'
+    temporary = output_dir / '.checkpoint.tmp'
+    torch.save(payload, temporary)
+    temporary.replace(path)
+
+    inference_payload = {
+        'format': 'mini-diffusion-lm-reasoning-inference-v1',
+        'step': step,
+        'config': experiment.to_dict(),
+        'tokenizer_sha256': payload['tokenizer_sha256'],
+        'model': _inference_state_dict(model),
+    }
+    inference_temporary = output_dir / '.inference.tmp'
+    torch.save(inference_payload, inference_temporary)
+    inference_temporary.replace(output_dir / 'inference-latest.pt')
+    return path
+
+
+def _model_architecture(config: ModelConfig) -> dict[str, Any]:
+    """Config payload compared on resume; pretrained_path is machine-local and may move."""
+
+    fields = asdict(config)
+    fields.pop('pretrained_path', None)
+    return fields
+
+
+def _load_init_weights(model: DiffusionTransformer, path: str | Path) -> None:
+    """Initialize model weights from any project checkpoint, ignoring optimizer state.
+
+    Enables the pretrain-AR-then-adapt recipe: architectures must match tensor-wise,
+    while objective-level settings (forbidden outputs, objective) may differ.
+    """
+
+    checkpoint = torch.load(path, map_location='cpu', weights_only=False)
+    state = checkpoint.get('model')
+    if state is None:
+        raise ValueError(f'{path} does not contain model weights')
+    model.load_state_dict(state)
+
+
+def train(
+    experiment: ReasoningExperiment,
+    resume: str | Path | None = None,
+    max_run_steps: int | None = None,
+    init_weights: str | Path | None = None,
+) -> Path:
+    config = experiment.training
+    device = resolve_device(config.device)
+    seed_everything(config.seed, device)
+    configure_cuda_backends(device, config.require_fused_attention)
+    precision = resolve_precision(config.precision, device)
+
+    is_packed = experiment.reasoning.objective in PACKED_OBJECTIVES
+
+    def _open_dataset(path: str):
+        if is_packed:
+            return load_packed_dataset(path, experiment.model.max_seq_len)
+        return ReasoningTokenDataset(path)
+
+    train_dataset = _open_dataset(config.train_data)
+    datasets = [train_dataset]
+    val_loader = None
+    if config.val_data is not None:
+        val_dataset = _open_dataset(config.val_data)
+        datasets.append(val_dataset)
+        val_loader = DataLoader(
+            val_dataset,
+            batch_size=config.batch_size,
+            shuffle=False,
+            num_workers=config.num_workers,
+            pin_memory=device.type == 'cuda',
+        )
+    _validate_inputs(experiment, datasets)
+
+    data_generator = torch.Generator().manual_seed(config.seed)
+    batch_sampler = DeterministicBatchSampler(
+        len(train_dataset), config.batch_size, seed=config.seed
+    )
+    train_loader = DataLoader(
+        train_dataset,
+        batch_sampler=batch_sampler,
+        num_workers=config.num_workers,
+        pin_memory=device.type == 'cuda',
+        generator=data_generator,
+    )
+
+    load_pretrained = resume is None and init_weights is None
+    # fp32 master weights regardless of the checkpoint's serialized dtype (transformers 5
+    # defaults to 'auto'/bf16); compute precision comes from autocast like every project run.
+    model = build_denoiser(
+        experiment.model, load_pretrained=load_pretrained, dtype=torch.float32
+    ).to(device)
+    if not is_packed:
+        model.reasoning_block = int(train_dataset.metadata['block'])
+    if experiment.reasoning.adaptive:
+        metadata = train_dataset.metadata
+        if 'size_token_ids' not in metadata:
+            raise ValueError(f'{config.train_data} lacks size_token_ids; rebuild it adaptively')
+        model.adaptive_size_ids = tuple(sorted(int(v) for v in metadata['size_token_ids'].values()))
+        model.adaptive_end_think_id = int(metadata['reasoning_token_ids']['end_think'])
+    if init_weights is not None:
+        if resume is not None:
+            raise ValueError('init_weights and resume are mutually exclusive')
+        _load_init_weights(model, init_weights)
+        print(json.dumps({'event': 'init_weights', 'path': str(init_weights)}))
+    optimizer = build_optimizer(model, config)
+    scaler = create_grad_scaler(device.type == 'cuda' and precision == 'float16')
+
+    start_step = 0
+    micro_batches_seen = 0
+    if resume is not None:
+        checkpoint = torch.load(resume, map_location='cpu', weights_only=False)
+        if checkpoint.get('format') != REASONING_CHECKPOINT_FORMAT:
+            raise ValueError('unsupported checkpoint format')
+        if _model_architecture(
+            ModelConfig(**checkpoint['config']['model'])
+        ) != _model_architecture(experiment.model):
+            raise ValueError('checkpoint model configuration does not match the config')
+        model.load_state_dict(checkpoint['model'])
+        optimizer.load_state_dict(checkpoint['optimizer'])
+        scaler.load_state_dict(checkpoint.get('scaler', {}))
+        data_generator.set_state(checkpoint['data_generator_state'].cpu())
+        restore_rng_state(checkpoint['rng_state'])
+        start_step = int(checkpoint['step']) + 1
+        micro_batches_seen = int(checkpoint['micro_batches_seen'])
+        del checkpoint
+
+    batch_sampler.start_batch = micro_batches_seen
+    train_iterator = iter(train_loader)
+
+    output_dir = Path(config.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    with (output_dir / 'config.json').open('w', encoding='utf-8') as handle:
+        json.dump(experiment.to_dict(), handle, indent=2)
+        handle.write('\n')
+
+    print(
+        json.dumps(
+            {
+                'event': 'start',
+                'objective': experiment.reasoning.objective,
+                'device': str(device),
+                'precision': precision,
+                'parameters': model.num_parameters,
+                'parameters_human': format_parameter_count(model.num_parameters),
+                'training_examples': len(train_dataset),
+                'start_step': start_step,
+            }
+        )
+    )
+
+    model.train()
+    log_started = time.perf_counter()
+    log_loss = 0.0
+    log_metrics: dict[str, float] = {}
+    log_count = 0
+    end_step = config.max_steps
+    if max_run_steps is not None:
+        end_step = min(end_step, start_step + max_run_steps)
+    last_checkpoint = output_dir / 'latest.pt'
+
+    for step in range(start_step, end_step):
+        lr = learning_rate(step, config)
+        for group in optimizer.param_groups:
+            group['lr'] = lr
+        optimizer.zero_grad(set_to_none=True)
+
+        for _ in range(config.gradient_accumulation_steps):
+            batch = next(train_iterator)
+            if is_packed:
+                tokens, regions = batch, None
+            else:
+                tokens, regions = batch
+                regions = regions.to(device, non_blocking=True)
+            tokens = tokens.to(device, non_blocking=True)
+            micro_batches_seen += 1
+            with autocast_context(device, precision):
+                loss, metrics = _objective_step(model, tokens, regions, experiment, step=step)
+                scaled = loss / config.gradient_accumulation_steps
+            scaler.scale(scaled).backward()
+            log_loss += float(loss) / config.gradient_accumulation_steps
+            for key, value in metrics.items():
+                log_metrics[key] = (
+                    log_metrics.get(key, 0.0) + value / config.gradient_accumulation_steps
+                )
+
+        scaler.unscale_(optimizer)
+        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_clip)
+        scaler.step(optimizer)
+        scaler.update()
+        log_count += 1
+
+        if (step + 1) % config.log_interval == 0:
+            elapsed = time.perf_counter() - log_started
+            payload = {
+                'event': 'train',
+                'step': step + 1,
+                'loss': log_loss / max(1, log_count),
+                'learning_rate': lr,
+                'grad_norm': float(grad_norm),
+                'steps_per_second': log_count / max(elapsed, 1e-9),
+            }
+            payload.update(
+                {key: value / max(1, log_count) for key, value in log_metrics.items()}
+            )
+            print(json.dumps(payload))
+            log_started = time.perf_counter()
+            log_loss = 0.0
+            log_metrics = {}
+            log_count = 0
+
+        if val_loader is not None and (step + 1) % config.eval_interval == 0:
+            metrics = _evaluate(model, val_loader, experiment, device, precision)
+            print(json.dumps({'event': 'validation', 'step': step + 1, **metrics}))
+
+        if (step + 1) % config.save_interval == 0:
+            last_checkpoint = _save_checkpoint(
+                output_dir,
+                model,
+                optimizer,
+                scaler,
+                experiment,
+                step,
+                micro_batches_seen,
+                data_generator,
+            )
+            print(json.dumps({'event': 'checkpoint', 'path': str(last_checkpoint)}))
+
+    final_step = end_step - 1
+    if not last_checkpoint.exists() or (final_step + 1) % config.save_interval != 0:
+        last_checkpoint = _save_checkpoint(
+            output_dir,
+            model,
+            optimizer,
+            scaler,
+            experiment,
+            final_step,
+            micro_batches_seen,
+            data_generator,
+        )
+    event = 'complete' if end_step == config.max_steps else 'paused'
+    print(json.dumps({'event': event, 'checkpoint': str(last_checkpoint)}))
+    return last_checkpoint
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--config', type=Path, required=True)
+    parser.add_argument('--resume', type=Path)
+    parser.add_argument('--init-weights', type=Path)
+    parser.add_argument('--max-run-steps', type=int)
+    parser.add_argument('--device')
+    parser.add_argument(
+        '--precision', choices=('auto', 'float32', 'bfloat16', 'float16')
+    )
+    args = parser.parse_args()
+
+    experiment = load_reasoning_config(args.config)
+    if args.device or args.precision:
+        from dataclasses import replace
+
+        training = replace(
+            experiment.training,
+            device=args.device or experiment.training.device,
+            precision=args.precision or experiment.training.precision,
+        )
+        experiment = ReasoningExperiment(
+            model=experiment.model, training=training, reasoning=experiment.reasoning
+        )
+    train(
+        experiment,
+        resume=args.resume,
+        max_run_steps=args.max_run_steps,
+        init_weights=args.init_weights,
+    )
+
+
+if __name__ == '__main__':
+    main()

src/diffusion_lm/sample.py

@@ -0,0 +1,153 @@
+"""Generate fixed-length text by iteratively unmasking tokens."""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import hashlib
+from pathlib import Path
+
+import torch
+
+from diffusion_lm.config import ModelConfig
+from diffusion_lm.diffusion import iterative_unmask
+from diffusion_lm.model import DiffusionTransformer
+from diffusion_lm.tokenizer import load_tokenizer, special_token_id, special_token_ids
+from diffusion_lm.train import resolve_device
+
+
+def load_model(checkpoint_path: str | Path, device: torch.device) -> DiffusionTransformer:
+    try:
+        # mmap keeps unused optimizer tensors in a full training checkpoint off
+        # resident RAM. The compact inference export remains the preferred input.
+        checkpoint = torch.load(
+            checkpoint_path,
+            map_location="cpu",
+            weights_only=False,
+            mmap=True,
+        )
+    except TypeError:  # PyTorch versions before mmap= support.
+        checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+    checkpoint_format = checkpoint.get("format")
+    if checkpoint_format not in {
+        "mini-diffusion-lm-checkpoint-v1",
+        "mini-diffusion-lm-inference-v1",
+    }:
+        raise ValueError("unsupported checkpoint format")
+    model_config = ModelConfig(**checkpoint["config"]["model"])
+    model = DiffusionTransformer(model_config)
+    if checkpoint_format == "mini-diffusion-lm-inference-v1" and device.type != "cpu":
+        # The weights-only export is BF16. Keep that dtype on accelerators instead
+        # of silently expanding a 1B model back to FP32 during load_state_dict.
+        first_weight = next(iter(checkpoint["model"].values()))
+        if first_weight.is_floating_point():
+            model = model.to(device=device, dtype=first_weight.dtype)
+    model.load_state_dict(checkpoint["model"])
+    model.tokenizer_sha256 = checkpoint.get("tokenizer_sha256")
+    del checkpoint
+    gc.collect()
+    return model.to(device).eval()
+
+
+def generate(
+    model: DiffusionTransformer,
+    tokenizer_path: str | Path,
+    *,
+    prompt: str = "",
+    generation_length: int = 128,
+    num_samples: int = 1,
+    steps: int = 64,
+    temperature: float = 1.0,
+    strategy: str = "ancestral",
+    seed: int = 1337,
+) -> list[str]:
+    tokenizer = load_tokenizer(tokenizer_path)
+    tokenizer_hash = hashlib.sha256(Path(tokenizer_path).read_bytes()).hexdigest()
+    if model.tokenizer_sha256 is not None and tokenizer_hash != model.tokenizer_sha256:
+        raise ValueError("tokenizer file does not match the tokenizer used for training")
+    if tokenizer.get_vocab_size(with_added_tokens=True) != model.config.vocab_size:
+        raise ValueError("tokenizer vocabulary does not match the checkpoint")
+    mask_id = special_token_id(tokenizer, "mask")
+    if mask_id != model.config.mask_token_id:
+        raise ValueError("tokenizer mask id does not match the checkpoint")
+    if generation_length <= 0 or num_samples <= 0:
+        raise ValueError("generation_length and num_samples must be positive")
+
+    prompt_ids = tokenizer.encode(prompt).ids if prompt else []
+    total_length = len(prompt_ids) + generation_length
+    if total_length > model.config.max_seq_len:
+        raise ValueError(
+            f"prompt plus generation uses {total_length} tokens, but model limit is "
+            f"{model.config.max_seq_len}"
+        )
+
+    device = next(model.parameters()).device
+    input_ids = torch.full(
+        (num_samples, total_length),
+        model.config.mask_token_id,
+        dtype=torch.long,
+        device=device,
+    )
+    if prompt_ids:
+        input_ids[:, : len(prompt_ids)] = torch.tensor(prompt_ids, device=device)
+
+    torch.manual_seed(seed)
+    if device.type == "cuda":
+        torch.cuda.manual_seed_all(seed)
+    elif device.type == "mps" and hasattr(torch.mps, "manual_seed"):
+        torch.mps.manual_seed(seed)
+    role_ids = special_token_ids(tokenizer)
+    blocked = tuple(role_ids[role] for role in ("pad", "unk", "bos", "mask"))
+    result = iterative_unmask(
+        model,
+        input_ids,
+        model.config.mask_token_id,
+        steps=steps,
+        temperature=temperature,
+        strategy=strategy,  # type: ignore[arg-type]
+        blocked_token_ids=blocked,
+    ).cpu()
+
+    eos_id = special_token_id(tokenizer, "eos")
+    texts: list[str] = []
+    for row in result.tolist():
+        if eos_id in row[len(prompt_ids) :]:
+            eos_position = row.index(eos_id, len(prompt_ids))
+            row = row[:eos_position]
+        texts.append(tokenizer.decode(row, skip_special_tokens=True))
+    return texts
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--checkpoint", type=Path, required=True)
+    parser.add_argument("--tokenizer", type=Path, required=True)
+    parser.add_argument("--prompt", default="")
+    parser.add_argument("--length", type=int, default=128, help="number of completion tokens")
+    parser.add_argument("--num-samples", type=int, default=1)
+    parser.add_argument("--steps", type=int, default=64)
+    parser.add_argument("--temperature", type=float, default=1.0)
+    parser.add_argument("--strategy", choices=("ancestral", "confidence"), default="ancestral")
+    parser.add_argument("--seed", type=int, default=1337)
+    parser.add_argument("--device", default="auto")
+    args = parser.parse_args()
+
+    device = resolve_device(args.device)
+    model = load_model(args.checkpoint, device)
+    texts = generate(
+        model,
+        args.tokenizer,
+        prompt=args.prompt,
+        generation_length=args.length,
+        num_samples=args.num_samples,
+        steps=args.steps,
+        temperature=args.temperature,
+        strategy=args.strategy,
+        seed=args.seed,
+    )
+    for index, text in enumerate(texts, start=1):
+        print(f"[{index}] {text}")
+
+
+if __name__ == "__main__":
+    main()

src/diffusion_lm/stories.py

@@ -0,0 +1,264 @@
+"""TinyStories preparation: AR-pretraining stream and instruct data with thought slots.
+
+Two products over one shared tokenizer:
+
+- ``prepare-pretrain``: TinyStoriesV2 stories packed as a continuous EOS-separated
+  uint16 stream for causal-LM pretraining.
+- ``prepare-instruct``: TinyStories-Instruct records rendered with the reasoning
+  slot geometry — prompt holds the writing instruction, thought slots hold the
+  requirement restatements plus the story plan (the summary, which never appears
+  in the prompt), and the answer region holds the story.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import re
+from pathlib import Path
+from typing import Iterator
+
+import numpy as np
+
+from diffusion_lm.reasoning import (
+    ExampleEncoder,
+    LayoutSpec,
+    ReasoningExample,
+    REASONING_SPECIAL_TOKENS,
+    _write_packed,
+)
+from diffusion_lm.tokenizer import (
+    load_tokenizer,
+    special_token_ids,
+    train_tokenizer_from_iterator,
+)
+
+_FIELD_RE = re.compile(r'^(Features|Words|Summary|Random sentence):\s*(.+)$', re.MULTILINE)
+_STORY_RE = re.compile(r'^Story:\s*$', re.MULTILINE)
+
+
+def iter_stories(path: Path) -> Iterator[str]:
+    """Yield individual stories from an ``<|endoftext|>``-separated text file."""
+
+    buffer = ''
+    with path.open('r', encoding='utf-8') as handle:
+        while True:
+            chunk = handle.read(1 << 24)
+            if not chunk:
+                break
+            buffer += chunk
+            *complete, buffer = buffer.split('<|endoftext|>')
+            for piece in complete:
+                story = piece.strip()
+                if story:
+                    yield story
+    tail = buffer.strip()
+    if tail:
+        yield tail
+
+
+def iter_instruct_records(path: Path) -> Iterator[tuple[dict[str, str], str]]:
+    """Yield ``(fields, story)`` pairs from a TinyStories-Instruct dump."""
+
+    for record in iter_stories(path):
+        story_match = _STORY_RE.search(record)
+        if not story_match:
+            continue
+        header = record[: story_match.start()]
+        story = record[story_match.end():].strip()
+        fields = dict(_FIELD_RE.findall(header))
+        if story and fields:
+            yield fields, story
+
+
+def instruct_example(fields: dict[str, str], story: str) -> ReasoningExample | None:
+    """Render one record as prompt, thought steps, and the story answer.
+
+    The summary becomes plan thoughts and is deliberately excluded from the
+    prompt, so planning is generative rather than copyable.
+    """
+
+    summary = ' '.join(fields.get('Summary', '').split())
+    if not summary:
+        return None
+    prompt_parts = ['Write a short story.']
+    thoughts: list[str] = []
+    features = ' '.join(fields.get('Features', '').split())
+    words = ' '.join(fields.get('Words', '').split())
+    sentence = ' '.join(fields.get('Random sentence', '').split())
+    if features:
+        prompt_parts.append(f'It should feature: {features}.')
+        thoughts.append(f'The story needs these elements: {features}.')
+    if words:
+        prompt_parts.append(f'Use the words: {words}.')
+        thoughts.append(f'I have to work in the words {words}.')
+    if sentence:
+        prompt_parts.append(f'Include the sentence: {sentence}')
+        thoughts.append(f'The sentence "{sentence}" must appear.')
+    thoughts.append(f'Plan: {summary}')
+    story = '\n'.join(line.strip() for line in story.splitlines() if line.strip())
+    return ReasoningExample(
+        problem=' '.join(prompt_parts),
+        steps=tuple(thoughts),
+        answer=story,
+        expected_answer='',
+    )
+
+
+def prepare_pretrain(args: argparse.Namespace) -> None:
+    tokenizer_path = Path(args.tokenizer)
+    if tokenizer_path.is_file():
+        tokenizer = load_tokenizer(tokenizer_path)
+        print(f'reusing tokenizer {tokenizer_path}')
+    else:
+        def _texts() -> Iterator[str]:
+            for index, story in enumerate(iter_stories(Path(args.train))):
+                if index >= args.tokenizer_sample:
+                    break
+                yield story
+
+        tokenizer = train_tokenizer_from_iterator(
+            _texts(),
+            tokenizer_path,
+            vocab_size=args.vocab_size,
+            min_frequency=4,
+            length=args.tokenizer_sample,
+            extra_special_tokens=REASONING_SPECIAL_TOKENS,
+        )
+        print(f'trained tokenizer {tokenizer_path}')
+
+    eos_id = special_token_ids(tokenizer)['eos']
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    for split, source in (('train', args.train), ('validation', args.validation)):
+        source_path = Path(source)
+        out_path = output_dir / f'{split}.bin'
+        token_count = 0
+        document_count = 0
+        batch: list[str] = []
+        with out_path.open('wb') as destination:
+            def _flush(batch: list[str]) -> tuple[int, int]:
+                encodings = tokenizer.encode_batch(batch, add_special_tokens=False)
+                tokens = 0
+                for encoding in encodings:
+                    ids = encoding.ids + [eos_id]
+                    np.asarray(ids, dtype=np.uint16).tofile(destination)
+                    tokens += len(ids)
+                return tokens, len(encodings)
+
+            for story in iter_stories(source_path):
+                batch.append(story)
+                if len(batch) >= 2048:
+                    tokens, docs = _flush(batch)
+                    token_count += tokens
+                    document_count += docs
+                    batch = []
+            if batch:
+                tokens, docs = _flush(batch)
+                token_count += tokens
+                document_count += docs
+
+        metadata = {
+            'format': 'mini-diffusion-lm-packed-tokens-v1',
+            'dtype': 'uint16',
+            'token_count': token_count,
+            'document_count': document_count,
+            'vocab_size': tokenizer.get_vocab_size(with_added_tokens=True),
+            'mask_token_id': special_token_ids(tokenizer)['mask'],
+            'eos_token_id': eos_id,
+            'special_token_ids': special_token_ids(tokenizer),
+            'tokenizer_sha256': hashlib.sha256(tokenizer_path.read_bytes()).hexdigest(),
+            'source_files': [str(source_path)],
+        }
+        with (out_path.parent / f'{out_path.name}.json').open('w', encoding='utf-8') as handle:
+            json.dump(metadata, handle, indent=2)
+            handle.write('\n')
+        print(f'{split}: {document_count:,} stories, {token_count:,} tokens -> {out_path}')
+
+
+def prepare_instruct(args: argparse.Namespace) -> None:
+    tokenizer = load_tokenizer(Path(args.tokenizer))
+    spec = LayoutSpec(seq_len=args.seq_len, block=args.block, max_slots=args.max_slots)
+    encoder = ExampleEncoder(tokenizer, spec)
+    output_dir = Path(args.output_dir)
+
+    for split, source in (('train', args.train), ('validation', args.validation)):
+        flat, flat_regions, slotted, slotted_regions = [], [], [], []
+        prompts: list[dict[str, str]] = []
+        dropped = 0
+        count = 0
+        for fields, story in iter_instruct_records(Path(source)):
+            if args.limit and count >= args.limit:
+                break
+            example = instruct_example(fields, story)
+            if example is None:
+                dropped += 1
+                continue
+            encoded = encoder.encode_example(example)
+            if encoded is None:
+                dropped += 1
+                continue
+            count += 1
+            flat.append(encoded.flat)
+            flat_regions.append(encoded.flat_regions)
+            slotted.append(encoded.slotted)
+            slotted_regions.append(encoded.slotted_regions)
+            if split == 'validation' and len(prompts) < 500:
+                prompts.append({'problem': example.problem, 'expected_answer': ''})
+        if not flat:
+            raise ValueError(f'no usable records in {source}')
+        for layout, tokens, regions in (
+            ('flat', flat, flat_regions),
+            ('slotted', slotted, slotted_regions),
+        ):
+            _write_packed(
+                output_dir / f'{split}-{layout}.bin',
+                np.stack(tokens),
+                np.stack(regions),
+                layout=layout,
+                spec=spec,
+                tokenizer_path=Path(args.tokenizer),
+                tokenizer=tokenizer,
+            )
+        print(f'{split}: {count:,} examples ({dropped:,} dropped) -> {output_dir}')
+        if split == 'validation':
+            with (output_dir / 'validation-problems.jsonl').open('w', encoding='utf-8') as fh:
+                for record in prompts:
+                    fh.write(json.dumps(record, ensure_ascii=False) + '\n')
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=__doc__)
+    subparsers = parser.add_subparsers(dest='command', required=True)
+
+    pretrain = subparsers.add_parser('prepare-pretrain')
+    pretrain.add_argument('--train', required=True)
+    pretrain.add_argument('--validation', required=True)
+    pretrain.add_argument('--output-dir', required=True)
+    pretrain.add_argument('--tokenizer', required=True)
+    pretrain.add_argument('--vocab-size', type=int, default=8192)
+    pretrain.add_argument('--tokenizer-sample', type=int, default=400_000)
+
+    instruct = subparsers.add_parser('prepare-instruct')
+    instruct.add_argument('--train', required=True)
+    instruct.add_argument('--validation', required=True)
+    instruct.add_argument('--output-dir', required=True)
+    instruct.add_argument('--tokenizer', required=True)
+    instruct.add_argument('--seq-len', type=int, default=768)
+    instruct.add_argument('--block', type=int, default=32)
+    instruct.add_argument('--max-slots', type=int, default=9)
+    instruct.add_argument('--limit', type=int, default=0)
+    return parser
+
+
+def main() -> None:
+    args = _build_parser().parse_args()
+    if args.command == 'prepare-pretrain':
+        prepare_pretrain(args)
+    elif args.command == 'prepare-instruct':
+        prepare_instruct(args)
+
+
+if __name__ == '__main__':
+    main()

src/diffusion_lm/tokenizer.py

@@ -0,0 +1,357 @@
+"""Train byte-level BPE tokenizers and encode packed token files."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+from pathlib import Path
+from typing import Iterable, Literal, Protocol, TypeAlias
+
+import numpy as np
+from tokenizers import Tokenizer, decoders, models, normalizers, pre_tokenizers, trainers
+from tqdm import tqdm
+
+
+SpecialTokenRole: TypeAlias = Literal["pad", "unk", "bos", "eos", "mask"]
+SPECIAL_TOKEN_ROLES: tuple[SpecialTokenRole, ...] = ("pad", "unk", "bos", "eos", "mask")
+
+# These sentinels deliberately include a project-specific random namespace. Common names such as
+# ``[MASK]`` occur in documentation and source code often enough to stop a web-scale encoder.
+_SPECIAL_TOKEN_NAMESPACE = "mdlm-v1-8f4e2a6c-917b-4d31-b3f7-6a2c8d5e0f19"
+SPECIAL_TOKENS = tuple(
+    f"<|{_SPECIAL_TOKEN_NAMESPACE}:{role}|>" for role in SPECIAL_TOKEN_ROLES
+)
+LEGACY_SPECIAL_TOKENS = ("[PAD]", "[UNK]", "[BOS]", "[EOS]", "[MASK]")
+# Native spellings honored on pretrained-backbone tokenizers (e.g. Qwen3). Only roles with a
+# trained native token map here; the rest are added as namespaced sentinels on free vocab rows.
+PRETRAINED_SPECIAL_TOKENS: tuple[str | None, ...] = (None, None, None, "<|endoftext|>", None)
+
+
+def _role_token_candidates(index: int) -> tuple[str, ...]:
+    candidates = [SPECIAL_TOKENS[index], LEGACY_SPECIAL_TOKENS[index]]
+    pretrained = PRETRAINED_SPECIAL_TOKENS[index]
+    if pretrained is not None:
+        candidates.append(pretrained)
+    return tuple(candidates)
+
+
+class _TokenizerLike(Protocol):
+    def token_to_id(self, token: str) -> int | None: ...
+
+
+def _raw_tokenizer(tokenizer: _TokenizerLike) -> _TokenizerLike:
+    return getattr(tokenizer, "raw_tokenizer", tokenizer)
+
+
+class RoleAwareTokenizer:
+    """Proxy a Tokenizer while keeping legacy role spellings usable by old callers.
+
+    Encoding is always delegated unchanged, so strings such as ``[MASK]`` remain ordinary text in
+    newly trained tokenizers. Only explicit ``token_to_id`` lookups receive alias compatibility.
+    """
+
+    def __init__(self, tokenizer: Tokenizer) -> None:
+        self.raw_tokenizer = tokenizer
+
+    def token_to_id(self, token: str) -> int | None:
+        token_id = self.raw_tokenizer.token_to_id(token)
+        if token_id is not None:
+            return token_id
+        for index in range(len(SPECIAL_TOKEN_ROLES)):
+            candidates = _role_token_candidates(index)
+            if token in candidates:
+                for concrete in candidates:
+                    token_id = self.raw_tokenizer.token_to_id(concrete)
+                    if token_id is not None:
+                        return token_id
+        return None
+
+    def __getattr__(self, name: str):
+        return getattr(self.raw_tokenizer, name)
+
+
+def _role_index(role: SpecialTokenRole | str) -> int:
+    try:
+        return SPECIAL_TOKEN_ROLES.index(role)  # type: ignore[arg-type]
+    except ValueError as exc:
+        choices = ", ".join(SPECIAL_TOKEN_ROLES)
+        raise ValueError(f"unknown special-token role {role!r}; expected one of {choices}") from exc
+
+
+def special_token_string(tokenizer: _TokenizerLike, role: SpecialTokenRole | str) -> str:
+    """Return the concrete token string used for ``role`` by a new or legacy tokenizer."""
+
+    index = _role_index(role)
+    raw = _raw_tokenizer(tokenizer)
+    for token in _role_token_candidates(index):
+        if raw.token_to_id(token) is not None:
+            return token
+    raise ValueError(f"tokenizer is missing the {role!r} special token")
+
+
+def special_token_id(tokenizer: _TokenizerLike, role: SpecialTokenRole | str) -> int:
+    """Resolve a special-token ID by semantic role, independent of its serialized spelling."""
+
+    token = special_token_string(tokenizer, role)
+    token_id = _raw_tokenizer(tokenizer).token_to_id(token)
+    if token_id is None:  # Kept defensive for non-Tokenizer protocol implementations.
+        raise ValueError(f"tokenizer is missing the {role!r} special token")
+    return token_id
+
+
+def special_token_ids(tokenizer: _TokenizerLike) -> dict[SpecialTokenRole, int]:
+    """Return all semantic special-token IDs for a new or legacy tokenizer."""
+
+    return {role: special_token_id(tokenizer, role) for role in SPECIAL_TOKEN_ROLES}
+
+
+def _build_tokenizer() -> Tokenizer:
+    tokenizer = Tokenizer(models.BPE(unk_token=SPECIAL_TOKENS[1]))
+    tokenizer.normalizer = normalizers.NFC()
+    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
+    tokenizer.decoder = decoders.ByteLevel()
+    return tokenizer
+
+
+def _build_trainer(
+    *,
+    vocab_size: int,
+    min_frequency: int,
+    max_token_length: int,
+    extra_special_tokens: tuple[str, ...] = (),
+) -> trainers.BpeTrainer:
+    special_tokens = list(SPECIAL_TOKENS) + list(extra_special_tokens)
+    if len(set(special_tokens)) != len(special_tokens):
+        raise ValueError('extra_special_tokens must not duplicate the role tokens or each other')
+    if vocab_size <= len(special_tokens) + 256:
+        raise ValueError("vocab_size must leave room for the byte alphabet and special tokens")
+    if min_frequency <= 0:
+        raise ValueError("min_frequency must be positive")
+    if max_token_length <= 0:
+        raise ValueError("max_token_length must be positive")
+    return trainers.BpeTrainer(
+        vocab_size=vocab_size,
+        min_frequency=min_frequency,
+        special_tokens=special_tokens,
+        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
+        max_token_length=max_token_length,
+        show_progress=True,
+    )
+
+
+def _validate_special_token_layout(tokenizer: Tokenizer) -> None:
+    for expected_id, role in enumerate(SPECIAL_TOKEN_ROLES):
+        actual_id = special_token_id(tokenizer, role)
+        if actual_id != expected_id:
+            raise ValueError(
+                f"expected the {role!r} special token at id {expected_id}, found {actual_id}"
+            )
+
+
+def _save_tokenizer(tokenizer: Tokenizer, output_path: str | Path) -> None:
+    output = Path(output_path)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    temporary = output.with_name(f".{output.name}.tmp")
+    temporary.unlink(missing_ok=True)
+    tokenizer.save(str(temporary), pretty=True)
+    temporary.replace(output)
+
+
+TokenizerTrainingItem: TypeAlias = str | list[str] | tuple[str, ...]
+
+
+def train_tokenizer_from_iterator(
+    texts: Iterable[TokenizerTrainingItem],
+    output_path: str | Path,
+    *,
+    vocab_size: int = 32_768,
+    min_frequency: int = 10,
+    max_token_length: int = 64,
+    length: int | None = None,
+    extra_special_tokens: tuple[str, ...] = (),
+) -> RoleAwareTokenizer:
+    """Train a collision-resistant byte-level BPE directly from a text iterator.
+
+    Iterator items may be individual strings or batches of strings. Batched iteration avoids
+    materializing a large text export and is the intended entry point for Parquet corpora.
+    ``extra_special_tokens`` are assigned the ids directly after the role tokens.
+    """
+
+    if length is not None and length < 0:
+        raise ValueError("length must be non-negative")
+    tokenizer = _build_tokenizer()
+    trainer = _build_trainer(
+        vocab_size=vocab_size,
+        min_frequency=min_frequency,
+        max_token_length=max_token_length,
+        extra_special_tokens=extra_special_tokens,
+    )
+    tokenizer.train_from_iterator(texts, trainer=trainer, length=length)
+    _validate_special_token_layout(tokenizer)
+    _save_tokenizer(tokenizer, output_path)
+    return RoleAwareTokenizer(tokenizer)
+
+
+def train_tokenizer(
+    input_paths: Iterable[str | Path],
+    output_path: str | Path,
+    *,
+    vocab_size: int = 32_768,
+    min_frequency: int = 10,
+    max_token_length: int = 64,
+) -> RoleAwareTokenizer:
+    """Train a byte-level BPE from local text files."""
+
+    paths = [str(Path(path)) for path in input_paths]
+    if not paths:
+        raise ValueError("at least one input text file is required")
+    missing = [path for path in paths if not Path(path).is_file()]
+    if missing:
+        raise FileNotFoundError(f"missing tokenizer inputs: {missing}")
+
+    tokenizer = _build_tokenizer()
+    trainer = _build_trainer(
+        vocab_size=vocab_size,
+        min_frequency=min_frequency,
+        max_token_length=max_token_length,
+    )
+    tokenizer.train(paths, trainer)
+    _validate_special_token_layout(tokenizer)
+    _save_tokenizer(tokenizer, output_path)
+    return RoleAwareTokenizer(tokenizer)
+
+
+def load_tokenizer(path: str | Path) -> RoleAwareTokenizer:
+    """Load and validate a project, legacy, or pretrained-backbone token layout.
+
+    Project-trained tokenizers place the role tokens at ids 0-4 and keep the strict layout
+    check. Pretrained-backbone tokenizers (role tokens appended on free vocab rows, ``pad``
+    far from 0) only need every role to resolve to some id.
+    """
+
+    tokenizer_path = Path(path)
+    if not tokenizer_path.is_file():
+        raise FileNotFoundError(f"tokenizer not found: {tokenizer_path}")
+    tokenizer = Tokenizer.from_file(str(tokenizer_path))
+    if special_token_id(tokenizer, "pad") == 0:
+        _validate_special_token_layout(tokenizer)
+    else:
+        special_token_ids(tokenizer)
+    return RoleAwareTokenizer(tokenizer)
+
+
+def token_metadata_path(token_path: str | Path) -> Path:
+    path = Path(token_path)
+    return path.with_suffix(path.suffix + ".json")
+
+
+def encode_files(
+    tokenizer_path: str | Path,
+    input_paths: Iterable[str | Path],
+    output_path: str | Path,
+) -> dict[str, object]:
+    """Encode newline-delimited documents to a compact, memory-mappable file.
+
+    This compatibility path remains useful for small corpora. Web-scale Parquet preparation lives
+    in :mod:`diffusion_lm.corpus` and preserves embedded newlines within each document.
+    """
+
+    tokenizer = load_tokenizer(tokenizer_path)
+    inputs = [Path(path) for path in input_paths]
+    if not inputs:
+        raise ValueError("at least one input text file is required")
+    missing = [str(path) for path in inputs if not path.is_file()]
+    if missing:
+        raise FileNotFoundError(f"missing corpus inputs: {missing}")
+
+    vocab_size = tokenizer.get_vocab_size(with_added_tokens=True)
+    dtype = np.dtype("uint16" if vocab_size <= np.iinfo(np.uint16).max else "uint32")
+    role_ids = special_token_ids(tokenizer)
+    eos_id = role_ids["eos"]
+    reserved_ids = set(role_ids.values())
+
+    output = Path(output_path)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    token_count = 0
+    document_count = 0
+    with output.open("wb") as destination:
+        for input_path in inputs:
+            with input_path.open("r", encoding="utf-8") as source:
+                for line_number, line in enumerate(
+                    tqdm(source, desc=f"encoding {input_path.name}", unit="docs"), start=1
+                ):
+                    text = line.rstrip("\r\n")
+                    if not text:
+                        continue
+                    token_ids = tokenizer.encode(text, add_special_tokens=False).ids
+                    encountered = reserved_ids.intersection(token_ids)
+                    if encountered:
+                        raise ValueError(
+                            f"{input_path}:{line_number} encodes reserved special-token ids "
+                            f"{sorted(encountered)}; remove literal special tokens from the corpus"
+                        )
+                    token_ids.append(eos_id)
+                    np.asarray(token_ids, dtype=dtype).tofile(destination)
+                    token_count += len(token_ids)
+                    document_count += 1
+
+    tokenizer_bytes = Path(tokenizer_path).read_bytes()
+    metadata: dict[str, object] = {
+        "format": "mini-diffusion-lm-packed-tokens-v1",
+        "dtype": dtype.name,
+        "token_count": token_count,
+        "document_count": document_count,
+        "vocab_size": vocab_size,
+        "mask_token_id": role_ids["mask"],
+        "eos_token_id": eos_id,
+        "special_token_ids": role_ids,
+        "tokenizer_sha256": hashlib.sha256(tokenizer_bytes).hexdigest(),
+        "source_files": [str(path) for path in inputs],
+    }
+    metadata_path = token_metadata_path(output)
+    with metadata_path.open("w", encoding="utf-8") as handle:
+        json.dump(metadata, handle, indent=2)
+        handle.write("\n")
+    return metadata
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=__doc__)
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    train_parser = subparsers.add_parser("train", help="train a byte-level BPE tokenizer")
+    train_parser.add_argument("--input", type=Path, nargs="+", required=True)
+    train_parser.add_argument("--output", type=Path, required=True)
+    train_parser.add_argument("--vocab-size", type=int, default=32_768)
+    train_parser.add_argument("--min-frequency", type=int, default=10)
+    train_parser.add_argument("--max-token-length", type=int, default=64)
+
+    encode_parser = subparsers.add_parser("encode", help="encode text into packed tokens")
+    encode_parser.add_argument("--tokenizer", type=Path, required=True)
+    encode_parser.add_argument("--input", type=Path, nargs="+", required=True)
+    encode_parser.add_argument("--output", type=Path, required=True)
+    return parser
+
+
+def main() -> None:
+    args = _build_parser().parse_args()
+    if args.command == "train":
+        tokenizer = train_tokenizer(
+            args.input,
+            args.output,
+            vocab_size=args.vocab_size,
+            min_frequency=args.min_frequency,
+            max_token_length=args.max_token_length,
+        )
+        print(f"saved {tokenizer.get_vocab_size():,}-token tokenizer to {args.output}")
+    elif args.command == "encode":
+        metadata = encode_files(args.tokenizer, args.input, args.output)
+        print(
+            f"wrote {metadata['token_count']:,} tokens from "
+            f"{metadata['document_count']:,} documents to {args.output}"
+        )
+
+
+if __name__ == "__main__":
+    main()

src/diffusion_lm/train.py

@@ -0,0 +1,666 @@
+"""Single-device training loop for the masked discrete diffusion LM."""
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import gc
+import hashlib
+import json
+import math
+import random
+import shutil
+import time
+from dataclasses import replace
+from pathlib import Path
+import numpy as np
+import torch
+from torch import Tensor
+from torch.optim import AdamW, Optimizer
+from torch.utils.data import DataLoader
+
+from diffusion_lm.config import ExperimentConfig, ModelConfig, TrainingConfig, load_config
+from diffusion_lm.data import DeterministicBatchSampler, load_packed_dataset
+from diffusion_lm.diffusion import corrupt_tokens, diffusion_cross_entropy
+from diffusion_lm.model import DiffusionTransformer, format_parameter_count
+from diffusion_lm.tokenizer import load_tokenizer, special_token_id
+
+
+def seed_everything(seed: int, device: torch.device) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if device.type == "cuda":
+        torch.cuda.manual_seed_all(seed)
+    if device.type == "mps" and hasattr(torch.mps, "manual_seed"):
+        # Materialize the lazy MPS runtime before setting its RNG. Seeding before
+        # the first allocation can otherwise be overwritten during initialization.
+        torch.empty((), device="mps")
+        torch.mps.synchronize()
+        torch.mps.manual_seed(seed)
+
+
+def capture_rng_state(device: torch.device) -> dict[str, object]:
+    state: dict[str, object] = {
+        "python": random.getstate(),
+        "numpy": np.random.get_state(),
+        "torch": torch.get_rng_state(),
+    }
+    if device.type == "cuda":
+        state["cuda"] = torch.cuda.get_rng_state_all()
+    if device.type == "mps" and hasattr(torch.mps, "get_rng_state"):
+        state["mps"] = torch.mps.get_rng_state()
+    return state
+
+
+def restore_rng_state(state: dict[str, object]) -> None:
+    random.setstate(state["python"])  # type: ignore[arg-type]
+    np.random.set_state(state["numpy"])  # type: ignore[arg-type]
+    torch.set_rng_state(state["torch"].cpu())  # type: ignore[union-attr]
+    if torch.cuda.is_available() and "cuda" in state:
+        torch.cuda.set_rng_state_all(  # type: ignore[arg-type]
+            [rng_state.cpu() for rng_state in state["cuda"]]  # type: ignore[union-attr]
+        )
+    if (
+        torch.backends.mps.is_available()
+        and "mps" in state
+        and hasattr(torch.mps, "set_rng_state")
+    ):
+        torch.mps.set_rng_state(state["mps"].cpu())  # type: ignore[union-attr]
+
+
+def resolve_device(requested: str) -> torch.device:
+    if requested != "auto":
+        device = torch.device(requested)
+        if device.type == "cuda" and not torch.cuda.is_available():
+            raise RuntimeError("CUDA was requested but is unavailable")
+        if device.type == "mps" and not torch.backends.mps.is_available():
+            raise RuntimeError("MPS was requested but is unavailable")
+        return device
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def resolve_precision(requested: str, device: torch.device) -> str:
+    if requested != "auto":
+        if device.type == "cpu" and requested == "float16":
+            raise ValueError("float16 training on CPU is unsupported; use float32 or bfloat16")
+        return requested
+    if device.type == "cuda":
+        return "bfloat16" if torch.cuda.is_bf16_supported() else "float16"
+    # Float32 is the most reliable default for CPU and Apple Silicon in this MVP.
+    return "float32"
+
+
+def autocast_context(device: torch.device, precision: str):
+    if precision == "float32":
+        return contextlib.nullcontext()
+    dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16}[precision]
+    return torch.autocast(device_type=device.type, dtype=dtype)
+
+
+def configure_cuda_backends(device: torch.device, require_fused_attention: bool) -> None:
+    """Enable Ampere-friendly kernels and optionally forbid quadratic math attention."""
+
+    if device.type != "cuda":
+        if require_fused_attention:
+            raise ValueError("fused attention can only be required on CUDA")
+        return
+
+    torch.set_float32_matmul_precision("high")
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cuda.enable_flash_sdp(True)
+    torch.backends.cuda.enable_mem_efficient_sdp(True)
+    # MultiheadAttention delegates to scaled_dot_product_attention when it is
+    # called with need_weights=False (as TransformerEncoderLayer does). When
+    # required, disabling math makes an unsupported fused path fail fast rather
+    # than silently allocating an L x L attention matrix. Set it in both branches
+    # so repeated train()/benchmark calls in one process cannot leak backend state.
+    torch.backends.cuda.enable_math_sdp(not require_fused_attention)
+
+
+def build_optimizer(model: DiffusionTransformer, config: TrainingConfig) -> Optimizer:
+    decay: list[Tensor] = []
+    no_decay: list[Tensor] = []
+    for parameter in model.parameters():
+        if not parameter.requires_grad:
+            continue
+        (decay if parameter.ndim >= 2 else no_decay).append(parameter)
+    parameter_groups = [
+        {"params": decay, "weight_decay": config.weight_decay},
+        {"params": no_decay, "weight_decay": 0.0},
+    ]
+    common = {
+        "lr": config.learning_rate,
+        "betas": (0.9, 0.95),
+        "eps": 1e-8,
+    }
+    if config.optimizer == "adamw8bit":
+        try:
+            import bitsandbytes as bnb
+        except ImportError as exc:
+            raise RuntimeError(
+                'optimizer=adamw8bit requires bitsandbytes; install the "gpu" extra'
+            ) from exc
+        if config.optimizer_embedding_32bit:
+            # Bitsandbytes recommends higher-precision optimizer state for NLP
+            # embeddings. The LM head is tied to this exact Parameter, so one
+            # override protects both roles at a modest memory cost.
+            manager = bnb.optim.GlobalOptimManager.get_instance()
+            manager.register_module_override(
+                model.token_embedding,
+                "weight",
+                {"optim_bits": 32},
+            )
+        return bnb.optim.AdamW8bit(
+            parameter_groups,
+            min_8bit_size=config.optimizer_min_8bit_size,
+            **common,
+        )
+    return AdamW(parameter_groups, foreach=False, **common)
+
+
+def accumulation_mask_probabilities(
+    batch_size: int,
+    micro_batch_index: int,
+    config: TrainingConfig,
+    offset: Tensor,
+    device: torch.device,
+) -> Tensor:
+    """Stratify diffusion times across a complete gradient-accumulation step."""
+
+    slots = config.batch_size * config.gradient_accumulation_steps
+    start = micro_batch_index * config.batch_size
+    indices = torch.arange(start, start + batch_size, device=device, dtype=torch.float32)
+    unit = (offset + indices / max(1, slots)) % 1.0
+    return config.mask_eps + (1.0 - config.mask_eps) * unit
+
+
+def learning_rate(step: int, config: TrainingConfig) -> float:
+    if step < config.warmup_steps:
+        return config.learning_rate * (step + 1) / max(1, config.warmup_steps)
+    progress = (step - config.warmup_steps) / max(1, config.max_steps - config.warmup_steps - 1)
+    cosine = 0.5 * (1.0 + math.cos(math.pi * min(progress, 1.0)))
+    return config.min_learning_rate + cosine * (
+        config.learning_rate - config.min_learning_rate
+    )
+
+
+def create_grad_scaler(enabled: bool):
+    """Use the unified API when available and retain PyTorch 2.2 support."""
+
+    unified_scaler = getattr(torch.amp, "GradScaler", None)
+    if unified_scaler is not None:
+        return unified_scaler("cuda", enabled=enabled)
+    return torch.cuda.amp.GradScaler(enabled=enabled)
+
+
+def validate_inputs(config: ExperimentConfig) -> None:
+    tokenizer = load_tokenizer(config.training.tokenizer)
+    actual_vocab = tokenizer.get_vocab_size(with_added_tokens=True)
+    actual_mask = special_token_id(tokenizer, "mask")
+    if actual_vocab != config.model.vocab_size:
+        raise ValueError(
+            f"config vocab_size is {config.model.vocab_size}, tokenizer has {actual_vocab}; "
+            "train the requested tokenizer or update the model budget"
+        )
+    if actual_mask != config.model.mask_token_id:
+        raise ValueError(
+            f"config mask_token_id is {config.model.mask_token_id}, tokenizer uses {actual_mask}"
+        )
+    tokenizer_hash = hashlib.sha256(Path(config.training.tokenizer).read_bytes()).hexdigest()
+    for data_path in (config.training.train_data, config.training.val_data):
+        if data_path is None:
+            continue
+        dataset = load_packed_dataset(data_path, config.model.max_seq_len)
+        metadata = dataset.metadata
+        if int(metadata["vocab_size"]) != config.model.vocab_size:
+            raise ValueError(f"{data_path} was encoded with a different vocabulary size")
+        if metadata["tokenizer_sha256"] != tokenizer_hash:
+            raise ValueError(f"{data_path} was encoded with a different tokenizer file")
+
+
+@torch.no_grad()
+def evaluate(
+    model: DiffusionTransformer,
+    loader: DataLoader[Tensor],
+    device: torch.device,
+    precision: str,
+    mask_eps: float,
+    max_batches: int,
+) -> dict[str, float]:
+    training_rng_state = capture_rng_state(device)
+    was_training = model.training
+    try:
+        # Fixed corruption masks make validation checkpoints directly comparable.
+        # The complete caller RNG state is restored below, so evaluation remains
+        # invisible to the subsequent training trajectory.
+        seed_everything(0, device)
+        model.eval()
+        losses: list[float] = []
+        correct_weighted = 0.0
+        masked_total = 0
+        for batch_index, clean_tokens in enumerate(loader):
+            if batch_index >= max_batches:
+                break
+            clean_tokens = clean_tokens.to(device, non_blocking=True)
+            # Cover the full noise range deterministically even when evaluation uses
+            # microbatch one. Random batch-1 evaluation can otherwise miss the hard
+            # near-fully-masked regime for many consecutive checkpoints.
+            level = mask_eps + (1.0 - mask_eps) * (batch_index + 0.5) / max_batches
+            mask_probability = torch.full(
+                (clean_tokens.shape[0],), level, device=device, dtype=torch.float32
+            )
+            corruption = corrupt_tokens(
+                clean_tokens,
+                model.config.mask_token_id,
+                mask_probability=mask_probability,
+                eps=mask_eps,
+            )
+            with autocast_context(device, precision):
+                logits = model(corruption.noisy_tokens, output_positions=corruption.mask)
+                output = diffusion_cross_entropy(logits, clean_tokens, corruption)
+            losses.append(float(output.loss))
+            correct_weighted += float(output.masked_accuracy) * output.masked_tokens
+            masked_total += output.masked_tokens
+        return {
+            "loss": sum(losses) / max(1, len(losses)),
+            "masked_accuracy": correct_weighted / max(1, masked_total),
+        }
+    finally:
+        restore_rng_state(training_rng_state)
+        model.train(was_training)
+
+
+def save_checkpoint(
+    output_dir: Path,
+    model: DiffusionTransformer,
+    optimizer: Optimizer,
+    scaler,
+    experiment: ExperimentConfig,
+    step: int,
+    tokens_seen: int,
+    keep_last_checkpoints: int,
+    data_generator: torch.Generator,
+    micro_batches_seen: int,
+) -> Path:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    checkpoint = {
+        "format": "mini-diffusion-lm-checkpoint-v1",
+        "step": step,
+        "tokens_seen": tokens_seen,
+        "config": experiment.to_dict(),
+        "tokenizer_sha256": hashlib.sha256(
+            Path(experiment.training.tokenizer).read_bytes()
+        ).hexdigest(),
+        "rng_state": capture_rng_state(next(model.parameters()).device),
+        "data_generator_state": data_generator.get_state(),
+        "micro_batches_seen": micro_batches_seen,
+        "model": model.state_dict(),
+        "optimizer": optimizer.state_dict(),
+        "scaler": scaler.state_dict(),
+    }
+    numbered_path = output_dir / f"step-{step:08d}.pt"
+    temporary_path = output_dir / ".checkpoint.tmp"
+    torch.save(checkpoint, temporary_path)
+    temporary_path.replace(numbered_path)
+
+    # Keep `latest.pt` as a hard link when possible so a large optimizer state is
+    # not stored twice. The temporary name makes replacement atomic.
+    latest_path = output_dir / "latest.pt"
+    latest_temporary = output_dir / ".latest.tmp"
+    latest_temporary.unlink(missing_ok=True)
+    try:
+        latest_temporary.hardlink_to(numbered_path)
+    except OSError:
+        shutil.copyfile(numbered_path, latest_temporary)
+    latest_temporary.replace(latest_path)
+
+    if keep_last_checkpoints:
+        numbered_checkpoints = sorted(output_dir.glob("step-*.pt"))
+        for old_checkpoint in numbered_checkpoints[:-keep_last_checkpoints]:
+            old_checkpoint.unlink()
+    if experiment.training.save_inference_checkpoint:
+        save_inference_checkpoint(output_dir, model, experiment, step, tokens_seen)
+    return numbered_path
+
+
+def _inference_state_dict(model: DiffusionTransformer) -> dict[str, Tensor]:
+    """Copy weights to CPU BF16 while preserving tied tensor storage."""
+
+    converted: dict[str, Tensor] = {}
+    shared: dict[tuple[object, ...], Tensor] = {}
+    for name, tensor in model.state_dict().items():
+        key = (
+            tensor.untyped_storage().data_ptr(),
+            tensor.storage_offset(),
+            tuple(tensor.shape),
+            tuple(tensor.stride()),
+        )
+        value = shared.get(key)
+        if value is None:
+            dtype = torch.bfloat16 if tensor.is_floating_point() else tensor.dtype
+            value = tensor.detach().to(device="cpu", dtype=dtype)
+            shared[key] = value
+        converted[name] = value
+    return converted
+
+
+def save_inference_checkpoint(
+    output_dir: Path,
+    model: DiffusionTransformer,
+    experiment: ExperimentConfig,
+    step: int,
+    tokens_seen: int,
+) -> Path:
+    """Write a compact weights-only checkpoint for sampling and the playground."""
+
+    path = output_dir / "inference-latest.pt"
+    temporary = output_dir / ".inference.tmp"
+    state = _inference_state_dict(model)
+    payload = {
+        "format": "mini-diffusion-lm-inference-v1",
+        "step": step,
+        "tokens_seen": tokens_seen,
+        "config": experiment.to_dict(),
+        "tokenizer_sha256": hashlib.sha256(
+            Path(experiment.training.tokenizer).read_bytes()
+        ).hexdigest(),
+        "model": state,
+    }
+    torch.save(payload, temporary)
+    temporary.replace(path)
+    del payload, state
+    gc.collect()
+    return path
+
+
+def train(
+    experiment: ExperimentConfig,
+    resume: str | Path | None = None,
+    max_run_steps: int | None = None,
+) -> Path:
+    config = experiment.training
+    if max_run_steps is not None and max_run_steps <= 0:
+        raise ValueError("max_run_steps must be positive")
+    device = resolve_device(config.device)
+    seed_everything(config.seed, device)
+    configure_cuda_backends(device, config.require_fused_attention)
+    validate_inputs(experiment)
+
+    precision = resolve_precision(config.precision, device)
+    train_dataset = load_packed_dataset(config.train_data, experiment.model.max_seq_len)
+    data_generator = torch.Generator().manual_seed(config.seed)
+    train_batch_sampler = DeterministicBatchSampler(
+        len(train_dataset), config.batch_size, seed=config.seed
+    )
+    train_loader = DataLoader(
+        train_dataset,
+        batch_sampler=train_batch_sampler,
+        num_workers=config.num_workers,
+        pin_memory=device.type == "cuda",
+        generator=data_generator,
+    )
+
+    val_loader = None
+    if config.val_data is not None:
+        val_dataset = load_packed_dataset(config.val_data, experiment.model.max_seq_len)
+        val_loader = DataLoader(
+            val_dataset,
+            batch_size=config.batch_size,
+            shuffle=False,
+            num_workers=config.num_workers,
+            pin_memory=device.type == "cuda",
+        )
+
+    model = DiffusionTransformer(experiment.model).to(device)
+    optimizer = build_optimizer(model, config)
+    scaler = create_grad_scaler(device.type == "cuda" and precision == "float16")
+    start_step = 0
+    tokens_seen = 0
+    micro_batches_seen = 0
+    if resume is not None:
+        checkpoint = torch.load(resume, map_location="cpu", weights_only=False)
+        if checkpoint.get("format") != "mini-diffusion-lm-checkpoint-v1":
+            raise ValueError("unsupported checkpoint format")
+        if ModelConfig(**checkpoint["config"]["model"]) != experiment.model:
+            raise ValueError("checkpoint model configuration does not match the requested config")
+        checkpoint_training = checkpoint["config"].get("training", {})
+        checkpoint_optimizer = checkpoint_training.get("optimizer", "adamw")
+        if checkpoint_optimizer != config.optimizer:
+            raise ValueError(
+                f"checkpoint optimizer is {checkpoint_optimizer}, requested {config.optimizer}"
+            )
+        if config.optimizer == "adamw8bit":
+            checkpoint_min_size = int(
+                checkpoint_training.get("optimizer_min_8bit_size", 4096)
+            )
+            checkpoint_embedding_32bit = bool(
+                # An absent legacy field means no explicit 32-bit override was
+                # guaranteed. Never silently reinterpret it as the safer setting.
+                checkpoint_training.get("optimizer_embedding_32bit", False)
+            )
+            if checkpoint_min_size != config.optimizer_min_8bit_size:
+                raise ValueError("checkpoint 8-bit optimizer minimum tensor size does not match")
+            if checkpoint_embedding_32bit != config.optimizer_embedding_32bit:
+                raise ValueError("checkpoint embedding optimizer precision does not match")
+        current_tokenizer_hash = hashlib.sha256(Path(config.tokenizer).read_bytes()).hexdigest()
+        if checkpoint.get("tokenizer_sha256") != current_tokenizer_hash:
+            raise ValueError("checkpoint was trained with a different tokenizer")
+        model.load_state_dict(checkpoint["model"])
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        optimizer.param_groups[0]["weight_decay"] = config.weight_decay
+        optimizer.param_groups[1]["weight_decay"] = 0.0
+        for group in optimizer.param_groups:
+            group["betas"] = (0.9, 0.95)
+            group["eps"] = 1e-8
+        scaler.load_state_dict(checkpoint.get("scaler", {}))
+        if "data_generator_state" in checkpoint:
+            data_generator.set_state(checkpoint["data_generator_state"].cpu())
+        if "rng_state" in checkpoint:
+            restore_rng_state(checkpoint["rng_state"])
+        start_step = int(checkpoint["step"]) + 1
+        tokens_seen = int(checkpoint.get("tokens_seen", 0))
+        micro_batches_seen = int(
+            checkpoint.get(
+                "micro_batches_seen",
+                start_step * int(checkpoint["config"]["training"]["gradient_accumulation_steps"]),
+            )
+        )
+        del checkpoint
+        gc.collect()
+
+    train_batch_sampler.start_batch = micro_batches_seen
+    train_iterator = iter(train_loader)
+
+    output_dir = Path(config.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    with (output_dir / "config.json").open("w", encoding="utf-8") as handle:
+        json.dump(experiment.to_dict(), handle, indent=2)
+        handle.write("\n")
+
+    print(
+        json.dumps(
+            {
+                "event": "start",
+                "device": str(device),
+                "precision": precision,
+                "parameters": model.num_parameters,
+                "parameters_human": format_parameter_count(model.num_parameters),
+                "optimizer": config.optimizer,
+                "activation_checkpointing": experiment.model.activation_checkpointing,
+                "fused_attention_required": config.require_fused_attention,
+                "training_blocks": len(train_dataset),
+                "start_step": start_step,
+            }
+        )
+    )
+
+    last_checkpoint = output_dir / "latest.pt"
+    model.train()
+    log_started = time.perf_counter()
+    log_loss = torch.zeros((), device=device)
+    log_accuracy = torch.zeros((), device=device)
+    log_tokens = 0
+    end_step = config.max_steps
+    if max_run_steps is not None:
+        end_step = min(end_step, start_step + max_run_steps)
+    for step in range(start_step, end_step):
+        lr = learning_rate(step, config)
+        for group in optimizer.param_groups:
+            group["lr"] = lr
+        optimizer.zero_grad(set_to_none=True)
+        step_loss = torch.zeros((), device=device)
+        step_accuracy = torch.zeros((), device=device)
+        noise_offset = torch.rand((), device=device)
+
+        for micro_batch_index in range(config.gradient_accumulation_steps):
+            clean_tokens = next(train_iterator).to(device, non_blocking=True)
+            micro_batches_seen += 1
+            mask_probability = accumulation_mask_probabilities(
+                clean_tokens.shape[0],
+                micro_batch_index,
+                config,
+                noise_offset,
+                device,
+            )
+            corruption = corrupt_tokens(
+                clean_tokens,
+                experiment.model.mask_token_id,
+                mask_probability=mask_probability,
+                eps=config.mask_eps,
+            )
+            with autocast_context(device, precision):
+                logits = model(corruption.noisy_tokens, output_positions=corruption.mask)
+                output = diffusion_cross_entropy(logits, clean_tokens, corruption)
+                scaled_loss = output.loss / config.gradient_accumulation_steps
+            scaler.scale(scaled_loss).backward()
+            step_loss += output.loss.detach() / config.gradient_accumulation_steps
+            step_accuracy += output.masked_accuracy.detach() / config.gradient_accumulation_steps
+            tokens_seen += clean_tokens.numel()
+            log_tokens += clean_tokens.numel()
+
+        scaler.unscale_(optimizer)
+        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_clip)
+        scaler.step(optimizer)
+        scaler.update()
+        log_loss += step_loss
+        log_accuracy += step_accuracy
+
+        if (step + 1) % config.log_interval == 0:
+            elapsed = time.perf_counter() - log_started
+            print(
+                json.dumps(
+                    {
+                        "event": "train",
+                        "step": step + 1,
+                        "loss": float(log_loss / config.log_interval),
+                        "masked_accuracy": float(log_accuracy / config.log_interval),
+                        "learning_rate": lr,
+                        "grad_norm": float(grad_norm),
+                        "tokens_seen": tokens_seen,
+                        "tokens_per_second": log_tokens / max(elapsed, 1e-9),
+                        "memory_allocated_gib": (
+                            torch.cuda.memory_allocated(device) / 1024**3
+                            if device.type == "cuda"
+                            else 0.0
+                        ),
+                        "memory_reserved_gib": (
+                            torch.cuda.memory_reserved(device) / 1024**3
+                            if device.type == "cuda"
+                            else 0.0
+                        ),
+                        "peak_memory_allocated_gib": (
+                            torch.cuda.max_memory_allocated(device) / 1024**3
+                            if device.type == "cuda"
+                            else 0.0
+                        ),
+                    }
+                )
+            )
+            log_started = time.perf_counter()
+            log_loss.zero_()
+            log_accuracy.zero_()
+            log_tokens = 0
+
+        if val_loader is not None and (step + 1) % config.eval_interval == 0:
+            metrics = evaluate(
+                model,
+                val_loader,
+                device,
+                precision,
+                config.mask_eps,
+                config.eval_batches,
+            )
+            print(json.dumps({"event": "validation", "step": step + 1, **metrics}))
+
+        if (step + 1) % config.save_interval == 0:
+            last_checkpoint = save_checkpoint(
+                output_dir,
+                model,
+                optimizer,
+                scaler,
+                experiment,
+                step,
+                tokens_seen,
+                config.keep_last_checkpoints,
+                data_generator,
+                micro_batches_seen,
+            )
+            print(json.dumps({"event": "checkpoint", "path": str(last_checkpoint)}))
+
+    final_step = end_step - 1
+    if final_step < start_step:
+        raise ValueError("checkpoint step is already at or beyond max_steps")
+    if not last_checkpoint.exists() or (final_step + 1) % config.save_interval != 0:
+        last_checkpoint = save_checkpoint(
+            output_dir,
+            model,
+            optimizer,
+            scaler,
+            experiment,
+            final_step,
+            tokens_seen,
+            config.keep_last_checkpoints,
+            data_generator,
+            micro_batches_seen,
+        )
+    event = "complete" if end_step == config.max_steps else "paused"
+    print(json.dumps({"event": event, "checkpoint": str(last_checkpoint)}))
+    return last_checkpoint
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", type=Path, required=True)
+    parser.add_argument("--resume", type=Path)
+    parser.add_argument(
+        "--max-run-steps",
+        type=int,
+        help="stop safely after this many optimizer steps (useful for scheduled jobs/tests)",
+    )
+    parser.add_argument("--device", help="override config device, e.g. cpu, mps, cuda")
+    parser.add_argument(
+        "--precision",
+        choices=("auto", "float32", "bfloat16", "float16"),
+        help="override config precision",
+    )
+    args = parser.parse_args()
+
+    experiment = load_config(args.config)
+    if args.device or args.precision:
+        training = replace(
+            experiment.training,
+            device=args.device or experiment.training.device,
+            precision=args.precision or experiment.training.precision,
+        )
+        experiment = replace(experiment, training=training)
+    train(experiment, resume=args.resume, max_run_steps=args.max_run_steps)
+
+
+if __name__ == "__main__":
+    main()

tokenizer-qwen3-adaptive.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:561587f831d1efad45896c86901cc142551f8f22c8ddbe097890f8f0f0c1bf15
+size 11424780