v2: training script, eval harness, eval set and baseline metrics

v2/eval.py

@@ -0,0 +1,471 @@
+# /// script
+# requires-python = ">=3.11,<3.14"
+# dependencies = [
+#   "torch>=2.6,<2.8",
+#   "datasets>=3.0",
+#   "numpy<3",
+#   "onnxruntime>=1.20",
+#   "transformers>=4.46",
+#   "tokenizers>=0.20",
+# ]
+# ///
+"""Scores autocomplete models the way people use them: type real text, pause, and see if the gray suggestion is right.
+
+Every model sees the same pauses in the same text and is decoded with the same rule as offline.page:
+follow the model's top pick while the running product of probabilities stays above a cutoff, and show
+the text up to the last finished word. A suggestion is right when the text really continues that way.
+
+Two scenarios:
+  typing   the cursor is at the end of what's written so far.
+  editing  the cursor is mid-sentence and real text follows a gap of 2-6 words. Only fill-in-the-middle
+           models are given that text; the rest see only what's before the cursor.
+
+Usage:
+  uv run eval.py --model tinynext-v1 --onnx model_quint8.onnx
+  uv run eval.py --model smollm2-135m --hf HuggingFaceTB/SmolLM2-135M
+  uv run eval.py --table                  # compare every metrics/*.json
+"""
+
+import argparse
+import hashlib
+import json
+import random
+import re
+import statistics
+import time
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import numpy as np
+
+HERE = Path(__file__).parent
+EVAL_SET = HERE / "eval_set.json"
+METRICS = HERE / "metrics"
+
+CUTOFFS = (0.3, 0.4, 0.5, 0.6, 0.7)
+MAIN_CUTOFF = 0.5  # what offline.page uses
+PAUSE_EVERY = 12  # characters between pauses
+MIN_SHOWN_CHARS = 4  # offline.page hides shorter suggestions
+MAX_WORDS = 6
+MAX_CHARS = 48
+DOMAINS = ("articles", "messages", "email")
+SCENARIOS = ("typing", "editing")
+PUNCTUATION = ".,!?;:"
+
+# ---- Eval text -------------------------------------------------------------------------------------------------
+# Kept apart from training: Wikipedia isn't in either model's data, OpenAssistant validation isn't in SmolTalk's
+# training split, and the emails come from the senders the original enron-autocomplete held out for its test set.
+
+CUT_REPLY = re.compile(
+    r"(?im)^\s*(?:-+\s*original message|-+\s*forwarded by|-+\s*inline attachment|from:\s|to:\s|sent:\s"
+    r"|subject:\s|cc:\s|.*\bwrote:\s*$|_{5,}|-{5,}|={5,}|\*{5,}|\d{1,2}/\d{1,2}/\d{2,4}\s+\d{1,2}:\d{2})"
+)
+SENSITIVE = re.compile(r"(?i)https?://|www\.|\S+@\S+|\(?\b\d{3}\)?[-. ]\d{3}[-. ]\d{4}\b|\b\d{3}-\d{2}-\d{4}\b")
+BULK_SENDER = re.compile(r"(?i)announce|newsletter|mailer|outlook\.team|^all[._]|general\.|chairman|enron\.(?:on|net)")
+SMART_CHARS = str.maketrans({"‘": "'", "’": "'", "“": '"', "”": '"', "–": "-", "—": "-", "\t": " "})
+
+
+def clean_email(body: str) -> list[str]:
+    """The same cleaning enron-autocomplete was trained with."""
+    text = body.replace("\r\n", "\n").replace("=20", " ").replace("=\n", "").translate(SMART_CHARS)
+    text = "\n".join(line for line in text.split("\n") if not line.lstrip().startswith(">"))
+    cut = CUT_REPLY.search(text)
+    if cut:
+        text = text[: cut.start()]
+    text = re.sub(r"(?m)^-- *\n[\s\S]*", "", text)
+    paragraphs = []
+    for block in re.split(r"\n\s*\n", text):
+        lines = [re.sub(r" +", " ", line).strip() for line in block.split("\n") if line.strip()]
+        if not lines:
+            continue
+        joined = lines[0]
+        for previous, line in zip(lines, lines[1:]):
+            joined += (" " if len(previous) >= 50 else "\n") + line
+        if SENSITIVE.search(joined) or sum(c.isalpha() for c in joined) < 0.6 * len(joined):
+            continue
+        paragraphs.append(joined)
+    return paragraphs
+
+
+def is_held_out_sender(sender: str) -> bool:
+    """enron-autocomplete put senders whose hash bucket is under 5 in its test split."""
+    sender = sender.lower()
+    if "@enron.com" not in sender or BULK_SENDER.search(sender):
+        return False
+    return int(hashlib.sha1(sender.encode()).hexdigest(), 16) % 100 < 5
+
+
+def ascii_paragraphs(text: str, min_chars: int = 40) -> list[str]:
+    text = text.translate(SMART_CHARS).encode("ascii", "ignore").decode()
+    return [p.strip() for p in re.split(r"\n\s*\n|\n", text) if len(p.strip()) >= min_chars]
+
+
+def build_eval_set(docs_per_domain: int = 80) -> dict[str, list[list[str]]]:
+    from datasets import load_dataset
+
+    sets: dict[str, list[list[str]]] = {domain: [] for domain in DOMAINS}
+    wiki = load_dataset(
+        "wikimedia/wikipedia", data_files="20231101.en/train-00040-of-00041.parquet", split="train", streaming=True
+    )
+    for row in wiki:
+        paragraphs = [p for p in ascii_paragraphs(row["text"], 80) if not p.endswith(":")][:6]
+        if len(paragraphs) >= 3:
+            sets["articles"].append(paragraphs)
+        if len(sets["articles"]) >= docs_per_domain:
+            break
+
+    for row in load_dataset("OpenAssistant/oasst1", split="validation"):
+        if row["lang"] == "en" and "```" not in row["text"]:
+            paragraphs = ascii_paragraphs(row["text"])[:6]
+            if paragraphs:
+                sets["messages"].append(paragraphs)
+        if len(sets["messages"]) >= docs_per_domain:
+            break
+
+    for row in load_dataset("corbt/enron-emails", split="train", streaming=True):
+        if is_held_out_sender(row["from"] or ""):
+            paragraphs = [p for p in clean_email(row["body"] or "") if len(p) >= 40][:6]
+            if paragraphs:
+                sets["email"].append(paragraphs)
+        if len(sets["email"]) >= docs_per_domain:
+            break
+    return sets
+
+
+def load_eval_set() -> dict[str, list[list[str]]]:
+    if not EVAL_SET.exists():
+        EVAL_SET.write_text(json.dumps(build_eval_set(), indent=1))
+    return json.loads(EVAL_SET.read_text())
+
+
+@dataclass
+class Pause:
+    domain: str
+    scenario: str
+    before: str  # the document up to the cursor
+    after: str  # text after the cursor that a model may see (empty when typing)
+    truth: str  # what the person actually typed next
+
+
+def pauses(eval_set: dict[str, list[list[str]]], max_per_domain: int) -> list[Pause]:
+    rng = random.Random(0)
+    found: list[Pause] = []
+    for domain in DOMAINS:
+        in_domain: list[Pause] = []
+        for paragraphs in eval_set[domain]:
+            for index, paragraph in enumerate(paragraphs):
+                earlier = "\n".join(paragraphs[:index])
+                for at in range(PAUSE_EVERY, len(paragraph) - 8, PAUSE_EVERY):
+                    before = (earlier + "\n" if earlier else "") + paragraph[:at]
+                    truth = paragraph[at:]
+                    in_domain.append(Pause(domain, "typing", before, "", truth))
+                    # Editing: a gap of 2-6 words, then the rest of the paragraph as it was written.
+                    gap_words = rng.randint(2, 6)
+                    gap = re.match(r"(\S*\s+){%d}" % gap_words, truth)
+                    if gap and len(truth) > gap.end() + 10:
+                        in_domain.append(Pause(domain, "editing", before, truth[gap.end():], truth))
+        rng.shuffle(in_domain)
+        found += [p for p in in_domain if p.scenario == "typing"][:max_per_domain]
+        found += [p for p in in_domain if p.scenario == "editing"][:max_per_domain]
+    return found
+
+
+# ---- Decoding -------------------------------------------------------------------------------------------------
+
+
+@dataclass
+class Trajectory:
+    """Suggestions that end at a word boundary, with the model's confidence in each. Greedy decoding gives the
+    same path for every cutoff, so one decode answers all of them."""
+
+    candidates: list[tuple[str, float]] = field(default_factory=list)
+    milliseconds: float = 0.0
+
+    def at(self, cutoff: float) -> str:
+        best = ""
+        for text, confidence in self.candidates:
+            if confidence >= cutoff:
+                best = text
+        return best.rstrip()
+
+
+class Model:
+    name: str
+    sees_after: bool = False
+
+    def suggest(self, before: str, after: str) -> Trajectory:
+        raise NotImplementedError
+
+
+class CharGRU(Model):
+    """enron-autocomplete / tinynext v1: one character per step, reading only the current paragraph."""
+
+    HIDDEN, END, NEWLINE, FIRST, OFFSET = 512, 1, 2, 3, 29
+
+    def __init__(self, name: str, onnx_path: str):
+        import onnxruntime as ort
+
+        self.name = name
+        self.session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
+        self.read_text: str | None = None
+        self.reading = None
+
+    def step(self, char_id: int, state):
+        logits, state = self.session.run(None, {"input_id": np.array([char_id], np.int64), "state": state})
+        return logits[0], state
+
+    def encode(self, char: str) -> int | None:
+        if char == "\n":
+            return self.NEWLINE
+        code = ord(char)
+        return code - self.OFFSET if 32 <= code <= 126 else None
+
+    def read(self, text: str):
+        if self.read_text is None or not text.startswith(self.read_text):
+            self.reading = self.step(self.END, np.zeros((1, self.HIDDEN), np.float32))
+            self.read_text = ""
+        for char in text[len(self.read_text):]:
+            char_id = self.encode(char)
+            if char_id is not None:
+                self.reading = self.step(char_id, self.reading[1])
+        self.read_text = text
+        return self.reading
+
+    def suggest(self, before: str, after: str) -> Trajectory:
+        began = time.perf_counter()
+        logits, state = self.read(before.split("\n")[-1])
+        trajectory, run, confidence, words = Trajectory(), "", 1.0, 0
+        while len(run) < MAX_CHARS and confidence >= min(CUTOFFS):
+            probabilities = np.exp(logits - logits.max())
+            probabilities /= probabilities.sum()
+            char_id = int(probabilities.argmax())
+            if char_id < self.FIRST:
+                if run:
+                    trajectory.candidates.append((run, confidence))
+                break
+            char = chr(char_id + self.OFFSET)
+            if char == " ":
+                trajectory.candidates.append((run, confidence))
+                words += 1
+                if words >= MAX_WORDS:
+                    break
+            confidence *= float(probabilities[char_id])
+            run += char
+            if char in PUNCTUATION:
+                trajectory.candidates.append((run, confidence))
+            logits, state = self.step(char_id, state)
+        trajectory.milliseconds = (time.perf_counter() - began) * 1000
+        return trajectory
+
+
+class TokenLM(Model):
+    """A word-piece model. `next_logits(ids, cache)` runs one forward pass and returns (last logits, cache).
+
+    The word being typed is "healed": the prompt stops before it, and the first generated pieces must spell it
+    out, so a half-typed word like "resig" can still become "resignation" in one piece.
+    """
+
+    def __init__(self, name: str, next_logits: Callable, encode: Callable[[str], list[int]], token_texts: list[str],
+                 stop_ids: set[int], prompt: Callable[[list[int], list[int]], list[int]] | None = None,
+                 context_chars: int = 1500, suffix_chars: int = 300):
+        self.name = name
+        self.next_logits = next_logits
+        self.encode = encode
+        self.token_texts = token_texts
+        self.stop_ids = stop_ids
+        self.prompt = prompt
+        self.sees_after = prompt is not None
+        self.context_chars = context_chars
+        self.suffix_chars = suffix_chars
+        lengths = np.array([len(t) for t in token_texts])
+        self.nonempty = lengths > 0
+
+    def allowed(self, need: str) -> np.ndarray:
+        """Pieces consistent with the rest of the half-typed word."""
+        return np.array([bool(t) and (need.startswith(t) or t.startswith(need)) for t in self.token_texts])
+
+    def suggest(self, before: str, after: str) -> Trajectory:
+        began = time.perf_counter()
+        before = before[-self.context_chars:]
+        # The half-typed word keeps its leading space, since word pieces start with one. A line break stays in the prompt.
+        space, newline = before.rfind(" "), before.rfind("\n")
+        if space > newline:
+            stem, need = before[:space], before[space:]
+        else:
+            stem, need = before[: newline + 1], before[newline + 1:]
+        prefix_ids = self.encode(stem)
+        if self.prompt:
+            ids = self.prompt(prefix_ids, self.encode(after[: self.suffix_chars]) if after else [])
+        else:
+            ids = prefix_ids
+        logits, cache = self.next_logits(ids, None)
+
+        trajectory, generated, confidence = Trajectory(), "", 1.0
+        for _ in range(16):
+            logits = logits.astype(np.float64)
+            if need:
+                mask = self.allowed(need)
+                if not mask.any():
+                    break
+                logits = np.where(mask, logits, -np.inf)
+            probabilities = np.exp(logits - logits.max())
+            probabilities /= probabilities.sum()
+            token = int(probabilities.argmax())
+            text = self.token_texts[token]
+            if token in self.stop_ids or "\n" in text:
+                if generated:
+                    trajectory.candidates.append((generated, confidence))
+                break
+            if need:
+                consumed = min(len(need), len(text))
+                text, need = text[consumed:], need[consumed:]
+            if text[:1].isspace() and generated:
+                trajectory.candidates.append((generated, confidence))
+            confidence *= float(probabilities[token])
+            generated += text
+            if generated.rstrip()[-1:] in PUNCTUATION and generated.strip():
+                trajectory.candidates.append((generated, confidence))
+            if confidence < min(CUTOFFS) or len(generated) > MAX_CHARS or len(generated.split()) > MAX_WORDS:
+                break
+            logits, cache = self.next_logits([token], cache)
+        trajectory.milliseconds = (time.perf_counter() - began) * 1000
+        return trajectory
+
+
+def hf_causal_lm(name: str, repo: str) -> TokenLM:
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    device = "mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu"
+    tokenizer = AutoTokenizer.from_pretrained(repo)
+    model = AutoModelForCausalLM.from_pretrained(repo, torch_dtype=torch.float32).to(device).eval()
+
+    @torch.no_grad()
+    def next_logits(ids, cache):
+        out = model(torch.tensor([ids], device=device), past_key_values=cache, use_cache=True)
+        return out.logits[0, -1].float().cpu().numpy(), out.past_key_values
+
+    texts = [tokenizer.decode([i]) for i in range(len(tokenizer))]
+    stops = {tokenizer.eos_token_id} if tokenizer.eos_token_id is not None else set()
+    encode = lambda text: tokenizer.encode(text, add_special_tokens=False) or [tokenizer.eos_token_id]
+    return TokenLM(name, next_logits, encode, texts, stops)
+
+
+# ---- Scoring ---------------------------------------------------------------------------------------------------
+
+
+def word_by_word(suggestion: str, truth: str) -> int:
+    """Characters you'd keep accepting a word at a time until a word is wrong."""
+    kept = 0
+    for match in re.finditer(r"\s*\S+", suggestion):
+        if truth.startswith(suggestion[: match.end()]):
+            kept = match.end()
+        else:
+            break
+    return kept
+
+
+def score(model: Model, all_pauses: list[Pause]) -> dict:
+    results: dict = {"model": model.name, "sees_after": model.sees_after, "cutoffs": {}, "examples": [], "latency_ms": {}}
+    latencies: list[float] = []
+    raw = []
+    for index, pause in enumerate(all_pauses):
+        trajectory = model.suggest(pause.before, pause.after if model.sees_after else "")
+        latencies.append(trajectory.milliseconds)
+        raw.append((pause, trajectory))
+        if index % 500 == 0:
+            print(f"  {model.name}: {index}/{len(all_pauses)} pauses, median {statistics.median(latencies):.1f} ms", flush=True)
+
+    for cutoff in CUTOFFS:
+        table: dict = {}
+        for domain in DOMAINS:
+            for scenario in SCENARIOS:
+                rows = [(p, t.at(cutoff)) for p, t in raw if p.domain == domain and p.scenario == scenario]
+                if not rows:
+                    continue
+                shown = [(p, s) for p, s in rows if len(s) >= MIN_SHOWN_CHARS]
+                right = [(p, s) for p, s in shown if p.truth.startswith(s)]
+                multi = [(p, s) for p, s in shown if len(s.split()) >= 2]
+                multi_right = [(p, s) for p, s in multi if p.truth.startswith(s)]
+                typed = len(rows) * PAUSE_EVERY
+                table[f"{domain}/{scenario}"] = {
+                    "pauses": len(rows),
+                    "shown": round(len(shown) / len(rows), 4),
+                    "right": round(len(right) / max(1, len(shown)), 4),
+                    "avg_right_chars": round(sum(len(s) for _, s in right) / max(1, len(right)), 2),
+                    "multiword_shown": round(len(multi) / len(rows), 4),
+                    "multiword_right": round(len(multi_right) / max(1, len(multi)), 4),
+                    "saved_with_tab": round(sum(len(s) for _, s in right) / typed, 4),
+                    "saved_word_by_word": round(sum(word_by_word(s, p.truth) for p, s in shown) / typed, 4),
+                }
+        results["cutoffs"][str(cutoff)] = table
+
+    for pause, trajectory in raw[:: max(1, len(raw) // 60)]:
+        suggestion = trajectory.at(MAIN_CUTOFF)
+        if len(suggestion) >= MIN_SHOWN_CHARS:
+            results["examples"].append({
+                "domain": pause.domain, "scenario": pause.scenario, "before": pause.before[-60:],
+                "suggestion": suggestion, "truth": pause.truth[: len(suggestion) + 12], "right": pause.truth.startswith(suggestion),
+            })
+    results["latency_ms"] = {
+        "p50": round(statistics.median(latencies), 2),
+        "p95": round(float(np.percentile(latencies, 95)), 2),
+        "note": "Measured in this eval process, not in a browser.",
+    }
+    return results
+
+
+# ---- Report ----------------------------------------------------------------------------------------------------
+
+
+def table(paths: list[Path]) -> str:
+    runs = [json.loads(p.read_text()) for p in paths]
+    lines = []
+    for scenario in SCENARIOS:
+        lines.append(f"\n### {scenario.capitalize()}, {int(MAIN_CUTOFF * 100)}% cutoff, suggestions of {MIN_SHOWN_CHARS}+ characters\n")
+        lines.append("| Model | Domain | Shown | Right | Avg right | 2+ words shown | 2+ words right | Saved (Tab) | Saved (word by word) |")
+        lines.append("|---|---|---|---|---|---|---|---|---|")
+        for run in runs:
+            for domain in DOMAINS:
+                m = run["cutoffs"][str(MAIN_CUTOFF)].get(f"{domain}/{scenario}")
+                if not m:
+                    continue
+                lines.append(
+                    f"| {run['model']} | {domain} | {m['shown']:.0%} | {m['right']:.0%} | {m['avg_right_chars']:.1f} ch "
+                    f"| {m['multiword_shown']:.0%} | {m['multiword_right']:.0%} | {m['saved_with_tab']:.1%} | {m['saved_word_by_word']:.1%} |"
+                )
+    lines.append("\n| Model | Latency p50 | p95 |\n|---|---|---|")
+    for run in runs:
+        lines.append(f"| {run['model']} | {run['latency_ms']['p50']} ms | {run['latency_ms']['p95']} ms |")
+    return "\n".join(lines)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", help="Name for the metrics file")
+    parser.add_argument("--onnx", help="A character-level GRU exported like enron-autocomplete")
+    parser.add_argument("--hf", help="A Hugging Face causal LM to evaluate, e.g. HuggingFaceTB/SmolLM2-135M")
+    parser.add_argument("--max-pauses", type=int, default=1000, help="Per domain and scenario")
+    parser.add_argument("--table", action="store_true", help="Print a comparison of every metrics/*.json")
+    args = parser.parse_args()
+
+    if args.table:
+        print(table(sorted(METRICS.glob("*.json"))))
+        return
+    if args.onnx:
+        model: Model = CharGRU(args.model, args.onnx)
+    elif args.hf:
+        model = hf_causal_lm(args.model, args.hf)
+    else:
+        parser.error("Pass --onnx, --hf or --table. train.py evaluates its own model through score().")
+    all_pauses = pauses(load_eval_set(), args.max_pauses)
+    print(f"{len(all_pauses)} pauses", flush=True)
+    results = score(model, all_pauses)
+    METRICS.mkdir(exist_ok=True)
+    (METRICS / f"{args.model}.json").write_text(json.dumps(results, indent=1))
+    print(table([METRICS / f"{args.model}.json"]))
+
+
+if __name__ == "__main__":
+    main()

v2/metrics/smollm2-135m.json

@@ -0,0 +1,427 @@
+{
+ "model": "smollm2-135m",
+ "sees_after": false,
+ "cutoffs": {
+  "0.3": {
+   "articles/typing": {
+    "pauses": 400,
+    "shown": 0.33,
+    "right": 0.6515,
+    "avg_right_chars": 7.33,
+    "multiword_shown": 0.175,
+    "multiword_right": 0.6286,
+    "saved_with_tab": 0.1313,
+    "saved_word_by_word": 0.1462
+   },
+   "articles/editing": {
+    "pauses": 400,
+    "shown": 0.3175,
+    "right": 0.685,
+    "avg_right_chars": 6.25,
+    "multiword_shown": 0.1475,
+    "multiword_right": 0.6271,
+    "saved_with_tab": 0.1133,
+    "saved_word_by_word": 0.1221
+   },
+   "messages/typing": {
+    "pauses": 400,
+    "shown": 0.4375,
+    "right": 0.6057,
+    "avg_right_chars": 9.35,
+    "multiword_shown": 0.3125,
+    "multiword_right": 0.6,
+    "saved_with_tab": 0.2065,
+    "saved_word_by_word": 0.25
+   },
+   "messages/editing": {
+    "pauses": 400,
+    "shown": 0.44,
+    "right": 0.6136,
+    "avg_right_chars": 9.51,
+    "multiword_shown": 0.2825,
+    "multiword_right": 0.5841,
+    "saved_with_tab": 0.214,
+    "saved_word_by_word": 0.2498
+   },
+   "email/typing": {
+    "pauses": 400,
+    "shown": 0.2725,
+    "right": 0.5321,
+    "avg_right_chars": 5.97,
+    "multiword_shown": 0.0875,
+    "multiword_right": 0.4286,
+    "saved_with_tab": 0.0721,
+    "saved_word_by_word": 0.0879
+   },
+   "email/editing": {
+    "pauses": 400,
+    "shown": 0.27,
+    "right": 0.6204,
+    "avg_right_chars": 6.22,
+    "multiword_shown": 0.1175,
+    "multiword_right": 0.6383,
+    "saved_with_tab": 0.0869,
+    "saved_word_by_word": 0.0942
+   }
+  },
+  "0.4": {
+   "articles/typing": {
+    "pauses": 400,
+    "shown": 0.25,
+    "right": 0.72,
+    "avg_right_chars": 7.04,
+    "multiword_shown": 0.115,
+    "multiword_right": 0.6739,
+    "saved_with_tab": 0.1056,
+    "saved_word_by_word": 0.1131
+   },
+   "articles/editing": {
+    "pauses": 400,
+    "shown": 0.2375,
+    "right": 0.7895,
+    "avg_right_chars": 6.17,
+    "multiword_shown": 0.1,
+    "multiword_right": 0.75,
+    "saved_with_tab": 0.0965,
+    "saved_word_by_word": 0.1008
+   },
+   "messages/typing": {
+    "pauses": 400,
+    "shown": 0.3625,
+    "right": 0.6759,
+    "avg_right_chars": 8.56,
+    "multiword_shown": 0.2375,
+    "multiword_right": 0.6421,
+    "saved_with_tab": 0.1748,
+    "saved_word_by_word": 0.2052
+   },
+   "messages/editing": {
+    "pauses": 400,
+    "shown": 0.3825,
+    "right": 0.6993,
+    "avg_right_chars": 8.83,
+    "multiword_shown": 0.2225,
+    "multiword_right": 0.6742,
+    "saved_with_tab": 0.1969,
+    "saved_word_by_word": 0.2167
+   },
+   "email/typing": {
+    "pauses": 400,
+    "shown": 0.2075,
+    "right": 0.6747,
+    "avg_right_chars": 5.62,
+    "multiword_shown": 0.0375,
+    "multiword_right": 0.6667,
+    "saved_with_tab": 0.0656,
+    "saved_word_by_word": 0.0688
+   },
+   "email/editing": {
+    "pauses": 400,
+    "shown": 0.2025,
+    "right": 0.6914,
+    "avg_right_chars": 5.43,
+    "multiword_shown": 0.065,
+    "multiword_right": 0.6538,
+    "saved_with_tab": 0.0633,
+    "saved_word_by_word": 0.0667
+   }
+  },
+  "0.5": {
+   "articles/typing": {
+    "pauses": 400,
+    "shown": 0.185,
+    "right": 0.7973,
+    "avg_right_chars": 6.61,
+    "multiword_shown": 0.0675,
+    "multiword_right": 0.8148,
+    "saved_with_tab": 0.0813,
+    "saved_word_by_word": 0.0829
+   },
+   "articles/editing": {
+    "pauses": 400,
+    "shown": 0.18,
+    "right": 0.8889,
+    "avg_right_chars": 5.92,
+    "multiword_shown": 0.07,
+    "multiword_right": 0.8214,
+    "saved_with_tab": 0.079,
+    "saved_word_by_word": 0.0819
+   },
+   "messages/typing": {
+    "pauses": 400,
+    "shown": 0.31,
+    "right": 0.7823,
+    "avg_right_chars": 8.3,
+    "multiword_shown": 0.185,
+    "multiword_right": 0.7703,
+    "saved_with_tab": 0.1677,
+    "saved_word_by_word": 0.1773
+   },
+   "messages/editing": {
+    "pauses": 400,
+    "shown": 0.32,
+    "right": 0.7969,
+    "avg_right_chars": 8.48,
+    "multiword_shown": 0.1675,
+    "multiword_right": 0.8209,
+    "saved_with_tab": 0.1802,
+    "saved_word_by_word": 0.1913
+   },
+   "email/typing": {
+    "pauses": 400,
+    "shown": 0.1625,
+    "right": 0.7538,
+    "avg_right_chars": 5.61,
+    "multiword_shown": 0.0275,
+    "multiword_right": 0.8182,
+    "saved_with_tab": 0.0573,
+    "saved_word_by_word": 0.0585
+   },
+   "email/editing": {
+    "pauses": 400,
+    "shown": 0.1575,
+    "right": 0.7302,
+    "avg_right_chars": 5.24,
+    "multiword_shown": 0.045,
+    "multiword_right": 0.7778,
+    "saved_with_tab": 0.0502,
+    "saved_word_by_word": 0.0521
+   }
+  },
+  "0.6": {
+   "articles/typing": {
+    "pauses": 400,
+    "shown": 0.15,
+    "right": 0.9167,
+    "avg_right_chars": 6.11,
+    "multiword_shown": 0.0475,
+    "multiword_right": 0.8947,
+    "saved_with_tab": 0.07,
+    "saved_word_by_word": 0.071
+   },
+   "articles/editing": {
+    "pauses": 400,
+    "shown": 0.1525,
+    "right": 0.9344,
+    "avg_right_chars": 5.46,
+    "multiword_shown": 0.055,
+    "multiword_right": 0.9545,
+    "saved_with_tab": 0.0648,
+    "saved_word_by_word": 0.066
+   },
+   "messages/typing": {
+    "pauses": 400,
+    "shown": 0.2625,
+    "right": 0.8476,
+    "avg_right_chars": 7.99,
+    "multiword_shown": 0.145,
+    "multiword_right": 0.8103,
+    "saved_with_tab": 0.1481,
+    "saved_word_by_word": 0.1554
+   },
+   "messages/editing": {
+    "pauses": 400,
+    "shown": 0.2575,
+    "right": 0.8641,
+    "avg_right_chars": 8.08,
+    "multiword_shown": 0.1325,
+    "multiword_right": 0.9057,
+    "saved_with_tab": 0.1498,
+    "saved_word_by_word": 0.1544
+   },
+   "email/typing": {
+    "pauses": 400,
+    "shown": 0.125,
+    "right": 0.82,
+    "avg_right_chars": 5.44,
+    "multiword_shown": 0.0175,
+    "multiword_right": 0.8571,
+    "saved_with_tab": 0.0465,
+    "saved_word_by_word": 0.0471
+   },
+   "email/editing": {
+    "pauses": 400,
+    "shown": 0.11,
+    "right": 0.7955,
+    "avg_right_chars": 5.23,
+    "multiword_shown": 0.035,
+    "multiword_right": 0.7857,
+    "saved_with_tab": 0.0381,
+    "saved_word_by_word": 0.039
+   }
+  },
+  "0.7": {
+   "articles/typing": {
+    "pauses": 400,
+    "shown": 0.1175,
+    "right": 0.9149,
+    "avg_right_chars": 6.09,
+    "multiword_shown": 0.03,
+    "multiword_right": 0.9167,
+    "saved_with_tab": 0.0546,
+    "saved_word_by_word": 0.0548
+   },
+   "articles/editing": {
+    "pauses": 400,
+    "shown": 0.0925,
+    "right": 0.9459,
+    "avg_right_chars": 5.51,
+    "multiword_shown": 0.03,
+    "multiword_right": 0.9167,
+    "saved_with_tab": 0.0402,
+    "saved_word_by_word": 0.0415
+   },
+   "messages/typing": {
+    "pauses": 400,
+    "shown": 0.1925,
+    "right": 0.9351,
+    "avg_right_chars": 7.86,
+    "multiword_shown": 0.0925,
+    "multiword_right": 0.8919,
+    "saved_with_tab": 0.1179,
+    "saved_word_by_word": 0.1206
+   },
+   "messages/editing": {
+    "pauses": 400,
+    "shown": 0.1975,
+    "right": 0.9367,
+    "avg_right_chars": 7.42,
+    "multiword_shown": 0.0875,
+    "multiword_right": 0.9143,
+    "saved_with_tab": 0.1144,
+    "saved_word_by_word": 0.1175
+   },
+   "email/typing": {
+    "pauses": 400,
+    "shown": 0.095,
+    "right": 0.8684,
+    "avg_right_chars": 5.55,
+    "multiword_shown": 0.0125,
+    "multiword_right": 1.0,
+    "saved_with_tab": 0.0381,
+    "saved_word_by_word": 0.0381
+   },
+   "email/editing": {
+    "pauses": 400,
+    "shown": 0.075,
+    "right": 0.8,
+    "avg_right_chars": 5.25,
+    "multiword_shown": 0.0225,
+    "multiword_right": 0.7778,
+    "saved_with_tab": 0.0262,
+    "saved_word_by_word": 0.0271
+   }
+  }
+ },
+ "examples": [
+  {
+   "domain": "articles",
+   "scenario": "typing",
+   "before": "ROME KAVALKADE\". The segment featured acrobats, bareback rid",
+   "suggestion": "ers,",
+   "truth": "ers, clowns, tig",
+   "right": true
+  },
+  {
+   "domain": "articles",
+   "scenario": "typing",
+   "before": "Street/Perimeter Road, just south of the Miami International",
+   "suggestion": " Airport.",
+   "truth": " Airport runway.",
+   "right": false
+  },
+  {
+   "domain": "articles",
+   "scenario": "editing",
+   "before": "le tree, with a number of ecotypes, and in forms closely res",
+   "suggestion": "embling",
+   "truth": "embling other speci",
+   "right": true
+  },
+  {
+   "domain": "messages",
+   "scenario": "typing",
+   "before": " L'Hpital's Rule can be used to analyze the behavior of thes",
+   "suggestion": "e systems",
+   "truth": "e systems under speci",
+   "right": true
+  },
+  {
+   "domain": "messages",
+   "scenario": "typing",
+   "before": "f you know the formula for calculating the coefficient of fr",
+   "suggestion": "iction,",
+   "truth": "iction as a functio",
+   "right": false
+  },
+  {
+   "domain": "messages",
+   "scenario": "editing",
+   "before": "eleration. Set the acceleration equation equal to zero and s",
+   "suggestion": "olve for",
+   "truth": "olve for time. More ",
+   "right": true
+  },
+  {
+   "domain": "messages",
+   "scenario": "editing",
+   "before": "tion. Social media has allowed people to connect with each o",
+   "suggestion": "ther",
+   "truth": "ther and share i",
+   "right": true
+  },
+  {
+   "domain": "messages",
+   "scenario": "editing",
+   "before": "day, it is the capital and largest city of Germany, and a gl",
+   "suggestion": "obal",
+   "truth": "obal center for ",
+   "right": true
+  },
+  {
+   "domain": "messages",
+   "scenario": "editing",
+   "before": "d applications include:\n1. Physics: Calculus plays a vital r",
+   "suggestion": "ole in",
+   "truth": "ole in physics, an",
+   "right": true
+  },
+  {
+   "domain": "email",
+   "scenario": "typing",
+   "before": "ay the course with Continental plus I am confident that Cont",
+   "suggestion": "inental",
+   "truth": "inental can either ",
+   "right": true
+  },
+  {
+   "domain": "email",
+   "scenario": "editing",
+   "before": "ured products and marketing of electricity for the United St",
+   "suggestion": "ates",
+   "truth": "ates ? Actively ",
+   "right": true
+  },
+  {
+   "domain": "email",
+   "scenario": "editing",
+   "before": " action is necessary as contractually we are coming to the e",
+   "suggestion": "nd of",
+   "truth": "nd of the accepta",
+   "right": true
+  },
+  {
+   "domain": "email",
+   "scenario": "editing",
+   "before": "we should get something by the end of the year.\nSimilarly, C",
+   "suggestion": "harlie",
+   "truth": "ontinental has not",
+   "right": false
+  }
+ ],
+ "latency_ms": {
+  "p50": 128.9,
+  "p95": 437.94,
+  "note": "Measured in this eval process, not in a browser."
+ }
+}

v2/metrics/tinynext-v1.json

@@ -0,0 +1,339 @@
+{
+ "model": "tinynext-v1",
+ "sees_after": false,
+ "cutoffs": {
+  "0.3": {
+   "articles/typing": {
+    "pauses": 1000,
+    "shown": 0.11,
+    "right": 0.4909,
+    "avg_right_chars": 5.24,
+    "multiword_shown": 0.019,
+    "multiword_right": 0.3684,
+    "saved_with_tab": 0.0236,
+    "saved_word_by_word": 0.0252
+   },
+   "articles/editing": {
+    "pauses": 1000,
+    "shown": 0.081,
+    "right": 0.5309,
+    "avg_right_chars": 5.37,
+    "multiword_shown": 0.018,
+    "multiword_right": 0.6111,
+    "saved_with_tab": 0.0192,
+    "saved_word_by_word": 0.0199
+   },
+   "messages/typing": {
+    "pauses": 1000,
+    "shown": 0.137,
+    "right": 0.4891,
+    "avg_right_chars": 5.25,
+    "multiword_shown": 0.042,
+    "multiword_right": 0.5238,
+    "saved_with_tab": 0.0293,
+    "saved_word_by_word": 0.0313
+   },
+   "messages/editing": {
+    "pauses": 1000,
+    "shown": 0.141,
+    "right": 0.4539,
+    "avg_right_chars": 5.39,
+    "multiword_shown": 0.043,
+    "multiword_right": 0.4419,
+    "saved_with_tab": 0.0288,
+    "saved_word_by_word": 0.0316
+   },
+   "email/typing": {
+    "pauses": 1000,
+    "shown": 0.182,
+    "right": 0.5165,
+    "avg_right_chars": 5.76,
+    "multiword_shown": 0.042,
+    "multiword_right": 0.3333,
+    "saved_with_tab": 0.0451,
+    "saved_word_by_word": 0.0483
+   },
+   "email/editing": {
+    "pauses": 1000,
+    "shown": 0.188,
+    "right": 0.5372,
+    "avg_right_chars": 5.47,
+    "multiword_shown": 0.041,
+    "multiword_right": 0.4878,
+    "saved_with_tab": 0.046,
+    "saved_word_by_word": 0.0486
+   }
+  },
+  "0.4": {
+   "articles/typing": {
+    "pauses": 1000,
+    "shown": 0.074,
+    "right": 0.6081,
+    "avg_right_chars": 5.07,
+    "multiword_shown": 0.007,
+    "multiword_right": 0.4286,
+    "saved_with_tab": 0.019,
+    "saved_word_by_word": 0.0192
+   },
+   "articles/editing": {
+    "pauses": 1000,
+    "shown": 0.057,
+    "right": 0.5614,
+    "avg_right_chars": 5.34,
+    "multiword_shown": 0.007,
+    "multiword_right": 0.8571,
+    "saved_with_tab": 0.0143,
+    "saved_word_by_word": 0.0143
+   },
+   "messages/typing": {
+    "pauses": 1000,
+    "shown": 0.092,
+    "right": 0.6304,
+    "avg_right_chars": 5.14,
+    "multiword_shown": 0.026,
+    "multiword_right": 0.5769,
+    "saved_with_tab": 0.0248,
+    "saved_word_by_word": 0.0261
+   },
+   "messages/editing": {
+    "pauses": 1000,
+    "shown": 0.094,
+    "right": 0.5638,
+    "avg_right_chars": 5.26,
+    "multiword_shown": 0.019,
+    "multiword_right": 0.5789,
+    "saved_with_tab": 0.0232,
+    "saved_word_by_word": 0.0239
+   },
+   "email/typing": {
+    "pauses": 1000,
+    "shown": 0.13,
+    "right": 0.6,
+    "avg_right_chars": 5.63,
+    "multiword_shown": 0.024,
+    "multiword_right": 0.4583,
+    "saved_with_tab": 0.0366,
+    "saved_word_by_word": 0.038
+   },
+   "email/editing": {
+    "pauses": 1000,
+    "shown": 0.143,
+    "right": 0.6224,
+    "avg_right_chars": 5.39,
+    "multiword_shown": 0.023,
+    "multiword_right": 0.6957,
+    "saved_with_tab": 0.04,
+    "saved_word_by_word": 0.0408
+   }
+  },
+  "0.5": {
+   "articles/typing": {
+    "pauses": 1000,
+    "shown": 0.051,
+    "right": 0.7059,
+    "avg_right_chars": 5.06,
+    "multiword_shown": 0.002,
+    "multiword_right": 0.5,
+    "saved_with_tab": 0.0152,
+    "saved_word_by_word": 0.0152
+   },
+   "articles/editing": {
+    "pauses": 1000,
+    "shown": 0.037,
+    "right": 0.6486,
+    "avg_right_chars": 5.12,
+    "multiword_shown": 0.005,
+    "multiword_right": 1.0,
+    "saved_with_tab": 0.0103,
+    "saved_word_by_word": 0.0103
+   },
+   "messages/typing": {
+    "pauses": 1000,
+    "shown": 0.062,
+    "right": 0.7581,
+    "avg_right_chars": 5.04,
+    "multiword_shown": 0.01,
+    "multiword_right": 0.8,
+    "saved_with_tab": 0.0198,
+    "saved_word_by_word": 0.0198
+   },
+   "messages/editing": {
+    "pauses": 1000,
+    "shown": 0.071,
+    "right": 0.6338,
+    "avg_right_chars": 5.16,
+    "multiword_shown": 0.009,
+    "multiword_right": 0.5556,
+    "saved_with_tab": 0.0193,
+    "saved_word_by_word": 0.0197
+   },
+   "email/typing": {
+    "pauses": 1000,
+    "shown": 0.093,
+    "right": 0.6774,
+    "avg_right_chars": 5.68,
+    "multiword_shown": 0.016,
+    "multiword_right": 0.4375,
+    "saved_with_tab": 0.0298,
+    "saved_word_by_word": 0.0305
+   },
+   "email/editing": {
+    "pauses": 1000,
+    "shown": 0.098,
+    "right": 0.6939,
+    "avg_right_chars": 5.22,
+    "multiword_shown": 0.009,
+    "multiword_right": 0.6667,
+    "saved_with_tab": 0.0296,
+    "saved_word_by_word": 0.0297
+   }
+  },
+  "0.6": {
+   "articles/typing": {
+    "pauses": 1000,
+    "shown": 0.036,
+    "right": 0.8056,
+    "avg_right_chars": 5.07,
+    "multiword_shown": 0.002,
+    "multiword_right": 0.5,
+    "saved_with_tab": 0.0123,
+    "saved_word_by_word": 0.0123
+   },
+   "articles/editing": {
+    "pauses": 1000,
+    "shown": 0.024,
+    "right": 0.6667,
+    "avg_right_chars": 5.0,
+    "multiword_shown": 0.003,
+    "multiword_right": 1.0,
+    "saved_with_tab": 0.0067,
+    "saved_word_by_word": 0.0067
+   },
+   "messages/typing": {
+    "pauses": 1000,
+    "shown": 0.041,
+    "right": 0.7805,
+    "avg_right_chars": 4.94,
+    "multiword_shown": 0.007,
+    "multiword_right": 0.8571,
+    "saved_with_tab": 0.0132,
+    "saved_word_by_word": 0.0132
+   },
+   "messages/editing": {
+    "pauses": 1000,
+    "shown": 0.045,
+    "right": 0.6667,
+    "avg_right_chars": 5.0,
+    "multiword_shown": 0.003,
+    "multiword_right": 0.3333,
+    "saved_with_tab": 0.0125,
+    "saved_word_by_word": 0.0125
+   },
+   "email/typing": {
+    "pauses": 1000,
+    "shown": 0.056,
+    "right": 0.75,
+    "avg_right_chars": 4.9,
+    "multiword_shown": 0.004,
+    "multiword_right": 0.0,
+    "saved_with_tab": 0.0172,
+    "saved_word_by_word": 0.0174
+   },
+   "email/editing": {
+    "pauses": 1000,
+    "shown": 0.067,
+    "right": 0.7313,
+    "avg_right_chars": 4.86,
+    "multiword_shown": 0.004,
+    "multiword_right": 0.75,
+    "saved_with_tab": 0.0198,
+    "saved_word_by_word": 0.0198
+   }
+  },
+  "0.7": {
+   "articles/typing": {
+    "pauses": 1000,
+    "shown": 0.025,
+    "right": 0.84,
+    "avg_right_chars": 5.1,
+    "multiword_shown": 0.001,
+    "multiword_right": 1.0,
+    "saved_with_tab": 0.0089,
+    "saved_word_by_word": 0.0089
+   },
+   "articles/editing": {
+    "pauses": 1000,
+    "shown": 0.018,
+    "right": 0.6667,
+    "avg_right_chars": 5.0,
+    "multiword_shown": 0.001,
+    "multiword_right": 1.0,
+    "saved_with_tab": 0.005,
+    "saved_word_by_word": 0.005
+   },
+   "messages/typing": {
+    "pauses": 1000,
+    "shown": 0.025,
+    "right": 0.84,
+    "avg_right_chars": 4.9,
+    "multiword_shown": 0.005,
+    "multiword_right": 0.8,
+    "saved_with_tab": 0.0086,
+    "saved_word_by_word": 0.0086
+   },
+   "messages/editing": {
+    "pauses": 1000,
+    "shown": 0.036,
+    "right": 0.6944,
+    "avg_right_chars": 4.96,
+    "multiword_shown": 0.002,
+    "multiword_right": 0.5,
+    "saved_with_tab": 0.0103,
+    "saved_word_by_word": 0.0103
+   },
+   "email/typing": {
+    "pauses": 1000,
+    "shown": 0.042,
+    "right": 0.8333,
+    "avg_right_chars": 4.89,
+    "multiword_shown": 0.001,
+    "multiword_right": 0.0,
+    "saved_with_tab": 0.0143,
+    "saved_word_by_word": 0.0143
+   },
+   "email/editing": {
+    "pauses": 1000,
+    "shown": 0.047,
+    "right": 0.8298,
+    "avg_right_chars": 4.95,
+    "multiword_shown": 0.002,
+    "multiword_right": 1.0,
+    "saved_with_tab": 0.0161,
+    "saved_word_by_word": 0.0161
+   }
+  }
+ },
+ "examples": [
+  {
+   "domain": "messages",
+   "scenario": "typing",
+   "before": "onsidered new. Those decisions make the graphs as neat as po",
+   "suggestion": "ssible",
+   "truth": "ssible, and it is ",
+   "right": true
+  },
+  {
+   "domain": "email",
+   "scenario": "editing",
+   "before": "ured products and marketing of electricity for the United St",
+   "suggestion": "ates",
+   "truth": "ates ? Actively ",
+   "right": true
+  }
+ ],
+ "latency_ms": {
+  "p50": 4.22,
+  "p95": 16.52,
+  "note": "Measured in this eval process, not in a browser."
+ }
+}

v2/train.py

@@ -0,0 +1,460 @@
+# /// script
+# requires-python = ">=3.11,<3.14"
+# dependencies = [
+#   "torch>=2.6,<2.8",
+#   "datasets>=3.0",
+#   "numpy<3",
+#   "tokenizers>=0.20",
+#   "transformers>=4.46",
+#   "onnx>=1.17",
+#   "onnxruntime>=1.20",
+#   "huggingface_hub>=0.30",
+# ]
+# ///
+"""tinynext v2: a small word-piece transformer for in-browser autocomplete that can also fill in the middle.
+
+One script does everything: it trains a tokenizer, builds a mix of general text, conversation and email,
+trains the model, scores it with eval.py (the same harness as tinynext v1 and SmolLM2), exports ONNX with a
+key/value cache for the browser, and uploads the lot to the model repo under v2/.
+
+  hf jobs uv run --flavor a100-large --secrets HF_TOKEN train.py
+  uv run train.py --smoke --no-push      # a few minutes on a laptop, to check the whole pipeline
+"""
+
+import argparse
+import importlib.util
+import json
+import math
+import os
+import random
+import shutil
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+HERE = Path(__file__).parent
+SPECIAL = ["<|endoftext|>", "<|pre|>", "<|suf|>", "<|mid|>"]
+EOT, PRE, SUF, MID = range(4)
+
+# ---- Shared eval code ------------------------------------------------------------------------------------------
+
+
+def load_eval_module(repo: str, folder: str):
+    """Uses eval.py and eval_set.json next to this script, or fetches them from the repo, so every model is scored
+    by one implementation on the same text."""
+    path = HERE / "eval.py"
+    if not path.exists():
+        from huggingface_hub import hf_hub_download
+
+        for name in ("eval.py", "eval_set.json"):
+            shutil.copy(hf_hub_download(repo, f"{folder}/{name}"), HERE / name)
+    spec = importlib.util.spec_from_file_location("tinynext_eval", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+# ---- Data ------------------------------------------------------------------------------------------------------
+
+
+def fineweb_docs():
+    from datasets import load_dataset
+
+    for row in load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT", split="train", streaming=True):
+        yield row["text"]
+
+
+def smoltalk_docs():
+    from datasets import load_dataset
+
+    for row in load_dataset("HuggingFaceTB/smoltalk", "all", split="train", streaming=True):
+        turns = [m["content"] for m in row["messages"] if m["role"] in ("user", "assistant")]
+        if turns:
+            yield "\n\n".join(turns)
+
+
+def enron_docs(ev):
+    """Training senders only: the eval's held-out senders (and enron-autocomplete's val split) are skipped."""
+    import hashlib
+
+    from datasets import load_dataset
+
+    for row in load_dataset("corbt/enron-emails", split="train", streaming=True):
+        sender = (row["from"] or "").lower()
+        if "@enron.com" not in sender or ev.BULK_SENDER.search(sender):
+            continue
+        if int(hashlib.sha1(sender.encode()).hexdigest(), 16) % 100 < 10:
+            continue
+        paragraphs = ev.clean_email(row["body"] or "")
+        if paragraphs:
+            yield "\n\n".join(paragraphs)
+
+
+def mixed_docs(ev, weights: dict[str, float], seed: int):
+    """Interleaves the sources by weight, forever (until the caller stops)."""
+    rng = random.Random(seed)
+    sources = {"fineweb": fineweb_docs(), "smoltalk": smoltalk_docs(), "enron": enron_docs(ev)}
+    names, probs = zip(*weights.items())
+    counts = dict.fromkeys(names, 0)
+    while True:
+        name = rng.choices(names, probs)[0]
+        try:
+            doc = next(sources[name])
+        except StopIteration:
+            sources[name] = {"fineweb": fineweb_docs, "smoltalk": smoltalk_docs}.get(name, lambda: enron_docs(ev))()
+            continue
+        counts[name] += 1
+        # Log the first document from each source and then every 50k, so a stalled download shows which one.
+        if counts[name] == 1 or counts[name] % 50_000 == 0:
+            print(f"  data seed {seed}: {name} doc {counts[name]:,}", flush=True)
+        yield doc
+
+
+def clean(text: str, ev) -> str:
+    """The browser model only sees keyboard characters, so train on the same."""
+    return text.translate(ev.SMART_CHARS).encode("ascii", "ignore").decode().strip()
+
+
+def train_tokenizer(docs, vocab: int, sample_docs: int):
+    from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers
+
+    tokenizer = Tokenizer(models.BPE())
+    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
+    tokenizer.decoder = decoders.ByteLevel()
+    trainer = trainers.BpeTrainer(vocab_size=vocab, special_tokens=SPECIAL, min_frequency=2,
+                                  initial_alphabet=pre_tokenizers.ByteLevel.alphabet())
+    tokenizer.train_from_iterator((next(docs) for _ in range(sample_docs)), trainer=trainer)
+    return tokenizer
+
+
+def fim(ids: list[int], rng: random.Random, max_middle: int, max_suffix: int, max_prefix: int) -> list[int]:
+    """Suffix first, then prefix, then the middle, so typing at the cursor only appends to the prompt."""
+    if len(ids) < 8:
+        return ids
+    start = rng.randint(1, len(ids) - 2)
+    end = min(len(ids), start + rng.randint(1, max_middle))
+    prefix, middle, suffix = ids[max(0, start - max_prefix):start], ids[start:end], ids[end:end + max_suffix]
+    if rng.random() < 0.25:
+        suffix = []  # also learn the plain "nothing after the cursor" case in this format
+    return [SUF, *suffix, PRE, *prefix, MID, *middle, EOT]
+
+
+def build_tokens(tokenizer, docs, ev, total: int, seq: int, fim_rate: float, seed: int) -> np.ndarray:
+    """Packs documents into rows of seq+1 tokens. A share of chunks become fill-in-the-middle examples."""
+    rng = random.Random(seed)
+    out = np.empty(total + seq + 1, dtype=np.uint16)
+    filled, started, next_report = 0, time.time(), 0
+    while filled < total:
+        batch = [clean(next(docs), ev) for _ in range(512)]
+        for ids in tokenizer.encode_batch([b for b in batch if b]):
+            ids = ids.ids
+            for at in range(0, len(ids), seq - 8):
+                chunk = ids[at: at + seq - 8]
+                piece = fim(chunk, rng, 24, 64, seq - 100) if rng.random() < fim_rate else chunk + [EOT]
+                take = min(len(piece), len(out) - filled)
+                out[filled: filled + take] = piece[:take]
+                filled += take
+                if filled >= total:
+                    break
+            if filled >= total:
+                break
+        if filled >= next_report:
+            print(f"  tokens {filled / 1e6:.0f}M / {total / 1e6:.0f}M ({time.time() - started:.0f}s)", flush=True)
+            next_report += 20_000_000
+    rows = total // (seq + 1)
+    return out[: rows * (seq + 1)].reshape(rows, seq + 1)
+
+
+# ---- Model -----------------------------------------------------------------------------------------------------
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * self.weight
+
+
+def rotate(x, cos, sin):
+    half = x.shape[-1] // 2
+    x1, x2 = x[..., :half], x[..., half:]
+    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
+
+
+class Block(nn.Module):
+    def __init__(self, dim: int, heads: int, ffn: int):
+        super().__init__()
+        self.heads, self.head_dim = heads, dim // heads
+        self.norm1, self.norm2 = RMSNorm(dim), RMSNorm(dim)
+        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
+        self.out = nn.Linear(dim, dim, bias=False)
+        self.gate_up = nn.Linear(dim, 2 * ffn, bias=False)
+        self.down = nn.Linear(ffn, dim, bias=False)
+
+    def forward(self, x, cos, sin, past_k, past_v, mask):
+        batch, length, _ = x.shape
+        q, k, v = self.qkv(self.norm1(x)).view(batch, length, 3, self.heads, self.head_dim).unbind(2)
+        q, k, v = rotate(q.transpose(1, 2), cos, sin), rotate(k.transpose(1, 2), cos, sin), v.transpose(1, 2)
+        if past_k is not None:
+            k, v = torch.cat([past_k, k], dim=2), torch.cat([past_v, v], dim=2)
+        if mask is None:
+            attended = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        else:
+            attended = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        x = x + self.out(attended.transpose(1, 2).reshape(batch, length, -1))
+        gate, up = self.gate_up(self.norm2(x)).chunk(2, dim=-1)
+        return x + self.down(F.silu(gate) * up), k, v
+
+
+class TinyNext(nn.Module):
+    def __init__(self, vocab: int, dim: int, layers: int, heads: int, ffn: int):
+        super().__init__()
+        self.config = {"vocab": vocab, "dim": dim, "layers": layers, "heads": heads, "ffn": ffn}
+        self.embed = nn.Embedding(vocab, dim)
+        self.blocks = nn.ModuleList(Block(dim, heads, ffn) for _ in range(layers))
+        self.norm = RMSNorm(dim)
+        self.head_dim = dim // heads
+        inverse = 1.0 / (10000 ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim))
+        self.register_buffer("inverse_frequency", inverse, persistent=False)
+        nn.init.normal_(self.embed.weight, std=0.02)
+
+    def forward(self, ids, past_k=None, past_v=None):
+        """ids: [batch, length]. past_k/past_v: [layers, batch, heads, past, head_dim] or None.
+        Returns logits for every position and the updated cache."""
+        past = 0 if past_k is None else past_k.shape[3]
+        length = ids.shape[1]
+        positions = torch.arange(length, device=ids.device) + past
+        angles = positions[:, None].float() * self.inverse_frequency[None, :]
+        cos, sin = angles.cos()[None, None], angles.sin()[None, None]
+        mask = None
+        if past_k is not None:
+            keys = torch.arange(past + length, device=ids.device)
+            mask = keys[None, :] <= positions[:, None]
+        x = self.embed(ids)
+        new_k, new_v = [], []
+        for index, block in enumerate(self.blocks):
+            x, k, v = block(x, cos, sin, None if past_k is None else past_k[index], None if past_v is None else past_v[index], mask)
+            new_k.append(k)
+            new_v.append(v)
+        logits = self.norm(x) @ self.embed.weight.T
+        return logits, torch.stack(new_k), torch.stack(new_v)
+
+
+class BrowserStep(nn.Module):
+    """The ONNX graph: always takes a cache (possibly empty) and returns only the last position's logits."""
+
+    def __init__(self, model: TinyNext):
+        super().__init__()
+        self.model = model
+
+    def forward(self, ids, past_k, past_v):
+        logits, k, v = self.model(ids, past_k, past_v)
+        return logits[:, -1], k, v
+
+
+# ---- Training --------------------------------------------------------------------------------------------------
+
+
+def evaluate_loss(model, rows: np.ndarray, device, batch: int) -> float:
+    model.eval()
+    losses = []
+    with torch.no_grad(), torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"):
+        for at in range(0, len(rows), batch):
+            chunk = torch.from_numpy(rows[at: at + batch].astype(np.int64)).to(device)
+            logits, _, _ = model(chunk[:, :-1])
+            losses.append(F.cross_entropy(logits.float().flatten(0, 1), chunk[:, 1:].flatten()).item())
+    model.train()
+    return float(np.mean(losses))
+
+
+def train(model, rows: np.ndarray, val_rows: np.ndarray, args, device) -> list[dict]:
+    steps = args.epochs * len(rows) // args.batch
+    decay = [p for n, p in model.named_parameters() if p.dim() >= 2]
+    no_decay = [p for n, p in model.named_parameters() if p.dim() < 2]
+    optimizer = torch.optim.AdamW([{"params": decay, "weight_decay": 0.1}, {"params": no_decay, "weight_decay": 0.0}],
+                                  lr=args.lr, betas=(0.9, 0.95), fused=device.type == "cuda")
+    warmup = max(1, steps // 50)
+    schedule = lambda step: min(1, step / warmup) * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * min(1, step / steps))))
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, schedule)
+    compiled = torch.compile(model) if device.type == "cuda" else model
+    history, rng, started = [], np.random.default_rng(0), time.time()
+    for step in range(steps):
+        index = rng.integers(0, len(rows), args.batch)
+        chunk = torch.from_numpy(rows[index].astype(np.int64)).to(device, non_blocking=True)
+        with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"):
+            logits, _, _ = compiled(chunk[:, :-1])
+            loss = F.cross_entropy(logits.float().flatten(0, 1), chunk[:, 1:].flatten())
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        optimizer.step()
+        scheduler.step()
+        if step % args.log_every == 0 or step == steps - 1:
+            entry = {"step": step, "train_loss": round(loss.item(), 4), "minutes": round((time.time() - started) / 60, 2)}
+            if step % (args.log_every * 10) == 0 or step == steps - 1:
+                entry["val_loss"] = round(evaluate_loss(model, val_rows, device, args.batch), 4)
+            history.append(entry)
+            print(json.dumps(entry), flush=True)
+    return history
+
+
+# ---- Eval, export --------------------------------------------------------------------------------------------
+
+
+def as_token_lm(ev, name: str, model: TinyNext, tokenizer, device):
+    texts = [tokenizer.decode([i], skip_special_tokens=False) if i >= len(SPECIAL) else "" for i in range(tokenizer.get_vocab_size())]
+
+    @torch.no_grad()
+    def next_logits(ids, cache):
+        past_k, past_v = cache if cache else (None, None)
+        logits, k, v = model(torch.tensor([ids], device=device), past_k, past_v)
+        return logits[0, -1].float().cpu().numpy(), (k, v)
+
+    encode = lambda text: tokenizer.encode(text).ids or [EOT]
+    prompt = lambda prefix, suffix: prefix if not suffix else [SUF, *suffix, PRE, *prefix, MID]
+    return ev.TokenLM(name, next_logits, encode, texts, set(range(len(SPECIAL))), prompt=prompt)
+
+
+def export(model: TinyNext, output: Path, device) -> dict:
+    import onnxruntime as ort
+    from onnxruntime.quantization import QuantType, quantize_dynamic
+
+    model = model.float().eval().to("cpu")
+    cfg = model.config
+    head_dim = cfg["dim"] // cfg["heads"]
+    ids = torch.tensor([[5, 6, 7]], dtype=torch.long)
+    empty = torch.zeros(cfg["layers"], 1, cfg["heads"], 0, head_dim)
+    # Trace with a non-empty cache so no shape gets fixed at zero; the checks below run with an empty one.
+    with torch.no_grad():
+        _, trace_k, trace_v = model(torch.tensor([[1, 2]]))
+    float_path, int8_path = output / "model.onnx", output / "model_int8.onnx"
+    torch.onnx.export(
+        BrowserStep(model), (ids, trace_k, trace_v), float_path, opset_version=17,
+        input_names=["input_ids", "past_k", "past_v"], output_names=["logits", "present_k", "present_v"],
+        dynamic_axes={"input_ids": {1: "length"}, "past_k": {3: "past"}, "past_v": {3: "past"},
+                      "present_k": {3: "total"}, "present_v": {3: "total"}},
+        dynamo=False,
+    )
+    quantize_dynamic(float_path, int8_path, weight_type=QuantType.QInt8, op_types_to_quantize=["MatMul"])
+
+    report = {}
+    with torch.no_grad():
+        expected, k, v = model(ids)
+        expected_next, _, _ = model(torch.tensor([[8]]), k, v)
+    for name, path in (("float", float_path), ("int8", int8_path)):
+        session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
+        logits, pk, pv = session.run(None, {"input_ids": ids.numpy(), "past_k": empty.numpy(), "past_v": empty.numpy()})
+        logits2, _, _ = session.run(None, {"input_ids": np.array([[8]]), "past_k": pk, "past_v": pv})
+        timings = []
+        for _ in range(200):
+            began = time.perf_counter()
+            session.run(None, {"input_ids": np.array([[8]]), "past_k": pk, "past_v": pv})
+            timings.append((time.perf_counter() - began) * 1000)
+        report[name] = {
+            "megabytes": round(path.stat().st_size / 1e6, 2),
+            "same_top_token": bool(logits.argmax() == expected[0, -1].argmax() and logits2.argmax() == expected_next[0, -1].argmax()),
+            "cached_step_ms_p50_cpu": round(float(np.median(timings)), 3),
+        }
+    model.to(device)
+    return report
+
+
+# ---- Main ------------------------------------------------------------------------------------------------------
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--vocab", type=int, default=8192)
+    parser.add_argument("--dim", type=int, default=320)
+    parser.add_argument("--layers", type=int, default=8)
+    parser.add_argument("--heads", type=int, default=5)
+    parser.add_argument("--ffn", type=int, default=864)
+    parser.add_argument("--seq", type=int, default=256)
+    parser.add_argument("--tokens", type=float, default=800e6, help="Training tokens to build")
+    parser.add_argument("--epochs", type=float, default=1.0)
+    parser.add_argument("--batch", type=int, default=256)
+    parser.add_argument("--lr", type=float, default=3e-3)
+    parser.add_argument("--fim-rate", type=float, default=0.5)
+    parser.add_argument("--mix", default="fineweb=0.6,smoltalk=0.3,enron=0.1")
+    parser.add_argument("--tokenizer-docs", type=int, default=200_000)
+    parser.add_argument("--eval-pauses", type=int, default=1000)
+    parser.add_argument("--log-every", type=int, default=100)
+    parser.add_argument("--repo", default="cfahlgren1/tinynext")
+    parser.add_argument("--folder", default="v2")
+    parser.add_argument("--name", default="tinynext-v2")
+    parser.add_argument("--smoke", action="store_true", help="Tiny settings to check the pipeline end to end")
+    parser.add_argument("--no-push", action="store_true")
+    args = parser.parse_args()
+    if args.smoke:
+        args.vocab, args.dim, args.layers, args.heads, args.ffn = 2048, 64, 2, 2, 128
+        args.tokens, args.batch, args.tokenizer_docs, args.eval_pauses, args.log_every = 3e6, 32, 3000, 20, 20
+        args.name = "tinynext-v2-smoke"
+    return args
+
+
+def main():
+    args = parse_args()
+    device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
+    torch.manual_seed(0)
+    output = HERE / "output"
+    output.mkdir(exist_ok=True)
+    ev = load_eval_module(args.repo, args.folder)
+    weights = {k: float(v) for k, v in (part.split("=") for part in args.mix.split(","))}
+    print(f"device {device}, mix {weights}", flush=True)
+
+    tokenizer = train_tokenizer((clean(d, ev) for d in mixed_docs(ev, weights, seed=1)), args.vocab, args.tokenizer_docs)
+    tokenizer.save(str(output / "tokenizer.json"))
+    # Validation takes the first documents of the stream and training continues after them, so they never overlap.
+    docs = mixed_docs(ev, weights, seed=2)
+    val_rows = build_tokens(tokenizer, docs, ev, int(min(2e6, args.tokens / 20)), args.seq, args.fim_rate, seed=5)
+    rows = build_tokens(tokenizer, docs, ev, int(args.tokens), args.seq, args.fim_rate, seed=3)
+
+    model = TinyNext(tokenizer.get_vocab_size(), args.dim, args.layers, args.heads, args.ffn).to(device)
+    parameters = sum(p.numel() for p in model.parameters())
+    print(f"{parameters / 1e6:.2f}M parameters, {rows.size / 1e6:.0f}M tokens", flush=True)
+    started = time.time()
+    history = train(model, rows, val_rows, args, device)
+    minutes = (time.time() - started) / 60
+
+    model.eval()
+    torch.save({"config": model.config, "state_dict": model.state_dict()}, output / "model.pt")
+    scored = ev.score(as_token_lm(ev, args.name, model, tokenizer, device), ev.pauses(ev.load_eval_set(), args.eval_pauses))
+    (output / "metrics").mkdir(exist_ok=True)
+    (output / "metrics" / f"{args.name}.json").write_text(json.dumps(scored, indent=1))
+    exported = export(model, output, device)
+
+    config = {
+        **model.config, "parameters": parameters, "seq": args.seq, "special_tokens": dict(zip(SPECIAL, range(len(SPECIAL)))),
+        "prompt": "No text after the cursor: the text before it. Otherwise: <|suf|> after <|pre|> before <|mid|>.",
+    }
+    (output / "config.json").write_text(json.dumps(config, indent=1))
+    report = {"args": vars(args), "parameters": parameters, "train_minutes": round(minutes, 1), "history": history, "exports": exported}
+    (output / "report.json").write_text(json.dumps(report, indent=1))
+    print(json.dumps(exported, indent=1))
+    print(ev.table([output / "metrics" / f"{args.name}.json"]))
+
+    for name in ("train.py", "eval.py", "eval_set.json"):
+        shutil.copy(HERE / name, output / name)
+    if (HERE / "metrics").exists():
+        for path in (HERE / "metrics").glob("*.json"):
+            shutil.copy(path, output / "metrics" / path.name)
+    if args.no_push:
+        print(f"Skipping upload; everything is in {output}")
+        return
+    from huggingface_hub import HfApi
+
+    HfApi(token=os.environ.get("HF_TOKEN")).upload_folder(
+        folder_path=str(output), repo_id=args.repo, path_in_repo=args.folder, repo_type="model",
+        commit_message=f"{args.name}: {parameters / 1e6:.1f}M params, trained {minutes:.0f} min",
+    )
+    print(f"Uploaded to https://huggingface.co/{args.repo}/tree/main/{args.folder}")
+
+
+if __name__ == "__main__":
+    main()