tinynext-v2-nofim: 12.5M params, 28 min on GPU

v2/experiments/tinynext-v2-nofim/config.json

@@ -0,0 +1,16 @@
+{
+ "vocab": 8192,
+ "dim": 320,
+ "layers": 8,
+ "heads": 5,
+ "ffn": 864,
+ "parameters": 12539200,
+ "seq": 512,
+ "special_tokens": {
+  "<|endoftext|>": 0,
+  "<|pre|>": 1,
+  "<|suf|>": 2,
+  "<|mid|>": 3
+ },
+ "prompt": "No text after the cursor: the text before it. Otherwise: <|suf|> after <|pre|> before <|mid|>."
+}

v2/experiments/tinynext-v2-nofim/eval.py

@@ -0,0 +1,499 @@
+# /// 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)
+
+
+def onnx_token_lm(name: str, folder: str, fim: bool) -> TokenLM:
+    """A tinynext v2 export: graph.onnx + weights.f32 (compressed, expanded) or model.onnx, plus tokenizer.json."""
+    import onnxruntime as ort
+    from tokenizers import Tokenizer
+
+    folder = Path(folder)
+    graph = folder / "graph.onnx" if (folder / "graph.onnx").exists() else folder / "model.onnx"
+    session = ort.InferenceSession(str(graph), providers=["CPUExecutionProvider"])
+    config = json.loads((folder / "config.json").read_text())
+    tokenizer = Tokenizer.from_file(str(folder / "tokenizer.json"))
+    empty = np.zeros((config["layers"], 1, config["heads"], 0, config["dim"] // config["heads"]), np.float32)
+
+    def next_logits(ids, cache):
+        past_k, past_v = cache if cache else (empty, empty)
+        logits, k, v = session.run(None, {"input_ids": np.array([ids], np.int64), "past_k": past_k, "past_v": past_v})
+        return logits[0], (k, v)
+
+    special = len(config["special_tokens"])
+    texts = [tokenizer.decode([i], skip_special_tokens=False) if i >= special else "" for i in range(config["vocab"])]
+    encode = lambda text: tokenizer.encode(text).ids or [0]
+    prompt = (lambda prefix, suffix: prefix if not suffix else [2, *suffix, 1, *prefix, 3]) if fim else None
+    return TokenLM(name, next_logits, encode, texts, set(range(special)), prompt=prompt)
+
+
+# ---- 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("--onnx-lm", help="Folder with a tinynext v2 ONNX export, config.json and tokenizer.json")
+    parser.add_argument("--no-fim", action="store_true", help="With --onnx-lm: the model wasn't trained to fill in the middle")
+    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)
+    elif args.onnx_lm:
+        model = onnx_token_lm(args.model, args.onnx_lm, fim=not args.no_fim)
+    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/experiments/tinynext-v2-nofim/metrics/tinynext-v2-nofim.json

@@ -0,0 +1,363 @@
+{
+ "model": "tinynext-v2-nofim",
+ "sees_after": false,
+ "cutoffs": {
+  "0.3": {
+   "articles/typing": {
+    "pauses": 1000,
+    "shown": 0.191,
+    "right": 0.5969,
+    "avg_right_chars": 5.94,
+    "multiword_shown": 0.069,
+    "multiword_right": 0.5217,
+    "saved_with_tab": 0.0564,
+    "saved_word_by_word": 0.064
+   },
+   "articles/editing": {
+    "pauses": 1000,
+    "shown": 0.163,
+    "right": 0.6074,
+    "avg_right_chars": 6.13,
+    "multiword_shown": 0.07,
+    "multiword_right": 0.6,
+    "saved_with_tab": 0.0506,
+    "saved_word_by_word": 0.0554
+   },
+   "messages/typing": {
+    "pauses": 1000,
+    "shown": 0.3,
+    "right": 0.58,
+    "avg_right_chars": 7.48,
+    "multiword_shown": 0.18,
+    "multiword_right": 0.5611,
+    "saved_with_tab": 0.1084,
+    "saved_word_by_word": 0.1323
+   },
+   "messages/editing": {
+    "pauses": 1000,
+    "shown": 0.3,
+    "right": 0.5767,
+    "avg_right_chars": 7.34,
+    "multiword_shown": 0.182,
+    "multiword_right": 0.5275,
+    "saved_with_tab": 0.1057,
+    "saved_word_by_word": 0.1288
+   },
+   "email/typing": {
+    "pauses": 1000,
+    "shown": 0.24,
+    "right": 0.5292,
+    "avg_right_chars": 6.03,
+    "multiword_shown": 0.08,
+    "multiword_right": 0.45,
+    "saved_with_tab": 0.0638,
+    "saved_word_by_word": 0.0718
+   },
+   "email/editing": {
+    "pauses": 1000,
+    "shown": 0.237,
+    "right": 0.557,
+    "avg_right_chars": 5.95,
+    "multiword_shown": 0.091,
+    "multiword_right": 0.5385,
+    "saved_with_tab": 0.0655,
+    "saved_word_by_word": 0.0717
+   }
+  },
+  "0.4": {
+   "articles/typing": {
+    "pauses": 1000,
+    "shown": 0.133,
+    "right": 0.7293,
+    "avg_right_chars": 5.86,
+    "multiword_shown": 0.043,
+    "multiword_right": 0.7209,
+    "saved_with_tab": 0.0473,
+    "saved_word_by_word": 0.0497
+   },
+   "articles/editing": {
+    "pauses": 1000,
+    "shown": 0.114,
+    "right": 0.7281,
+    "avg_right_chars": 5.9,
+    "multiword_shown": 0.042,
+    "multiword_right": 0.7143,
+    "saved_with_tab": 0.0408,
+    "saved_word_by_word": 0.0436
+   },
+   "messages/typing": {
+    "pauses": 1000,
+    "shown": 0.239,
+    "right": 0.7071,
+    "avg_right_chars": 7.33,
+    "multiword_shown": 0.125,
+    "multiword_right": 0.712,
+    "saved_with_tab": 0.1032,
+    "saved_word_by_word": 0.111
+   },
+   "messages/editing": {
+    "pauses": 1000,
+    "shown": 0.233,
+    "right": 0.6567,
+    "avg_right_chars": 7.03,
+    "multiword_shown": 0.126,
+    "multiword_right": 0.6032,
+    "saved_with_tab": 0.0896,
+    "saved_word_by_word": 0.1005
+   },
+   "email/typing": {
+    "pauses": 1000,
+    "shown": 0.158,
+    "right": 0.6392,
+    "avg_right_chars": 6.06,
+    "multiword_shown": 0.043,
+    "multiword_right": 0.5581,
+    "saved_with_tab": 0.051,
+    "saved_word_by_word": 0.0548
+   },
+   "email/editing": {
+    "pauses": 1000,
+    "shown": 0.166,
+    "right": 0.6807,
+    "avg_right_chars": 5.77,
+    "multiword_shown": 0.046,
+    "multiword_right": 0.6957,
+    "saved_with_tab": 0.0543,
+    "saved_word_by_word": 0.0568
+   }
+  },
+  "0.5": {
+   "articles/typing": {
+    "pauses": 1000,
+    "shown": 0.1,
+    "right": 0.81,
+    "avg_right_chars": 5.74,
+    "multiword_shown": 0.03,
+    "multiword_right": 0.7667,
+    "saved_with_tab": 0.0387,
+    "saved_word_by_word": 0.0402
+   },
+   "articles/editing": {
+    "pauses": 1000,
+    "shown": 0.084,
+    "right": 0.7738,
+    "avg_right_chars": 5.6,
+    "multiword_shown": 0.028,
+    "multiword_right": 0.7143,
+    "saved_with_tab": 0.0303,
+    "saved_word_by_word": 0.0322
+   },
+   "messages/typing": {
+    "pauses": 1000,
+    "shown": 0.182,
+    "right": 0.8022,
+    "avg_right_chars": 7.25,
+    "multiword_shown": 0.084,
+    "multiword_right": 0.7976,
+    "saved_with_tab": 0.0882,
+    "saved_word_by_word": 0.0916
+   },
+   "messages/editing": {
+    "pauses": 1000,
+    "shown": 0.17,
+    "right": 0.7176,
+    "avg_right_chars": 6.76,
+    "multiword_shown": 0.088,
+    "multiword_right": 0.6364,
+    "saved_with_tab": 0.0688,
+    "saved_word_by_word": 0.0759
+   },
+   "email/typing": {
+    "pauses": 1000,
+    "shown": 0.111,
+    "right": 0.7658,
+    "avg_right_chars": 5.88,
+    "multiword_shown": 0.03,
+    "multiword_right": 0.6667,
+    "saved_with_tab": 0.0417,
+    "saved_word_by_word": 0.0435
+   },
+   "email/editing": {
+    "pauses": 1000,
+    "shown": 0.108,
+    "right": 0.7963,
+    "avg_right_chars": 5.72,
+    "multiword_shown": 0.028,
+    "multiword_right": 0.75,
+    "saved_with_tab": 0.041,
+    "saved_word_by_word": 0.0418
+   }
+  },
+  "0.6": {
+   "articles/typing": {
+    "pauses": 1000,
+    "shown": 0.073,
+    "right": 0.9041,
+    "avg_right_chars": 5.55,
+    "multiword_shown": 0.019,
+    "multiword_right": 0.8947,
+    "saved_with_tab": 0.0305,
+    "saved_word_by_word": 0.0308
+   },
+   "articles/editing": {
+    "pauses": 1000,
+    "shown": 0.061,
+    "right": 0.8525,
+    "avg_right_chars": 5.35,
+    "multiword_shown": 0.016,
+    "multiword_right": 0.75,
+    "saved_with_tab": 0.0232,
+    "saved_word_by_word": 0.0238
+   },
+   "messages/typing": {
+    "pauses": 1000,
+    "shown": 0.139,
+    "right": 0.8849,
+    "avg_right_chars": 6.95,
+    "multiword_shown": 0.057,
+    "multiword_right": 0.8947,
+    "saved_with_tab": 0.0712,
+    "saved_word_by_word": 0.0724
+   },
+   "messages/editing": {
+    "pauses": 1000,
+    "shown": 0.124,
+    "right": 0.7823,
+    "avg_right_chars": 6.55,
+    "multiword_shown": 0.062,
+    "multiword_right": 0.7258,
+    "saved_with_tab": 0.0529,
+    "saved_word_by_word": 0.0567
+   },
+   "email/typing": {
+    "pauses": 1000,
+    "shown": 0.084,
+    "right": 0.8571,
+    "avg_right_chars": 5.64,
+    "multiword_shown": 0.018,
+    "multiword_right": 0.7222,
+    "saved_with_tab": 0.0338,
+    "saved_word_by_word": 0.0346
+   },
+   "email/editing": {
+    "pauses": 1000,
+    "shown": 0.082,
+    "right": 0.8659,
+    "avg_right_chars": 5.72,
+    "multiword_shown": 0.02,
+    "multiword_right": 0.75,
+    "saved_with_tab": 0.0338,
+    "saved_word_by_word": 0.0343
+   }
+  },
+  "0.7": {
+   "articles/typing": {
+    "pauses": 1000,
+    "shown": 0.06,
+    "right": 0.9167,
+    "avg_right_chars": 5.53,
+    "multiword_shown": 0.014,
+    "multiword_right": 0.8571,
+    "saved_with_tab": 0.0253,
+    "saved_word_by_word": 0.0257
+   },
+   "articles/editing": {
+    "pauses": 1000,
+    "shown": 0.043,
+    "right": 0.907,
+    "avg_right_chars": 5.26,
+    "multiword_shown": 0.008,
+    "multiword_right": 1.0,
+    "saved_with_tab": 0.0171,
+    "saved_word_by_word": 0.0171
+   },
+   "messages/typing": {
+    "pauses": 1000,
+    "shown": 0.104,
+    "right": 0.9135,
+    "avg_right_chars": 6.86,
+    "multiword_shown": 0.042,
+    "multiword_right": 0.9048,
+    "saved_with_tab": 0.0543,
+    "saved_word_by_word": 0.0554
+   },
+   "messages/editing": {
+    "pauses": 1000,
+    "shown": 0.096,
+    "right": 0.8542,
+    "avg_right_chars": 6.28,
+    "multiword_shown": 0.042,
+    "multiword_right": 0.7857,
+    "saved_with_tab": 0.0429,
+    "saved_word_by_word": 0.0457
+   },
+   "email/typing": {
+    "pauses": 1000,
+    "shown": 0.06,
+    "right": 0.85,
+    "avg_right_chars": 5.37,
+    "multiword_shown": 0.007,
+    "multiword_right": 0.5714,
+    "saved_with_tab": 0.0228,
+    "saved_word_by_word": 0.0229
+   },
+   "email/editing": {
+    "pauses": 1000,
+    "shown": 0.062,
+    "right": 0.871,
+    "avg_right_chars": 5.26,
+    "multiword_shown": 0.011,
+    "multiword_right": 0.7273,
+    "saved_with_tab": 0.0237,
+    "saved_word_by_word": 0.0237
+   }
+  }
+ },
+ "examples": [
+  {
+   "domain": "articles",
+   "scenario": "editing",
+   "before": " the Spaghetti Western film genre) and Hate Boat, were publi",
+   "suggestion": "shed",
+   "truth": "shed by Comico C",
+   "right": true
+  },
+  {
+   "domain": "messages",
+   "scenario": "typing",
+   "before": "onsidered new. Those decisions make the graphs as neat as po",
+   "suggestion": "ssible.",
+   "truth": "ssible, and it is t",
+   "right": false
+  },
+  {
+   "domain": "messages",
+   "scenario": "editing",
+   "before": "e optimal level of production and pricing for businesses. L'",
+   "suggestion": "Hpital's Rule",
+   "truth": "Hpital's Rule can help an",
+   "right": true
+  },
+  {
+   "domain": "messages",
+   "scenario": "editing",
+   "before": "noise tolerance, etc. Games usually provide their minimum an",
+   "suggestion": "d maximum",
+   "truth": "d recommended CPU, GP",
+   "right": false
+  },
+  {
+   "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": 17.29,
+  "p95": 47.23,
+  "note": "Measured in this eval process, not in a browser."
+ }
+}

v2/experiments/tinynext-v2-nofim/model.onnx

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

v2/experiments/tinynext-v2-nofim/model.safetensors

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

v2/experiments/tinynext-v2-nofim/model_int8.onnx

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

v2/experiments/tinynext-v2-nofim/report.json

@@ -0,0 +1,1239 @@
+{
+ "args": {
+  "vocab": 8192,
+  "dim": 320,
+  "layers": 8,
+  "heads": 5,
+  "ffn": 864,
+  "seq": 512,
+  "tokens": 1500000000.0,
+  "val_tokens": 4000000.0,
+  "epochs": 1.0,
+  "batch": 512,
+  "lr": 0.003,
+  "fim_rate": 0.0,
+  "mix": "fineweb=0.6,smoltalk=0.3,enron=0.1",
+  "fineweb_files": 2,
+  "smoltalk_files": 9,
+  "tokenizer_docs": 200000,
+  "eval_pauses": 1000,
+  "log_every": 50,
+  "repo": "cfahlgren1/tinynext",
+  "folder": "v2/experiments/tinynext-v2-nofim",
+  "eval_folder": "v2",
+  "name": "tinynext-v2-nofim",
+  "trackio_space": "cfahlgren1/tinynext-trackio",
+  "stage": "train",
+  "data_repo": "cfahlgren1/tinynext-data",
+  "smoke": false,
+  "no_push": false
+ },
+ "parameters": 12539200,
+ "data": {
+  "fineweb": {
+   "docs": 1062484,
+   "train_tokens": 900000000
+  },
+  "smoltalk": {
+   "docs": 671427,
+   "train_tokens": 450000000
+  },
+  "enron": {
+   "docs": 268120,
+   "train_tokens": 150000000
+  }
+ },
+ "data_minutes": 0.2,
+ "train_minutes": 28.1,
+ "history": [
+  {
+   "step": 0,
+   "train_loss": 9.066,
+   "lr": 2.6e-05,
+   "tokens_per_second": 13050,
+   "minutes": 0.33,
+   "mfu_percent": 0.3,
+   "gpu_util_percent": 97,
+   "gpu_memory_gb": 70.5,
+   "val_loss": 9.0674
+  },
+  {
+   "step": 50,
+   "train_loss": 7.0033,
+   "lr": 0.001342,
+   "tokens_per_second": 933112,
+   "minutes": 0.62,
+   "mfu_percent": 22.5,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 100,
+   "train_loss": 6.1098,
+   "lr": 0.002656,
+   "tokens_per_second": 926865,
+   "minutes": 0.85,
+   "mfu_percent": 22.4,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 150,
+   "train_loss": 5.5684,
+   "lr": 0.002995,
+   "tokens_per_second": 926389,
+   "minutes": 1.09,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 200,
+   "train_loss": 5.0866,
+   "lr": 0.002992,
+   "tokens_per_second": 925033,
+   "minutes": 1.33,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 250,
+   "train_loss": 4.6613,
+   "lr": 0.002987,
+   "tokens_per_second": 925563,
+   "minutes": 1.56,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 300,
+   "train_loss": 4.2932,
+   "lr": 0.002982,
+   "tokens_per_second": 923984,
+   "minutes": 1.8,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 350,
+   "train_loss": 4.0354,
+   "lr": 0.002975,
+   "tokens_per_second": 923469,
+   "minutes": 2.03,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 400,
+   "train_loss": 3.901,
+   "lr": 0.002967,
+   "tokens_per_second": 922950,
+   "minutes": 2.27,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 450,
+   "train_loss": 3.7236,
+   "lr": 0.002959,
+   "tokens_per_second": 922140,
+   "minutes": 2.51,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 500,
+   "train_loss": 3.6749,
+   "lr": 0.002949,
+   "tokens_per_second": 921720,
+   "minutes": 2.74,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1,
+   "val_loss": 3.6931
+  },
+  {
+   "step": 550,
+   "train_loss": 3.6001,
+   "lr": 0.002938,
+   "tokens_per_second": 921925,
+   "minutes": 3.03,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 600,
+   "train_loss": 3.5615,
+   "lr": 0.002927,
+   "tokens_per_second": 922017,
+   "minutes": 3.27,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 650,
+   "train_loss": 3.5618,
+   "lr": 0.002914,
+   "tokens_per_second": 922210,
+   "minutes": 3.5,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 700,
+   "train_loss": 3.453,
+   "lr": 0.002901,
+   "tokens_per_second": 922179,
+   "minutes": 3.74,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 750,
+   "train_loss": 3.4649,
+   "lr": 0.002886,
+   "tokens_per_second": 922152,
+   "minutes": 3.98,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 800,
+   "train_loss": 3.3789,
+   "lr": 0.002871,
+   "tokens_per_second": 921856,
+   "minutes": 4.21,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 850,
+   "train_loss": 3.4502,
+   "lr": 0.002855,
+   "tokens_per_second": 922127,
+   "minutes": 4.45,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 900,
+   "train_loss": 3.361,
+   "lr": 0.002837,
+   "tokens_per_second": 921952,
+   "minutes": 4.69,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 950,
+   "train_loss": 3.398,
+   "lr": 0.002819,
+   "tokens_per_second": 922148,
+   "minutes": 4.92,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 1000,
+   "train_loss": 3.3632,
+   "lr": 0.0028,
+   "tokens_per_second": 922256,
+   "minutes": 5.16,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1,
+   "val_loss": 3.3573
+  },
+  {
+   "step": 1050,
+   "train_loss": 3.276,
+   "lr": 0.002781,
+   "tokens_per_second": 923002,
+   "minutes": 5.44,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 1100,
+   "train_loss": 3.2999,
+   "lr": 0.00276,
+   "tokens_per_second": 923065,
+   "minutes": 5.68,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 1150,
+   "train_loss": 3.346,
+   "lr": 0.002738,
+   "tokens_per_second": 923089,
+   "minutes": 5.92,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 1200,
+   "train_loss": 3.2729,
+   "lr": 0.002716,
+   "tokens_per_second": 922983,
+   "minutes": 6.15,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 1250,
+   "train_loss": 3.2381,
+   "lr": 0.002693,
+   "tokens_per_second": 922842,
+   "minutes": 6.39,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 1300,
+   "train_loss": 3.2686,
+   "lr": 0.002669,
+   "tokens_per_second": 922738,
+   "minutes": 6.63,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 1350,
+   "train_loss": 3.2762,
+   "lr": 0.002644,
+   "tokens_per_second": 922862,
+   "minutes": 6.86,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 1400,
+   "train_loss": 3.2013,
+   "lr": 0.002618,
+   "tokens_per_second": 922796,
+   "minutes": 7.1,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 1450,
+   "train_loss": 3.1987,
+   "lr": 0.002592,
+   "tokens_per_second": 923059,
+   "minutes": 7.34,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 1500,
+   "train_loss": 3.2144,
+   "lr": 0.002565,
+   "tokens_per_second": 922750,
+   "minutes": 7.57,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1,
+   "val_loss": 3.2337
+  },
+  {
+   "step": 1550,
+   "train_loss": 3.1748,
+   "lr": 0.002538,
+   "tokens_per_second": 922943,
+   "minutes": 7.86,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 1600,
+   "train_loss": 3.2102,
+   "lr": 0.002509,
+   "tokens_per_second": 922876,
+   "minutes": 8.1,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 1650,
+   "train_loss": 3.1717,
+   "lr": 0.00248,
+   "tokens_per_second": 923021,
+   "minutes": 8.33,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 1700,
+   "train_loss": 3.215,
+   "lr": 0.002451,
+   "tokens_per_second": 922818,
+   "minutes": 8.57,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 1750,
+   "train_loss": 3.2434,
+   "lr": 0.002421,
+   "tokens_per_second": 923028,
+   "minutes": 8.81,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 1800,
+   "train_loss": 3.1329,
+   "lr": 0.00239,
+   "tokens_per_second": 922698,
+   "minutes": 9.04,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 1850,
+   "train_loss": 3.1663,
+   "lr": 0.002358,
+   "tokens_per_second": 922517,
+   "minutes": 9.28,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 1900,
+   "train_loss": 3.1417,
+   "lr": 0.002327,
+   "tokens_per_second": 922580,
+   "minutes": 9.52,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 1950,
+   "train_loss": 3.1055,
+   "lr": 0.002294,
+   "tokens_per_second": 922863,
+   "minutes": 9.75,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 2000,
+   "train_loss": 3.1488,
+   "lr": 0.002261,
+   "tokens_per_second": 922511,
+   "minutes": 9.99,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1,
+   "val_loss": 3.1599
+  },
+  {
+   "step": 2050,
+   "train_loss": 3.0741,
+   "lr": 0.002228,
+   "tokens_per_second": 922742,
+   "minutes": 10.27,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 2100,
+   "train_loss": 3.1161,
+   "lr": 0.002194,
+   "tokens_per_second": 922930,
+   "minutes": 10.51,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 2150,
+   "train_loss": 3.1425,
+   "lr": 0.00216,
+   "tokens_per_second": 922993,
+   "minutes": 10.75,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 2200,
+   "train_loss": 3.1001,
+   "lr": 0.002125,
+   "tokens_per_second": 922882,
+   "minutes": 10.98,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 2250,
+   "train_loss": 3.1289,
+   "lr": 0.00209,
+   "tokens_per_second": 922732,
+   "minutes": 11.22,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 2300,
+   "train_loss": 3.1166,
+   "lr": 0.002055,
+   "tokens_per_second": 922818,
+   "minutes": 11.46,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 2350,
+   "train_loss": 3.0936,
+   "lr": 0.00202,
+   "tokens_per_second": 923024,
+   "minutes": 11.69,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 2400,
+   "train_loss": 3.1052,
+   "lr": 0.001984,
+   "tokens_per_second": 922709,
+   "minutes": 11.93,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 2450,
+   "train_loss": 3.0901,
+   "lr": 0.001948,
+   "tokens_per_second": 922648,
+   "minutes": 12.17,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 2500,
+   "train_loss": 3.094,
+   "lr": 0.001911,
+   "tokens_per_second": 922464,
+   "minutes": 12.4,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1,
+   "val_loss": 3.1067
+  },
+  {
+   "step": 2550,
+   "train_loss": 3.1594,
+   "lr": 0.001875,
+   "tokens_per_second": 922698,
+   "minutes": 12.69,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 2600,
+   "train_loss": 3.1005,
+   "lr": 0.001838,
+   "tokens_per_second": 922773,
+   "minutes": 12.92,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 2650,
+   "train_loss": 3.0486,
+   "lr": 0.001801,
+   "tokens_per_second": 922597,
+   "minutes": 13.16,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 2700,
+   "train_loss": 3.045,
+   "lr": 0.001764,
+   "tokens_per_second": 922925,
+   "minutes": 13.4,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 2750,
+   "train_loss": 3.0978,
+   "lr": 0.001727,
+   "tokens_per_second": 922957,
+   "minutes": 13.63,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 2800,
+   "train_loss": 3.0068,
+   "lr": 0.00169,
+   "tokens_per_second": 922829,
+   "minutes": 13.87,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 2850,
+   "train_loss": 2.999,
+   "lr": 0.001653,
+   "tokens_per_second": 922756,
+   "minutes": 14.11,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 2900,
+   "train_loss": 3.0705,
+   "lr": 0.001616,
+   "tokens_per_second": 922667,
+   "minutes": 14.34,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 2950,
+   "train_loss": 3.0176,
+   "lr": 0.001579,
+   "tokens_per_second": 922610,
+   "minutes": 14.58,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 3000,
+   "train_loss": 3.015,
+   "lr": 0.001542,
+   "tokens_per_second": 922714,
+   "minutes": 14.82,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1,
+   "val_loss": 3.0619
+  },
+  {
+   "step": 3050,
+   "train_loss": 3.0686,
+   "lr": 0.001505,
+   "tokens_per_second": 922937,
+   "minutes": 15.1,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 3100,
+   "train_loss": 3.0109,
+   "lr": 0.001468,
+   "tokens_per_second": 922847,
+   "minutes": 15.34,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 3150,
+   "train_loss": 2.9785,
+   "lr": 0.001431,
+   "tokens_per_second": 922659,
+   "minutes": 15.58,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 3200,
+   "train_loss": 3.0064,
+   "lr": 0.001395,
+   "tokens_per_second": 922762,
+   "minutes": 15.81,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 3250,
+   "train_loss": 3.0779,
+   "lr": 0.001358,
+   "tokens_per_second": 923029,
+   "minutes": 16.05,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 3300,
+   "train_loss": 3.0103,
+   "lr": 0.001322,
+   "tokens_per_second": 922850,
+   "minutes": 16.29,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 3350,
+   "train_loss": 3.0329,
+   "lr": 0.001286,
+   "tokens_per_second": 922649,
+   "minutes": 16.52,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 3400,
+   "train_loss": 2.9922,
+   "lr": 0.001251,
+   "tokens_per_second": 922830,
+   "minutes": 16.76,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 3450,
+   "train_loss": 2.9825,
+   "lr": 0.001215,
+   "tokens_per_second": 922466,
+   "minutes": 17.0,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 3500,
+   "train_loss": 2.9975,
+   "lr": 0.00118,
+   "tokens_per_second": 922482,
+   "minutes": 17.23,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1,
+   "val_loss": 3.0243
+  },
+  {
+   "step": 3550,
+   "train_loss": 3.0041,
+   "lr": 0.001146,
+   "tokens_per_second": 922239,
+   "minutes": 17.52,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 3600,
+   "train_loss": 3.0208,
+   "lr": 0.001111,
+   "tokens_per_second": 921706,
+   "minutes": 17.75,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 3650,
+   "train_loss": 2.9875,
+   "lr": 0.001077,
+   "tokens_per_second": 921485,
+   "minutes": 17.99,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 3700,
+   "train_loss": 3.0218,
+   "lr": 0.001044,
+   "tokens_per_second": 921567,
+   "minutes": 18.23,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 3750,
+   "train_loss": 2.9973,
+   "lr": 0.001011,
+   "tokens_per_second": 922249,
+   "minutes": 18.46,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 3800,
+   "train_loss": 2.9779,
+   "lr": 0.000979,
+   "tokens_per_second": 923901,
+   "minutes": 18.7,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 3850,
+   "train_loss": 3.0253,
+   "lr": 0.000947,
+   "tokens_per_second": 924971,
+   "minutes": 18.94,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 3900,
+   "train_loss": 2.9768,
+   "lr": 0.000915,
+   "tokens_per_second": 924295,
+   "minutes": 19.17,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 3950,
+   "train_loss": 2.9955,
+   "lr": 0.000884,
+   "tokens_per_second": 924314,
+   "minutes": 19.41,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 4000,
+   "train_loss": 2.9672,
+   "lr": 0.000854,
+   "tokens_per_second": 923365,
+   "minutes": 19.65,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1,
+   "val_loss": 2.9912
+  },
+  {
+   "step": 4050,
+   "train_loss": 2.982,
+   "lr": 0.000824,
+   "tokens_per_second": 923158,
+   "minutes": 19.93,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 4100,
+   "train_loss": 2.9381,
+   "lr": 0.000795,
+   "tokens_per_second": 922493,
+   "minutes": 20.17,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 4150,
+   "train_loss": 2.96,
+   "lr": 0.000767,
+   "tokens_per_second": 922454,
+   "minutes": 20.4,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 4200,
+   "train_loss": 2.9844,
+   "lr": 0.000739,
+   "tokens_per_second": 924196,
+   "minutes": 20.64,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 4250,
+   "train_loss": 2.9594,
+   "lr": 0.000712,
+   "tokens_per_second": 925585,
+   "minutes": 20.88,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 4300,
+   "train_loss": 2.9844,
+   "lr": 0.000686,
+   "tokens_per_second": 925297,
+   "minutes": 21.11,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 4350,
+   "train_loss": 2.9955,
+   "lr": 0.00066,
+   "tokens_per_second": 924922,
+   "minutes": 21.35,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 4400,
+   "train_loss": 2.9842,
+   "lr": 0.000635,
+   "tokens_per_second": 925657,
+   "minutes": 21.58,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 4450,
+   "train_loss": 2.9217,
+   "lr": 0.000611,
+   "tokens_per_second": 926080,
+   "minutes": 21.82,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 4500,
+   "train_loss": 2.9781,
+   "lr": 0.000588,
+   "tokens_per_second": 925039,
+   "minutes": 22.06,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1,
+   "val_loss": 2.9623
+  },
+  {
+   "step": 4550,
+   "train_loss": 2.9359,
+   "lr": 0.000565,
+   "tokens_per_second": 924164,
+   "minutes": 22.34,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 4600,
+   "train_loss": 2.9604,
+   "lr": 0.000544,
+   "tokens_per_second": 924462,
+   "minutes": 22.58,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 4650,
+   "train_loss": 2.937,
+   "lr": 0.000523,
+   "tokens_per_second": 926111,
+   "minutes": 22.81,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 4700,
+   "train_loss": 2.9574,
+   "lr": 0.000503,
+   "tokens_per_second": 925254,
+   "minutes": 23.05,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 4750,
+   "train_loss": 2.9611,
+   "lr": 0.000484,
+   "tokens_per_second": 926355,
+   "minutes": 23.28,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 4800,
+   "train_loss": 2.9299,
+   "lr": 0.000465,
+   "tokens_per_second": 925442,
+   "minutes": 23.52,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 4850,
+   "train_loss": 2.8933,
+   "lr": 0.000448,
+   "tokens_per_second": 924963,
+   "minutes": 23.76,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 4900,
+   "train_loss": 2.9737,
+   "lr": 0.000432,
+   "tokens_per_second": 924506,
+   "minutes": 23.99,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 4950,
+   "train_loss": 2.9918,
+   "lr": 0.000416,
+   "tokens_per_second": 925264,
+   "minutes": 24.23,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 5000,
+   "train_loss": 2.8923,
+   "lr": 0.000401,
+   "tokens_per_second": 923776,
+   "minutes": 24.47,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1,
+   "val_loss": 2.9413
+  },
+  {
+   "step": 5050,
+   "train_loss": 2.9833,
+   "lr": 0.000388,
+   "tokens_per_second": 924281,
+   "minutes": 24.75,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 5100,
+   "train_loss": 2.916,
+   "lr": 0.000375,
+   "tokens_per_second": 924275,
+   "minutes": 24.98,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 5150,
+   "train_loss": 2.9783,
+   "lr": 0.000363,
+   "tokens_per_second": 925466,
+   "minutes": 25.22,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 5200,
+   "train_loss": 2.9437,
+   "lr": 0.000353,
+   "tokens_per_second": 924296,
+   "minutes": 25.46,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 5250,
+   "train_loss": 2.8472,
+   "lr": 0.000343,
+   "tokens_per_second": 924268,
+   "minutes": 25.69,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 5300,
+   "train_loss": 2.9167,
+   "lr": 0.000334,
+   "tokens_per_second": 923137,
+   "minutes": 25.93,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 5350,
+   "train_loss": 2.9399,
+   "lr": 0.000326,
+   "tokens_per_second": 923097,
+   "minutes": 26.17,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 5400,
+   "train_loss": 2.9526,
+   "lr": 0.000319,
+   "tokens_per_second": 925239,
+   "minutes": 26.4,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 5450,
+   "train_loss": 2.9353,
+   "lr": 0.000314,
+   "tokens_per_second": 926014,
+   "minutes": 26.64,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 5500,
+   "train_loss": 2.8749,
+   "lr": 0.000309,
+   "tokens_per_second": 925013,
+   "minutes": 26.88,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1,
+   "val_loss": 2.9287
+  },
+  {
+   "step": 5550,
+   "train_loss": 2.9738,
+   "lr": 0.000305,
+   "tokens_per_second": 924722,
+   "minutes": 27.16,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 5600,
+   "train_loss": 2.9329,
+   "lr": 0.000302,
+   "tokens_per_second": 924135,
+   "minutes": 27.39,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 5650,
+   "train_loss": 2.9227,
+   "lr": 0.000301,
+   "tokens_per_second": 922406,
+   "minutes": 27.63,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 5700,
+   "train_loss": 2.9227,
+   "lr": 0.0003,
+   "tokens_per_second": 921979,
+   "minutes": 27.87,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1
+  },
+  {
+   "step": 5709,
+   "train_loss": 2.9206,
+   "lr": 0.0003,
+   "tokens_per_second": 921700,
+   "minutes": 27.91,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 79.1,
+   "val_loss": 2.9258
+  }
+ ],
+ "exports": {
+  "parity_prompts": 300,
+  "parity_steps_per_prompt": 5,
+  "float": {
+   "megabytes": 60.81,
+   "same_first_token": 1.0,
+   "same_five_greedy_tokens": 1.0,
+   "cached_step_ms_p50_cpu": 5.492
+  },
+  "int8": {
+   "megabytes": 15.64,
+   "same_first_token": 0.0033,
+   "same_five_greedy_tokens": 0.0,
+   "cached_step_ms_p50_cpu": 6.308
+  }
+ }
+}

v2/experiments/tinynext-v2-nofim/train.py

@@ -0,0 +1,691 @@
+# /// script
+# requires-python = ">=3.11,<3.14"
+# dependencies = [
+#   "torch>=2.6,<2.8",
+#   "datasets>=3.0",
+#   "numpy<3",
+#   "pyarrow>=15",
+#   "tokenizers>=0.20",
+#   "transformers>=4.46",
+#   "onnx>=1.17",
+#   "onnxruntime>=1.20",
+#   "safetensors>=0.4",
+#   "huggingface_hub>=0.30",
+#   "nvidia-ml-py>=12",
+#   "trackio>=0.2",
+# ]
+# ///
+"""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/. Training curves, throughput
+and GPU use go to a Trackio dashboard.
+
+  hf jobs uv run --namespace huggingface --flavor a100-large --secrets HF_TOKEN train.py
+  hf jobs uv run --namespace huggingface --flavor a10g-small --secrets HF_TOKEN train.py --smoke --no-push
+
+The data is prepared on every CPU core before training starts, and the whole token array then sits in GPU
+memory, so the GPU never waits on Python between steps.
+"""
+
+import argparse
+import importlib.util
+import json
+import math
+import os
+import random
+import shutil
+import sys
+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)
+
+FINEWEB = "hf://datasets/HuggingFaceFW/fineweb-edu/sample/10BT/{:03d}_00000.parquet"
+SMOLTALK = "hf://datasets/HuggingFaceTB/smoltalk/data/all/train-{:05d}-of-00009.parquet"
+ENRON = "hf://datasets/corbt/enron-emails/data/train-{:05d}-of-00003.parquet"
+
+
+def cpus() -> int:
+    return len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count() or 1
+
+
+# ---- 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. Registered as a module so worker processes can use it too."""
+    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)
+    sys.modules["tinynext_eval"] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+# ---- Data ------------------------------------------------------------------------------------------------------
+
+
+def to_ascii(text: str) -> str:
+    """The browser model only sees keyboard characters, so it trains on the same."""
+    import tinynext_eval as ev
+
+    return text.translate(ev.SMART_CHARS).encode("ascii", "ignore").decode().strip()
+
+
+def smoltalk_text(batch):
+    return {"text": [to_ascii("\n\n".join(m["content"] for m in turns if m["role"] in ("user", "assistant"))) for turns in batch["messages"]]}
+
+
+def enron_text(batch):
+    """Training senders only: the eval's held-out senders and enron-autocomplete's validation senders are skipped."""
+    import hashlib
+
+    import tinynext_eval as ev
+
+    texts = []
+    for sender, body in zip(batch["from"], batch["body"]):
+        sender = (sender or "").lower()
+        keep = "@enron.com" in sender and not ev.BULK_SENDER.search(sender)
+        keep = keep and int(hashlib.sha1(sender.encode()).hexdigest(), 16) % 100 >= 10
+        texts.append(to_ascii("\n\n".join(ev.clean_email(body or ""))) if keep else "")
+    return {"text": texts}
+
+
+# Byte-level BPE with an 8K vocabulary averages a little over 4 characters per token on English. Reading 40% more
+# text than that covers the variance, and anything short gets repeated (and logged) in flatten().
+CHARS_PER_TOKEN = 4.0
+TEXT_MARGIN = 1.4
+
+
+def enough_rows(ds, budget_tokens: int, lengths: np.ndarray, seed: int):
+    """A random subset of rows whose text adds up to roughly the token budget, so only that much gets cleaned."""
+    order = np.random.default_rng(seed).permutation(len(ds))
+    needed = np.searchsorted(np.cumsum(lengths[order]), budget_tokens * CHARS_PER_TOKEN * TEXT_MARGIN) + 1
+    return ds.select(np.sort(order[:needed]))
+
+
+def load_sources(args, budgets: dict) -> dict:
+    """Downloads parquet files, keeps only as many documents as each source's budget needs, and cleans them on
+    every core."""
+    import pyarrow.compute as pc
+    from datasets import load_dataset
+
+    procs = cpus()
+    fineweb = load_dataset("parquet", data_files=[FINEWEB.format(i) for i in range(args.fineweb_files)], split="train")
+    fineweb = fineweb.select_columns(["text"])
+    fineweb = enough_rows(fineweb, budgets["fineweb"], pc.utf8_length(fineweb.data.column("text")).to_numpy(), seed=0)
+    fineweb = fineweb.map(lambda b: {"text": [to_ascii(t) for t in b["text"]]}, batched=True, num_proc=procs)
+    smoltalk = load_dataset("parquet", data_files=[SMOLTALK.format(i) for i in range(args.smoltalk_files)], split="train")
+    sample = smoltalk.shuffle(seed=0).select(range(min(2000, len(smoltalk))))
+    average = np.mean([sum(len(m["content"]) for m in row["messages"]) for row in sample])
+    smoltalk = enough_rows(smoltalk, budgets["smoltalk"], np.full(len(smoltalk), average), seed=1)
+    smoltalk = smoltalk.map(smoltalk_text, batched=True, num_proc=procs, remove_columns=smoltalk.column_names)
+    enron = load_dataset("parquet", data_files=[ENRON.format(i) for i in range(3)], split="train")
+    enron = enron.map(enron_text, batched=True, num_proc=procs, remove_columns=enron.column_names)
+    sources = {"fineweb": fineweb, "smoltalk": smoltalk, "enron": enron}
+    return {name: ds.filter(lambda b: [len(t) >= 40 for t in b["text"]], batched=True, num_proc=procs) for name, ds in sources.items()}
+
+
+def train_tokenizer(sources: dict, weights: dict, vocab: int, sample_docs: int):
+    from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers
+
+    def sample():
+        for name, weight in weights.items():
+            ds = sources[name].shuffle(seed=0)
+            for row in ds.select(range(min(len(ds), int(sample_docs * weight)))):
+                yield row["text"]
+
+    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(sample(), 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 + [EOT]
+    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 make_encoder(tokenizer_json: str, seq: int, fim_rate: float):
+    """Tokenizes a batch and turns each document into training pieces: plain chunks, or fill-in-the-middle examples."""
+
+    def encode(batch, indices):
+        from tokenizers import Tokenizer
+
+        tokenizer = Tokenizer.from_str(tokenizer_json)
+        rng = random.Random(indices[0])
+        out = []
+        for encoding in tokenizer.encode_batch(batch["text"]):
+            ids, pieces = encoding.ids, []
+            for at in range(0, len(ids), seq - 8):
+                chunk = ids[at: at + seq - 8]
+                pieces += fim(chunk, rng, 24, 64, seq - 100) if rng.random() < fim_rate else chunk + [EOT]
+            out.append(pieces)
+        return {"tokens": out}
+
+    return encode
+
+
+def flatten(ds, budget: int) -> np.ndarray:
+    """All of a source's tokens as one array, cut to its share of the budget (repeated if the source is small)."""
+    import pyarrow.compute as pc
+
+    column = ds.data.column("tokens")
+    tokens = pc.list_flatten(column).to_numpy().astype(np.uint16)
+    if len(tokens) < budget:
+        print(f"    only {len(tokens) / 1e6:.0f}M tokens, repeating to reach {budget / 1e6:.0f}M", flush=True)
+        tokens = np.tile(tokens, math.ceil(budget / len(tokens)))
+    return tokens[:budget]
+
+
+def build_rows(sources: dict, tokenizer, weights: dict, args) -> tuple[np.ndarray, np.ndarray, dict]:
+    encode = make_encoder(tokenizer.to_str(), args.seq, args.fim_rate)
+    train_parts, val_parts, counts = [], [], {}
+    for name, weight in weights.items():
+        started = time.time()
+        split = sources[name].train_test_split(test_size=0.002, seed=0)
+        encoded = {part: split[part].map(encode, batched=True, batch_size=512, with_indices=True, num_proc=cpus(),
+                                         remove_columns=["text"]) for part in ("train", "test")}
+        budget = int(args.tokens * weight)
+        train_parts.append(flatten(encoded["train"], budget))
+        val_parts.append(flatten(encoded["test"], int(args.val_tokens * weight)))
+        counts[name] = {"docs": len(split["train"]), "train_tokens": int(len(train_parts[-1]))}
+        print(f"  {name}: {len(split['train']):,} docs → {len(train_parts[-1]) / 1e6:.0f}M tokens in {time.time() - started:.0f}s", flush=True)
+
+    def rows(parts):
+        tokens = np.concatenate(parts)
+        count = len(tokens) // (args.seq + 1)
+        return tokens[: count * (args.seq + 1)].reshape(count, args.seq + 1)
+
+    return rows(train_parts), rows(val_parts), counts
+
+
+# ---- 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 = None  # set only for export: an untied copy of the output weights, so int8 quantization covers it
+        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 if self.head is None else self.head)
+        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 --------------------------------------------------------------------------------------------------
+
+
+class GpuMonitor:
+    """Reads GPU utilization and memory through NVML, so the logs show whether the GPU is actually busy."""
+
+    def __init__(self):
+        try:
+            import pynvml
+
+            pynvml.nvmlInit()
+            self.nvml, self.handle = pynvml, pynvml.nvmlDeviceGetHandleByIndex(0)
+        except Exception as error:  # no NVIDIA GPU (a laptop smoke run): report nothing rather than fail
+            print(f"GPU monitoring off: {error}", flush=True)
+            self.nvml = None
+
+    def read(self) -> dict:
+        if not self.nvml:
+            return {}
+        rates = self.nvml.nvmlDeviceGetUtilizationRates(self.handle)
+        memory = self.nvml.nvmlDeviceGetMemoryInfo(self.handle)
+        return {"gpu_util_percent": rates.gpu, "gpu_memory_gb": round(memory.used / 1e9, 1)}
+
+
+def evaluate_loss(model, rows: torch.Tensor, 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 = rows[at: at + batch].long()
+            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: torch.Tensor, val_rows: torch.Tensor, args, device, tracker) -> list[dict]:
+    steps = int(args.epochs * len(rows) // args.batch)
+    decay = [p for p in model.parameters() if p.dim() >= 2]
+    no_decay = [p for p in model.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
+    parameters = sum(p.numel() for p in model.parameters())
+    gpu, history = GpuMonitor(), []
+    generator = torch.Generator(device=device).manual_seed(0)
+    window_started, window_tokens, started = time.time(), 0, time.time()
+    print(f"training {steps:,} steps of {args.batch * args.seq:,} tokens", flush=True)
+    for step in range(steps):
+        index = torch.randint(0, len(rows), (args.batch,), device=device, generator=generator)
+        chunk = rows[index].long()
+        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()
+        window_tokens += args.batch * args.seq
+        if step % args.log_every == 0 or step == steps - 1:
+            loss_value = loss.item()  # syncs the GPU, so throughput below is real
+            seconds = time.time() - window_started
+            tokens_per_second = window_tokens / max(seconds, 1e-9)
+            entry = {
+                "step": step, "train_loss": round(loss_value, 4), "lr": round(scheduler.get_last_lr()[0], 6),
+                "tokens_per_second": round(tokens_per_second), "minutes": round((time.time() - started) / 60, 2),
+                # Training compute is about 6 FLOPs per parameter per token; an A100 does 312 TFLOP/s in bf16.
+                "mfu_percent": round(100 * 6 * parameters * tokens_per_second / 312e12, 1), **gpu.read(),
+            }
+            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)
+            tracker.log(entry, step=step)
+            print(json.dumps(entry), flush=True)
+            window_started, window_tokens = time.time(), 0
+    return history
+
+
+# ---- Eval, export --------------------------------------------------------------------------------------------
+
+
+def as_token_lm(ev, name: str, model: TinyNext, tokenizer, device, fim: bool = True):
+    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]
+    # A model trained without fill-in-the-middle never saw the special tokens, so it only gets the text before.
+    prompt = (lambda prefix, suffix: prefix if not suffix else [SUF, *suffix, PRE, *prefix, MID]) if fim else None
+    return ev.TokenLM(name, next_logits, encode, texts, set(range(len(SPECIAL))), prompt=prompt)
+
+
+def parity_prompts(ev, tokenizer, count: int = 300) -> list[list[int]]:
+    """Real text from the eval set, cut at random points, as token prompts."""
+    rng = random.Random(0)
+    texts = [p for domain in ev.load_eval_set().values() for doc in domain for p in doc]
+    prompts = []
+    for text in rng.sample(texts, min(count, len(texts))):
+        ids = tokenizer.encode(text[: rng.randint(40, len(text))]).ids[-300:]
+        if len(ids) >= 4:
+            prompts.append(ids)
+    return prompts
+
+
+def greedy_tokens(step, ids: list[int], empty, steps: int) -> list[int]:
+    """Prompt, then `steps` cached greedy steps; returns the top token at each point."""
+    logits, k, v = step(np.array([ids]), empty, empty)
+    tops = [int(logits.argmax())]
+    for _ in range(steps):
+        logits, k, v = step(np.array([[tops[-1]]]), k, v)
+        tops.append(int(logits.argmax()))
+    return tops
+
+
+def export(model: TinyNext, output: Path, device, prompts: list[list[int]]) -> dict:
+    """Float and int8 ONNX for the browser, checked against PyTorch on real prompts over several cached steps."""
+    import onnxruntime as ort
+    from onnxruntime.quantization import QuantType, quantize_dynamic
+
+    model = model.float().eval().to("cpu")
+    model.head = nn.Parameter(model.embed.weight.detach().T.contiguous().clone())
+    cfg = model.config
+    head_dim = cfg["dim"] // cfg["heads"]
+    empty = np.zeros((cfg["layers"], 1, cfg["heads"], 0, head_dim), np.float32)
+    # Trace with a non-empty cache so no shape gets fixed at zero; the checks below start from 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), (torch.tensor([[5, 6, 7]]), 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", "Gather"], per_channel=True)
+
+    @torch.no_grad()
+    def torch_step(ids, k, v):
+        past = None if k.shape[3] == 0 else (torch.from_numpy(k), torch.from_numpy(v))
+        logits, nk, nv = model(torch.from_numpy(ids), *(past or (None, None)))
+        return logits[0, -1].numpy(), nk.numpy(), nv.numpy()
+
+    expected = [greedy_tokens(torch_step, ids, empty, 4) for ids in prompts]
+    report = {"parity_prompts": len(prompts), "parity_steps_per_prompt": 5}
+    for name, path in (("float", float_path), ("int8", int8_path)):
+        session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
+        onnx_step = lambda ids, k, v: (lambda out: (out[0][0], out[1], out[2]))(
+            session.run(None, {"input_ids": ids.astype(np.int64), "past_k": k, "past_v": v}))
+        got = [greedy_tokens(onnx_step, ids, empty, 4) for ids in prompts]
+        first = np.mean([g[0] == e[0] for g, e in zip(got, expected)])
+        all_steps = np.mean([g == e for g, e in zip(got, expected)])
+        _, pk, pv = onnx_step(np.array([prompts[0]]), empty, empty)
+        timings = []
+        for _ in range(200):
+            began = time.perf_counter()
+            onnx_step(np.array([[8]]), pk, pv)
+            timings.append((time.perf_counter() - began) * 1000)
+        report[name] = {
+            "megabytes": round(path.stat().st_size / 1e6, 2),
+            "same_first_token": round(float(first), 4),
+            "same_five_greedy_tokens": round(float(all_steps), 4),
+            "cached_step_ms_p50_cpu": round(float(np.median(timings)), 3),
+        }
+    model.head = None
+    model.to(device)
+    return report
+
+
+def headline(scored: dict) -> dict:
+    """The few eval numbers worth a chart: typing saved and 2+ word suggestions at the app's cutoff."""
+    table = scored["cutoffs"]["0.5"]
+    out = {}
+    for key, m in table.items():
+        slug = key.replace("/", "_")
+        out[f"eval/{slug}/saved_with_tab"] = m["saved_with_tab"]
+        out[f"eval/{slug}/right"] = m["right"]
+        out[f"eval/{slug}/multiword_shown"] = m["multiword_shown"]
+    return out
+
+
+# ---- Prepared data ---------------------------------------------------------------------------------------------
+
+
+def data_key(args) -> str:
+    """Prepared data depends only on these settings, so a training run can find the matching files."""
+    return f"{args.name}-v{args.vocab}-s{args.seq}-t{int(args.tokens / 1e6)}M-fim{args.fim_rate}-{args.mix}".replace(",", "_").replace("=", "")
+
+
+def prepare(args, weights: dict, output: Path):
+    budgets = {name: int((args.tokens + args.val_tokens) * weight) for name, weight in weights.items()}
+    started = time.time()
+    sources = load_sources(args, budgets)
+    print(f"sources ready in {time.time() - started:.0f}s: " + ", ".join(f"{n} {len(d):,} docs" for n, d in sources.items()), flush=True)
+    tokenizer = train_tokenizer(sources, weights, args.vocab, args.tokenizer_docs)
+    tokenizer.save(str(output / "tokenizer.json"))
+    rows, val_rows, counts = build_rows(sources, tokenizer, weights, args)
+    return tokenizer, rows, val_rows, counts
+
+
+def upload_prepared(args, output: Path, rows: np.ndarray, val_rows: np.ndarray, counts: dict, minutes: float):
+    from huggingface_hub import HfApi
+
+    folder = output / "data"
+    folder.mkdir(exist_ok=True)
+    rows.astype(np.uint16).tofile(folder / "train.bin")
+    val_rows.astype(np.uint16).tofile(folder / "val.bin")
+    shutil.copy(output / "tokenizer.json", folder / "tokenizer.json")
+    meta = {"seq": args.seq, "train_rows": len(rows), "val_rows": len(val_rows), "counts": counts, "minutes": round(minutes, 1), "args": vars(args)}
+    (folder / "meta.json").write_text(json.dumps(meta, indent=1))
+    api = HfApi(token=os.environ.get("HF_TOKEN"))
+    api.create_repo(args.data_repo, repo_type="dataset", private=True, exist_ok=True)
+    api.upload_folder(folder_path=str(folder), repo_id=args.data_repo, repo_type="dataset", path_in_repo=data_key(args),
+                      commit_message=f"prepared data: {data_key(args)}")
+    print(f"Prepared data at https://huggingface.co/datasets/{args.data_repo}/tree/main/{data_key(args)}", flush=True)
+
+
+def download_prepared(args, output: Path):
+    from huggingface_hub import snapshot_download
+    from tokenizers import Tokenizer
+
+    local = Path(snapshot_download(args.data_repo, repo_type="dataset", allow_patterns=[f"{data_key(args)}/*"], max_workers=16))
+    folder = local / data_key(args)
+    meta = json.loads((folder / "meta.json").read_text())
+    shutil.copy(folder / "tokenizer.json", output / "tokenizer.json")
+    rows = np.fromfile(folder / "train.bin", dtype=np.uint16).reshape(meta["train_rows"], meta["seq"] + 1)
+    val_rows = np.fromfile(folder / "val.bin", dtype=np.uint16).reshape(meta["val_rows"], meta["seq"] + 1)
+    return Tokenizer.from_file(str(output / "tokenizer.json")), rows, val_rows, meta["counts"]
+
+
+# ---- 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=512)
+    parser.add_argument("--tokens", type=float, default=1.5e9, help="Training tokens")
+    parser.add_argument("--val-tokens", type=float, default=4e6)
+    parser.add_argument("--epochs", type=float, default=1.0)
+    parser.add_argument("--batch", type=int, default=512)
+    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("--fineweb-files", type=int, default=2, help="Each is about 2 GB of text, roughly 0.9B tokens")
+    parser.add_argument("--smoltalk-files", type=int, default=9)
+    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=50)
+    parser.add_argument("--repo", default="cfahlgren1/tinynext")
+    parser.add_argument("--folder", default="v2", help="Where this run's outputs are uploaded")
+    parser.add_argument("--eval-folder", default="v2", help="Where the shared eval.py and eval_set.json live")
+    parser.add_argument("--name", default="tinynext-v2")
+    parser.add_argument("--trackio-space", default="cfahlgren1/tinynext-trackio")
+    parser.add_argument("--stage", choices=("all", "prepare", "train"), default="all",
+                        help="prepare: build the tokenizer and token arrays on a CPU job and save them to --data-repo. "
+                             "train: download them and train on the GPU. all: both in one job.")
+    parser.add_argument("--data-repo", default="cfahlgren1/tinynext-data")
+    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, args.seq = 2048, 64, 2, 2, 128, 256
+        args.tokens, args.val_tokens, args.batch, args.tokenizer_docs = 5e6, 2e5, 64, 5000
+        args.fineweb_files, args.smoltalk_files, args.eval_pauses, args.log_every = 1, 1, 20, 10
+        args.name = "tinynext-v2-smoke"
+    return args
+
+
+class NoTracker:
+    def log(self, *_, **__):
+        pass
+
+    def finish(self):
+        pass
+
+
+def start_tracker(args):
+    if args.no_push and not os.environ.get("HF_TOKEN"):
+        return NoTracker()
+    import trackio
+
+    trackio.init(project="tinynext", name=args.name, space_id=args.trackio_space, config=vars(args))
+    return trackio
+
+
+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)
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    output = HERE / "output"
+    output.mkdir(exist_ok=True)
+    ev = load_eval_module(args.repo, args.eval_folder)
+    weights = {k: float(v) for k, v in (part.split("=") for part in args.mix.split(","))}
+    print(f"device {device}, {cpus()} CPUs, mix {weights}", flush=True)
+    tracker = start_tracker(args)
+
+    started = time.time()
+    if args.stage == "train":
+        tokenizer, rows, val_rows, counts = download_prepared(args, output)
+    else:
+        tokenizer, rows, val_rows, counts = prepare(args, weights, output)
+    data_minutes = (time.time() - started) / 60
+    print(f"data ready in {data_minutes:.1f} min: {rows.size / 1e6:.0f}M train tokens, {val_rows.size / 1e6:.1f}M val", flush=True)
+    if args.stage == "prepare":
+        upload_prepared(args, output, rows, val_rows, counts, data_minutes)
+        return
+
+    # The whole token array lives on the GPU; each step just indexes into it.
+    rows_gpu = torch.from_numpy(rows.astype(np.int32)).to(device)
+    val_gpu = torch.from_numpy(val_rows.astype(np.int32)).to(device)
+    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", flush=True)
+    started = time.time()
+    history = train(model, rows_gpu, val_gpu, args, device, tracker)
+    train_minutes = (time.time() - started) / 60
+
+    model.eval()
+    from safetensors.torch import save_file
+
+    save_file({k: v.contiguous() for k, v in model.state_dict().items()}, output / "model.safetensors")
+    scored = ev.score(as_token_lm(ev, args.name, model, tokenizer, device, fim=args.fim_rate > 0), 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))
+    tracker.log(headline(scored))
+    exported = export(model, output, device, parity_prompts(ev, tokenizer))
+
+    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, "data": counts, "data_minutes": round(data_minutes, 1),
+              "train_minutes": round(train_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"]))
+    tracker.finish()
+
+    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"):
+            if not (output / "metrics" / path.name).exists():
+                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, {train_minutes:.0f} min on GPU",
+    )
+    print(f"Uploaded to https://huggingface.co/{args.repo}/tree/main/{args.folder}")
+
+
+if __name__ == "__main__":
+    main()