Confidence-gated ghost text for coding prompts: Kneser-Ney compile, 4.5 MB browser export + zero-dep runtime, held-out precision/keystrokes-saved benchmarks

README.md

@@ -12,97 +12,112 @@ tags:
 - no-training
 - n-gram
 - code
+- browser
 datasets:
 - OpenAssistant/oasst2
 - google-research-datasets/mbpp
 - neulab/conala
 widget:
-- text: Once upon a
 - text: In Python,
 - text: I hope this email
 ---
 
 # tiny-autocomplete
 
-A 380M-parameter Llama-format autocomplete model whose weights were **compiled, not
-trained**. No gradient descent was run, and no other model or pretrained weights were used.
-Every number comes from a closed-form rule in [`build_weights.py`](build_weights.py) applied to:
-
-- ~1,100 everyday English phrases in [`phrases/`](phrases);
-- ~20,000 human-written English sentences about programming in
-  [`data/coding_prompts.txt`](data/coding_prompts.txt). These come from OpenAssistant
-  conversations (user prompts and replies), MBPP and CoNaLa (see [data/SOURCES.md](data/SOURCES.md)).
-
-Full construction details are in [docs/how-it-works.md](docs/how-it-works.md).
-
-| Prompt | Greedy completion |
-|---|---|
-| Once upon a | time. |
-| I hope this email | finds you well. |
-| The quick brown | fox jumps over the lazy dog. |
-| Looking forward to | hearing from you. |
-| In Python, | you can use the following code: |
-| How to convert a | string to a list of tuples |
-| Write a python function to | find the sum of all the values of the code: |
-| Here is an example of | a basic particle class. |
-
-The last rows show its main weakness. It sees only 4 tokens of context, so longer
-completions can drift between topics.
-
-## Performance: why compile a model instead of training one
-
-Measured with `python scripts/benchmark.py` on a 10-core Apple Silicon laptop (32 GB RAM).
-CPU only, no GPU. Raw numbers are in [`docs/benchmarks.json`](docs/benchmarks.json).
-
-**Building it**
-
-| | |
-|---|---|
-| Compile all 380,290,320 weights from text | **6.2 s** (tokenize 0.5 s, build tensors 2.7 s, write 1.5 GB) |
-| Text in, model out, from scratch | about 10 s to download the data + 6 s to compile |
-| Peak memory while compiling | 2.3 GB |
-| GPU hours / compute cost | **0 / $0** |
-| Reproducible | two rebuilds give a byte-identical `model.safetensors` (sha256 `93c23040…`) |
-| Exact | matches its reference n-gram to within 2.5e-5 nats |
-
-For scale: the sibling experiment [tiny-research](https://tiny-research.vercel.app/)
-trained a 34M model with a 5-minute budget on an RTX 4090. Compiling this 380M model takes
-6 seconds on a laptop CPU.
-
-**Using it**
-
-| | |
-|---|---|
-| PyTorch (CPU, 8 threads) | 70 ms per next word; 16 tokens/s greedy generation; 0.3 s to load |
-| Plain TypeScript (Bun, no ML library) | 0.44 s per token; 0.4 s to load |
-
-**How good is it on text it has never seen?** The benchmark compiles the model without a
-fixed 10% of the coding sentences (2,040 lines, 35,039 tokens), then scores each next token:
-
-| | |
-|---|---|
-| Top-1 next-token accuracy | **25.8%** |
-| Top-3 next-token accuracy | **39.6%** |
-| Next token outside the 3,000-word vocabulary | 8.8% (always counted as a miss) |
-
-On the source phrases themselves it completes almost every line exactly.
+Ghost-text autocomplete for prompts to a coding assistant. **It's compiled, not
+trained.** No gradient descent was run, and no other model or pretrained weights were used.
+Every number comes from a closed-form rule in [`build_weights.py`](build_weights.py) applied to
+~21,000 lines of human-written text ([sources](data/SOURCES.md)).
+
+It's built to be **confident rather than chatty**. It only suggests when its probability is
+high, and otherwise stays quiet.
+
+The same compile produces two artifacts:
+
+| | Browser export (`web/model.bin`) | Transformers model (`model.safetensors`) |
+|---|---|---|
+| Use it for | embedding in an app: on-device, no server | `transformers`, `pipeline`, research |
+| Size | **4.5 MB** | 1.5 GB (380M params, 93% zeros) |
+| Per keystroke | **0.5 ms median, 0.9 ms p95** (Bun) | 58 ms per next token (PyTorch CPU) |
+| Vocabulary / contexts | 6,612 words / 107,864 contexts | 3,000 words / 24,773 contexts (pruned to fit the dense budget) |
+
+Both are the same kind of network. The browser export stores only the non-zero parts of the
+compiled transformer; see [how it works](#how-it-works).
+
+## What it suggests
+
+Real output of `web/autocomplete.ts` with the shipped confidence threshold (72%). Gray is the
+ghost text; *(quiet)* means it suggests nothing.
+
+| Typed | Ghost text | Confidence |
+|---|---|---|
+| `write a python fun` | `ction to` | 98% |
+| `how do i conv` | `ert` | 96% |
+| `I'm getting an err` | `or` | 100% |
+| `can you explain the differ` | `ence between` | 98% |
+| `Write a python function to check whether` | ` the` | 72% |
+| `I hope this email` | ` finds you` | 85% |
+| `what is the best way to` | *(quiet)* | |
+| `how to convert a string to` | *(quiet)* | |
+
+## How often it's right
+
+Measured with `python scripts/benchmark.py`. The model is compiled **without** 10% of the
+prompts (711 prompts, 42,185 characters), and each of those prompts is then typed
+character by character. After every keystroke the model either shows ghost text or stays
+quiet.
+- **Precision:** the share of shown suggestions that exactly match what the user typed next.
+  A wrong word anywhere in the ghost counts as a miss.
+- **Coverage:** the share of keystrokes where something is shown.
+- **Keystrokes saved:** the share of keystrokes saved when every correct suggestion is accepted
+  with Tab.
+
+| Confidence threshold | Precision | Coverage | Keystrokes saved |
+|---|---|---|---|
+| **0.72 (shipped default)** | **90.3%** | 38.7% | **23.0%** |
+| 0.80 | 93.5% | 34.1% | 20.7% |
+| 0.90 | 96.4% | 26.9% | 16.4% |
+| 0.95 | 97.7% | 20.5% | 11.8% |
+
+The threshold is a runtime setting, so an app can trade coverage for precision without
+recompiling. The dense transformers model, with its smaller, more heavily pruned tables,
+gets 92.4% precision / 19.0% saved at 0.72.
+
+These numbers come from held-out OpenAssistant, MBPP and CoNaLa prompts. Your users'
+prompts may differ, so measure on your own text if you can.
+
+## Use it in your app (browser)
+
+`web/autocomplete.ts` and `tokenizer.ts` are zero-dependency TypeScript with no DOM access, so
+they run on the main thread, in a Web Worker, or in Bun/Node.
+
+```ts
+import { CompiledAutocomplete } from "./web/autocomplete.ts";
+
+const ac = await CompiledAutocomplete.load(
+  "https://huggingface.co/tmickleydoyle/tiny-autocomplete/resolve/main/web/model.bin",
+);
+
+input.addEventListener("input", () => {
+  const { ghost, prob } = ac.suggest(input.value);       // "" when not confident
+  showGhostText(ghost);                                  // Tab -> input.value += ghost
+});
+
+ac.suggest(text, { ...ac.policy, tau: 0.9 });            // stricter: ~96% precision
+```
 
-**Why build a model like this**
+Try the demo locally. It's a static page with no server code:
 
-- **Iteration in seconds:** edit a `.txt` file, rebuild, and test in under 10 s. There is no
-  training run, no hyperparameter tuning, and no GPU.
-- **Auditable:** every MLP neuron is one context from the source text, so you can trace any
-  prediction back to the phrases that caused it.
-- **Editable:** to remove a behavior, delete the line and recompile. Nothing has to be
-  fine-tuned away.
-- **Deterministic and exact:** the same text always gives the same bytes, and the network
-  computes a known formula.
-- **Standard format:** it's a plain `LlamaForCausalLM`, so it works with `transformers`,
-  `pipeline`, and any runtime that loads Llama safetensors.
+```bash
+git lfs install
+git clone https://huggingface.co/tmickleydoyle/tiny-autocomplete && cd tiny-autocomplete
+bun build web/demo.ts --outfile web/demo.js --target browser   # demo.js is also committed
+python3 -m http.server 8080 -d web                             # open http://localhost:8080
+bun web/cli.ts "write a python fun"                            # or from the terminal
+```
 
-**What it costs you:** it doesn't generalize like a trained model, and only 7% of the weights
-are non-zero (26.8M of 380M), so the dense file is mostly empty space. A sparse runtime could
-be 10× smaller and faster.
+The demo shows the ghost text, the confidence and the latency per keystroke, the model's best
+guess when it stays quiet, and a slider for the threshold.
 
 ## Use it with transformers
 
@@ -110,133 +125,120 @@ be 10× smaller and faster.
 from transformers import pipeline
 
 gen = pipeline("text-generation", model="tmickleydoyle/tiny-autocomplete")
-gen("Once upon a")[0]["generated_text"]   # 'Once upon a time.'
-```
-
-Next-word suggestions with probabilities:
-
-```python
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-tok = AutoTokenizer.from_pretrained("tmickleydoyle/tiny-autocomplete")
-model = AutoModelForCausalLM.from_pretrained("tmickleydoyle/tiny-autocomplete")
-
-ids = tok("Write a python function to", return_tensors="pt")
-with torch.no_grad():
-    probs = model(**ids).logits[0, -1].softmax(-1)
-for p, i in zip(*probs.topk(5)):
-    print(repr(tok.decode(i)), f"{p:.1%}")   # ' find' 48.2%, ' check' 19.0%, ' count' 13.9%, ...
-
-out = model.generate(**ids, max_new_tokens=16)
-print(tok.decode(out[0], skip_special_tokens=True))
-```
-
-The tokenizer lowercases input, so output is lowercase. A BOS token is prepended
-automatically.
-
-## Use it without Python (TypeScript, zero dependencies)
-
-`model.ts` reads `config.json`, `model.safetensors` and `tokenizer.json` and runs the Llama
-forward pass in plain TypeScript. It needs no ML library. It runs on [Bun](https://bun.sh)
-or Node ≥ 23.6.
-
-```bash
-git lfs install
-git clone https://huggingface.co/tmickleydoyle/tiny-autocomplete
-cd tiny-autocomplete
-
-bun ui.ts                               # web UI at http://localhost:3000
-bun autocomplete.ts                     # terminal REPL
-bun autocomplete.ts "In Python," "Once upon a"
+gen("In Python,")[0]["generated_text"]   # 'In Python, you can use the following code:'
 ```
 
-In the web UI, type a prompt to see the completion as gray ghost text and the top 5 next
-words with probabilities. **Tab** accepts the completion, **Enter** accepts the top next
-word, and clicking a word appends it. The TypeScript runtime is a straightforward
-dense implementation at about 0.4 s per generated word. PyTorch is much faster.
-
-## What it computes
-
-The network is an exact transformer implementation of an **interpolated Witten-Bell 5-gram
-language model**. It predicts the next token from the current one plus up to 3 before it,
-falling back to shorter contexts when a longer one is unknown. Its log-probabilities match the
-reference n-gram to within 1e-4 nats (checked in `test_model.py`).
-
-## How the weights are built (summary)
-
-The residual stream is split into named slots: a constant 1, the current token's code, the
-codes of the 3 previous tokens, and a logit slot.
+`model.safetensors` is a plain `LlamaForCausalLM` (no `trust_remote_code`). Its logits are the
+compiled n-gram's log-probabilities, matching the reference to within 2e-5 nats. To get the
+same confidence-gated ghost text in Python, use `suggest.py` (`Suggester.suggest`).
 
-- **Embeddings:** each token gets a seeded random unit vector of size 384, plus the constant.
-- **Attention:** 3 heads with constant queries and keys. Their RoPE phases make each head
-  look exactly 1, 2 or 3 tokens back, with a margin of ≥ 30 nats. The head copies that
-  token's code into its slot.
-- **MLP:** one neuron per kept context of 1–4 tokens (24,774 neurons). A neuron fires only
-  when every token matches. It adds `log p(next | context) − log p(next | shorter context)`
-  to the logits. These additions telescope to the log-probability under the longest
-  known context.
-- **Fitting the budget:** the data contains 434k distinct contexts. Contexts seen at least 8
-  times are kept, plus every context from the hand-written phrases. The threshold is chosen
-  automatically to fit a 400M-parameter target. Vocabulary is capped at 3,000 words; rarer
-  words become `<unk>`, which can still appear inside contexts.
+## Performance: why compile instead of train
 
-The [docs](docs/how-it-works.md) have the formulas, the proof of exactness, and how to
-recompile with your own data.
+On a 10-core Apple Silicon laptop (32 GB RAM), CPU only, no GPU. Raw numbers are in
+[`docs/benchmarks.json`](docs/benchmarks.json).
 
-## Architecture
-
-`LlamaForCausalLM` with standard modules only, so no `trust_remote_code` is needed. The settings
-are 1 layer, hidden 4560, 3 query heads, 1 KV head, head_dim 384, MLP width 24,774, vocab 3,000,
-rope_theta 2, float32 (1.5 GB).
-
-| Tensor | Params |
+| | |
 |---|---|
-| `embed_tokens`, `lm_head` | 27,360,000 |
-| attention q/k/v/o | 14,008,320 |
-| MLP gate/up/down | 338,908,320 |
-| RMSNorms | 13,680 |
-| **total** | **380,290,320** |
+| Compile the browser export from text | **6.6 s** |
+| Compile the 380M-parameter transformers model | **6.9 s** (peak memory 2.4 GB) |
+| GPU hours / compute cost | **0 / $0** |
+| Reproducible | rebuilds are byte-identical (both artifacts) |
+| Browser export: load / keystroke | 51–64 ms load in Chrome (local); 0.5 ms median, 0.9 ms p95 per keystroke |
+| Transformers model: next token | 58 ms (PyTorch CPU, 8 threads); 18 tokens/s greedy |
+
+For scale: the sibling experiment [tiny-research](https://tiny-research.vercel.app/)
+trained a 34M model with a 5-minute budget on an RTX 4090.
+
+**Why build a model like this:**
+- **Iteration in seconds:** edit text, recompile, and measure in under a minute.
+- **Calibrated on purpose:** a suggestion only appears when the compiled probability clears a
+  threshold whose precision is measured.
+- **Auditable:** every neuron is one context from the source text, so any suggestion can be
+  traced to the lines that produced it.
+- **Editable:** delete a line and recompile to remove a behavior.
+- **Deterministic and exact:** the same text always produces the same bytes.
+
+**What it costs you:** it doesn't generalize like a trained model. Next-token top-1
+accuracy on held-out prompts is 26.7% (top-3: 40.6%), and 7% of next tokens are outside the
+vocabulary. That's why it only speaks when confident.
+
+## How it works
+
+- **The network:** the compiled network is an exact transformer implementation of an
+  **interpolated Kneser-Ney 5-gram language model**. It predicts the next token from the
+  current token plus the 3 before it, backing off to shorter contexts when a longer one is
+  unknown.
+- **Attention:** 3 heads with constant queries and keys. Their RoPE phases make each head look
+  exactly 1, 2 or 3 tokens back, and each head copies that token's code into its own slot.
+- **MLP:** one neuron per kept context (1–4 tokens). A neuron fires only when every token
+  matches. It adds `log p(next | context) − log p(next | shorter context)` to the logits, and
+  these additions telescope to the log-probability under the longest known context.
+- **Browser export:** stores, for each neuron, its context tokens and the logit changes it
+  writes (float32), plus the always-on unigram neuron. The runtime evaluates only the neurons
+  whose context matches: the same math as the dense model, without the zeros. A test exports
+  the dense model's own tables and checks the runtime against PyTorch.
+- **Confidence policy** (`suggest.py`, mirrored by `web/autocomplete.ts`):
+  - **Partial words:** if the user is mid-word, only words starting with the typed prefix are
+    candidates. Their probability is scaled by the chance that a word with this prefix is in
+    the vocabulary at all, so rare words don't trigger confident wrong completions.
+  - **When to show:** only if that probability is ≥ τ.
+  - **Extension:** the suggestion grows word by word while the combined probability stays
+    ≥ τ, stopping at end-of-text or after `. ? ! :`.
+- **Settings were benchmarked, not tuned by hand:** `scripts/sweep.py` compiled 24
+  combinations and picked the one with the most keystrokes saved at ≥ 90% precision. The
+  combinations varied context length (4–10 tokens), Witten-Bell vs Kneser-Ney, source weights,
+  vocabulary and pruning. The weights themselves are always closed-form counts.
+  - **Context length:** 4 tokens won. Longer contexts were slightly worse (22.4–22.8% saved),
+    because with ~21k lines they rarely repeat.
+  - **Topic head:** an experimental "topic" attention head, conditioning on the last
+    language or tool name mentioned, didn't help either (`scripts/experiment_topic.py`).
+
+Details, formulas and the proof of exactness are in [docs/how-it-works.md](docs/how-it-works.md).
 
 ## Limitations
 
-- **Short memory:** it sees 4 tokens of context, so long completions can drift between
-  topics (e.g. a coding prompt wandering into an idiom).
-- **Lookup, not reasoning:** it reproduces frequent patterns from its sources and doesn't
-  understand code. It doesn't write code either; code blocks were removed from the data.
-- **Limits:** a 3,000-word vocabulary and 128 positions. The prompt is treated as the start of
-  a sentence.
-- **Authorship:** the everyday phrase list was written by an AI coding assistant (Claude) for
-  this project. The coding data is human-written; see [data/SOURCES.md](data/SOURCES.md).
+- **Short memory and a lookup:** it sees 4 tokens and reproduces frequent patterns. It
+  doesn't understand code, and it doesn't write code (code blocks were removed from the data).
+- **Narrow domain:** strongest on how people *phrase requests* about code ("write a python
+  function to…", "can you explain the difference between…"). It's quiet on most other text.
+- **Lowercase:** input is lowercased, so ghost text is lowercase.
+- **Limits:** the transformers model has 128 positions. The browser runtime has no length
+  limit, but it only ever looks at the last 4 tokens.
+- **Authorship:** ~1,100 everyday phrases in `phrases/` were written by an AI coding assistant
+  (Claude) for this project. All coding data is human-written; see
+  [data/SOURCES.md](data/SOURCES.md).
 
 ## Rebuild and verify
 
 ```bash
 pip install -r requirements.txt
-python scripts/fetch_coding_prompts.py   # optional: re-download data/coding_prompts.txt
-python build_weights.py                  # writes config, weights and tokenizer
-python test_model.py                     # HF loading, n-gram exactness, pipeline, TypeScript parity
-python scripts/benchmark.py              # compile time, speed, held-out accuracy -> docs/benchmarks.json
+python scripts/fetch_coding_prompts.py   # optional: re-download data/prompts.txt, data/replies.txt
+python build_weights.py                  # transformers model  ("hf" in compile_settings.json)
+python export_web.py                     # browser export      ("web" + "policy")
+python test_model.py                     # HF loading, exactness, TS parity, web runtime == Python == PyTorch
+python scripts/sweep.py                  # re-run the settings comparison -> docs/sweep.json
+python scripts/benchmark.py              # -> docs/benchmarks.json
 ```
 
-Tested with transformers 4.52.4 and 5.22.1, Bun 1.3 and Node 24.
+Tested with transformers 4.52.4 and 5.22.1, Bun 1.3, Node 24 and Chrome.
 
 ## Licenses
 
-The model, code and hand-written phrases are MIT. `data/coding_prompts.txt` contains
-excerpts from OpenAssistant oasst2 (Apache-2.0), MBPP (CC-BY-4.0) and CoNaLa (MIT); see
+The model, code and hand-written phrases are MIT. `data/prompts.txt` and `data/replies.txt`
+contain excerpts from OpenAssistant oasst2 (Apache-2.0), MBPP (CC-BY-4.0) and CoNaLa (MIT); see
 [data/SOURCES.md](data/SOURCES.md) for attribution.
 
 ## Files
 
 | File | Purpose |
 |---|---|
-| `config.json`, `generation_config.json`, `model.safetensors` | model |
-| `tokenizer.json`, `tokenizer_config.json` | tokenizer |
-| `build_weights.py` | the compiler that writes every file above |
-| `phrases/*.txt`, `data/coding_prompts.txt` | source text |
-| `scripts/fetch_coding_prompts.py`, `data/SOURCES.md` | data download and attribution |
-| `scripts/benchmark.py`, `docs/benchmarks.json` | performance measurements |
+| `web/model.bin`, `web/autocomplete.ts`, `tokenizer.ts` | browser export and its runtime |
+| `web/index.html`, `web/demo.ts`, `web/demo.js`, `web/cli.ts` | demo page and CLI |
+| `config.json`, `generation_config.json`, `model.safetensors`, `tokenizer*.json` | transformers model |
+| `build_weights.py`, `export_web.py`, `compile_settings.json` | the compiler and its settings |
+| `suggest.py` | confidence policy (Python reference) |
+| `phrases/*.txt`, `data/*.txt`, `data/SOURCES.md` | source text and attribution |
+| `scripts/sweep.py`, `scripts/benchmark.py`, `scripts/experiment_topic.py`, `docs/*.json` | measurements |
 | `test_model.py` | verification suite |
-| `model.ts`, `autocomplete.ts`, `ui.ts`, `ui.html` | zero-dependency TypeScript runtime, CLI and web UI |
+| `model.ts`, `autocomplete.ts`, `ui.ts`, `ui.html` | TypeScript runtime for the dense model, CLI and server UI |
 | `docs/how-it-works.md` | full construction details |

build_weights.py

@@ -1,34 +1,40 @@
-"""Compile phrases/*.txt into a Llama-format autocomplete model. No training.
+"""Compile text into a Llama-format autocomplete model. No training.
 
-Every weight is set by a closed-form rule; nothing is fit by gradient descent
-and no other model is used. The result is an exact transformer implementation
-of an interpolated Witten-Bell 5-gram language model over the phrase list.
+Every weight is set by a closed-form rule from n-gram counts; nothing is fit by
+gradient descent and no other model is used. The result is an exact transformer
+implementation of an interpolated n-gram language model (Witten-Bell or
+Kneser-Ney) over the source text.
 
-Residual stream layout (hidden_size = 1 + 4*CODE_DIM + vocab):
-    [0]                     constant 1 (acts as a bias, Llama has none)
-    cur slot   (CODE_DIM)   code of the current token        <- embedding
-    prev1..3   (CODE_DIM)   codes of the 1..3 previous tokens <- attention
-    out slot   (vocab)      next-token logits, scaled by OUT_SCALE <- MLP
+Residual stream layout (hidden_size = 1 + (n_prev + 1) * CODE_DIM + vocab):
+    [0]                       constant 1 (acts as a bias, Llama has none)
+    cur slot     (CODE_DIM)   code of the current token          <- embedding
+    prev1..n     (CODE_DIM)   codes of the 1..n previous tokens   <- attention
+    out slot     (vocab)      next-token logits, scaled by OUT_SCALE <- MLP
 
 Token codes are seeded random unit vectors (nearly orthogonal). Attention head
-a (a = 1..3) has constant query/key vectors whose RoPE phases make the score
-peak at exactly a positions back; its value copies that token's code into
+a (a = 1..n_prev) has constant query/key vectors whose RoPE phases make the
+score peak at exactly a positions back; its value copies that token's code into
 slot prev_a. Positions before the start fall back to position 0 (BOS).
 
-The MLP has one neuron per context seen in the phrases (1-4 tokens, plus one
-always-on unigram neuron). A neuron's gate sums code matches across slots
-minus (k - 0.5), so it fires only when all k tokens match. It writes
-log p(next | ctx) - log p(next | shorter ctx) into the out slot. All suffixes
-of the true context fire together, so the sum telescopes to the log
-probability under the longest matching context.
+The MLP has one neuron per kept context (1..n_prev+1 tokens) plus one always-on
+unigram neuron. A neuron's gate sums code matches across slots minus (k - 0.5),
+so it fires only when all k tokens match. It writes
+log p(next | ctx) - log p(next | shorter ctx) into the out slot. All suffixes of
+the true context fire together, so the sum telescopes to the log probability
+under the longest kept context.
+
+Compile settings live in compile_settings.json ("hf" for this dense model,
+"web" for the sparse browser export written by export_web.py).
 """
 
 import json
 import math
+import zlib
 from collections import Counter, defaultdict
-from functools import lru_cache
+from dataclasses import dataclass
 from pathlib import Path
 
+import numpy as np
 import torch
 from safetensors.torch import save_file
 from tokenizers import Tokenizer, decoders, models, normalizers, pre_tokenizers, processors
@@ -36,18 +42,13 @@ from transformers import GenerationConfig, LlamaConfig, PreTrainedTokenizerFast
 
 ROOT = Path(__file__).parent
 PHRASE_DIR = ROOT / "phrases"  # hand-written everyday phrases
-DATA_DIR = ROOT / "data"  # downloaded human-written prompts (scripts/fetch_coding_prompts.py)
+DATA_DIR = ROOT / "data"  # human-written prompts/replies (scripts/fetch_coding_prompts.py)
+SETTINGS_PATH = ROOT / "compile_settings.json"
 EOS, UNK = "<|endoftext|>", "<unk>"  # EOS doubles as BOS and pad, like GPT-2
-
-PHRASE_WEIGHT = 5  # each hand-written phrase counts as 5 occurrences of data lines
-MIN_WORD_COUNT = 3  # rarer words map to <unk> (which may still appear inside contexts)
-MAX_VOCAB = 3000  # every hand-phrase word, then the most frequent data words; each word costs a hidden dim
-PARAM_TARGET = 400_000_000  # contexts below an auto-chosen count are pruned to fit this
+EOS_ID, UNK_ID = 0, 1
 
 SEED = 1234
 CODE_DIM = 384
-N_PREV = 3  # previous tokens attended to; context = N_PREV + current
-ORDER = N_PREV + 1
 MAX_POS = 128
 ATTN_MARGIN = 30.0  # min attention logit gap (nats) between target and any other position
 BETA = 100.0  # gate sharpness; matched neurons sit at z = BETA / 2
@@ -59,16 +60,33 @@ NORMALIZER = normalizers.Sequence([normalizers.NFKC(), normalizers.Lowercase()])
 PRE_TOKENIZER = pre_tokenizers.ByteLevel(add_prefix_space=True)
 
 
-def pieces(text: str) -> list[str]:
-    return [p for p, _ in PRE_TOKENIZER.pre_tokenize_str(NORMALIZER.normalize_str(text))]
+@dataclass(frozen=True)
+class Settings:
+    n_prev: int = 3  # previous tokens attended to; context = n_prev + current token
+    smoothing: str = "wb"  # "wb" (Witten-Bell) or "kn" (interpolated Kneser-Ney)
+    discount: float = 0.75  # Kneser-Ney absolute discount
+    phrase_weight: float = 5.0  # weight of one hand-written phrase (phrases/*.txt)
+    prompt_weight: float = 1.0  # weight of one user prompt (data/prompts.txt)
+    reply_weight: float = 1.0  # weight of one assistant reply sentence (data/replies.txt)
+    max_vocab: int = 3000
+    min_word_count: float = 3
+    min_count: float | None = None  # None: smallest context count that fits param_target
+    param_target: int = 400_000_000
 
+    @property
+    def order(self) -> int:
+        return self.n_prev + 1
 
-def load_phrases() -> list[list[str]]:
-    return load_lines(sorted(PHRASE_DIR.glob("*.txt")))
 
+def load_settings(name: str) -> Settings:
+    return Settings(**json.loads(SETTINGS_PATH.read_text())[name])
 
-def load_data() -> list[list[str]]:
-    return load_lines(sorted(DATA_DIR.glob("*.txt")))
+
+# ---------- text ----------
+
+
+def pieces(text: str) -> list[str]:
+    return [p for p, _ in PRE_TOKENIZER.pre_tokenize_str(NORMALIZER.normalize_str(text))]
 
 
 def load_lines(paths: list[Path]) -> list[list[str]]:
@@ -85,92 +103,222 @@ def load_lines(paths: list[Path]) -> list[list[str]]:
     return out
 
 
-def padded_events(seq: list[int]):
-    """(contexts of length 0..ORDER, next token) for every position of a sequence."""
-    # ORDER BOS tokens: the last one is the real BOS position the model sees.
-    padded = [0] * ORDER + seq + [0]
-    for i in range(ORDER - 1, len(padded) - 1):
-        yield [tuple(padded[i - k + 1 : i + 1]) if k else () for k in range(ORDER + 1)], padded[i + 1]
+def load_phrases() -> list[list[str]]:
+    return load_lines(sorted(PHRASE_DIR.glob("*.txt")))
+
 
+def load_sources() -> dict[str, list[list[str]]]:
+    return {
+        "phrases": load_phrases(),
+        "prompts": load_lines([DATA_DIR / "prompts.txt"]),
+        "replies": load_lines([DATA_DIR / "replies.txt"]),
+    }
 
-def count_ngrams(weighted: list[tuple[list[int], int]]) -> list[defaultdict]:
-    counts = [defaultdict(Counter) for _ in range(ORDER + 1)]
+
+def heldout_split(lines: list[list[str]], pct: int = 10) -> tuple[list, list]:
+    """Deterministic split by content hash: (train, test)."""
+    test = [t for t in lines if zlib.crc32(" ".join(t).encode()) % 100 < pct]
+    train = [t for t in lines if zlib.crc32(" ".join(t).encode()) % 100 >= pct]
+    return train, test
+
+
+# ---------- counting ----------
+
+
+def weighted_sources(s: Settings, sources: dict) -> list[tuple[list[list[str]], float]]:
+    weights = {"phrases": s.phrase_weight, "prompts": s.prompt_weight, "replies": s.reply_weight}
+    return [(sources[name], w) for name, w in weights.items() if w > 0 and name in sources]
+
+
+def build_vocab(s: Settings, sources: dict) -> list[str]:
+    freq = Counter()
+    for lines, w in weighted_sources(s, sources):
+        for toks in lines:
+            for p in toks:
+                freq[p] += w
+    by_freq = lambda p: (-freq[p], p)  # noqa: E731
+    hand = {p for t in sources["phrases"] for p in t} if s.phrase_weight > 0 else set()
+    rest = sorted((p for p in freq if p not in hand and freq[p] >= s.min_word_count), key=by_freq)
+    words = hand | set(rest[: max(0, s.max_vocab - 2 - len(hand))])
+    return [EOS, UNK] + sorted(words, key=by_freq)
+
+
+def padded_events(seq: list[int], order: int):
+    """(contexts of length 0..order, next token) for every position of a sequence."""
+    # `order` BOS tokens: the last one is the real BOS position the model sees.
+    padded = [EOS_ID] * order + seq + [EOS_ID]
+    for i in range(order - 1, len(padded) - 1):
+        yield [tuple(padded[i - k + 1 : i + 1]) if k else () for k in range(order + 1)], padded[i + 1]
+
+
+def count_ngrams(weighted: list[tuple[list[int], float]], order: int) -> list[defaultdict]:
+    counts = [defaultdict(Counter) for _ in range(order + 1)]
     for seq, weight in weighted:
-        for ctxs, nxt in padded_events(seq):
-            if nxt == 1:  # never learn to predict <unk>
+        for ctxs, nxt in padded_events(seq, order):
+            if nxt == UNK_ID:  # never learn to predict <unk>
                 continue
             for k, ctx in enumerate(ctxs):
                 counts[k][ctx][nxt] += weight
     return counts
 
 
-def kept_contexts(counts: list[defaultdict], min_count: int, forced: set) -> list[dict]:
-    """Contexts that get a neuron: seen >= min_count times, or from the hand-written phrases.
-
-    A suffix is seen at least as often as any context ending in it, and suffixes of
-    hand-written contexts are hand-written, so every kept context's backoff chain is kept.
-    """
-    return [counts[0]] + [
-        {ctx: t for ctx, t in counts[k].items() if ctx in forced or sum(t.values()) >= min_count}
-        for k in range(1, ORDER + 1)
-    ]
-
-
 class NGram:
-    """Interpolated Witten-Bell n-gram over the kept contexts; the reference the weights must reproduce."""
+    """Interpolated n-gram over the kept contexts; the reference the weights must reproduce.
 
-    def __init__(self, counts: list[dict], vocab_size: int):
-        self.V = vocab_size
-        self.counts = counts
-        for k in range(2, ORDER + 1):
-            assert all(ctx[1:] in counts[k - 1] for ctx in counts[k]), "backoff chain broken"
+    Witten-Bell: p_k(j) = (c_j + u * p_{k-1}(j)) / (c + u)
+    Kneser-Ney:  p_k(j) = (c_j - d_j) / c + (sum_j d_j / c) * p_{k-1}(j),  d_j = min(D, c_j),
+                 with continuation counts below the top order (KenLM-style)
+    Both have the form p_k(j) = A_j + gamma * p_{k-1}(j), so each context's neuron can write
+    log p_k(j) - log(gamma * p_{k-1}(j)) for its seen tokens and 0 elsewhere.
+    """
 
-        uni = torch.zeros(vocab_size, dtype=torch.float64)
-        for j, c in self.counts[0][()].items():
+    def __init__(self, raw: list[dict], keep: list[set], vocab_size: int, s: Settings):
+        self.s, self.V, self.order = s, vocab_size, s.order
+        self.support = {ctx: sum(raw[k][ctx].values()) for k in range(1, self.order + 1) for ctx in keep[k]}
+        if s.smoothing == "kn":
+            cont = [defaultdict(Counter) for _ in range(self.order)]
+            for k in range(self.order):
+                for ctx, table in raw[k + 1].items():
+                    for j in table:
+                        cont[k][ctx[1:]][j] += 1
+            self.tables = [cont[0][()]] + [{c: cont[k][c] for c in keep[k]} for k in range(1, self.order)]
+            self.tables.append({c: raw[self.order][c] for c in keep[self.order]})
+        elif s.smoothing == "wb":
+            self.tables = [raw[0][()]] + [{c: raw[k][c] for c in keep[k]} for k in range(1, self.order + 1)]
+        else:
+            raise ValueError(s.smoothing)
+        for k in range(2, self.order + 1):
+            assert all(c[1:] in self.tables[k - 1] for c in self.tables[k]), "backoff chain broken"
+
+        uni = np.zeros(vocab_size)
+        for j, c in self.tables[0].items():
             uni[j] = c
-        logits = torch.log(uni + 1.0)  # add-one smoothing
-        logits[1] = UNK_LOGIT
+        logits = np.log(uni + 1.0)  # add-one smoothing
+        logits[UNK_ID] = UNK_LOGIT
         self.base_logits = logits - logits[2:].mean()  # constant shift; softmax-invariant
-        self.p0 = torch.softmax(self.base_logits, 0)
+        p0 = np.exp(self.base_logits - self.base_logits.max())
+        self.p0 = p0 / p0.sum()
+        self._prob, self._delta, self._stats = {}, {}, {}
+
+    def _table_stats(self, ctx: tuple):
+        if ctx not in self._stats:
+            table = self.tables[len(ctx)][ctx]
+            total = sum(table.values())
+            if self.s.smoothing == "wb":
+                u = len(table)
+                self._stats[ctx] = (table, {j: c / (total + u) for j, c in table.items()}, u / (total + u))
+            else:
+                d = {j: min(self.s.discount, c) for j, c in table.items()}
+                self._stats[ctx] = (table, {j: (c - d[j]) / total for j, c in table.items()}, sum(d.values()) / total)
+        return self._stats[ctx]
 
-    @lru_cache(maxsize=None)
     def prob(self, j: int, ctx: tuple) -> float:
-        if not ctx:
-            return self.p0[j].item()
-        lower = self.prob(j, ctx[1:])
-        table = self.counts[len(ctx)].get(ctx)
-        if table is None:
-            return lower
-        total, uniq = sum(table.values()), len(table)
-        return (table.get(j, 0) + uniq * lower) / (total + uniq)
+        key = (j, ctx)
+        if key not in self._prob:
+            if not ctx:
+                p = self.p0[j]
+            else:
+                lower = self.prob(j, ctx[1:])
+                if ctx in self.tables[len(ctx)]:
+                    _, a, gamma = self._table_stats(ctx)
+                    p = a.get(j, 0.0) + gamma * lower
+                else:
+                    p = lower
+            self._prob[key] = p
+        return self._prob[key]
 
     def delta(self, ctx: tuple) -> dict[int, float]:
-        """log p_k - log p_{k-1}, minus the constant log(u/(c+u)) shared by unseen tokens."""
-        table = self.counts[len(ctx)][ctx]
-        uniq = len(table)
-        out = {}
-        for j, c in table.items():
-            q = self.prob(j, ctx[1:])
-            out[j] = math.log(c + uniq * q) - math.log(uniq * q)
-        return out
+        """log p_k - log p_{k-1}, minus the constant log(gamma) shared by unseen tokens."""
+        if ctx not in self._delta:
+            table, a, gamma = self._table_stats(ctx)
+            out = {}
+            for j in table:
+                q = gamma * self.prob(j, ctx[1:])
+                out[j] = math.log(a[j] + q) - math.log(q)
+            self._delta[ctx] = out
+        return self._delta[ctx]
+
+    def delta_arrays(self, ctx: tuple) -> tuple[np.ndarray, np.ndarray]:
+        d = self.delta(ctx)
+        return np.fromiter(d.keys(), dtype=np.int64, count=len(d)), np.fromiter(d.values(), dtype=np.float64, count=len(d))
 
     def contexts(self) -> list[tuple]:
-        return [ctx for k in range(1, ORDER + 1) for ctx in sorted(self.counts[k])]
+        return [ctx for k in range(1, self.order + 1) for ctx in sorted(self.tables[k])]
 
-    def reference_logits(self, ids: list[int]) -> torch.Tensor:
-        """Logits (up to a constant) the model should emit after the last token of ids."""
+    def window(self, ids: list[int]) -> tuple:
+        """Last order tokens (oldest .. current); positions before 0 read BOS like the attention heads."""
         t = len(ids) - 1
-        window = [ids[max(t - a, 0)] for a in range(N_PREV, -1, -1)]  # oldest .. current
-        logits = self.base_logits.clone()
-        for k in range(1, ORDER + 1):
-            ctx = tuple(window[-k:])
-            if ctx in self.counts[k]:
-                for j, d in self.delta(ctx).items():
-                    logits[j] += d
-        return logits
+        return tuple(ids[max(t - a, 0)] for a in range(self.s.n_prev, -1, -1))
+
+    def predict(self, ids: list[int]) -> tuple[np.ndarray, float, int]:
+        """(logits up to a constant, support of the longest matched context, its length)."""
+        window = self.window(ids)
+        logits = self.base_logits.copy()
+        support, length = 0.0, 0
+        for k in range(1, self.order + 1):
+            ctx = window[-k:]
+            if ctx in self.tables[k]:
+                idx, vals = self.delta_arrays(ctx)
+                logits[idx] += vals
+                support, length = self.support[ctx], k
+        return logits, support, length
+
+    def reference_logits(self, ids: list[int]) -> torch.Tensor:
+        return torch.from_numpy(self.predict(ids)[0])
 
 
-def rope_attention_scale(head_dim: int) -> tuple[float, torch.Tensor, float]:
+def compile_tables(s: Settings, sources: dict):
+    """vocab, NGram and stats for settings s (min_count=None fits the dense param budget)."""
+    vocab = build_vocab(s, sources)
+    index = {p: i for i, p in enumerate(vocab)}
+    weighted = [([index.get(p, UNK_ID) for p in toks], w) for lines, w in weighted_sources(s, sources) for toks in lines]
+    raw = count_ngrams(weighted, s.order)
+    forced = set()
+    if s.phrase_weight > 0:
+        for toks in sources["phrases"]:
+            for ctxs, _ in padded_events([index[p] for p in toks], s.order):
+                forced.update(ctxs[1:])
+
+    totals = {ctx: sum(t.values()) for k in range(1, s.order + 1) for ctx, t in raw[k].items()}
+    min_count = s.min_count
+    if min_count is None:
+        # Smallest count threshold whose kept neurons (forced + total >= threshold) fit the budget.
+        hist = Counter(0 if ctx in forced else c for ctx, c in totals.items())
+        thresholds = sorted(c for c in hist if c > 0)
+        running, kept_at = hist[0], {}
+        for c in reversed(thresholds):
+            running += hist[c]
+            kept_at[c] = running
+        fits = [c for c in thresholds if dense_params(len(vocab), s.n_prev, 1 + kept_at[c]) <= s.param_target]
+        if not fits:
+            raise SystemExit("hand-written contexts alone exceed param_target; lower max_vocab")
+        min_count = fits[0]
+    keep = [set()] + [{c for c in raw[k] if c in forced or totals[c] >= min_count} for k in range(1, s.order + 1)]
+    ngram = NGram(raw, keep, len(vocab), s)
+    stats = {
+        "vocab": len(vocab),
+        "min_context_count": min_count,
+        "contexts_seen": len(totals),
+        "contexts_kept": sum(len(k) for k in keep),
+    }
+    return vocab, ngram, stats
+
+
+# ---------- dense transformer ----------
+
+
+def hidden_size(V: int, n_prev: int) -> int:
+    # transformers validates hidden_size % num_heads == 0; pad with unused zero dims.
+    raw = 1 + (n_prev + 1) * CODE_DIM + V
+    return -(-raw // (8 * n_prev)) * (8 * n_prev)
+
+
+def dense_params(V: int, n_prev: int, neurons: int) -> int:
+    D, h = hidden_size(V, n_prev), CODE_DIM
+    return 2 * V * D + (2 * n_prev + 2) * h * D + 3 * neurons * D + 3 * D
+
+
+def rope_attention_scale(head_dim: int, n_prev: int) -> tuple[float, torch.Tensor, float]:
     """Pick rope_theta maximizing the worst-case attention margin; return (theta, inv_freq, q_amp)."""
     best = None
     for theta in ROPE_THETA_CANDIDATES:
@@ -178,7 +326,7 @@ def rope_attention_scale(head_dim: int) -> tuple[float, torch.Tensor, float]:
         offsets = torch.arange(-MAX_POS, MAX_POS + 1, dtype=torch.float64)
         kernel = torch.cos(offsets[:, None] * inv_freq[None, :]).sum(1)  # K(x), index x + MAX_POS
         worst = math.inf
-        for lag in range(1, N_PREV + 1):
+        for lag in range(1, n_prev + 1):
             for t in range(MAX_POS):
                 s_target = max(t - lag, 0)
                 xs = torch.tensor([t - s - lag for s in range(t + 1)]) + MAX_POS
@@ -195,51 +343,15 @@ def rope_attention_scale(head_dim: int) -> tuple[float, torch.Tensor, float]:
     return theta, inv_freq, q_amp
 
 
-def build(phrases: list[list[str]], data: list[list[str]]):
-    freq = Counter()
-    for toks, w in [(t, PHRASE_WEIGHT) for t in phrases] + [(t, 1) for t in data]:
-        for p in toks:
-            freq[p] += w
-    by_freq = lambda p: (-freq[p], p)  # noqa: E731
-    hand_words = {p for t in phrases for p in t}
-    data_words = sorted((p for p in freq if p not in hand_words and freq[p] >= MIN_WORD_COUNT), key=by_freq)
-    words = hand_words | set(data_words[: max(0, MAX_VOCAB - 2 - len(hand_words))])
-    vocab = [EOS, UNK] + sorted(words, key=by_freq)
-    index = {p: i for i, p in enumerate(vocab)}
-    V, m = len(vocab), CODE_DIM
-
-    def ids(toks: list[str]) -> list[int]:
-        return [index.get(p, 1) for p in toks]
-
-    counts = count_ngrams([(ids(t), PHRASE_WEIGHT) for t in phrases] + [(ids(t), 1) for t in data])
-    forced = {ctx for t in phrases for ctxs, _ in padded_events(ids(t)) for ctx in ctxs[1:]}
-
-    CONST, CUR = 0, 1
-    prev_slot = [1 + a * m for a in range(N_PREV + 1)]  # prev_slot[0] is the cur slot
-    OUT = 1 + (N_PREV + 1) * m
-    n_heads, h = N_PREV, m
-    # transformers validates hidden_size % num_heads == 0; pad with unused zero dims.
-    D = -(-(OUT + V) // (8 * n_heads)) * (8 * n_heads)
-
-    def n_params(neurons: int) -> int:
-        return 2 * V * D + (2 * n_heads + 2) * h * D + 3 * neurons * D + 3 * D
-
-    totals = [Counter() for _ in range(ORDER + 1)]  # how many contexts have each total count
-    for k in range(1, ORDER + 1):
-        for ctx, t in counts[k].items():
-            totals[k][0 if ctx in forced else sum(t.values())] += 1
-    max_total = max(c for k in range(1, ORDER + 1) for c in totals[k])
-    min_count = 1
-    while True:
-        N = 1 + sum(n for k in range(1, ORDER + 1) for c, n in totals[k].items() if c == 0 or c >= min_count)
-        if n_params(N) <= PARAM_TARGET:
-            break
-        if min_count > max_total:
-            raise SystemExit(f"hand-written contexts alone need {n_params(N):,} params > PARAM_TARGET; lower MAX_VOCAB")
-        min_count += 1
-    ngram = NGram(kept_contexts(counts, min_count, forced), V)
+def build_dense(vocab: list[str], ngram: NGram, s: Settings):
+    V, m, n_prev = len(vocab), CODE_DIM, s.n_prev
     contexts = ngram.contexts()
-    assert 1 + len(contexts) == N
+    N = 1 + len(contexts)
+    CONST, CUR = 0, 1
+    prev_slot = [1 + a * m for a in range(n_prev + 1)]  # prev_slot[0] is the cur slot
+    OUT = 1 + (n_prev + 1) * m
+    n_heads, h = n_prev, m
+    D = hidden_size(V, n_prev)
 
     gen = torch.Generator().manual_seed(SEED)
     codes = torch.randn(V, m, generator=gen, dtype=torch.float64)
@@ -253,7 +365,7 @@ def build(phrases: list[list[str]], data: list[list[str]]):
     embed[:, CONST] = 1.0
     embed[:, CUR : CUR + m] = codes
 
-    theta, inv_freq, q_amp = rope_attention_scale(h)
+    theta, inv_freq, q_amp = rope_attention_scale(h, n_prev)
     half = h // 2
     q_proj = torch.zeros(n_heads * h, D, dtype=torch.float64)
     for a in range(n_heads):
@@ -275,7 +387,7 @@ def build(phrases: list[list[str]], data: list[list[str]]):
     on = (BETA / 2) * torch.sigmoid(torch.tensor(BETA / 2, dtype=torch.float64)).item()  # silu(BETA/2)
     up[:, CONST] = 1.0
     gate[0, CONST] = BETA / 2  # always-on unigram neuron
-    down[OUT : OUT + V, 0] = OUT_SCALE * ngram.base_logits / on
+    down[OUT : OUT + V, 0] = torch.from_numpy(OUT_SCALE * ngram.base_logits / on).float()
     for n, ctx in enumerate(contexts, start=1):
         k = len(ctx)
         for a, tok in enumerate(reversed(ctx)):  # a = 0 current, a = 1 previous, ...
@@ -288,9 +400,9 @@ def build(phrases: list[list[str]], data: list[list[str]]):
     lm_head[torch.arange(V), OUT + torch.arange(V)] = 1.0 / OUT_SCALE
 
     # RMSNorm weights undo the norm for the (fixed) input norms: |x|^2 = 2 before
-    # attention, 1 + 1 + N_PREV after it.
+    # attention, 1 + 1 + n_prev after it.
     attn_norm = math.sqrt(2 / D)
-    mlp_norm = math.sqrt((2 + N_PREV) / D)
+    mlp_norm = math.sqrt((2 + n_prev) / D)
     state = {
         "model.embed_tokens.weight": embed,
         "model.layers.0.input_layernorm.weight": torch.full((D,), attn_norm),
@@ -325,16 +437,13 @@ def build(phrases: list[list[str]], data: list[list[str]]):
         pad_token_id=0,
     )
     info = {
-        "vocab": V,
         "hidden": D,
         "neurons": N,
-        "min_context_count": min_count,
-        "contexts_seen": sum(len(counts[k]) for k in range(1, ORDER + 1)),
         "rope_theta": theta,
         "max_code_coherence": round(max_coherence, 3),
         "params": sum(t.numel() for t in state.values()),
     }
-    return vocab, ngram, state, config, info
+    return state, config, info
 
 
 def build_tokenizer(vocab: list[str]) -> PreTrainedTokenizerFast:
@@ -355,13 +464,17 @@ def build_tokenizer(vocab: list[str]) -> PreTrainedTokenizerFast:
 
 
 def compile_model():
-    """Build everything in memory (used by the tests for the n-gram reference)."""
-    phrases, data = load_phrases(), load_data()
-    return ({"phrases": len(phrases), "data_lines": len(data)}, *build(phrases, data))
+    """Build the dense ("hf") model in memory (used by the tests for the n-gram reference)."""
+    s = load_settings("hf")
+    sources = load_sources()
+    vocab, ngram, stats = compile_tables(s, sources)
+    state, config, info = build_dense(vocab, ngram, s)
+    lines = {name: len(v) for name, v in sources.items()}
+    return lines, vocab, ngram, state, config, {**stats, **info}
 
 
 def main() -> None:
-    sources, vocab, _, state, config, info = compile_model()
+    lines, vocab, _, state, config, info = compile_model()
     assert info["params"] <= 1_000_000_000, info
 
     save_file(state, ROOT / "model.safetensors", metadata={"format": "pt"})
@@ -386,7 +499,7 @@ def main() -> None:
     tok_cfg["model_input_names"] = ["input_ids", "attention_mask"]
     tok_cfg_path.write_text(json.dumps(tok_cfg, indent=2) + "\n")
 
-    print(json.dumps({**sources, **info}, indent=2))
+    print(json.dumps({"lines": lines, **info}, indent=2))
 
 
 if __name__ == "__main__":

compile_settings.json

@@ -0,0 +1,27 @@
+{
+  "_comment": "Chosen by scripts/sweep.py (docs/sweep.json): best keystrokes saved at >= 90% precision on held-out prompts. hf = dense transformers model (vocab capped for width); web = sparse browser export.",
+  "hf": {
+    "n_prev": 3,
+    "smoothing": "kn",
+    "phrase_weight": 1.0,
+    "prompt_weight": 1.0,
+    "reply_weight": 1.0,
+    "max_vocab": 3000,
+    "min_count": null,
+    "param_target": 400000000
+  },
+  "web": {
+    "n_prev": 3,
+    "smoothing": "kn",
+    "phrase_weight": 1.0,
+    "prompt_weight": 1.0,
+    "reply_weight": 1.0,
+    "max_vocab": 8000,
+    "min_count": 2
+  },
+  "policy": {
+    "tau": 0.72,
+    "min_support": 0,
+    "max_words": 8
+  }
+}

data/SOURCES.md

@@ -1,19 +1,30 @@
 # Data sources
 
-`coding_prompts.txt` is produced by `scripts/fetch_coding_prompts.py`. Every line is
-human-written English prose taken from the datasets below. LLM-generated messages are
-excluded (OASST `synthetic` flag). The script keeps only prose: it removes code blocks,
-code-looking lines and markdown, keeps up to 2 sentences per prompt and 5 per reply, and
-keeps sentences of 3–24 words.
+`prompts.txt` and `replies.txt` are produced by `scripts/fetch_coding_prompts.py`. Every line is
+human-written English prose taken from the datasets below. LLM-generated messages are excluded
+(OASST `synthetic` flag). The script keeps only prose: it removes code blocks, code-looking
+lines and markdown, keeps up to 2 sentences per prompt and 5 per reply, and keeps sentences
+of 3–24 words.
 
-| Dataset | What is used | License | Attribution |
-|---|---|---|---|
-| [OpenAssistant/oasst2](https://huggingface.co/datasets/OpenAssistant/oasst2) | English prompter **and** assistant messages (all turns, reviewed, not deleted, not synthetic) from conversations whose opening prompt is about code | Apache-2.0 | Köpf et al., *OpenAssistant Conversations: Democratizing Large Language Model Alignment*, 2023 |
-| [google-research-datasets/mbpp](https://huggingface.co/datasets/google-research-datasets/mbpp) | task descriptions (`text`), all splits | CC-BY-4.0 | Austin et al., *Program Synthesis with Large Language Models*, 2021 |
-| [neulab/conala](https://huggingface.co/datasets/neulab/conala) | curated intents (`intent`), paired train + test | MIT | Yin et al., *Learning to Mine Aligned Code and Natural Language Pairs from Stack Overflow*, MSR 2018 |
+| File | Dataset | What is used | License | Attribution |
+|---|---|---|---|---|
+| `prompts.txt` | [OpenAssistant/oasst2](https://huggingface.co/datasets/OpenAssistant/oasst2) | English **user** messages (all turns, reviewed, not deleted, not synthetic) from conversations whose opening prompt is about code | Apache-2.0 | Köpf et al., *OpenAssistant Conversations: Democratizing Large Language Model Alignment*, 2023 |
+| `prompts.txt` | [google-research-datasets/mbpp](https://huggingface.co/datasets/google-research-datasets/mbpp) | task descriptions (`text`), all splits | CC-BY-4.0 | Austin et al., *Program Synthesis with Large Language Models*, 2021 |
+| `prompts.txt` | [neulab/conala](https://huggingface.co/datasets/neulab/conala) | curated intents (`intent`), paired train + test | MIT | Yin et al., *Learning to Mine Aligned Code and Natural Language Pairs from Stack Overflow*, MSR 2018 |
+| `replies.txt` | [OpenAssistant/oasst2](https://huggingface.co/datasets/OpenAssistant/oasst2) | English **assistant** messages, written by human volunteers, from the same conversations | Apache-2.0 | as above |
 
-Counts from the last fetch: 17,352 OASST sentences, 938 MBPP, 2,820 CoNaLa, and 19,974
-unique lines after removing duplicates.
+Counts from the last fetch:
 
-The model weights contain only n-gram statistics derived from these lines, plus the
-hand-written phrases in `../phrases/`. No dataset was used to train anything.
+| Source | Sentences |
+|---|---|
+| OASST user turns | 4,169 |
+| MBPP | 938 |
+| CoNaLa | 2,820 |
+| OASST assistant turns | 13,183 |
+
+After removing duplicates: **7,133** lines in `prompts.txt` and **12,841** in `replies.txt`.
+
+The weights contain only n-gram statistics derived from these lines, plus the hand-written
+phrases in `../phrases/`. No dataset was used to train anything. 10% of `prompts.txt` (chosen
+by content hash) is held out when measuring precision and keystrokes saved
+(`scripts/sweep.py`, `scripts/benchmark.py`).

docs/benchmarks.json

@@ -7,41 +7,123 @@
     "python": "3.13.4",
     "torch": "2.15.0"
   },
-  "compile": {
-    "compile_seconds": 6.2,
-    "compile_runs": [
-      6.4,
-      6.2
+  "compile_dense": {
+    "seconds": 6.9,
+    "runs": [
+      6.9,
+      7.0
     ],
-    "peak_memory_gb": 2.3,
     "byte_identical_rebuild": true,
-    "model_sha256": "93c23040ccf6393710b0ee1d2fff98d57dc4193a27c4911cac7374e2c28f50f6",
-    "file_gb": 1.52,
-    "tokenize_seconds": 0.49,
-    "build_tensors_seconds": 2.74,
+    "sha256": "ae857db3e321c35782192c198a7778987964e9cc147d2c67428fea9bf5777100",
+    "mb": 1521.16,
+    "peak_memory_gb": 2.4,
+    "tokenize_seconds": 0.48,
+    "count_seconds": 2.37,
+    "build_tensors_seconds": 1.07,
     "source_lines": 21106,
     "params": 380290320,
     "nonzero_params": 26792230,
+    "vocab": 3000,
+    "min_context_count": 8.0,
     "contexts_seen": 434284,
-    "neurons": 24774
+    "contexts_kept": 24773
+  },
+  "compile_web": {
+    "seconds": 6.6,
+    "runs": [
+      7.2,
+      6.6
+    ],
+    "byte_identical_rebuild": true,
+    "sha256": "d770bf4ec6644ccab5646e2b1244b58b185378b371c52b1c263b5a0ebf3f8499",
+    "mb": 4.49,
+    "contexts": 107864,
+    "vocab": 6612
   },
   "pytorch": {
-    "load_seconds": 0.3,
-    "next_word_ms": 70.3,
-    "generate_tokens_per_second": 16.4,
+    "load_seconds": 0.1,
+    "next_word_ms": 57.9,
+    "generate_tokens_per_second": 18.0,
     "threads": 8
   },
-  "typescript": {
-    "runtime": "bun, plain TypeScript, no ML library",
-    "load_seconds": 0.4,
-    "seconds_per_token": 0.44
+  "web_keystroke_latency": {
+    "runtime": "web/autocomplete.ts on Bun",
+    "median_ms": 0.5,
+    "p95_ms": 0.92,
+    "max_ms": 3.48
+  },
+  "max_logprob_error_vs_ngram": 1.6793374655321713e-05,
+  "heldout_typing_web": {
+    "policy": {
+      "tau": 0.72,
+      "min_support": 0,
+      "max_words": 8
+    },
+    "heldout_prompts": 711,
+    "characters": 42185,
+    "tau=0.72": {
+      "precision": 0.9029,
+      "coverage": 0.3866,
+      "keystrokes_saved": 0.2296,
+      "shown": 16034
+    },
+    "tau=0.8": {
+      "precision": 0.9351,
+      "coverage": 0.3412,
+      "keystrokes_saved": 0.2066,
+      "shown": 14153
+    },
+    "tau=0.9": {
+      "precision": 0.9644,
+      "coverage": 0.2692,
+      "keystrokes_saved": 0.1642,
+      "shown": 11165
+    },
+    "tau=0.95": {
+      "precision": 0.9773,
+      "coverage": 0.2049,
+      "keystrokes_saved": 0.1176,
+      "shown": 8497
+    }
+  },
+  "heldout_typing_dense_tables": {
+    "policy": {
+      "tau": 0.72,
+      "min_support": 0,
+      "max_words": 8
+    },
+    "heldout_prompts": 711,
+    "characters": 42185,
+    "tau=0.72": {
+      "precision": 0.9243,
+      "coverage": 0.3245,
+      "keystrokes_saved": 0.1902,
+      "shown": 13458
+    },
+    "tau=0.8": {
+      "precision": 0.9516,
+      "coverage": 0.2698,
+      "keystrokes_saved": 0.1574,
+      "shown": 11191
+    },
+    "tau=0.9": {
+      "precision": 0.9736,
+      "coverage": 0.1921,
+      "keystrokes_saved": 0.1057,
+      "shown": 7966
+    },
+    "tau=0.95": {
+      "precision": 0.9843,
+      "coverage": 0.1322,
+      "keystrokes_saved": 0.0698,
+      "shown": 5481
+    }
   },
-  "max_logprob_error_vs_ngram": 2.4948840973593178e-05,
-  "heldout": {
-    "heldout_lines": 2040,
-    "heldout_tokens": 35039,
-    "next_token_top1": 0.258,
-    "next_token_top3": 0.396,
-    "oov_rate": 0.088
+  "heldout_next_token_dense": {
+    "heldout_lines": 711,
+    "heldout_tokens": 9466,
+    "next_token_top1": 0.267,
+    "next_token_top3": 0.406,
+    "oov_rate": 0.07
   }
 }

docs/experiment_topic.json

@@ -0,0 +1,41 @@
+{
+  "topic_words": 48,
+  "results": [
+    {
+      "lambda": 0.0,
+      "tau": 0.72,
+      "min_support": 0,
+      "precision": 0.9029,
+      "coverage": 0.3866,
+      "keystrokes_saved": 0.2296,
+      "shown": 16034
+    },
+    {
+      "lambda": 0.25,
+      "tau": 0.72,
+      "min_support": 0,
+      "precision": 0.9032,
+      "coverage": 0.3864,
+      "keystrokes_saved": 0.2297,
+      "shown": 16026
+    },
+    {
+      "lambda": 0.5,
+      "tau": 0.72,
+      "min_support": 0,
+      "precision": 0.9031,
+      "coverage": 0.386,
+      "keystrokes_saved": 0.2293,
+      "shown": 16007
+    },
+    {
+      "lambda": 1.0,
+      "tau": 0.73,
+      "min_support": 0,
+      "precision": 0.904,
+      "coverage": 0.3811,
+      "keystrokes_saved": 0.2281,
+      "shown": 15804
+    }
+  ]
+}

docs/how-it-works.md

@@ -1,60 +1,62 @@
 # How tiny-autocomplete was built
 
-This model was **compiled, not trained**. No gradient descent, no optimizer, no
-pretrained weights, and no other model's outputs were used. `build_weights.py` turns
-text into a standard `LlamaForCausalLM` checkpoint using closed-form rules.
-Running it twice gives byte-identical weights.
+This model was **compiled, not trained**. No gradient descent, no optimizer, no pretrained
+weights, and no other model's outputs were used.
+- `build_weights.py` turns text into a standard `LlamaForCausalLM` checkpoint using
+  closed-form rules.
+- `export_web.py` writes the same kind of network in sparse form for the browser.
+- Rebuilds are byte-identical.
 
-The network computes an **interpolated Witten-Bell 5-gram language model** exactly. It
-predicts the next token from the current token and up to 3 before it.
+The network computes an **interpolated Kneser-Ney 5-gram language model** exactly. It predicts
+the next token from the current token and the 3 before it. Witten-Bell smoothing and other
+context lengths are supported too; the choice was made by measurement (section 9).
 
 ## 1. Source text and tokenizer
 
-There are two sources:
+| Source | Lines | Weight |
+|---|---|---|
+| `phrases/*.txt`: everyday English, written for this project by an AI coding assistant (Claude) | 1,139 | 1 |
+| `data/prompts.txt`: what people type to an assistant (OASST user turns, MBPP, CoNaLa) | 7,133 | 1 |
+| `data/replies.txt`: human-written assistant replies from OASST coding conversations | 12,834 | 1 |
 
-- **`phrases/*.txt`:** 1,139 common English sentences, one per line, in four files: social
-  phrases, everyday requests, idioms and proverbs, and opinions and work phrases. This
-  list was written for this project by an AI coding assistant (Claude). Each phrase counts
-  as **5** occurrences, and all of its contexts are always kept (section 5).
-- **`data/coding_prompts.txt`:** 19,967 human-written English sentences about programming.
-  `scripts/fetch_coding_prompts.py` downloads them from OpenAssistant oasst2 (prompts and
-  replies in coding conversations, with synthetic messages excluded), MBPP and CoNaLa. It
-  strips code and markdown and keeps short prose sentences. See `data/SOURCES.md` for
-  licenses and attribution. Each line counts once.
+`scripts/fetch_coding_prompts.py` downloads the data, removes code blocks, code-looking lines and
+markdown, and keeps short prose sentences. Licenses and attribution are in `data/SOURCES.md`.
 
 The tokenizer is word-level with GPT-2 style byte-level pre-tokenization:
 
 - NFKC normalization, then lowercase.
 - Split on the GPT-2 regex. Words keep their leading space as `Ġ`, and punctuation and
   contractions (`'s`, `'t`, ...) become their own tokens.
-- The vocabulary is every word in the hand-written phrases, then the most frequent data
-  words seen ≥ 3 times, up to 3,000 entries in total. Each word costs one hidden dimension
-  (the logit slot), so the cap matters. The two specials are `<|endoftext|>` (id 0, used as
-  BOS, EOS and pad) and `<unk>` (id 1). Rarer words become `<unk>`. `<unk>` can appear
-  inside contexts ("convert <unk> to a") but is never predicted.
+- Vocabulary:
+  - every hand-written word, then the most frequent data words seen ≥ 3 times;
+  - capped at **3,000** for the dense model, where each word costs one hidden dimension;
+  - capped at 8,000 for the browser export, where **6,612** words qualify;
+  - two specials: `<|endoftext|>` (id 0, used as BOS, EOS and pad) and `<unk>` (id 1);
+  - rarer words become `<unk>`, which can appear inside contexts but is never predicted.
 - A post-processor prepends BOS, so every prompt starts at position 0 with `<|endoftext|>`.
 
 ## 2. The reference n-gram
 
-Each line becomes `[BOS, BOS, BOS, BOS, w1, ..., wn, EOS]`. For every position, the
-compiler counts which token follows each context of length k = 0..4 that ends there
-(weighted by source).
+Each line becomes `[BOS × 4, w1, ..., wn, EOS]`. For every position, the compiler counts which
+token follows each context of length k = 0..4 that ends there, weighted by source.
 
 - **Order 0:** add-one-smoothed unigram, with `<unk>` pinned at logit -30.
-- **Order k ≥ 1:** Witten-Bell interpolation with the next-shorter context (drop the
-  oldest token):
+- **Order k ≥ 1:** interpolate with the next-shorter context `ctx[1:]`. Both smoothing methods
+  have the form `p_k(j) = A_j + γ · p_{k-1}(j)`:
 
   ```
-  p_k(j | ctx) = (c(ctx, j) + u(ctx) * p_{k-1}(j | ctx[1:])) / (c(ctx) + u(ctx))
+  Witten-Bell:  A_j = c_j / (c + u)                  γ = u / (c + u)
+  Kneser-Ney:   A_j = (c_j - d_j) / c,  d_j = min(0.75, c_j)    γ = Σ d_j / c
   ```
 
-  Here `c` is a count and `u(ctx)` is the number of distinct tokens seen after `ctx`.
-  An unseen context backs off: `p_k = p_{k-1}`.
+  Here `c_j` is the count of token j after the context, `c` the total, and `u` the number of
+  distinct next tokens. Kneser-Ney uses **continuation counts** (how many distinct tokens
+  precede `ctx + j`) for every order below the top one, as KenLM does. An unseen or pruned
+  context backs off: `p_k = p_{k-1}`.
 
-The model must reproduce `p_K(· | longest kept context)`, where K is the length of the
-longest context that has a neuron (section 5).
+The model must reproduce `p_K(· | longest kept context)`.
 
-## 3. The residual stream as named slots
+## 3. The residual stream as named slots (dense model)
 
 | Dims | Slot | Written by |
 |---|---|---|
@@ -66,20 +68,19 @@ longest context that has a neuron (section 5).
 | 1537–4536 | **logits**, one dim per vocab entry, scaled by 1e-5 | MLP |
 | 4537–4559 | unused padding (transformers wants hidden % heads == 0) | — |
 
-**Token codes** are seeded Gaussian vectors in 384 dimensions, normalized to unit
-length. With 3,000 of them, the largest |cosine| between two different codes is 0.256.
-That's far enough below 1 that a "match" (dot = 1) and a "non-match" (dot ≤ 0.26) can't
-be confused.
+**Token codes** are seeded Gaussian vectors in 384 dimensions, normalized to unit length. With
+3,000 of them, the largest |cosine| between two different codes is 0.256. That's far enough
+below 1 that a "match" (dot = 1) and a "non-match" (dot ≤ 0.26) can't be confused.
 
 ## 4. Attention: previous-token heads from RoPE alone
 
-Llama has no absolute positions. The only position signal is RoPE, which rotates query
-and key pairs by `position * inv_freq[i]`.
+Llama has no absolute positions. The only position signal is RoPE, which rotates query and key
+pairs by `position * inv_freq[i]`.
 
-- The **key** for every token is the same constant (read from the constant dim): 1 on
-  the first half of each RoPE pair.
+- The **key** for every token is the same constant (read from the constant dim): 1 on the
+  first half of each RoPE pair.
 - The **query** for head *a* is also constant, with each pair's phase set to
-  `-a * inv_freq[i]`. The attention score between positions t and s is then
+  `-a * inv_freq[i]`. Then
 
   ```
   score(t, s) = A * Σ_i cos((t - s - a) * inv_freq[i]) / sqrt(head_dim)
@@ -88,91 +89,154 @@ and key pairs by `position * inv_freq[i]`.
   which peaks exactly when `s = t - a`.
 - When `t < a` there is no position `a` back, and the closest position is 0, which is BOS.
   The model then "sees" BOS in that slot, matching the BOS padding used when counting.
-- `rope_theta` is picked from a candidate list to maximize the worst-case gap between the
-  target position and every other position, for all t < 128 and a ∈ {1, 2, 3}. It came out
-  to 2. The amplitude `A` then scales that worst gap to **30 nats**, so softmax puts
-  > 0.9999999 weight on the target.
+- `rope_theta` is picked to maximize the worst-case gap between the target position and every
+  other position, for all t < 128 and every head. It came out to 2. `A` scales that gap to
+  **30 nats**.
 - **Value** = the current-token code. **Output projection** = identity into slot *a*.
 
-All 3 heads share one key/value head (`num_key_value_heads = 1`).
+All heads share one key/value head (`num_key_value_heads = 1`).
 
 ## 5. MLP: one neuron per kept context
 
-Llama's MLP is `down(silu(gate(x)) * up(x))`. The compiler creates one neuron for each
-**kept** context (lengths 1–4) and one always-on unigram neuron.
-
-**Pruning to a budget.** The sources contain 434,284 distinct contexts, and each neuron
-costs 3 × hidden ≈ 13.7k parameters. A context is kept if it was seen at least
-`min_count` times or comes from the hand-written phrases. `min_count` is raised from 1
-until the model fits `PARAM_TARGET` (400M). It came out to **8**, keeping 24,773
-contexts. Pruning never breaks a backoff chain: a context's shorter suffix is always seen
-at least as often, and suffixes of hand-written contexts are hand-written. The reference
-n-gram uses the same kept set, so the model still matches it exactly.
-
-For a context `(c_{k-1}, ..., c_1, c_0)` of length k (c_0 = current token):
+Llama's MLP is `down(silu(gate(x)) * up(x))`. The compiler creates one neuron for each **kept**
+context (lengths 1–4) and one always-on unigram neuron. For a context `(c_{k-1}, ..., c_0)`
+(c_0 = current token):
 
 - **gate** = `100 * (code(c_0)·cur + code(c_1)·prev1 + ... - (k - 0.5) * const)`.
-  If all k tokens match, this is `+50`. If any one fails (dot ≤ 0.26), it's at most
-  `100 * (0.26 - 0.5) = -24`, and `silu(-24)` ≈ -1e-9, effectively off.
+  - If all k tokens match, this is `+50`.
+  - If any one fails, it's at most `100 * (0.26 - 0.5) = -24`, and `silu(-24)` ≈ -1e-9,
+    effectively off.
 - **up** = 1 (the constant dim).
-- **down** writes, for each token j seen after this context,
+- **down** writes, for each token j seen after the context, `log p_k(j) − log(γ · p_{k-1}(j))`.
+  Every unseen token would get the same `log γ`, which softmax ignores, so the column is sparse.
+
+**Why this is exact.** If a context of length k is kept, so are its shorter suffixes, and
+their neurons fire too. The unigram neuron writes `log p_0`, and the neurons for lengths 1..K
+add their deltas. The sum telescopes to `log p_K(· | longest kept context)` plus a constant.
+
+**Pruning.** A context is kept if it was seen at least `min_count` times (weighted) or comes
+from the hand-written phrases. A suffix is always seen at least as often as the contexts
+ending in it, so pruning never breaks a backoff chain. The reference n-gram uses the same kept
+set, so the model still matches it exactly.
+- **Dense model:** `min_count` is the smallest value that fits 400M parameters. That's **8**,
+  keeping 24,773 of 434,284 contexts. Each neuron costs 3 × hidden ≈ 13.7k parameters.
+- **Browser export:** `min_count` = **2**, keeping 107,864 of 520,490 contexts, because a
+  sparse neuron costs only a few bytes.
+
+## 6. Norms and output head (dense model)
+
+- The residual has a known fixed norm at each point: |x|² = 2 before attention and 5 after.
+  Each RMSNorm weight is `sqrt(norm² / hidden)`, so the norms leave the stream unchanged.
+- The MLP writes logits scaled by 1e-5, so they barely change the final norm (< 0.001%). At
+  1e-4, the 3,000-entry logit slot shifted the norm enough to cause 1e-2-nat errors on extreme
+  logits.
+- `lm_head` is `1e5` times the identity on the logit slot.
 
-  ```
-  delta_j = log p_k(j | ctx) - log p_{k-1}(j | ctx[1:])
-  ```
+## 7. Browser export (`web/model.bin`, `web/autocomplete.ts`)
 
-  into the logit slot. The constant `log(u/(c+u))` shared by every unseen token is
-  dropped, since softmax ignores constant shifts. That makes the column sparse.
+93% of the dense weights are zero. The attention and embedding weights are fixed by
+construction, and each MLP neuron is fully described by:
+- which tokens its gate matches (its context);
+- which logits its down-projection changes, and by how much.
 
-**Why this is exact.** If a context of length k is kept, so are all its shorter
-suffixes, and their neurons fire too. The unigram neuron writes `log p_0`, and the
-neurons for lengths 1..K add their deltas. The sum telescopes to `log p_K(· | longest
-known context)` plus a constant. No neuron has to know which context is the longest.
+`export_web.py` stores exactly that, plus the unigram neuron and the vocabulary:
+- **contents:** per context its token ids, its count (support) and the logit changes it
+  writes (as float32);
+- **size:** 4.5 MB for 107,864 contexts.
 
-## 6. Norms and output head
+`web/autocomplete.ts` looks up the neurons whose context matches the last 1–4 tokens (hash
+maps per context length), adds their deltas to the unigram logits, and applies softmax.
+- **Speed:** 0.5 ms median per keystroke.
+- **Exactness:** a test exports the dense model's own tables and checks that the runtime's
+  top-5 next tokens and log-probs equal PyTorch's.
 
-- The residual has a known fixed norm at each point: |x|² = 2 before attention and 5
-  after it. Each RMSNorm weight is set to `sqrt(norm² / hidden)`, so the norms leave the
-  stream unchanged.
-- The MLP writes logits scaled by 1e-5, so they barely change the final norm (< 0.001%).
-  At 1e-4, the 3,000-entry logit slot shifted the norm enough to cause 1e-2-nat errors on
-  extreme logits.
-- `lm_head` is `1e5` times the identity on the logit slot.
+## 8. Confidence policy (`suggest.py`, mirrored by `web/autocomplete.ts`)
+
+Nothing here is learned; the policy reads the compiled probabilities.
+
+1. **Candidates.** The last pre-token is the word being typed (just `Ġ` after a space). Only
+   vocabulary tokens starting with it are candidates (`fun` → `function`).
+2. **Probability.** Candidate probabilities are renormalized over the candidates, then scaled
+   by `1 − P(OOV | prefix)`: the weighted share of source words with this prefix that are
+   outside the vocabulary. Without that, a prefix whose only in-vocabulary completion is
+   `function` would get 100% even when the user is typing `funky`.
+3. **Show** the top candidate only if its probability ≥ τ (and the matched context's support
+   ≥ `min_support`; the sweep found 0 best).
+4. **Extend** greedily word by word while the product of probabilities stays ≥ τ, up to 8
+   words, stopping at end-of-text or after `. ? ! :`.
 
-## 7. Verification (`test_model.py`)
+τ = 0.72 by default. It's a runtime argument (`ac.suggest(text, {...ac.policy, tau})`).
 
-The tests check the following:
-- **Hugging Face loading:** the files load with `AutoModelForCausalLM`, `AutoTokenizer` and
-  `pipeline("text-generation")` on transformers 4.52 and 5.22.
-- **Exactness:** next-token log-probs match the reference n-gram within 1e-3 nats (observed
-  ~3e-5). This is checked on 200 phrase prefixes and 50 random word-salad sequences of up to
-  80 tokens.
-- **Completions:** a set of prompts produce the expected greedy completions.
-- **Unknown words:** out-of-vocabulary words back off cleanly.
-- **TypeScript parity:** the zero-dependency TypeScript runtime (`model.ts`) gives the same
-  top-3 tokens and log-probs as PyTorch.
+## 9. Choosing the settings (`scripts/sweep.py`)
 
-## 8. Changing it
+The compile rules are fixed. A handful of compile settings and the threshold were picked by
+benchmarking:
 
-Edit or add lines in `phrases/*.txt`, or add any `.txt` file to `data/` (one sentence per
-line). Then run `python build_weights.py` and `python test_model.py`. Adding data never
-increases the size; it raises the pruning threshold instead, so only the most frequent
-patterns survive. Settings at the top of `build_weights.py`:
+1. Hold out 10% of `data/prompts.txt` (711 prompts) and compile without them.
+2. Type each held-out prompt one character at a time. After each keystroke, record the
+   suggestion chain.
+3. Score every threshold: **precision** (shown suggestions that exactly match the next
+   characters), **coverage**, and **keystrokes saved** (accepting each correct suggestion with
+   one Tab).
+4. Keep the setting with the most keystrokes saved at precision ≥ 90%, and an export ≤ 50 MB.
 
-| Constant | Meaning |
+Highlights from `docs/sweep.json` (keystrokes saved at ≥ 90% precision, 0.01 threshold grid):
+
+| Setting | Keystrokes saved |
+|---|---|
+| Witten-Bell, 4-token context (phrases ×5) | 21.7% |
+| **Kneser-Ney, 4-token context (phrases ×5)** | **22.5%** |
+| Kneser-Ney, 5 / 6 / 8 / 10-token context (phrases ×5) | 22.2% / 22.0% / 21.6% / 21.6% |
+| no assistant replies (prompts + phrases only) | 20.1–20.9% |
+| assistant replies ×0.3 | 21.8–22.4% |
+| **final: KN, 4 tokens, all sources ×1, vocab 8,000, min_count 2** | **23.0%** |
+| final settings with 5 / 6 / 8 / 10-token context | 22.8% / 22.7% / 22.4% / 22.4% |
+
+- **Longer contexts:** they lose slightly. With ~21k lines, 5+-token contexts rarely repeat,
+  and when they don't match the model is back to the same 4-token prediction.
+- **Topic head** (`scripts/experiment_topic.py`): a content-addressed head that would carry
+  the most recent language or tool name (python, sql, react, …) from anywhere earlier in the
+  prompt.
+  - Tested as a closed-form log-ratio bias `λ · log p(w | topic) / p(w)`.
+  - At best it left keystrokes saved at 23.0% (22.97% at λ = 0.25), and λ = 1 was worse.
+  - It wasn't compiled.
+
+## 10. Verification (`test_model.py`)
+
+- **Hugging Face loading** with `AutoModelForCausalLM`, `AutoTokenizer` and `pipeline` on
+  transformers 4.52 and 5.22.
+- **Exactness:** dense log-probs match the reference n-gram within 1e-3 nats (observed ~2e-5)
+  on 200 phrase prefixes and 50 random sequences of up to 80 tokens.
+- **Completions:** a set of greedy completions; unknown words back off cleanly.
+- **Dense TypeScript runtime** (`model.ts`): same top-3 tokens and log-probs as PyTorch.
+- **Browser runtime:**
+  - the same suggestion chains (text, probability, support) as `suggest.py` on 300 random
+    prefixes;
+  - the same distribution as PyTorch when given the dense model's tables;
+  - `model.bin` ≤ 50 MB.
+
+## 11. Changing it
+
+Edit `phrases/*.txt` or `data/*.txt` (one sentence per line), or the settings in
+`compile_settings.json`. Then run `python build_weights.py && python export_web.py &&
+python test_model.py`. Run `python scripts/sweep.py` to re-check the settings.
+
+| Setting | Meaning |
 |---|---|
-| `PARAM_TARGET` | parameter budget the pruning threshold is fitted to (400M) |
-| `MAX_VOCAB`, `MIN_WORD_COUNT` | vocabulary cap and minimum word frequency |
-| `PHRASE_WEIGHT` | how many data occurrences one hand-written phrase counts as |
-| `CODE_DIM` | token-code size (bigger means fewer near-collisions and a larger model) |
-| `N_PREV` | how many previous tokens are attended to (n-gram order = N_PREV + 2) |
-| `MAX_POS` | longest sequence the attention margin is guaranteed for |
-| `ATTN_MARGIN`, `BETA` | sharpness of attention and of the neuron gates |
+| `n_prev` | previous tokens attended to (context = n_prev + 1) |
+| `smoothing`, `discount` | `kn` or `wb`; Kneser-Ney discount |
+| `phrase_weight`, `prompt_weight`, `reply_weight` | how much each source counts |
+| `max_vocab`, `min_word_count` | vocabulary cap and minimum word frequency |
+| `min_count` / `param_target` | pruning threshold, or the dense budget to fit it to |
+| `policy.tau`, `policy.min_support`, `policy.max_words` | confidence policy shipped in `model.bin` |
 
 ## History
 
 1. **120 parameters** (hidden size 2): one fixed next word for each of 22 words, with no
    context.
-2. **91M parameters:** the compiled 5-gram over the 1,139 hand-written phrases.
-3. **380M parameters (current):** adds ~20k human-written coding sentences, plus a
-   vocabulary cap and count pruning to stay within the budget.
+2. **91M parameters:** the compiled 5-gram over 1,139 hand-written phrases.
+3. **380M parameters:** adds ~20k human-written coding sentences, with a vocabulary cap and
+   count pruning to fit the budget.
+4. **Current:** Kneser-Ney smoothing; a confidence-gated ghost-text policy; settings chosen by
+   held-out precision and keystrokes saved; a 4.5 MB browser export with a zero-dependency
+   runtime.

docs/sweep.json

@@ -0,0 +1,2181 @@
+{
+  "target_precision": 0.9,
+  "max_export_mb": 50,
+  "heldout_prompts": 711,
+  "best": {
+    "settings": {
+      "n_prev": 3,
+      "smoothing": "kn",
+      "discount": 0.75,
+      "phrase_weight": 1.0,
+      "prompt_weight": 1,
+      "reply_weight": 1.0,
+      "max_vocab": 8000,
+      "min_word_count": 3,
+      "min_count": 2,
+      "param_target": 400000000
+    },
+    "vocab": 6550,
+    "min_context_count": 2,
+    "contexts_seen": 509522,
+    "contexts_kept": 105467,
+    "export_mb": 4.2,
+    "best": {
+      "tau": 0.72,
+      "min_support": 0,
+      "precision": 0.9029,
+      "coverage": 0.3866,
+      "keystrokes_saved": 0.2296,
+      "shown": 16034
+    },
+    "seconds": 14.7
+  },
+  "runs": [
+    {
+      "settings": {
+        "n_prev": 3,
+        "smoothing": "wb",
+        "discount": 0.75,
+        "phrase_weight": 5,
+        "prompt_weight": 1,
+        "reply_weight": 1,
+        "max_vocab": 8000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 6550,
+      "min_context_count": 2,
+      "contexts_seen": 509522,
+      "contexts_kept": 105467,
+      "export_mb": 4.2,
+      "best": {
+        "tau": 0.81,
+        "min_support": 0,
+        "precision": 0.9023,
+        "coverage": 0.3647,
+        "keystrokes_saved": 0.2173,
+        "shown": 15125
+      },
+      "seconds": 13.3
+    },
+    {
+      "settings": {
+        "n_prev": 4,
+        "smoothing": "wb",
+        "discount": 0.75,
+        "phrase_weight": 5,
+        "prompt_weight": 1,
+        "reply_weight": 1,
+        "max_vocab": 8000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 6550,
+      "min_context_count": 2,
+      "contexts_seen": 754606,
+      "contexts_kept": 125943,
+      "export_mb": 4.9,
+      "best": {
+        "tau": 0.83,
+        "min_support": 0,
+        "precision": 0.9021,
+        "coverage": 0.3555,
+        "keystrokes_saved": 0.2116,
+        "shown": 14744
+      },
+      "seconds": 15.0
+    },
+    {
+      "settings": {
+        "n_prev": 5,
+        "smoothing": "wb",
+        "discount": 0.75,
+        "phrase_weight": 5,
+        "prompt_weight": 1,
+        "reply_weight": 1,
+        "max_vocab": 8000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 6550,
+      "min_context_count": 2,
+      "contexts_seen": 1005811,
+      "contexts_kept": 143202,
+      "export_mb": 5.4,
+      "best": {
+        "tau": 0.84,
+        "min_support": 0,
+        "precision": 0.9019,
+        "coverage": 0.3499,
+        "keystrokes_saved": 0.2084,
+        "shown": 14512
+      },
+      "seconds": 16.2
+    },
+    {
+      "settings": {
+        "n_prev": 7,
+        "smoothing": "wb",
+        "discount": 0.75,
+        "phrase_weight": 5,
+        "prompt_weight": 1,
+        "reply_weight": 1,
+        "max_vocab": 8000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 6550,
+      "min_context_count": 2,
+      "contexts_seen": 1514926,
+      "contexts_kept": 173445,
+      "export_mb": 6.6,
+      "best": {
+        "tau": 0.84,
+        "min_support": 0,
+        "precision": 0.9012,
+        "coverage": 0.3502,
+        "keystrokes_saved": 0.2079,
+        "shown": 14524
+      },
+      "seconds": 18.4
+    },
+    {
+      "settings": {
+        "n_prev": 9,
+        "smoothing": "wb",
+        "discount": 0.75,
+        "phrase_weight": 5,
+        "prompt_weight": 1,
+        "reply_weight": 1,
+        "max_vocab": 8000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 6550,
+      "min_context_count": 2,
+      "contexts_seen": 2027338,
+      "contexts_kept": 201278,
+      "export_mb": 7.7,
+      "best": {
+        "tau": 0.84,
+        "min_support": 0,
+        "precision": 0.9012,
+        "coverage": 0.3502,
+        "keystrokes_saved": 0.2079,
+        "shown": 14524
+      },
+      "seconds": 20.2
+    },
+    {
+      "settings": {
+        "n_prev": 3,
+        "smoothing": "kn",
+        "discount": 0.75,
+        "phrase_weight": 5,
+        "prompt_weight": 1,
+        "reply_weight": 1,
+        "max_vocab": 8000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 6550,
+      "min_context_count": 2,
+      "contexts_seen": 509522,
+      "contexts_kept": 105467,
+      "export_mb": 4.2,
+      "best": {
+        "tau": 0.73,
+        "min_support": 0,
+        "precision": 0.9033,
+        "coverage": 0.3819,
+        "keystrokes_saved": 0.2251,
+        "shown": 15837
+      },
+      "seconds": 14.1
+    },
+    {
+      "settings": {
+        "n_prev": 4,
+        "smoothing": "kn",
+        "discount": 0.75,
+        "phrase_weight": 5,
+        "prompt_weight": 1,
+        "reply_weight": 1,
+        "max_vocab": 8000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 6550,
+      "min_context_count": 2,
+      "contexts_seen": 754606,
+      "contexts_kept": 125943,
+      "export_mb": 4.9,
+      "best": {
+        "tau": 0.74,
+        "min_support": 0,
+        "precision": 0.9023,
+        "coverage": 0.3779,
+        "keystrokes_saved": 0.2224,
+        "shown": 15673
+      },
+      "seconds": 15.7
+    },
+    {
+      "settings": {
+        "n_prev": 5,
+        "smoothing": "kn",
+        "discount": 0.75,
+        "phrase_weight": 5,
+        "prompt_weight": 1,
+        "reply_weight": 1,
+        "max_vocab": 8000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 6550,
+      "min_context_count": 2,
+      "contexts_seen": 1005811,
+      "contexts_kept": 143202,
+      "export_mb": 5.4,
+      "best": {
+        "tau": 0.75,
+        "min_support": 0,
+        "precision": 0.9029,
+        "coverage": 0.3739,
+        "keystrokes_saved": 0.2196,
+        "shown": 15506
+      },
+      "seconds": 17.7
+    },
+    {
+      "settings": {
+        "n_prev": 7,
+        "smoothing": "kn",
+        "discount": 0.75,
+        "phrase_weight": 5,
+        "prompt_weight": 1,
+        "reply_weight": 1,
+        "max_vocab": 8000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 6550,
+      "min_context_count": 2,
+      "contexts_seen": 1514926,
+      "contexts_kept": 173445,
+      "export_mb": 6.6,
+      "best": {
+        "tau": 0.76,
+        "min_support": 0,
+        "precision": 0.9026,
+        "coverage": 0.3693,
+        "keystrokes_saved": 0.2162,
+        "shown": 15315
+      },
+      "seconds": 21.5
+    },
+    {
+      "settings": {
+        "n_prev": 9,
+        "smoothing": "kn",
+        "discount": 0.75,
+        "phrase_weight": 5,
+        "prompt_weight": 1,
+        "reply_weight": 1,
+        "max_vocab": 8000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 6550,
+      "min_context_count": 2,
+      "contexts_seen": 2027338,
+      "contexts_kept": 201278,
+      "export_mb": 7.7,
+      "best": {
+        "tau": 0.76,
+        "min_support": 0,
+        "precision": 0.9017,
+        "coverage": 0.3693,
+        "keystrokes_saved": 0.216,
+        "shown": 15316
+      },
+      "seconds": 24.4
+    },
+    {
+      "settings": {
+        "n_prev": 3,
+        "smoothing": "kn",
+        "discount": 0.75,
+        "phrase_weight": 1.0,
+        "prompt_weight": 1,
+        "reply_weight": 0.0,
+        "max_vocab": 8000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 2830,
+      "min_context_count": 2,
+      "contexts_seen": 133364,
+      "contexts_kept": 33000,
+      "export_mb": 1.3,
+      "best": {
+        "tau": 0.7,
+        "min_support": 0,
+        "precision": 0.901,
+        "coverage": 0.3569,
+        "keystrokes_saved": 0.2089,
+        "shown": 14803
+      },
+      "seconds": 8.8
+    },
+    {
+      "settings": {
+        "n_prev": 3,
+        "smoothing": "kn",
+        "discount": 0.75,
+        "phrase_weight": 1.0,
+        "prompt_weight": 1,
+        "reply_weight": 0.3,
+        "max_vocab": 8000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 3841,
+      "min_context_count": 2,
+      "contexts_seen": 454608,
+      "contexts_kept": 40744,
+      "export_mb": 2.6,
+      "best": {
+        "tau": 0.69,
+        "min_support": 0,
+        "precision": 0.9026,
+        "coverage": 0.3784,
+        "keystrokes_saved": 0.2238,
+        "shown": 15694
+      },
+      "seconds": 12.2
+    },
+    {
+      "settings": {
+        "n_prev": 3,
+        "smoothing": "kn",
+        "discount": 0.75,
+        "phrase_weight": 1.0,
+        "prompt_weight": 1,
+        "reply_weight": 1.0,
+        "max_vocab": 8000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 6550,
+      "min_context_count": 2,
+      "contexts_seen": 509522,
+      "contexts_kept": 105467,
+      "export_mb": 4.2,
+      "best": {
+        "tau": 0.72,
+        "min_support": 0,
+        "precision": 0.9029,
+        "coverage": 0.3866,
+        "keystrokes_saved": 0.2296,
+        "shown": 16034
+      },
+      "seconds": 14.7
+    },
+    {
+      "settings": {
+        "n_prev": 3,
+        "smoothing": "kn",
+        "discount": 0.75,
+        "phrase_weight": 5.0,
+        "prompt_weight": 1,
+        "reply_weight": 0.0,
+        "max_vocab": 8000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 2830,
+      "min_context_count": 2,
+      "contexts_seen": 133364,
+      "contexts_kept": 33000,
+      "export_mb": 1.3,
+      "best": {
+        "tau": 0.72,
+        "min_support": 0,
+        "precision": 0.9034,
+        "coverage": 0.3489,
+        "keystrokes_saved": 0.2007,
+        "shown": 14472
+      },
+      "seconds": 8.4
+    },
+    {
+      "settings": {
+        "n_prev": 3,
+        "smoothing": "kn",
+        "discount": 0.75,
+        "phrase_weight": 5.0,
+        "prompt_weight": 1,
+        "reply_weight": 0.3,
+        "max_vocab": 8000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 3841,
+      "min_context_count": 2,
+      "contexts_seen": 454608,
+      "contexts_kept": 40744,
+      "export_mb": 2.6,
+      "best": {
+        "tau": 0.7,
+        "min_support": 0,
+        "precision": 0.9006,
+        "coverage": 0.3758,
+        "keystrokes_saved": 0.2178,
+        "shown": 15588
+      },
+      "seconds": 11.9
+    },
+    {
+      "settings": {
+        "n_prev": 3,
+        "smoothing": "kn",
+        "discount": 0.75,
+        "phrase_weight": 1.0,
+        "prompt_weight": 1,
+        "reply_weight": 1.0,
+        "max_vocab": 3000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 3000,
+      "min_context_count": 2,
+      "contexts_seen": 425514,
+      "contexts_kept": 96265,
+      "export_mb": 4.0,
+      "best": {
+        "tau": 0.7,
+        "min_support": 2,
+        "precision": 0.9,
+        "coverage": 0.3601,
+        "keystrokes_saved": 0.2123,
+        "shown": 14933
+      },
+      "seconds": 11.2
+    },
+    {
+      "settings": {
+        "n_prev": 3,
+        "smoothing": "kn",
+        "discount": 0.75,
+        "phrase_weight": 1.0,
+        "prompt_weight": 1,
+        "reply_weight": 1.0,
+        "max_vocab": 3000,
+        "min_word_count": 3,
+        "min_count": 3,
+        "param_target": 400000000
+      },
+      "vocab": 3000,
+      "min_context_count": 3,
+      "contexts_seen": 425514,
+      "contexts_kept": 54867,
+      "export_mb": 3.0,
+      "best": {
+        "tau": 0.69,
+        "min_support": 0,
+        "precision": 0.9016,
+        "coverage": 0.358,
+        "keystrokes_saved": 0.2091,
+        "shown": 14849
+      },
+      "seconds": 11.1
+    },
+    {
+      "settings": {
+        "n_prev": 3,
+        "smoothing": "kn",
+        "discount": 0.75,
+        "phrase_weight": 1.0,
+        "prompt_weight": 1,
+        "reply_weight": 1.0,
+        "max_vocab": 8000,
+        "min_word_count": 3,
+        "min_count": 3,
+        "param_target": 400000000
+      },
+      "vocab": 6550,
+      "min_context_count": 3,
+      "contexts_seen": 509522,
+      "contexts_kept": 58669,
+      "export_mb": 3.1,
+      "best": {
+        "tau": 0.7,
+        "min_support": 0,
+        "precision": 0.9007,
+        "coverage": 0.3888,
+        "keystrokes_saved": 0.2294,
+        "shown": 16126
+      },
+      "seconds": 13.5
+    },
+    {
+      "settings": {
+        "n_prev": 3,
+        "smoothing": "kn",
+        "discount": 0.75,
+        "phrase_weight": 1.0,
+        "prompt_weight": 1,
+        "reply_weight": 1.0,
+        "max_vocab": 16000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 6550,
+      "min_context_count": 2,
+      "contexts_seen": 509522,
+      "contexts_kept": 105467,
+      "export_mb": 4.2,
+      "best": {
+        "tau": 0.72,
+        "min_support": 0,
+        "precision": 0.9029,
+        "coverage": 0.3866,
+        "keystrokes_saved": 0.2296,
+        "shown": 16034
+      },
+      "seconds": 13.2
+    },
+    {
+      "settings": {
+        "n_prev": 3,
+        "smoothing": "kn",
+        "discount": 0.75,
+        "phrase_weight": 1.0,
+        "prompt_weight": 1,
+        "reply_weight": 1.0,
+        "max_vocab": 16000,
+        "min_word_count": 3,
+        "min_count": 3,
+        "param_target": 400000000
+      },
+      "vocab": 6550,
+      "min_context_count": 3,
+      "contexts_seen": 509522,
+      "contexts_kept": 58669,
+      "export_mb": 3.1,
+      "best": {
+        "tau": 0.7,
+        "min_support": 0,
+        "precision": 0.9007,
+        "coverage": 0.3888,
+        "keystrokes_saved": 0.2294,
+        "shown": 16126
+      },
+      "seconds": 12.9
+    },
+    {
+      "settings": {
+        "n_prev": 4,
+        "smoothing": "kn",
+        "discount": 0.75,
+        "phrase_weight": 1.0,
+        "prompt_weight": 1,
+        "reply_weight": 1.0,
+        "max_vocab": 8000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 6550,
+      "min_context_count": 2,
+      "contexts_seen": 754606,
+      "contexts_kept": 125943,
+      "export_mb": 4.9,
+      "best": {
+        "tau": 0.73,
+        "min_support": 0,
+        "precision": 0.9025,
+        "coverage": 0.382,
+        "keystrokes_saved": 0.2279,
+        "shown": 15841
+      },
+      "seconds": 14.6
+    },
+    {
+      "settings": {
+        "n_prev": 5,
+        "smoothing": "kn",
+        "discount": 0.75,
+        "phrase_weight": 1.0,
+        "prompt_weight": 1,
+        "reply_weight": 1.0,
+        "max_vocab": 8000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 6550,
+      "min_context_count": 2,
+      "contexts_seen": 1005811,
+      "contexts_kept": 143202,
+      "export_mb": 5.4,
+      "best": {
+        "tau": 0.73,
+        "min_support": 2,
+        "precision": 0.9003,
+        "coverage": 0.3816,
+        "keystrokes_saved": 0.2273,
+        "shown": 15826
+      },
+      "seconds": 16.2
+    },
+    {
+      "settings": {
+        "n_prev": 7,
+        "smoothing": "kn",
+        "discount": 0.75,
+        "phrase_weight": 1.0,
+        "prompt_weight": 1,
+        "reply_weight": 1.0,
+        "max_vocab": 8000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 6550,
+      "min_context_count": 2,
+      "contexts_seen": 1514926,
+      "contexts_kept": 173445,
+      "export_mb": 6.6,
+      "best": {
+        "tau": 0.74,
+        "min_support": 0,
+        "precision": 0.9001,
+        "coverage": 0.3788,
+        "keystrokes_saved": 0.2244,
+        "shown": 15710
+      },
+      "seconds": 19.1
+    },
+    {
+      "settings": {
+        "n_prev": 9,
+        "smoothing": "kn",
+        "discount": 0.75,
+        "phrase_weight": 1.0,
+        "prompt_weight": 1,
+        "reply_weight": 1.0,
+        "max_vocab": 8000,
+        "min_word_count": 3,
+        "min_count": 2,
+        "param_target": 400000000
+      },
+      "vocab": 6550,
+      "min_context_count": 2,
+      "contexts_seen": 2027338,
+      "contexts_kept": 201278,
+      "export_mb": 7.7,
+      "best": {
+        "tau": 0.74,
+        "min_support": 2,
+        "precision": 0.9008,
+        "coverage": 0.378,
+        "keystrokes_saved": 0.2244,
+        "shown": 15678
+      },
+      "seconds": 22.9
+    }
+  ],
+  "best_grid": [
+    {
+      "tau": 0.5,
+      "min_support": 0,
+      "precision": 0.7865,
+      "coverage": 0.5115,
+      "keystrokes_saved": 0.2789,
+      "shown": 21215
+    },
+    {
+      "tau": 0.5,
+      "min_support": 2,
+      "precision": 0.7871,
+      "coverage": 0.5107,
+      "keystrokes_saved": 0.2786,
+      "shown": 21180
+    },
+    {
+      "tau": 0.5,
+      "min_support": 5,
+      "precision": 0.8083,
+      "coverage": 0.3863,
+      "keystrokes_saved": 0.2119,
+      "shown": 16021
+    },
+    {
+      "tau": 0.5,
+      "min_support": 10,
+      "precision": 0.8177,
+      "coverage": 0.3212,
+      "keystrokes_saved": 0.1777,
+      "shown": 13322
+    },
+    {
+      "tau": 0.51,
+      "min_support": 0,
+      "precision": 0.7919,
+      "coverage": 0.5048,
+      "keystrokes_saved": 0.2761,
+      "shown": 20937
+    },
+    {
+      "tau": 0.51,
+      "min_support": 2,
+      "precision": 0.7925,
+      "coverage": 0.504,
+      "keystrokes_saved": 0.2759,
+      "shown": 20902
+    },
+    {
+      "tau": 0.51,
+      "min_support": 5,
+      "precision": 0.8132,
+      "coverage": 0.3811,
+      "keystrokes_saved": 0.2099,
+      "shown": 15805
+    },
+    {
+      "tau": 0.51,
+      "min_support": 10,
+      "precision": 0.8217,
+      "coverage": 0.3173,
+      "keystrokes_saved": 0.1761,
+      "shown": 13161
+    },
+    {
+      "tau": 0.52,
+      "min_support": 0,
+      "precision": 0.797,
+      "coverage": 0.4991,
+      "keystrokes_saved": 0.2743,
+      "shown": 20699
+    },
+    {
+      "tau": 0.52,
+      "min_support": 2,
+      "precision": 0.7976,
+      "coverage": 0.4983,
+      "keystrokes_saved": 0.2741,
+      "shown": 20666
+    },
+    {
+      "tau": 0.52,
+      "min_support": 5,
+      "precision": 0.818,
+      "coverage": 0.3768,
+      "keystrokes_saved": 0.2085,
+      "shown": 15628
+    },
+    {
+      "tau": 0.52,
+      "min_support": 10,
+      "precision": 0.8261,
+      "coverage": 0.3142,
+      "keystrokes_saved": 0.1752,
+      "shown": 13030
+    },
+    {
+      "tau": 0.53,
+      "min_support": 0,
+      "precision": 0.803,
+      "coverage": 0.4929,
+      "keystrokes_saved": 0.2724,
+      "shown": 20441
+    },
+    {
+      "tau": 0.53,
+      "min_support": 2,
+      "precision": 0.8037,
+      "coverage": 0.4921,
+      "keystrokes_saved": 0.2722,
+      "shown": 20408
+    },
+    {
+      "tau": 0.53,
+      "min_support": 5,
+      "precision": 0.8236,
+      "coverage": 0.372,
+      "keystrokes_saved": 0.2073,
+      "shown": 15430
+    },
+    {
+      "tau": 0.53,
+      "min_support": 10,
+      "precision": 0.8309,
+      "coverage": 0.3104,
+      "keystrokes_saved": 0.1742,
+      "shown": 12873
+    },
+    {
+      "tau": 0.54,
+      "min_support": 0,
+      "precision": 0.8066,
+      "coverage": 0.4881,
+      "keystrokes_saved": 0.2709,
+      "shown": 20242
+    },
+    {
+      "tau": 0.54,
+      "min_support": 2,
+      "precision": 0.8072,
+      "coverage": 0.4873,
+      "keystrokes_saved": 0.2707,
+      "shown": 20209
+    },
+    {
+      "tau": 0.54,
+      "min_support": 5,
+      "precision": 0.8263,
+      "coverage": 0.3687,
+      "keystrokes_saved": 0.2062,
+      "shown": 15290
+    },
+    {
+      "tau": 0.54,
+      "min_support": 10,
+      "precision": 0.8336,
+      "coverage": 0.3077,
+      "keystrokes_saved": 0.1733,
+      "shown": 12763
+    },
+    {
+      "tau": 0.55,
+      "min_support": 0,
+      "precision": 0.813,
+      "coverage": 0.4822,
+      "keystrokes_saved": 0.2698,
+      "shown": 20000
+    },
+    {
+      "tau": 0.55,
+      "min_support": 2,
+      "precision": 0.8137,
+      "coverage": 0.4815,
+      "keystrokes_saved": 0.2696,
+      "shown": 19968
+    },
+    {
+      "tau": 0.55,
+      "min_support": 5,
+      "precision": 0.8324,
+      "coverage": 0.3641,
+      "keystrokes_saved": 0.2051,
+      "shown": 15101
+    },
+    {
+      "tau": 0.55,
+      "min_support": 10,
+      "precision": 0.8393,
+      "coverage": 0.3039,
+      "keystrokes_saved": 0.1725,
+      "shown": 12605
+    },
+    {
+      "tau": 0.56,
+      "min_support": 0,
+      "precision": 0.8183,
+      "coverage": 0.4766,
+      "keystrokes_saved": 0.2652,
+      "shown": 19765
+    },
+    {
+      "tau": 0.56,
+      "min_support": 2,
+      "precision": 0.8189,
+      "coverage": 0.4758,
+      "keystrokes_saved": 0.265,
+      "shown": 19734
+    },
+    {
+      "tau": 0.56,
+      "min_support": 5,
+      "precision": 0.8372,
+      "coverage": 0.3597,
+      "keystrokes_saved": 0.2007,
+      "shown": 14918
+    },
+    {
+      "tau": 0.56,
+      "min_support": 10,
+      "precision": 0.8438,
+      "coverage": 0.3003,
+      "keystrokes_saved": 0.1683,
+      "shown": 12455
+    },
+    {
+      "tau": 0.57,
+      "min_support": 0,
+      "precision": 0.8259,
+      "coverage": 0.4716,
+      "keystrokes_saved": 0.2665,
+      "shown": 19559
+    },
+    {
+      "tau": 0.57,
+      "min_support": 2,
+      "precision": 0.8265,
+      "coverage": 0.4709,
+      "keystrokes_saved": 0.2663,
+      "shown": 19530
+    },
+    {
+      "tau": 0.57,
+      "min_support": 5,
+      "precision": 0.8444,
+      "coverage": 0.356,
+      "keystrokes_saved": 0.2018,
+      "shown": 14766
+    },
+    {
+      "tau": 0.57,
+      "min_support": 10,
+      "precision": 0.8518,
+      "coverage": 0.2973,
+      "keystrokes_saved": 0.1697,
+      "shown": 12331
+    },
+    {
+      "tau": 0.58,
+      "min_support": 0,
+      "precision": 0.8361,
+      "coverage": 0.4668,
+      "keystrokes_saved": 0.265,
+      "shown": 19359
+    },
+    {
+      "tau": 0.58,
+      "min_support": 2,
+      "precision": 0.8366,
+      "coverage": 0.4661,
+      "keystrokes_saved": 0.2648,
+      "shown": 19332
+    },
+    {
+      "tau": 0.58,
+      "min_support": 5,
+      "precision": 0.8549,
+      "coverage": 0.3527,
+      "keystrokes_saved": 0.2004,
+      "shown": 14629
+    },
+    {
+      "tau": 0.58,
+      "min_support": 10,
+      "precision": 0.863,
+      "coverage": 0.2947,
+      "keystrokes_saved": 0.1685,
+      "shown": 12222
+    },
+    {
+      "tau": 0.59,
+      "min_support": 0,
+      "precision": 0.8403,
+      "coverage": 0.4613,
+      "keystrokes_saved": 0.2622,
+      "shown": 19132
+    },
+    {
+      "tau": 0.59,
+      "min_support": 2,
+      "precision": 0.8407,
+      "coverage": 0.4607,
+      "keystrokes_saved": 0.262,
+      "shown": 19106
+    },
+    {
+      "tau": 0.59,
+      "min_support": 5,
+      "precision": 0.8582,
+      "coverage": 0.3483,
+      "keystrokes_saved": 0.198,
+      "shown": 14445
+    },
+    {
+      "tau": 0.59,
+      "min_support": 10,
+      "precision": 0.8661,
+      "coverage": 0.2909,
+      "keystrokes_saved": 0.1663,
+      "shown": 12065
+    },
+    {
+      "tau": 0.6,
+      "min_support": 0,
+      "precision": 0.844,
+      "coverage": 0.4557,
+      "keystrokes_saved": 0.2575,
+      "shown": 18899
+    },
+    {
+      "tau": 0.6,
+      "min_support": 2,
+      "precision": 0.8444,
+      "coverage": 0.4551,
+      "keystrokes_saved": 0.2572,
+      "shown": 18873
+    },
+    {
+      "tau": 0.6,
+      "min_support": 5,
+      "precision": 0.862,
+      "coverage": 0.3438,
+      "keystrokes_saved": 0.1937,
+      "shown": 14259
+    },
+    {
+      "tau": 0.6,
+      "min_support": 10,
+      "precision": 0.8701,
+      "coverage": 0.2872,
+      "keystrokes_saved": 0.1626,
+      "shown": 11910
+    },
+    {
+      "tau": 0.61,
+      "min_support": 0,
+      "precision": 0.8482,
+      "coverage": 0.4501,
+      "keystrokes_saved": 0.2555,
+      "shown": 18667
+    },
+    {
+      "tau": 0.61,
+      "min_support": 2,
+      "precision": 0.8485,
+      "coverage": 0.4495,
+      "keystrokes_saved": 0.2553,
+      "shown": 18643
+    },
+    {
+      "tau": 0.61,
+      "min_support": 5,
+      "precision": 0.8655,
+      "coverage": 0.3396,
+      "keystrokes_saved": 0.1921,
+      "shown": 14085
+    },
+    {
+      "tau": 0.61,
+      "min_support": 10,
+      "precision": 0.8729,
+      "coverage": 0.2842,
+      "keystrokes_saved": 0.1613,
+      "shown": 11787
+    },
+    {
+      "tau": 0.62,
+      "min_support": 0,
+      "precision": 0.8551,
+      "coverage": 0.4405,
+      "keystrokes_saved": 0.2542,
+      "shown": 18271
+    },
+    {
+      "tau": 0.62,
+      "min_support": 2,
+      "precision": 0.8554,
+      "coverage": 0.44,
+      "keystrokes_saved": 0.254,
+      "shown": 18247
+    },
+    {
+      "tau": 0.62,
+      "min_support": 5,
+      "precision": 0.8729,
+      "coverage": 0.3316,
+      "keystrokes_saved": 0.1915,
+      "shown": 13753
+    },
+    {
+      "tau": 0.62,
+      "min_support": 10,
+      "precision": 0.8805,
+      "coverage": 0.2772,
+      "keystrokes_saved": 0.1614,
+      "shown": 11497
+    },
+    {
+      "tau": 0.63,
+      "min_support": 0,
+      "precision": 0.8595,
+      "coverage": 0.4349,
+      "keystrokes_saved": 0.2518,
+      "shown": 18035
+    },
+    {
+      "tau": 0.63,
+      "min_support": 2,
+      "precision": 0.8598,
+      "coverage": 0.4343,
+      "keystrokes_saved": 0.2516,
+      "shown": 18012
+    },
+    {
+      "tau": 0.63,
+      "min_support": 5,
+      "precision": 0.8767,
+      "coverage": 0.3276,
+      "keystrokes_saved": 0.1898,
+      "shown": 13587
+    },
+    {
+      "tau": 0.63,
+      "min_support": 10,
+      "precision": 0.884,
+      "coverage": 0.274,
+      "keystrokes_saved": 0.1601,
+      "shown": 11364
+    },
+    {
+      "tau": 0.64,
+      "min_support": 0,
+      "precision": 0.8637,
+      "coverage": 0.4308,
+      "keystrokes_saved": 0.2502,
+      "shown": 17865
+    },
+    {
+      "tau": 0.64,
+      "min_support": 2,
+      "precision": 0.864,
+      "coverage": 0.4302,
+      "keystrokes_saved": 0.2501,
+      "shown": 17843
+    },
+    {
+      "tau": 0.64,
+      "min_support": 5,
+      "precision": 0.8798,
+      "coverage": 0.3249,
+      "keystrokes_saved": 0.1886,
+      "shown": 13474
+    },
+    {
+      "tau": 0.64,
+      "min_support": 10,
+      "precision": 0.8871,
+      "coverage": 0.2715,
+      "keystrokes_saved": 0.1592,
+      "shown": 11260
+    },
+    {
+      "tau": 0.65,
+      "min_support": 0,
+      "precision": 0.8697,
+      "coverage": 0.4247,
+      "keystrokes_saved": 0.2479,
+      "shown": 17616
+    },
+    {
+      "tau": 0.65,
+      "min_support": 2,
+      "precision": 0.87,
+      "coverage": 0.4242,
+      "keystrokes_saved": 0.2477,
+      "shown": 17594
+    },
+    {
+      "tau": 0.65,
+      "min_support": 5,
+      "precision": 0.8849,
+      "coverage": 0.3206,
+      "keystrokes_saved": 0.1867,
+      "shown": 13296
+    },
+    {
+      "tau": 0.65,
+      "min_support": 10,
+      "precision": 0.8917,
+      "coverage": 0.268,
+      "keystrokes_saved": 0.1576,
+      "shown": 11113
+    },
+    {
+      "tau": 0.66,
+      "min_support": 0,
+      "precision": 0.8748,
+      "coverage": 0.4193,
+      "keystrokes_saved": 0.2456,
+      "shown": 17391
+    },
+    {
+      "tau": 0.66,
+      "min_support": 2,
+      "precision": 0.875,
+      "coverage": 0.4188,
+      "keystrokes_saved": 0.2454,
+      "shown": 17369
+    },
+    {
+      "tau": 0.66,
+      "min_support": 5,
+      "precision": 0.8883,
+      "coverage": 0.317,
+      "keystrokes_saved": 0.1852,
+      "shown": 13149
+    },
+    {
+      "tau": 0.66,
+      "min_support": 10,
+      "precision": 0.8943,
+      "coverage": 0.2651,
+      "keystrokes_saved": 0.1563,
+      "shown": 10994
+    },
+    {
+      "tau": 0.67,
+      "min_support": 0,
+      "precision": 0.8797,
+      "coverage": 0.4133,
+      "keystrokes_saved": 0.2425,
+      "shown": 17141
+    },
+    {
+      "tau": 0.67,
+      "min_support": 2,
+      "precision": 0.8799,
+      "coverage": 0.4128,
+      "keystrokes_saved": 0.2424,
+      "shown": 17122
+    },
+    {
+      "tau": 0.67,
+      "min_support": 5,
+      "precision": 0.8929,
+      "coverage": 0.3124,
+      "keystrokes_saved": 0.183,
+      "shown": 12958
+    },
+    {
+      "tau": 0.67,
+      "min_support": 10,
+      "precision": 0.8989,
+      "coverage": 0.2614,
+      "keystrokes_saved": 0.1542,
+      "shown": 10840
+    },
+    {
+      "tau": 0.68,
+      "min_support": 0,
+      "precision": 0.8852,
+      "coverage": 0.4085,
+      "keystrokes_saved": 0.2409,
+      "shown": 16943
+    },
+    {
+      "tau": 0.68,
+      "min_support": 2,
+      "precision": 0.8854,
+      "coverage": 0.4081,
+      "keystrokes_saved": 0.2407,
+      "shown": 16925
+    },
+    {
+      "tau": 0.68,
+      "min_support": 5,
+      "precision": 0.8966,
+      "coverage": 0.3091,
+      "keystrokes_saved": 0.1819,
+      "shown": 12820
+    },
+    {
+      "tau": 0.68,
+      "min_support": 10,
+      "precision": 0.9024,
+      "coverage": 0.2587,
+      "keystrokes_saved": 0.1531,
+      "shown": 10728
+    },
+    {
+      "tau": 0.69,
+      "min_support": 0,
+      "precision": 0.8894,
+      "coverage": 0.4034,
+      "keystrokes_saved": 0.2382,
+      "shown": 16731
+    },
+    {
+      "tau": 0.69,
+      "min_support": 2,
+      "precision": 0.8895,
+      "coverage": 0.403,
+      "keystrokes_saved": 0.2381,
+      "shown": 16713
+    },
+    {
+      "tau": 0.69,
+      "min_support": 5,
+      "precision": 0.899,
+      "coverage": 0.3057,
+      "keystrokes_saved": 0.1798,
+      "shown": 12680
+    },
+    {
+      "tau": 0.69,
+      "min_support": 10,
+      "precision": 0.9044,
+      "coverage": 0.2561,
+      "keystrokes_saved": 0.1512,
+      "shown": 10620
+    },
+    {
+      "tau": 0.7,
+      "min_support": 0,
+      "precision": 0.8942,
+      "coverage": 0.3982,
+      "keystrokes_saved": 0.2359,
+      "shown": 16516
+    },
+    {
+      "tau": 0.7,
+      "min_support": 2,
+      "precision": 0.8944,
+      "coverage": 0.3978,
+      "keystrokes_saved": 0.2358,
+      "shown": 16498
+    },
+    {
+      "tau": 0.7,
+      "min_support": 5,
+      "precision": 0.904,
+      "coverage": 0.3017,
+      "keystrokes_saved": 0.1779,
+      "shown": 12512
+    },
+    {
+      "tau": 0.7,
+      "min_support": 10,
+      "precision": 0.9094,
+      "coverage": 0.2529,
+      "keystrokes_saved": 0.1496,
+      "shown": 10490
+    },
+    {
+      "tau": 0.71,
+      "min_support": 0,
+      "precision": 0.8976,
+      "coverage": 0.3929,
+      "keystrokes_saved": 0.233,
+      "shown": 16295
+    },
+    {
+      "tau": 0.71,
+      "min_support": 2,
+      "precision": 0.8978,
+      "coverage": 0.3925,
+      "keystrokes_saved": 0.2329,
+      "shown": 16277
+    },
+    {
+      "tau": 0.71,
+      "min_support": 5,
+      "precision": 0.9062,
+      "coverage": 0.2981,
+      "keystrokes_saved": 0.1758,
+      "shown": 12365
+    },
+    {
+      "tau": 0.71,
+      "min_support": 10,
+      "precision": 0.9116,
+      "coverage": 0.2499,
+      "keystrokes_saved": 0.1477,
+      "shown": 10363
+    },
+    {
+      "tau": 0.72,
+      "min_support": 0,
+      "precision": 0.9029,
+      "coverage": 0.3866,
+      "keystrokes_saved": 0.2296,
+      "shown": 16034
+    },
+    {
+      "tau": 0.72,
+      "min_support": 2,
+      "precision": 0.9031,
+      "coverage": 0.3862,
+      "keystrokes_saved": 0.2295,
+      "shown": 16016
+    },
+    {
+      "tau": 0.72,
+      "min_support": 5,
+      "precision": 0.9113,
+      "coverage": 0.2934,
+      "keystrokes_saved": 0.173,
+      "shown": 12167
+    },
+    {
+      "tau": 0.72,
+      "min_support": 10,
+      "precision": 0.9172,
+      "coverage": 0.2457,
+      "keystrokes_saved": 0.1451,
+      "shown": 10191
+    },
+    {
+      "tau": 0.73,
+      "min_support": 0,
+      "precision": 0.9074,
+      "coverage": 0.381,
+      "keystrokes_saved": 0.2279,
+      "shown": 15800
+    },
+    {
+      "tau": 0.73,
+      "min_support": 2,
+      "precision": 0.9076,
+      "coverage": 0.3805,
+      "keystrokes_saved": 0.2278,
+      "shown": 15782
+    },
+    {
+      "tau": 0.73,
+      "min_support": 5,
+      "precision": 0.9149,
+      "coverage": 0.2894,
+      "keystrokes_saved": 0.1718,
+      "shown": 12004
+    },
+    {
+      "tau": 0.73,
+      "min_support": 10,
+      "precision": 0.9209,
+      "coverage": 0.2427,
+      "keystrokes_saved": 0.1443,
+      "shown": 10065
+    },
+    {
+      "tau": 0.74,
+      "min_support": 0,
+      "precision": 0.9114,
+      "coverage": 0.3765,
+      "keystrokes_saved": 0.2254,
+      "shown": 15615
+    },
+    {
+      "tau": 0.74,
+      "min_support": 2,
+      "precision": 0.9116,
+      "coverage": 0.3761,
+      "keystrokes_saved": 0.2252,
+      "shown": 15597
+    },
+    {
+      "tau": 0.74,
+      "min_support": 5,
+      "precision": 0.9183,
+      "coverage": 0.2863,
+      "keystrokes_saved": 0.17,
+      "shown": 11876
+    },
+    {
+      "tau": 0.74,
+      "min_support": 10,
+      "precision": 0.9242,
+      "coverage": 0.2402,
+      "keystrokes_saved": 0.1428,
+      "shown": 9964
+    },
+    {
+      "tau": 0.75,
+      "min_support": 0,
+      "precision": 0.9141,
+      "coverage": 0.3718,
+      "keystrokes_saved": 0.2234,
+      "shown": 15420
+    },
+    {
+      "tau": 0.75,
+      "min_support": 2,
+      "precision": 0.9143,
+      "coverage": 0.3714,
+      "keystrokes_saved": 0.2233,
+      "shown": 15403
+    },
+    {
+      "tau": 0.75,
+      "min_support": 5,
+      "precision": 0.9204,
+      "coverage": 0.2826,
+      "keystrokes_saved": 0.1686,
+      "shown": 11720
+    },
+    {
+      "tau": 0.75,
+      "min_support": 10,
+      "precision": 0.9261,
+      "coverage": 0.2373,
+      "keystrokes_saved": 0.1418,
+      "shown": 9843
+    },
+    {
+      "tau": 0.76,
+      "min_support": 0,
+      "precision": 0.918,
+      "coverage": 0.3657,
+      "keystrokes_saved": 0.2199,
+      "shown": 15166
+    },
+    {
+      "tau": 0.76,
+      "min_support": 2,
+      "precision": 0.9182,
+      "coverage": 0.3653,
+      "keystrokes_saved": 0.2197,
+      "shown": 15150
+    },
+    {
+      "tau": 0.76,
+      "min_support": 5,
+      "precision": 0.9246,
+      "coverage": 0.2782,
+      "keystrokes_saved": 0.1665,
+      "shown": 11538
+    },
+    {
+      "tau": 0.76,
+      "min_support": 10,
+      "precision": 0.9299,
+      "coverage": 0.2337,
+      "keystrokes_saved": 0.1403,
+      "shown": 9694
+    },
+    {
+      "tau": 0.77,
+      "min_support": 0,
+      "precision": 0.9224,
+      "coverage": 0.3601,
+      "keystrokes_saved": 0.2169,
+      "shown": 14934
+    },
+    {
+      "tau": 0.77,
+      "min_support": 2,
+      "precision": 0.9226,
+      "coverage": 0.3597,
+      "keystrokes_saved": 0.2168,
+      "shown": 14918
+    },
+    {
+      "tau": 0.77,
+      "min_support": 5,
+      "precision": 0.928,
+      "coverage": 0.2741,
+      "keystrokes_saved": 0.1643,
+      "shown": 11369
+    },
+    {
+      "tau": 0.77,
+      "min_support": 10,
+      "precision": 0.9325,
+      "coverage": 0.2304,
+      "keystrokes_saved": 0.1385,
+      "shown": 9557
+    },
+    {
+      "tau": 0.78,
+      "min_support": 0,
+      "precision": 0.9258,
+      "coverage": 0.3547,
+      "keystrokes_saved": 0.2143,
+      "shown": 14710
+    },
+    {
+      "tau": 0.78,
+      "min_support": 2,
+      "precision": 0.926,
+      "coverage": 0.3543,
+      "keystrokes_saved": 0.2142,
+      "shown": 14694
+    },
+    {
+      "tau": 0.78,
+      "min_support": 5,
+      "precision": 0.9311,
+      "coverage": 0.2708,
+      "keystrokes_saved": 0.1626,
+      "shown": 11230
+    },
+    {
+      "tau": 0.78,
+      "min_support": 10,
+      "precision": 0.9351,
+      "coverage": 0.2281,
+      "keystrokes_saved": 0.1373,
+      "shown": 9461
+    },
+    {
+      "tau": 0.79,
+      "min_support": 0,
+      "precision": 0.9307,
+      "coverage": 0.3489,
+      "keystrokes_saved": 0.2115,
+      "shown": 14469
+    },
+    {
+      "tau": 0.79,
+      "min_support": 2,
+      "precision": 0.931,
+      "coverage": 0.3485,
+      "keystrokes_saved": 0.2114,
+      "shown": 14454
+    },
+    {
+      "tau": 0.79,
+      "min_support": 5,
+      "precision": 0.9356,
+      "coverage": 0.2667,
+      "keystrokes_saved": 0.1606,
+      "shown": 11062
+    },
+    {
+      "tau": 0.79,
+      "min_support": 10,
+      "precision": 0.9391,
+      "coverage": 0.2247,
+      "keystrokes_saved": 0.1355,
+      "shown": 9320
+    },
+    {
+      "tau": 0.8,
+      "min_support": 0,
+      "precision": 0.9351,
+      "coverage": 0.3412,
+      "keystrokes_saved": 0.2066,
+      "shown": 14153
+    },
+    {
+      "tau": 0.8,
+      "min_support": 2,
+      "precision": 0.9352,
+      "coverage": 0.3409,
+      "keystrokes_saved": 0.2065,
+      "shown": 14139
+    },
+    {
+      "tau": 0.8,
+      "min_support": 5,
+      "precision": 0.9396,
+      "coverage": 0.2606,
+      "keystrokes_saved": 0.1567,
+      "shown": 10810
+    },
+    {
+      "tau": 0.8,
+      "min_support": 10,
+      "precision": 0.9424,
+      "coverage": 0.2192,
+      "keystrokes_saved": 0.1318,
+      "shown": 9093
+    },
+    {
+      "tau": 0.81,
+      "min_support": 0,
+      "precision": 0.9379,
+      "coverage": 0.335,
+      "keystrokes_saved": 0.2032,
+      "shown": 13894
+    },
+    {
+      "tau": 0.81,
+      "min_support": 2,
+      "precision": 0.938,
+      "coverage": 0.3347,
+      "keystrokes_saved": 0.2031,
+      "shown": 13881
+    },
+    {
+      "tau": 0.81,
+      "min_support": 5,
+      "precision": 0.9423,
+      "coverage": 0.2559,
+      "keystrokes_saved": 0.154,
+      "shown": 10615
+    },
+    {
+      "tau": 0.81,
+      "min_support": 10,
+      "precision": 0.945,
+      "coverage": 0.2153,
+      "keystrokes_saved": 0.1299,
+      "shown": 8931
+    },
+    {
+      "tau": 0.82,
+      "min_support": 0,
+      "precision": 0.9402,
+      "coverage": 0.3296,
+      "keystrokes_saved": 0.1999,
+      "shown": 13671
+    },
+    {
+      "tau": 0.82,
+      "min_support": 2,
+      "precision": 0.9403,
+      "coverage": 0.3293,
+      "keystrokes_saved": 0.1998,
+      "shown": 13658
+    },
+    {
+      "tau": 0.82,
+      "min_support": 5,
+      "precision": 0.9439,
+      "coverage": 0.2523,
+      "keystrokes_saved": 0.1517,
+      "shown": 10465
+    },
+    {
+      "tau": 0.82,
+      "min_support": 10,
+      "precision": 0.9465,
+      "coverage": 0.2124,
+      "keystrokes_saved": 0.1281,
+      "shown": 8808
+    },
+    {
+      "tau": 0.83,
+      "min_support": 0,
+      "precision": 0.9433,
+      "coverage": 0.323,
+      "keystrokes_saved": 0.1966,
+      "shown": 13396
+    },
+    {
+      "tau": 0.83,
+      "min_support": 2,
+      "precision": 0.9434,
+      "coverage": 0.3227,
+      "keystrokes_saved": 0.1965,
+      "shown": 13384
+    },
+    {
+      "tau": 0.83,
+      "min_support": 5,
+      "precision": 0.9473,
+      "coverage": 0.2474,
+      "keystrokes_saved": 0.1494,
+      "shown": 10260
+    },
+    {
+      "tau": 0.83,
+      "min_support": 10,
+      "precision": 0.9495,
+      "coverage": 0.2085,
+      "keystrokes_saved": 0.1263,
+      "shown": 8646
+    },
+    {
+      "tau": 0.84,
+      "min_support": 0,
+      "precision": 0.9472,
+      "coverage": 0.3166,
+      "keystrokes_saved": 0.1932,
+      "shown": 13130
+    },
+    {
+      "tau": 0.84,
+      "min_support": 2,
+      "precision": 0.9472,
+      "coverage": 0.3163,
+      "keystrokes_saved": 0.1931,
+      "shown": 13120
+    },
+    {
+      "tau": 0.84,
+      "min_support": 5,
+      "precision": 0.951,
+      "coverage": 0.2423,
+      "keystrokes_saved": 0.1469,
+      "shown": 10051
+    },
+    {
+      "tau": 0.84,
+      "min_support": 10,
+      "precision": 0.9519,
+      "coverage": 0.2047,
+      "keystrokes_saved": 0.1247,
+      "shown": 8490
+    },
+    {
+      "tau": 0.85,
+      "min_support": 0,
+      "precision": 0.9507,
+      "coverage": 0.3089,
+      "keystrokes_saved": 0.1886,
+      "shown": 12810
+    },
+    {
+      "tau": 0.85,
+      "min_support": 2,
+      "precision": 0.9507,
+      "coverage": 0.3086,
+      "keystrokes_saved": 0.1886,
+      "shown": 12800
+    },
+    {
+      "tau": 0.85,
+      "min_support": 5,
+      "precision": 0.9542,
+      "coverage": 0.2364,
+      "keystrokes_saved": 0.1436,
+      "shown": 9806
+    },
+    {
+      "tau": 0.85,
+      "min_support": 10,
+      "precision": 0.9551,
+      "coverage": 0.1999,
+      "keystrokes_saved": 0.1218,
+      "shown": 8292
+    },
+    {
+      "tau": 0.86,
+      "min_support": 0,
+      "precision": 0.9535,
+      "coverage": 0.3026,
+      "keystrokes_saved": 0.183,
+      "shown": 12550
+    },
+    {
+      "tau": 0.86,
+      "min_support": 2,
+      "precision": 0.9534,
+      "coverage": 0.3024,
+      "keystrokes_saved": 0.1829,
+      "shown": 12540
+    },
+    {
+      "tau": 0.86,
+      "min_support": 5,
+      "precision": 0.9569,
+      "coverage": 0.232,
+      "keystrokes_saved": 0.1391,
+      "shown": 9623
+    },
+    {
+      "tau": 0.86,
+      "min_support": 10,
+      "precision": 0.9574,
+      "coverage": 0.1965,
+      "keystrokes_saved": 0.1179,
+      "shown": 8149
+    },
+    {
+      "tau": 0.87,
+      "min_support": 0,
+      "precision": 0.9565,
+      "coverage": 0.2951,
+      "keystrokes_saved": 0.1793,
+      "shown": 12241
+    },
+    {
+      "tau": 0.87,
+      "min_support": 2,
+      "precision": 0.9565,
+      "coverage": 0.2949,
+      "keystrokes_saved": 0.1792,
+      "shown": 12232
+    },
+    {
+      "tau": 0.87,
+      "min_support": 5,
+      "precision": 0.9593,
+      "coverage": 0.2264,
+      "keystrokes_saved": 0.1366,
+      "shown": 9390
+    },
+    {
+      "tau": 0.87,
+      "min_support": 10,
+      "precision": 0.9594,
+      "coverage": 0.1919,
+      "keystrokes_saved": 0.1159,
+      "shown": 7958
+    },
+    {
+      "tau": 0.88,
+      "min_support": 0,
+      "precision": 0.9588,
+      "coverage": 0.2883,
+      "keystrokes_saved": 0.1753,
+      "shown": 11958
+    },
+    {
+      "tau": 0.88,
+      "min_support": 2,
+      "precision": 0.9587,
+      "coverage": 0.2881,
+      "keystrokes_saved": 0.1752,
+      "shown": 11949
+    },
+    {
+      "tau": 0.88,
+      "min_support": 5,
+      "precision": 0.9607,
+      "coverage": 0.2215,
+      "keystrokes_saved": 0.1336,
+      "shown": 9187
+    },
+    {
+      "tau": 0.88,
+      "min_support": 10,
+      "precision": 0.9609,
+      "coverage": 0.188,
+      "keystrokes_saved": 0.1135,
+      "shown": 7797
+    },
+    {
+      "tau": 0.89,
+      "min_support": 0,
+      "precision": 0.9614,
+      "coverage": 0.2802,
+      "keystrokes_saved": 0.1691,
+      "shown": 11621
+    },
+    {
+      "tau": 0.89,
+      "min_support": 2,
+      "precision": 0.9613,
+      "coverage": 0.28,
+      "keystrokes_saved": 0.1691,
+      "shown": 11612
+    },
+    {
+      "tau": 0.89,
+      "min_support": 5,
+      "precision": 0.9631,
+      "coverage": 0.2154,
+      "keystrokes_saved": 0.1285,
+      "shown": 8932
+    },
+    {
+      "tau": 0.89,
+      "min_support": 10,
+      "precision": 0.9631,
+      "coverage": 0.1829,
+      "keystrokes_saved": 0.1092,
+      "shown": 7587
+    },
+    {
+      "tau": 0.9,
+      "min_support": 0,
+      "precision": 0.9644,
+      "coverage": 0.2692,
+      "keystrokes_saved": 0.1642,
+      "shown": 11165
+    },
+    {
+      "tau": 0.9,
+      "min_support": 2,
+      "precision": 0.9643,
+      "coverage": 0.269,
+      "keystrokes_saved": 0.1642,
+      "shown": 11156
+    },
+    {
+      "tau": 0.9,
+      "min_support": 5,
+      "precision": 0.9666,
+      "coverage": 0.2069,
+      "keystrokes_saved": 0.1256,
+      "shown": 8580
+    },
+    {
+      "tau": 0.9,
+      "min_support": 10,
+      "precision": 0.9666,
+      "coverage": 0.1757,
+      "keystrokes_saved": 0.1072,
+      "shown": 7286
+    },
+    {
+      "tau": 0.91,
+      "min_support": 0,
+      "precision": 0.9673,
+      "coverage": 0.2609,
+      "keystrokes_saved": 0.1552,
+      "shown": 10820
+    },
+    {
+      "tau": 0.91,
+      "min_support": 2,
+      "precision": 0.9673,
+      "coverage": 0.2607,
+      "keystrokes_saved": 0.1551,
+      "shown": 10811
+    },
+    {
+      "tau": 0.91,
+      "min_support": 5,
+      "precision": 0.9696,
+      "coverage": 0.2008,
+      "keystrokes_saved": 0.1177,
+      "shown": 8328
+    },
+    {
+      "tau": 0.91,
+      "min_support": 10,
+      "precision": 0.9698,
+      "coverage": 0.171,
+      "keystrokes_saved": 0.1002,
+      "shown": 7093
+    },
+    {
+      "tau": 0.92,
+      "min_support": 0,
+      "precision": 0.9693,
+      "coverage": 0.2474,
+      "keystrokes_saved": 0.1485,
+      "shown": 10260
+    },
+    {
+      "tau": 0.92,
+      "min_support": 2,
+      "precision": 0.9693,
+      "coverage": 0.2472,
+      "keystrokes_saved": 0.1484,
+      "shown": 10253
+    },
+    {
+      "tau": 0.92,
+      "min_support": 5,
+      "precision": 0.9717,
+      "coverage": 0.19,
+      "keystrokes_saved": 0.1126,
+      "shown": 7878
+    },
+    {
+      "tau": 0.92,
+      "min_support": 10,
+      "precision": 0.9714,
+      "coverage": 0.1611,
+      "keystrokes_saved": 0.0954,
+      "shown": 6680
+    },
+    {
+      "tau": 0.93,
+      "min_support": 0,
+      "precision": 0.9722,
+      "coverage": 0.2361,
+      "keystrokes_saved": 0.1417,
+      "shown": 9791
+    },
+    {
+      "tau": 0.93,
+      "min_support": 2,
+      "precision": 0.9722,
+      "coverage": 0.2359,
+      "keystrokes_saved": 0.1417,
+      "shown": 9784
+    },
+    {
+      "tau": 0.93,
+      "min_support": 5,
+      "precision": 0.9746,
+      "coverage": 0.1815,
+      "keystrokes_saved": 0.1077,
+      "shown": 7526
+    },
+    {
+      "tau": 0.93,
+      "min_support": 10,
+      "precision": 0.9746,
+      "coverage": 0.154,
+      "keystrokes_saved": 0.0913,
+      "shown": 6388
+    },
+    {
+      "tau": 0.94,
+      "min_support": 0,
+      "precision": 0.9748,
+      "coverage": 0.2226,
+      "keystrokes_saved": 0.1317,
+      "shown": 9234
+    },
+    {
+      "tau": 0.94,
+      "min_support": 2,
+      "precision": 0.9748,
+      "coverage": 0.2225,
+      "keystrokes_saved": 0.1316,
+      "shown": 9228
+    },
+    {
+      "tau": 0.94,
+      "min_support": 5,
+      "precision": 0.9774,
+      "coverage": 0.1719,
+      "keystrokes_saved": 0.1007,
+      "shown": 7130
+    },
+    {
+      "tau": 0.94,
+      "min_support": 10,
+      "precision": 0.9775,
+      "coverage": 0.1459,
+      "keystrokes_saved": 0.0852,
+      "shown": 6052
+    },
+    {
+      "tau": 0.95,
+      "min_support": 0,
+      "precision": 0.9773,
+      "coverage": 0.2049,
+      "keystrokes_saved": 0.1176,
+      "shown": 8497
+    },
+    {
+      "tau": 0.95,
+      "min_support": 2,
+      "precision": 0.9773,
+      "coverage": 0.2047,
+      "keystrokes_saved": 0.1176,
+      "shown": 8491
+    },
+    {
+      "tau": 0.95,
+      "min_support": 5,
+      "precision": 0.9806,
+      "coverage": 0.1582,
+      "keystrokes_saved": 0.0897,
+      "shown": 6560
+    },
+    {
+      "tau": 0.95,
+      "min_support": 10,
+      "precision": 0.9806,
+      "coverage": 0.134,
+      "keystrokes_saved": 0.0756,
+      "shown": 5559
+    }
+  ]
+}

export_web.py

@@ -0,0 +1,102 @@
+"""Write web/model.bin: the compiled model in sparse form for the browser.
+
+The dense model.safetensors is ~93% zeros. Each MLP neuron of the compiled
+transformer is fully described by its context tokens (which slots its gate
+matches) and the logit changes its down-projection writes; the attention and
+embedding weights are fixed by construction. This file stores exactly that, so
+web/autocomplete.ts can evaluate only the neurons whose context matches: the
+same computation as the transformer, without the zeros.
+
+Layout (little-endian):
+    "TACB" | u32 version | u32 header bytes | header JSON | pad to 8
+    sections (each padded to 8), offsets relative to the first section:
+      base                 f32[V]           unigram neuron (always on)
+      ctx{k}   k=1..order  u16/u32[n_k * k] context token ids, oldest first, sorted
+      support{k}           f32[n_k]         weighted count of the context
+      offsets{k}           u32[n_k + 1]     slice of this neuron's deltas
+      dids{k}, dvals{k}    u16/u32, f32     next-token ids and logit changes
+
+    python export_web.py      # uses the "web" settings and "policy" in compile_settings.json
+"""
+
+import json
+from dataclasses import asdict
+from pathlib import Path
+
+import numpy as np
+
+import build_weights as bw
+import suggest as sg
+
+ROOT = Path(__file__).parent
+OUT = ROOT / "web" / "model.bin"
+MAGIC, VERSION = b"TACB", 1
+
+
+def load_policy() -> sg.Policy:
+    return sg.Policy(**json.loads(bw.SETTINGS_PATH.read_text())["policy"])
+
+
+def export(vocab: list[str], ngram: bw.NGram, oov: dict, policy: sg.Policy, path: Path) -> dict:
+    id_dtype = np.uint16 if len(vocab) < 65536 else np.uint32
+    sections, orders = [], []
+
+    def add(name: str, arr: np.ndarray) -> None:
+        sections.append((name, np.ascontiguousarray(arr).tobytes()))
+
+    add("base", ngram.base_logits.astype(np.float32))
+    for k in range(1, ngram.order + 1):
+        ctxs = sorted(ngram.tables[k])
+        offsets, dids, dvals = [0], [], []
+        for ctx in ctxs:
+            d = ngram.delta(ctx)
+            dids.extend(d.keys())
+            dvals.extend(d.values())
+            offsets.append(len(dids))
+        add(f"ctx{k}", np.array(ctxs, dtype=id_dtype).reshape(-1))
+        add(f"support{k}", np.array([ngram.support[c] for c in ctxs], dtype=np.float32))
+        add(f"offsets{k}", np.array(offsets, dtype=np.uint32))
+        add(f"dids{k}", np.array(dids, dtype=id_dtype))
+        add(f"dvals{k}", np.array(dvals, dtype=np.float32))
+        orders.append({"k": k, "contexts": len(ctxs), "pairs": len(dids)})
+
+    table, offset = [], 0
+    for name, data in sections:
+        table.append({"name": name, "offset": offset, "bytes": len(data)})
+        offset += -(-len(data) // 8) * 8
+    header = {
+        "n_prev": ngram.s.n_prev,
+        "vocab": vocab,
+        "eos_id": bw.EOS_ID,
+        "unk_id": bw.UNK_ID,
+        "id_bytes": np.dtype(id_dtype).itemsize,
+        "orders": orders,
+        "oov_prefix_rates": {p: round(r, 4) for p, r in oov.items()},
+        "policy": asdict(policy),
+        "stop_after": sorted(sg.STOP_AFTER),
+        "settings": asdict(ngram.s),
+        "sections": table,
+    }
+    head = json.dumps(header, separators=(",", ":")).encode()
+    path.parent.mkdir(exist_ok=True)
+    with open(path, "wb") as f:
+        f.write(MAGIC + np.array([VERSION, len(head)], dtype="<u4").tobytes() + head)
+        f.write(b"\0" * (-(12 + len(head)) % 8))
+        for _, data in sections:
+            f.write(data + b"\0" * (-len(data) % 8))
+    return {"bytes": path.stat().st_size, "contexts": sum(o["contexts"] for o in orders), "pairs": sum(o["pairs"] for o in orders)}
+
+
+def build(settings: bw.Settings, sources: dict, policy: sg.Policy, path: Path = OUT) -> dict:
+    vocab, ngram, stats = bw.compile_tables(settings, sources)
+    oov = sg.oov_prefix_rates(settings, sources, vocab)
+    return {**stats, **export(vocab, ngram, oov, policy, path)}
+
+
+def main() -> None:
+    info = build(bw.load_settings("web"), bw.load_sources(), load_policy())
+    print(json.dumps({**info, "mb": round(info["bytes"] / 1e6, 1), "out": str(OUT)}, indent=2))
+
+
+if __name__ == "__main__":
+    main()

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:93c23040ccf6393710b0ee1d2fff98d57dc4193a27c4911cac7374e2c28f50f6
+oid sha256:ae857db3e321c35782192c198a7778987964e9cc147d2c67428fea9bf5777100
 size 1521162624

model.ts

@@ -4,6 +4,7 @@
 
 import { readFileSync } from "node:fs";
 import { join } from "node:path";
+import { WordTokenizer } from "./tokenizer.ts";
 
 export type Config = {
   vocab_size: number;
@@ -44,63 +45,6 @@ function loadSafetensors(path: string): Map<string, Float32Array> {
   return tensors;
 }
 
-// ---------- tokenizer (WordLevel + ByteLevel, as in tokenizer.json) ----------
-
-function bytesToUnicode(): Map<number, string> {
-  const bs: number[] = [];
-  for (let b = 33; b <= 126; b++) bs.push(b);
-  for (let b = 161; b <= 172; b++) bs.push(b);
-  for (let b = 174; b <= 255; b++) bs.push(b);
-  const cs = [...bs];
-  let n = 0;
-  for (let b = 0; b < 256; b++) {
-    if (!bs.includes(b)) {
-      bs.push(b);
-      cs.push(256 + n++);
-    }
-  }
-  return new Map(bs.map((b, i) => [b, String.fromCodePoint(cs[i])]));
-}
-
-const GPT2_SPLIT = /'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+/gu;
-
-export class Tokenizer {
-  vocab: Map<string, number>;
-  inverse: string[];
-  special = new Set<number>();
-  unkId: number;
-  bosId: number;
-  private byteEnc = bytesToUnicode();
-  private byteDec = new Map([...this.byteEnc].map(([b, c]) => [c, b]));
-
-  constructor(path: string) {
-    const json = JSON.parse(readFileSync(path, "utf8"));
-    this.vocab = new Map(Object.entries<number>(json.model.vocab));
-    this.inverse = [];
-    for (const [tok, id] of this.vocab) this.inverse[id] = tok;
-    for (const t of json.added_tokens ?? []) if (t.special) this.special.add(t.id);
-    this.unkId = this.vocab.get(json.model.unk_token)!;
-    const bos = json.post_processor?.single?.find((p: any) => p.SpecialToken)?.SpecialToken?.id;
-    this.bosId = this.vocab.get(bos)!;
-  }
-
-  encode(text: string): number[] {
-    let norm = text.normalize("NFKC").toLowerCase();
-    if (!/^\s/.test(norm)) norm = " " + norm; // ByteLevel add_prefix_space
-    const ids = [this.bosId];
-    for (const piece of norm.match(GPT2_SPLIT) ?? []) {
-      const mapped = [...new TextEncoder().encode(piece)].map((b) => this.byteEnc.get(b)!).join("");
-      ids.push(this.vocab.get(mapped) ?? this.unkId);
-    }
-    return ids;
-  }
-
-  decode(ids: number[]): string {
-    const chars = ids.filter((id) => !this.special.has(id)).map((id) => this.inverse[id]).join("");
-    return new TextDecoder().decode(new Uint8Array([...chars].map((c) => this.byteDec.get(c) ?? 63)));
-  }
-}
-
 // ---------- Llama forward pass ----------
 
 function matvec(w: Float32Array, x: Float64Array, rows: number, cols: number): Float64Array {
@@ -228,12 +172,12 @@ export type Completion = {
 };
 
 export class Autocompleter {
-  tok: Tokenizer;
+  tok: WordTokenizer;
   model: Llama;
   maxNew: number;
 
   constructor(dir: string) {
-    this.tok = new Tokenizer(join(dir, "tokenizer.json"));
+    this.tok = WordTokenizer.fromJson(JSON.parse(readFileSync(join(dir, "tokenizer.json"), "utf8")));
     this.model = new Llama(dir);
     const gen = JSON.parse(readFileSync(join(dir, "generation_config.json"), "utf8"));
     this.maxNew = gen.max_new_tokens ?? 16;

scripts/benchmark.py

@@ -1,6 +1,6 @@
-"""Measure compile cost, inference speed, exactness and held-out accuracy.
+"""Measure compile cost, inference speed, exactness and held-out quality.
 
-Writes docs/benchmarks.json. Run from the repo root after build_weights.py:
+Writes docs/benchmarks.json. Run from the repo root after build_weights.py and export_web.py:
     python scripts/benchmark.py
 """
 
@@ -11,17 +11,21 @@ import platform
 import random
 import resource
 import shutil
+import statistics
 import subprocess
 import sys
 import time
-import zlib
 from pathlib import Path
 
 import torch
 
 ROOT = Path(__file__).resolve().parent.parent
 sys.path.insert(0, str(ROOT))
+sys.path.insert(0, str(ROOT / "scripts"))
 import build_weights as bw  # noqa: E402
+import export_web  # noqa: E402
+import suggest as sg  # noqa: E402
+import sweep  # noqa: E402
 
 PROMPTS = [
     "Once upon a", "I hope this email", "The quick brown", "Thank you for", "Looking forward to",
@@ -50,44 +54,53 @@ def hardware() -> dict:
             "python": platform.python_version(), "torch": torch.__version__}
 
 
-def compile_runs(n: int = 2) -> dict:
-    """Full `python build_weights.py` runs: wall time, peak memory, determinism."""
+def timed_runs(cmd: list[str], out: Path, n: int = 2) -> dict:
     times, hashes = [], []
     for _ in range(n):
         t0 = time.perf_counter()
-        subprocess.run([sys.executable, "build_weights.py"], cwd=ROOT, check=True, capture_output=True)
+        subprocess.run(cmd, cwd=ROOT, check=True, capture_output=True)
         times.append(time.perf_counter() - t0)
-        hashes.append(sha256(ROOT / "model.safetensors"))
-    peak = resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss
-    peak_gb = peak / 2**30 if sys.platform == "darwin" else peak / 2**20
+        hashes.append(sha256(out))
     return {
-        "compile_seconds": round(min(times), 1),
-        "compile_runs": [round(t, 1) for t in times],
-        "peak_memory_gb": round(peak_gb, 1),
+        "seconds": round(min(times), 1),
+        "runs": [round(t, 1) for t in times],
         "byte_identical_rebuild": len(set(hashes)) == 1,
-        "model_sha256": hashes[0],
-        "file_gb": round((ROOT / "model.safetensors").stat().st_size / 1e9, 2),
+        "sha256": hashes[0],
+        "mb": round(out.stat().st_size / 1e6, 2),
     }
 
 
-def compile_stages() -> dict:
+def compile_dense() -> dict:
+    runs = timed_runs([sys.executable, "build_weights.py"], ROOT / "model.safetensors")
+    peak = resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss
     t0 = time.perf_counter()
-    phrases, data = bw.load_phrases(), bw.load_data()
+    sources = bw.load_sources()
     t1 = time.perf_counter()
-    _, ngram, state, _, info = bw.build(phrases, data)
+    s = bw.load_settings("hf")
+    vocab, ngram, stats = bw.compile_tables(s, sources)
     t2 = time.perf_counter()
-    nonzero = sum(int((t != 0).sum()) for t in state.values())
+    state, _, info = bw.build_dense(vocab, ngram, s)
+    t3 = time.perf_counter()
     return {
+        **runs,
+        "peak_memory_gb": round(peak / 2**30 if sys.platform == "darwin" else peak / 2**20, 1),
         "tokenize_seconds": round(t1 - t0, 2),
-        "build_tensors_seconds": round(t2 - t1, 2),
-        "source_lines": len(phrases) + len(data),
+        "count_seconds": round(t2 - t1, 2),
+        "build_tensors_seconds": round(t3 - t2, 2),
+        "source_lines": sum(len(v) for v in sources.values()),
         "params": info["params"],
-        "nonzero_params": nonzero,
-        "contexts_seen": info["contexts_seen"],
-        "neurons": info["neurons"],
+        "nonzero_params": sum(int((t != 0).sum()) for t in state.values()),
+        **stats,
     }
 
 
+def compile_web() -> dict:
+    runs = timed_runs([sys.executable, "export_web.py"], ROOT / "web" / "model.bin")
+    header_len = int.from_bytes((ROOT / "web" / "model.bin").read_bytes()[8:12], "little")
+    header = json.loads((ROOT / "web" / "model.bin").read_bytes()[12 : 12 + header_len])
+    return {**runs, "contexts": sum(o["contexts"] for o in header["orders"]), "vocab": len(header["vocab"])}
+
+
 def torch_inference() -> dict:
     from transformers import AutoModelForCausalLM, AutoTokenizer
 
@@ -95,7 +108,6 @@ def torch_inference() -> dict:
     tok = AutoTokenizer.from_pretrained(ROOT)
     model = AutoModelForCausalLM.from_pretrained(ROOT).eval()
     load = time.perf_counter() - t0
-
     with torch.no_grad():
         ids = tok("Write a python function to", return_tensors="pt")
         model(**ids)  # warm-up
@@ -103,33 +115,27 @@ def torch_inference() -> dict:
         for _ in range(20):
             model(**ids)
         forward_ms = (time.perf_counter() - t0) / 20 * 1000
-
         new_tokens, t0 = 0, time.perf_counter()
         for p in PROMPTS:
             ids = tok(p, return_tensors="pt")
             out = model.generate(**ids, max_new_tokens=16)
             new_tokens += out.shape[1] - ids["input_ids"].shape[1]
         gen = time.perf_counter() - t0
-    return {
-        "load_seconds": round(load, 1),
-        "next_word_ms": round(forward_ms, 1),
-        "generate_tokens_per_second": round(new_tokens / gen, 1),
-        "threads": torch.get_num_threads(),
-    }
+    return {"load_seconds": round(load, 1), "next_word_ms": round(forward_ms, 1),
+            "generate_tokens_per_second": round(new_tokens / gen, 1), "threads": torch.get_num_threads()}
 
 
-def ts_inference() -> dict | None:
-    bun = shutil.which("bun")
-    if not bun:
+def web_latency(texts: list[str]) -> dict | None:
+    """Per-keystroke suggestion time of web/autocomplete.ts on held-out prefixes (Bun)."""
+    if not shutil.which("bun"):
         return None
-    t0 = time.perf_counter()
-    subprocess.run([bun, "autocomplete.ts", "--json"], cwd=ROOT, check=True, capture_output=True)
-    load = time.perf_counter() - t0
-    t0 = time.perf_counter()
-    out = subprocess.run([bun, "autocomplete.ts", "--json", *PROMPTS], cwd=ROOT, check=True, capture_output=True, text=True)
-    total = time.perf_counter() - t0 - load
-    words = sum(len(bw.pieces(r["completion"])) + 1 for r in json.loads(out.stdout))  # +1: the final EOS step
-    return {"runtime": "bun, plain TypeScript, no ML library", "load_seconds": round(load, 1), "seconds_per_token": round(total / words, 2)}
+    rng = random.Random(0)
+    prefixes = [t[: rng.randint(1, len(t) - 1)] for t in rng.sample(texts, 500)]
+    out = json.loads(subprocess.run(["bun", "web/cli.ts", "--json", *prefixes], cwd=ROOT, check=True,
+                                    capture_output=True, text=True).stdout)
+    ms = sorted(r["ms"] for r in out[50:])  # skip JIT warm-up
+    return {"runtime": "web/autocomplete.ts on Bun", "median_ms": round(statistics.median(ms), 2),
+            "p95_ms": round(ms[int(0.95 * len(ms))], 2), "max_ms": round(ms[-1], 2)}
 
 
 def exactness() -> float:
@@ -146,12 +152,23 @@ def exactness() -> float:
     return worst
 
 
-def heldout(frac: float = 0.1) -> dict:
-    """Compile without 10% of the data lines, then score next-word prediction on them."""
-    phrases, data = bw.load_phrases(), bw.load_data()
-    test = [t for t in data if zlib.crc32(" ".join(t).encode()) % 100 < frac * 100]
-    train = [t for t in data if zlib.crc32(" ".join(t).encode()) % 100 >= frac * 100]
-    vocab, ngram, _, _, _ = bw.build(phrases, train)
+def heldout_typing(name: str, texts: list[str], src: dict) -> dict:
+    """Simulated typing on held-out prompts with the shipped policy, compiled without them."""
+    s = bw.load_settings(name)
+    vocab, ngram, _ = bw.compile_tables(s, src)
+    sugg = sg.Suggester(vocab, ngram, sg.oov_prefix_rates(s, src, vocab))
+    recs = sweep.record(sugg, texts)
+    policy = export_web.load_policy()
+    curve = {f"tau={t}": sweep.evaluate(recs, sg.Policy(tau=t, min_support=policy.min_support))
+             for t in [policy.tau, 0.8, 0.9, 0.95]}
+    return {"policy": policy.__dict__, "heldout_prompts": len(texts), "characters": sum(map(len, texts)), **curve}
+
+
+def heldout_next_token(texts_src: dict) -> dict:
+    """Top-1/top-3 next-token accuracy of the dense model's tables on held-out prompts."""
+    s = bw.load_settings("hf")
+    _, test = bw.heldout_split(bw.load_sources()["prompts"])
+    vocab, ngram, _ = bw.compile_tables(s, texts_src)
     index = {p: i for i, p in enumerate(vocab)}
     top1 = top3 = total = oov = 0
     for toks in test:
@@ -165,27 +182,26 @@ def heldout(frac: float = 0.1) -> dict:
             ranked = ngram.reference_logits(ids[: t + 1]).topk(3).indices.tolist()
             top1 += ranked[0] == target
             top3 += target in ranked
-    return {
-        "heldout_lines": len(test),
-        "heldout_tokens": total,
-        "next_token_top1": round(top1 / total, 3),
-        "next_token_top3": round(top3 / total, 3),
-        "oov_rate": round(oov / total, 3),
-    }
+    return {"heldout_lines": len(test), "heldout_tokens": total, "next_token_top1": round(top1 / total, 3),
+            "next_token_top3": round(top3 / total, 3), "oov_rate": round(oov / total, 3)}
 
 
 def main() -> None:
     results = {"hardware": hardware()}
-    print("compiling (2 full runs)...", flush=True)
-    results["compile"] = {**compile_runs(), **compile_stages()}
+    print("compiling dense model (2 runs)...", flush=True)
+    results["compile_dense"] = compile_dense()
+    print("compiling web export (2 runs)...", flush=True)
+    results["compile_web"] = compile_web()
     print("inference...", flush=True)
     results["pytorch"] = torch_inference()
-    results["typescript"] = ts_inference()
+    texts, src = sweep.heldout_texts(), sweep.train_sources()
+    results["web_keystroke_latency"] = web_latency(texts)
     results["max_logprob_error_vs_ngram"] = exactness()
-    print("held-out accuracy...", flush=True)
-    results["heldout"] = heldout()
-    out = ROOT / "docs" / "benchmarks.json"
-    out.write_text(json.dumps(results, indent=2) + "\n")
+    print("held-out quality...", flush=True)
+    results["heldout_typing_web"] = heldout_typing("web", texts, src)
+    results["heldout_typing_dense_tables"] = heldout_typing("hf", texts, src)
+    results["heldout_next_token_dense"] = heldout_next_token(src)
+    (ROOT / "docs" / "benchmarks.json").write_text(json.dumps(results, indent=2) + "\n")
     print(json.dumps(results, indent=2))
 
 

scripts/experiment_topic.py

@@ -0,0 +1,112 @@
+"""Experiment: does long-range "topic" context help, before compiling it into attention?
+
+Hypothesis: a head that attends to the most recent topic word anywhere earlier in the
+prompt (python, sql, react, ...) gives context the 4-token n-gram window can't see.
+
+Closed-form trigger model (no fitting): for each topic word t, count next-word
+frequencies at positions where t is the most recent topic word earlier in the line.
+The bias added to the n-gram logits is
+    lambda * log( p(w | t) / p(w) ),   p(w | t) = (c(t, w) + 1) / (c(t) + V)
+kept only where c(t, w) >= MIN_COUNT. In a compiled transformer this would be one
+content-addressed attention head (keys = "is topic word" + code, recency via RoPE) and
+one MLP neuron per topic word writing these biases.
+
+Scored with the same held-out typing simulation as scripts/sweep.py.
+"""
+
+import json
+import sys
+from collections import Counter, defaultdict
+from pathlib import Path
+
+import numpy as np
+
+ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(ROOT))
+sys.path.insert(0, str(ROOT / "scripts"))
+import build_weights as bw  # noqa: E402
+import suggest as sg  # noqa: E402
+import sweep  # noqa: E402
+
+TOPIC_WORDS = [
+    "python", "javascript", "js", "typescript", "java", "c", "c++", "c#", "rust", "golang", "go", "ruby", "php",
+    "kotlin", "swift", "sql", "mysql", "postgresql", "sqlite", "bash", "shell", "powershell", "html", "css", "react",
+    "vue", "angular", "node", "django", "flask", "pandas", "numpy", "pytorch", "tensorflow", "excel", "linux", "windows",
+    "git", "github", "docker", "kubernetes", "regex", "json", "api", "unity", "blender", "matlab", "r", "lua", "haskell",
+]
+MIN_COUNT = 3
+
+
+class TopicSuggester(sg.Suggester):
+    def __init__(self, vocab, ngram, oov, topics: set[int], bias: dict[int, tuple[np.ndarray, np.ndarray]], lam: float):
+        super().__init__(vocab, ngram, oov)
+        self.topics, self.bias, self.lam = topics, bias, lam
+
+    def topic_of(self, ids: list[int]) -> int | None:
+        for i in reversed(ids):
+            if i in self.topics:
+                return i
+        return None
+
+    def dist(self, ids):
+        t = self.topic_of(ids)
+        key = (self.ngram.window(ids), t)
+        if key not in self._dist:
+            logits, support, _ = self.ngram.predict(ids)
+            if t is not None and t in self.bias:
+                idx, vals = self.bias[t]
+                logits[idx] += self.lam * vals
+            p = np.exp(logits - logits.max())
+            self._dist[key] = (p / p.sum(), support)
+        return self._dist[key]
+
+
+def topic_bias(vocab: list[str], src: dict, s: bw.Settings, topics: set[int]):
+    index = {p: i for i, p in enumerate(vocab)}
+    V = len(vocab)
+    uni, by_topic = Counter(), defaultdict(Counter)
+    for lines, w in bw.weighted_sources(s, src):
+        for toks in lines:
+            ids = [index.get(p, bw.UNK_ID) for p in toks] + [bw.EOS_ID]
+            topic = None
+            for i, nxt in enumerate(ids):
+                if nxt != bw.UNK_ID:
+                    uni[nxt] += w
+                    if topic is not None:
+                        by_topic[topic][nxt] += w
+                if i < len(ids) - 1 and ids[i] in topics:
+                    topic = ids[i]
+    total = sum(uni.values())
+    out = {}
+    for t, cnt in by_topic.items():
+        ct = sum(cnt.values())
+        keep = [(w, c) for w, c in cnt.items() if c >= MIN_COUNT]
+        if not keep:
+            continue
+        idx = np.array([w for w, _ in keep])
+        vals = np.array([np.log((c + 1) / (ct + V)) - np.log((uni[w] + 1) / (total + V)) for w, c in keep])
+        out[t] = (idx, vals)
+    return out
+
+
+def main() -> None:
+    s = bw.load_settings("web")
+    src, texts = sweep.train_sources(), sweep.heldout_texts()
+    vocab, ngram, _ = bw.compile_tables(s, src)
+    oov = sg.oov_prefix_rates(s, src, vocab)
+    index = {p: i for i, p in enumerate(vocab)}
+    topics = {index[p] for w in TOPIC_WORDS for p in ["Ġ" + w] if p in index}
+    bias = topic_bias(vocab, src, s, topics)
+    with_topic = sum(1 for t in texts if any(("Ġ" + w) in bw.pieces(t) for w in TOPIC_WORDS))
+    print(f"topic words in vocab: {len(topics)}, with biases: {len(bias)}; held-out prompts mentioning one: {with_topic}/{len(texts)}")
+    results = []
+    for lam in [0.0, 0.25, 0.5, 1.0]:
+        sugg = TopicSuggester(vocab, ngram, oov, topics, bias, lam)
+        best = sweep.best_policy(sweep.record(sugg, texts))["best"]
+        results.append({"lambda": lam, **best})
+        print(f"lambda={lam}: saved={best['keystrokes_saved']} prec={best['precision']} cov={best['coverage']} tau={best['tau']}", flush=True)
+    (ROOT / "docs" / "experiment_topic.json").write_text(json.dumps({"topic_words": len(topics), "results": results}, indent=2) + "\n")
+
+
+if __name__ == "__main__":
+    main()

scripts/fetch_coding_prompts.py

@@ -1,4 +1,7 @@
-"""Download human-written coding prompts and write data/coding_prompts.txt.
+"""Download human-written coding text into data/prompts.txt and data/replies.txt.
+
+prompts.txt: what people type to an assistant (OASST user turns, MBPP, CoNaLa).
+replies.txt: OASST assistant turns (human-written replies).
 
 Sources (all human-written, permissively licensed; see data/SOURCES.md):
   - OpenAssistant/oasst2  (Apache-2.0)  English prompter and assistant messages from
@@ -18,7 +21,7 @@ from pathlib import Path
 import pyarrow.parquet as pq
 from huggingface_hub import hf_hub_download
 
-OUT = Path(__file__).resolve().parent.parent / "data" / "coding_prompts.txt"
+DATA = Path(__file__).resolve().parent.parent / "data"
 MAX_WORDS = 24
 MAX_PROMPT_SENTENCES = 2
 MAX_REPLY_SENTENCES = 5
@@ -61,7 +64,7 @@ def clean(text: str, max_sentences: int = MAX_PROMPT_SENTENCES) -> list[str]:
     return out
 
 
-def oasst() -> list[str]:
+def oasst() -> tuple[list[str], list[str]]:
     path = hf_hub_download("OpenAssistant/oasst2", "2023-11-05_oasst2_ready.messages.jsonl.gz", repo_type="dataset")
     rows = [json.loads(line) for line in gzip.open(path, "rt")]
     rows = [
@@ -69,12 +72,14 @@ def oasst() -> list[str]:
         if r["lang"] == "en" and not r["deleted"] and r["review_result"] and not r["synthetic"]
     ]
     coding_trees = {r["message_tree_id"] for r in rows if r["parent_id"] is None and CODE_RE.search(r["text"])}
-    out = []
+    prompts, replies = [], []
     for r in rows:
         if r["message_tree_id"] in coding_trees:
-            limit = MAX_PROMPT_SENTENCES if r["role"] == "prompter" else MAX_REPLY_SENTENCES
-            out += clean(r["text"], limit)
-    return out
+            if r["role"] == "prompter":
+                prompts += clean(r["text"], MAX_PROMPT_SENTENCES)
+            else:
+                replies += clean(r["text"], MAX_REPLY_SENTENCES)
+    return prompts, replies
 
 
 def mbpp() -> list[str]:
@@ -97,19 +102,24 @@ def conala() -> list[str]:
 
 
 def main() -> None:
-    seen, lines = set(), []
-    counts = {}
-    for name, fetch in [("oasst2", oasst), ("mbpp", mbpp), ("conala", conala)]:
-        got = fetch()
-        counts[name] = len(got)
-        for s in got:
-            key = s.lower()
-            if key not in seen:
-                seen.add(key)
-                lines.append(s)
-    OUT.parent.mkdir(exist_ok=True)
-    OUT.write_text("\n".join(lines) + "\n")
-    print(json.dumps({**counts, "unique_lines": len(lines), "out": str(OUT)}, indent=2))
+    oasst_prompts, oasst_replies = oasst()
+    sources = {
+        "prompts": [("oasst2_user", oasst_prompts), ("mbpp", mbpp()), ("conala", conala())],
+        "replies": [("oasst2_assistant", oasst_replies)],
+    }
+    seen, report = set(), {}
+    DATA.mkdir(exist_ok=True)
+    for kind, parts in sources.items():
+        lines = []
+        for name, got in parts:
+            report[name] = len(got)
+            for s in got:
+                if s.lower() not in seen:
+                    seen.add(s.lower())
+                    lines.append(s)
+        (DATA / f"{kind}.txt").write_text("\n".join(lines) + "\n")
+        report[f"{kind}_unique"] = len(lines)
+    print(json.dumps(report, indent=2))
 
 
 if __name__ == "__main__":

scripts/sweep.py

@@ -0,0 +1,185 @@
+"""Choose compile settings and the confidence policy by simulated typing on held-out prompts.
+
+The weights are always produced by the same closed-form compile; this script only
+compares a handful of compile settings (context length, smoothing, source weights,
+vocabulary, pruning) and policy thresholds, and reports which one saves the most
+keystrokes while keeping precision >= TARGET_PRECISION.
+
+Simulation: each held-out prompt (10% of data/prompts.txt, never compiled in) is
+"typed" one character at a time. After every keystroke the policy shows ghost text
+or nothing. A shown suggestion is correct only if the user's next characters start
+with the whole ghost text. Keystrokes saved replays the prompt accepting every
+correct suggestion with one Tab press.
+
+    python scripts/sweep.py            # writes docs/sweep.json
+"""
+
+import json
+import sys
+import time
+import zlib
+from dataclasses import asdict, replace
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(ROOT))
+import build_weights as bw  # noqa: E402
+import suggest as sg  # noqa: E402
+
+TARGET_PRECISION = 0.90
+MAX_EXPORT_MB = 50
+TAUS = [round(0.5 + 0.01 * i, 2) for i in range(46)]  # 0.50 .. 0.95
+SUPPORTS = [0, 2, 5, 10]
+BASE = bw.Settings(n_prev=3, smoothing="wb", phrase_weight=5, prompt_weight=1, reply_weight=1, max_vocab=8000, min_count=2)
+
+
+def heldout_texts() -> list[str]:
+    seen, out = set(), []
+    for line in (bw.DATA_DIR / "prompts.txt").read_text().splitlines():
+        toks = bw.pieces(line)
+        if toks and zlib.crc32(" ".join(toks).encode()) % 100 < 10:
+            text = bw.NORMALIZER.normalize_str(line.strip())
+            if text not in seen:
+                seen.add(text)
+                out.append(text)
+    return out
+
+
+def train_sources() -> dict:
+    src = bw.load_sources()
+    src["prompts"], _ = bw.heldout_split(src["prompts"])
+    return src
+
+
+def record(sugg: sg.Suggester, texts: list[str]) -> list[tuple[str, list]]:
+    """For every prompt and every keystroke position, the unthresholded suggestion chain."""
+    return [(t, [None] + [sugg.steps(t[:i]) for i in range(1, len(t))]) for t in texts]
+
+
+def evaluate(records, policy: sg.Policy) -> dict:
+    shown = correct = positions = chars = strokes = 0
+    for text, steps in records:
+        ghosts = [None] + [sg.Suggester.apply(st, policy) for st in steps[1:]]
+        for i in range(1, len(text)):
+            positions += 1
+            if ghosts[i]:
+                shown += 1
+                correct += text.startswith(ghosts[i], i)
+        pos, n = 0, 0
+        while pos < len(text):
+            pos, n = pos + 1, n + 1
+            g = ghosts[pos] if pos < len(text) else ""
+            if g and text.startswith(g, pos):
+                pos, n = pos + len(g), n + 1
+        chars, strokes = chars + len(text), strokes + n
+    return {
+        "precision": round(correct / shown, 4) if shown else None,
+        "coverage": round(shown / positions, 4),
+        "keystrokes_saved": round(1 - strokes / chars, 4),
+        "shown": shown,
+    }
+
+
+def best_policy(records) -> dict:
+    results = []
+    for tau in TAUS:
+        for sup in SUPPORTS:
+            m = evaluate(records, sg.Policy(tau=tau, min_support=sup))
+            results.append({"tau": tau, "min_support": sup, **m})
+    ok = [r for r in results if r["precision"] is not None and r["precision"] >= TARGET_PRECISION]
+    best = max(ok, key=lambda r: r["keystrokes_saved"]) if ok else None
+    return {"best": best, "grid": results}
+
+
+def export_mb(ngram: bw.NGram, vocab: list[str], oov: dict) -> float:
+    """Size of web/model.bin for these tables (see export_web.py)."""
+    size = sum(len(p.encode()) + 1 for p in vocab) + 8 * len(vocab) + sum(len(p.encode()) + 5 for p in oov)
+    for k in range(1, ngram.order + 1):
+        for table in ngram.tables[k].values():
+            size += 2 * k + 8 + 6 * len(table)
+    return round(size / 1e6, 1)
+
+
+def run(s: bw.Settings, src: dict, texts: list[str]) -> dict:
+    t0 = time.time()
+    vocab, ngram, stats = bw.compile_tables(s, src)
+    oov = sg.oov_prefix_rates(s, src, vocab)
+    recs = record(sg.Suggester(vocab, ngram, oov), texts)
+    pol = best_policy(recs)
+    out = {
+        "settings": asdict(s),
+        **stats,
+        "export_mb": export_mb(ngram, vocab, oov),
+        "best": pol["best"],
+        "seconds": round(time.time() - t0, 1),
+        "grid": pol["grid"],
+    }
+    b = out["best"] or {}
+    print(
+        f"n_prev={s.n_prev} {s.smoothing} w=({s.phrase_weight},{s.prompt_weight},{s.reply_weight}) "
+        f"V={s.max_vocab} min={s.min_count} | {out['export_mb']} MB | "
+        f"saved={b.get('keystrokes_saved')} prec={b.get('precision')} cov={b.get('coverage')} "
+        f"tau={b.get('tau')} sup={b.get('min_support')} | {out['seconds']}s",
+        flush=True,
+    )
+    return out
+
+
+def score(r: dict) -> float:
+    if r["best"] is None or r["export_mb"] > MAX_EXPORT_MB:
+        return -1
+    return r["best"]["keystrokes_saved"]
+
+
+def main() -> None:
+    src, texts = train_sources(), heldout_texts()
+    print(f"held-out prompts: {len(texts)}, characters: {sum(map(len, texts))}", flush=True)
+    runs = []
+
+    print("== stage A: context length x smoothing", flush=True)
+    stage = [run(replace(BASE, n_prev=n, smoothing=sm), src, texts) for sm in ["wb", "kn"] for n in [3, 4, 5, 7, 9]]
+    runs += stage
+    best = max(stage, key=score)
+
+    print("== stage B: source weights", flush=True)
+    base = bw.Settings(**best["settings"])
+    stage = [
+        run(replace(base, phrase_weight=pw, reply_weight=rw), src, texts)
+        for pw in [1.0, 5.0]
+        for rw in [0.0, 0.3, 1.0]
+        if (pw, rw) != (base.phrase_weight, base.reply_weight)
+    ]
+    runs += stage
+    best = max(runs, key=score)
+
+    print("== stage C: vocabulary x pruning", flush=True)
+    base = bw.Settings(**best["settings"])
+    stage = [
+        run(replace(base, max_vocab=v, min_count=mc), src, texts)
+        for v in [3000, 8000, 16000]
+        for mc in [2, 3]
+        if (v, mc) != (base.max_vocab, base.min_count)
+    ]
+    runs += stage
+    best = max(runs, key=score)
+
+    print("== stage D: context length again, with the final settings", flush=True)
+    base = bw.Settings(**best["settings"])
+    stage = [run(replace(base, n_prev=n), src, texts) for n in [3, 4, 5, 7, 9] if n != base.n_prev]
+    runs += stage
+    best = max(runs, key=score)
+
+    out = ROOT / "docs" / "sweep.json"
+    out.write_text(json.dumps({
+        "target_precision": TARGET_PRECISION,
+        "max_export_mb": MAX_EXPORT_MB,
+        "heldout_prompts": len(texts),
+        "best": {k: v for k, v in best.items() if k != "grid"},
+        "runs": [{k: v for k, v in r.items() if k != "grid"} for r in runs],
+        "best_grid": best["grid"],
+    }, indent=2) + "\n")
+    print("BEST", json.dumps({k: v for k, v in best.items() if k != "grid"}, indent=2))
+
+
+if __name__ == "__main__":
+    main()

suggest.py

@@ -0,0 +1,139 @@
+"""Confidence-gated ghost-text suggestions (reference implementation).
+
+web/autocomplete.ts implements the same policy for the browser; the tests check
+they agree. Nothing here is learned: the policy reads the compiled model's
+next-token probabilities and only suggests when they clear fixed thresholds.
+
+Given the text typed so far:
+  1. The last pre-token is the word being typed (just "Ġ" after a space). Only
+     vocabulary tokens that start with it are candidates ("fun" -> "function").
+  2. Candidate probabilities are renormalized over the candidates and scaled by
+     the chance that a word with this prefix is in the vocabulary at all
+     (counted from the source text), so rare out-of-vocabulary words don't
+     produce confident wrong completions.
+  3. The suggestion is shown only if that probability >= tau and the longest
+     matched context was seen >= min_support times. It then extends greedily,
+     word by word, while the product of probabilities stays >= tau, stopping
+     at end-of-text or after . ? ! :
+"""
+
+from collections import Counter
+from dataclasses import dataclass
+
+import numpy as np
+
+from build_weights import EOS_ID, UNK_ID, NGram, Settings, pieces, weighted_sources
+
+STOP_AFTER = {".", "?", "!", ":"}
+MAX_WORDS = 8
+
+
+@dataclass(frozen=True)
+class Policy:
+    tau: float = 0.6  # min probability of the whole suggestion
+    min_support: float = 2  # min (weighted) count of the longest matched context
+    max_words: int = MAX_WORDS
+
+
+@dataclass(frozen=True)
+class Step:
+    text: str  # characters this step adds to the ghost text
+    prob: float  # probability of the suggestion up to and including this step
+    support: float  # count of the longest context matched for this step
+
+
+def _byte_decoder() -> dict[str, int]:
+    bs = list(range(33, 127)) + list(range(161, 173)) + list(range(174, 256))
+    cs, n = bs[:], 0
+    for b in range(256):
+        if b not in bs:
+            bs.append(b)
+            cs.append(256 + n)
+            n += 1
+    return {chr(c): b for b, c in zip(bs, cs)}
+
+
+BYTE_DECODER = _byte_decoder()
+
+
+def decode_piece(p: str) -> str:
+    return bytes(BYTE_DECODER[c] for c in p).decode("utf-8", errors="replace")
+
+
+def oov_prefix_rates(s: Settings, sources: dict, vocab: list[str]) -> dict[str, float]:
+    """P(word is out of vocabulary | word starts with prefix), from weighted source counts.
+
+    Only prefixes with a non-zero rate are stored; prefixes of in-vocabulary words that
+    never start an OOV word are omitted (rate 0).
+    """
+    in_vocab = set(vocab)
+    total, oov = Counter(), Counter()
+    for lines, w in weighted_sources(s, sources):
+        for toks in lines:
+            for p in toks:
+                for i in range(1, len(p) + 1):
+                    total[p[:i]] += w
+                    if p not in in_vocab:
+                        oov[p[:i]] += w
+    return {pre: oov[pre] / total[pre] for pre in oov}
+
+
+class Suggester:
+    def __init__(self, vocab: list[str], ngram: NGram, oov_rates: dict[str, float]):
+        self.vocab, self.ngram, self.oov = vocab, ngram, oov_rates
+        self.index = {p: i for i, p in enumerate(vocab)}
+        self.text = [decode_piece(p) for p in vocab]
+        self._dist, self._cands = {}, {}
+
+    def dist(self, ids: list[int]) -> tuple[np.ndarray, float]:
+        key = self.ngram.window(ids)
+        if key not in self._dist:
+            logits, support, _ = self.ngram.predict(ids)
+            p = np.exp(logits - logits.max())
+            self._dist[key] = (p / p.sum(), support)
+        return self._dist[key]
+
+    def candidates(self, partial: str) -> np.ndarray:
+        if partial not in self._cands:
+            self._cands[partial] = np.array(
+                [i for i, p in enumerate(self.vocab) if i > UNK_ID and p.startswith(partial)], dtype=np.int64
+            )
+        return self._cands[partial]
+
+    def steps(self, text: str, max_words: int = MAX_WORDS) -> list[Step]:
+        """Greedy suggestion chain with cumulative probabilities (before any thresholding)."""
+        pcs = pieces(text)
+        if not pcs:
+            return []
+        partial = pcs[-1]
+        cands = self.candidates(partial)
+        if not len(cands):
+            return []
+        ids = [EOS_ID] + [self.index.get(p, UNK_ID) for p in pcs[:-1]]
+        probs, support = self.dist(ids)
+        pc = probs[cands]
+        best = int(cands[pc.argmax()])
+        prob = (1.0 - self.oov.get(partial, 0.0)) * pc.max() / pc.sum()
+        out = [Step(self.text[best][len(decode_piece(partial)) :], prob, support)]
+        ids.append(best)
+        while self.vocab[ids[-1]] not in STOP_AFTER and len(out) < max_words:
+            probs, support = self.dist(ids)
+            nxt = int(probs.argmax())
+            if nxt == EOS_ID:
+                break
+            prob *= probs[nxt]
+            out.append(Step(self.text[nxt], prob, support))
+            ids.append(nxt)
+        return out
+
+    @staticmethod
+    def apply(steps: list[Step], policy: Policy) -> str:
+        ghost = ""
+        for i, st in enumerate(steps[: policy.max_words]):
+            if st.prob < policy.tau or st.support < policy.min_support:
+                break
+            ghost += st.text
+        return ghost
+
+    def suggest(self, text: str, policy: Policy = Policy()) -> str:
+        return self.apply(self.steps(text, policy.max_words), policy)

test_model.py

@@ -24,15 +24,15 @@ CASES = {
     "Thank you for your": "help.",
     "I hope this email": "finds you well.",
     "Can I have the": "menu, please?",
-    "Looking forward to": "hearing from you.",
+    "Looking forward to": "working with you.",
     "Actions speak": "louder than words.",
     "Where is the": "bathroom?",
     "Happy birthday": ".",
     "I'm so proud": "of you.",
-    # coding prompts (data/coding_prompts.txt)
+    # coding prompts and replies (data/*.txt)
     "In Python,": "you can use the following code:",
-    "How to convert a": "string to a list of tuples",
-    "Here is an example of": "a basic particle class.",
+    "You can use the": "following code:",
+    "How to convert a": "string to a file in python",
 }
 
 
@@ -103,5 +103,66 @@ class TestModel(unittest.TestCase):
                     self.assertEqual(res["completion"].strip(), self.complete(prompt))
 
 
+@unittest.skipUnless(shutil.which("bun"), "bun not installed")
+class TestWebExport(unittest.TestCase):
+    """web/model.bin + web/autocomplete.ts (run `python export_web.py` first)."""
+
+    @classmethod
+    def setUpClass(cls):
+        import export_web
+        import suggest
+
+        cls.export_web, cls.suggest = export_web, suggest
+        cls.settings, cls.sources = build_weights.load_settings("web"), build_weights.load_sources()
+        cls.policy = export_web.load_policy()
+        vocab, ngram, _ = build_weights.compile_tables(cls.settings, cls.sources)
+        oov = suggest.oov_prefix_rates(cls.settings, cls.sources, vocab)
+        cls.suggester = suggest.Suggester(vocab, ngram, oov)
+
+    def cli(self, *args: str) -> list:
+        out = subprocess.run(["bun", "web/cli.ts", *args], cwd=MODEL_DIR, capture_output=True, text=True, check=True)
+        return json.loads(out.stdout)
+
+    def test_size(self):
+        self.assertLessEqual((MODEL_DIR / "web" / "model.bin").stat().st_size, 50_000_000)
+
+    def test_matches_python_policy(self):
+        """Same suggestion chains (text, probability, support) as suggest.py on random prefixes."""
+        rng = random.Random(1)
+        lines = [ln for ln in (MODEL_DIR / "data" / "prompts.txt").read_text().splitlines() if len(ln) > 2]
+        prefixes = [t[: rng.randint(1, len(t) - 1)] for t in rng.sample(lines, 300)]
+        for p, res in zip(prefixes, self.cli("--json", *prefixes)):
+            with self.subTest(prefix=p):
+                py = self.suggester.steps(p, self.policy.max_words)
+                self.assertEqual([s.text for s in py], [s["text"] for s in res["steps"]])
+                for a, b in zip(py, res["steps"]):
+                    self.assertAlmostEqual(a.prob, b["prob"], places=3)
+                    self.assertAlmostEqual(a.support, b["support"], places=3)
+                if all(abs(s.prob - self.policy.tau) > 1e-3 for s in py):  # skip float32 ties at the threshold
+                    self.assertEqual(self.suggest.Suggester.apply(py, self.policy), res["ghost"])
+
+    def test_runtime_matches_dense_transformer(self):
+        """The sparse runtime, run on the dense model's own tables, reproduces PyTorch."""
+        import tempfile
+
+        s = build_weights.load_settings("hf")
+        vocab, ngram, _ = build_weights.compile_tables(s, self.sources)
+        oov = self.suggest.oov_prefix_rates(s, self.sources, vocab)
+        texts = ["Once upon a", "write a python function to", "zebra", "how do i", "thank you for your"]
+        tok = AutoTokenizer.from_pretrained(MODEL_DIR)
+        model = AutoModelForCausalLM.from_pretrained(MODEL_DIR).eval()
+        with tempfile.TemporaryDirectory() as tmp:
+            path = Path(tmp) / "model.bin"
+            self.export_web.export(vocab, ngram, oov, self.policy, path)
+            results = self.cli("--model", str(path), "--dist", *texts)
+        with torch.no_grad():
+            for t, res in zip(texts, results):
+                with self.subTest(text=t):
+                    lp = model(**tok(t, return_tensors="pt")).logits[0, -1].double().log_softmax(-1)
+                    self.assertEqual([i for i, _ in res["top"]], lp.topk(5).indices.tolist())
+                    for i, v in res["top"]:
+                        self.assertAlmostEqual(v, lp[i].item(), places=3)
+
+
 if __name__ == "__main__":
     unittest.main()

tokenizer.ts

@@ -0,0 +1,74 @@
+// Word-level tokenizer with GPT-2 byte-level pre-tokenization, matching tokenizer.json
+// (NFKC + lowercase, add_prefix_space, GPT-2 split regex). No Node or DOM APIs, so it
+// is shared by model.ts (Node/Bun) and web/autocomplete.ts (browser).
+
+function bytesToUnicode(): Map<number, string> {
+  const bs: number[] = [];
+  for (let b = 33; b <= 126; b++) bs.push(b);
+  for (let b = 161; b <= 172; b++) bs.push(b);
+  for (let b = 174; b <= 255; b++) bs.push(b);
+  const cs = [...bs];
+  let n = 0;
+  for (let b = 0; b < 256; b++) {
+    if (!bs.includes(b)) {
+      bs.push(b);
+      cs.push(256 + n++);
+    }
+  }
+  return new Map(bs.map((b, i) => [b, String.fromCodePoint(cs[i])]));
+}
+
+const BYTE_ENC = bytesToUnicode();
+const BYTE_DEC = new Map([...BYTE_ENC].map(([b, c]) => [c, b]));
+const GPT2_SPLIT = /'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+/gu;
+const encoder = new TextEncoder();
+const decoder = new TextDecoder();
+
+/** Byte-level pre-tokens ("Ġword" style), exactly as the Python tokenizer splits text. */
+export function preTokenize(text: string): string[] {
+  let norm = text.normalize("NFKC").toLowerCase();
+  if (norm && !/^\s/.test(norm)) norm = " " + norm; // ByteLevel add_prefix_space
+  return (norm.match(GPT2_SPLIT) ?? []).map((piece) => [...encoder.encode(piece)].map((b) => BYTE_ENC.get(b)!).join(""));
+}
+
+/** Byte-level piece(s) back to text. */
+export function decodePieces(pieces: string): string {
+  return decoder.decode(new Uint8Array([...pieces].map((c) => BYTE_DEC.get(c) ?? 63)));
+}
+
+export class WordTokenizer {
+  vocab: Map<string, number>;
+  inverse: string[];
+  special: Set<number>;
+  unkId: number;
+  bosId: number;
+
+  constructor(pieces: string[], unkId: number, bosId: number, special: Iterable<number>) {
+    this.unkId = unkId;
+    this.bosId = bosId;
+    this.inverse = pieces;
+    this.vocab = new Map(pieces.map((p, i) => [p, i]));
+    this.special = new Set(special);
+  }
+
+  /** tokenizer.json (WordLevel model, BOS template post-processor) -> tokenizer. */
+  static fromJson(json: any): WordTokenizer {
+    const pieces: string[] = [];
+    for (const [tok, id] of Object.entries<number>(json.model.vocab)) pieces[id] = tok;
+    const bos = json.post_processor?.single?.find((p: any) => p.SpecialToken)?.SpecialToken?.id;
+    const special = (json.added_tokens ?? []).filter((t: any) => t.special).map((t: any) => t.id);
+    return new WordTokenizer(pieces, json.model.vocab[json.model.unk_token], json.model.vocab[bos], special);
+  }
+
+  idOf(piece: string): number {
+    return this.vocab.get(piece) ?? this.unkId;
+  }
+
+  encode(text: string): number[] {
+    return [this.bosId, ...preTokenize(text).map((p) => this.idOf(p))];
+  }
+
+  decode(ids: number[]): string {
+    return decodePieces(ids.filter((id) => !this.special.has(id)).map((id) => this.inverse[id]).join(""));
+  }
+}

web/autocomplete.ts

@@ -0,0 +1,212 @@
+// Browser runtime for the compiled autocomplete model (zero dependencies, no DOM;
+// works on the main thread, in a Web Worker, or in Bun/Node).
+//
+// web/model.bin holds the compiled transformer in sparse form (see export_web.py):
+// the always-on unigram neuron plus, for every MLP neuron, its context tokens and
+// the logit changes it writes. Evaluating only the neurons whose context matches
+// gives the same next-token distribution as the dense model, without the zeros.
+// The confidence policy mirrors suggest.py.
+//
+//   const ac = await CompiledAutocomplete.load("/model.bin");
+//   ac.suggest("write a python fun").ghost   // "ction to"
+
+import { decodePieces, preTokenize } from "../tokenizer.ts";
+
+export type Policy = { tau: number; min_support: number; max_words: number };
+export type Step = { text: string; prob: number; support: number };
+export type Suggestion = { ghost: string; prob: number; steps: Step[] };
+
+type IdArray = Uint16Array | Uint32Array;
+type Order = {
+  k: number;
+  ctx: IdArray;
+  support: Float32Array;
+  offsets: Uint32Array;
+  dids: IdArray;
+  dvals: Float32Array;
+  index: Map<string, number>;
+};
+
+const MAGIC = "TACB";
+
+export class CompiledAutocomplete {
+  readonly vocab: string[];
+  readonly nPrev: number;
+  readonly order: number;
+  readonly eosId: number;
+  readonly unkId: number;
+  readonly policy: Policy;
+  readonly bytes: number;
+  readonly contexts: number;
+  private text: string[];
+  private pieceIndex: Map<string, number>;
+  private base: Float32Array;
+  private orders: Order[];
+  private oov: Record<string, number>;
+  private stopAfter: Set<string>;
+  private candCache = new Map<string, Int32Array>();
+  private distCache = new Map<string, { probs: Float64Array; support: number }>();
+
+  /** Fetch (with optional progress callback) or wrap an already-loaded buffer. */
+  static async load(source: string | ArrayBuffer, onProgress?: (fraction: number) => void): Promise<CompiledAutocomplete> {
+    if (source instanceof ArrayBuffer) return new CompiledAutocomplete(source);
+    const resp = await fetch(source);
+    if (!resp.ok || !resp.body) throw new Error(`failed to fetch ${source}: ${resp.status}`);
+    const total = Number(resp.headers.get("content-length")) || 0;
+    const reader = resp.body.getReader();
+    const chunks: Uint8Array[] = [];
+    let received = 0;
+    for (;;) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      chunks.push(value);
+      received += value.length;
+      if (total) onProgress?.(received / total);
+    }
+    const bytes = new Uint8Array(received);
+    let off = 0;
+    for (const c of chunks) {
+      bytes.set(c, off);
+      off += c.length;
+    }
+    onProgress?.(1);
+    return new CompiledAutocomplete(bytes.buffer);
+  }
+
+  constructor(buf: ArrayBuffer) {
+    const view = new DataView(buf);
+    const magic = new TextDecoder().decode(new Uint8Array(buf, 0, 4));
+    if (magic !== MAGIC) throw new Error("not a tiny-autocomplete model.bin");
+    const headerLen = view.getUint32(8, true);
+    const header = JSON.parse(new TextDecoder().decode(new Uint8Array(buf, 12, headerLen)));
+    const dataStart = 12 + headerLen + ((8 - ((12 + headerLen) % 8)) % 8);
+    const sections = new Map<string, { offset: number; bytes: number }>();
+    for (const s of header.sections) sections.set(s.name, s);
+    const Ids = header.id_bytes === 2 ? Uint16Array : Uint32Array;
+    const f32 = (name: string) => {
+      const s = sections.get(name)!;
+      return new Float32Array(buf, dataStart + s.offset, s.bytes / 4);
+    };
+    const u32 = (name: string) => {
+      const s = sections.get(name)!;
+      return new Uint32Array(buf, dataStart + s.offset, s.bytes / 4);
+    };
+    const ids = (name: string) => {
+      const s = sections.get(name)!;
+      return new Ids(buf, dataStart + s.offset, s.bytes / header.id_bytes);
+    };
+
+    this.vocab = header.vocab;
+    this.text = this.vocab.map((p) => decodePieces(p));
+    this.pieceIndex = new Map(this.vocab.map((p, i) => [p, i]));
+    this.nPrev = header.n_prev;
+    this.order = this.nPrev + 1;
+    this.eosId = header.eos_id;
+    this.unkId = header.unk_id;
+    this.oov = header.oov_prefix_rates;
+    this.policy = header.policy;
+    this.stopAfter = new Set(header.stop_after);
+    this.bytes = buf.byteLength;
+    this.base = f32("base");
+    this.orders = header.orders.map(({ k, contexts }: { k: number; contexts: number }) => {
+      const ctx = ids(`ctx${k}`);
+      const index = new Map<string, number>();
+      for (let i = 0; i < contexts; i++) index.set(ctx.subarray(i * k, (i + 1) * k).join(","), i);
+      return { k, ctx, support: f32(`support${k}`), offsets: u32(`offsets${k}`), dids: ids(`dids${k}`), dvals: f32(`dvals${k}`), index };
+    });
+    this.contexts = this.orders.reduce((n, o) => n + o.index.size, 0);
+  }
+
+  /** Last `order` token ids, oldest first; positions before 0 read BOS like the attention heads. */
+  private window(ids: number[]): number[] {
+    const t = ids.length - 1;
+    const w: number[] = [];
+    for (let a = this.nPrev; a >= 0; a--) w.push(ids[Math.max(t - a, 0)]);
+    return w;
+  }
+
+  /** Next-token distribution and support of the longest matched context. */
+  dist(ids: number[]): { probs: Float64Array; support: number } {
+    const window = this.window(ids);
+    const key = window.join(",");
+    const hit = this.distCache.get(key);
+    if (hit) return hit;
+    const logits = Float64Array.from(this.base);
+    let support = 0;
+    for (const o of this.orders) {
+      const i = o.index.get(window.slice(-o.k).join(","));
+      if (i === undefined) continue;
+      for (let j = o.offsets[i]; j < o.offsets[i + 1]; j++) logits[o.dids[j]] += o.dvals[j];
+      support = o.support[i];
+    }
+    let max = -Infinity;
+    for (const v of logits) if (v > max) max = v;
+    let z = 0;
+    const probs = new Float64Array(logits.length);
+    for (let i = 0; i < logits.length; i++) z += probs[i] = Math.exp(logits[i] - max);
+    for (let i = 0; i < probs.length; i++) probs[i] /= z;
+    if (this.distCache.size > 2000) this.distCache.clear();
+    const out = { probs, support };
+    this.distCache.set(key, out);
+    return out;
+  }
+
+  private candidates(partial: string): Int32Array {
+    let c = this.candCache.get(partial);
+    if (!c) {
+      const out: number[] = [];
+      for (let i = this.unkId + 1; i < this.vocab.length; i++) if (this.vocab[i].startsWith(partial)) out.push(i);
+      c = Int32Array.from(out);
+      if (this.candCache.size > 5000) this.candCache.clear();
+      this.candCache.set(partial, c);
+    }
+    return c;
+  }
+
+  /** Greedy suggestion chain with cumulative probabilities, before thresholding. */
+  steps(text: string, maxWords = this.policy.max_words): Step[] {
+    const pcs = preTokenize(text);
+    if (!pcs.length) return [];
+    const partial = pcs[pcs.length - 1];
+    const cands = this.candidates(partial);
+    if (!cands.length) return [];
+    const ids = [this.eosId, ...pcs.slice(0, -1).map((p) => this.pieceIndex.get(p) ?? this.unkId)];
+    let { probs, support } = this.dist(ids);
+    let best = cands[0];
+    let sum = 0;
+    for (const c of cands) {
+      sum += probs[c];
+      if (probs[c] > probs[best]) best = c;
+    }
+    let prob = (1 - (this.oov[partial] ?? 0)) * (probs[best] / sum);
+    const out: Step[] = [{ text: this.text[best].slice(decodePieces(partial).length), prob, support }];
+    ids.push(best);
+    while (!this.stopAfter.has(this.vocab[ids[ids.length - 1]]) && out.length < maxWords) {
+      ({ probs, support } = this.dist(ids));
+      let nxt = 0;
+      for (let i = 1; i < probs.length; i++) if (probs[i] > probs[nxt]) nxt = i;
+      if (nxt === this.eosId) break;
+      prob *= probs[nxt];
+      out.push({ text: this.text[nxt], prob, support });
+      ids.push(nxt);
+    }
+    return out;
+  }
+
+  static apply(steps: Step[], policy: Policy): { ghost: string; prob: number } {
+    let ghost = "";
+    let prob = 0;
+    for (const st of steps.slice(0, policy.max_words)) {
+      if (st.prob < policy.tau || st.support < policy.min_support) break;
+      ghost += st.text;
+      prob = st.prob;
+    }
+    return { ghost, prob };
+  }
+
+  /** Ghost text to show after `text`, or "" when the model isn't confident. */
+  suggest(text: string, policy: Policy = this.policy): Suggestion {
+    const steps = this.steps(text, policy.max_words);
+    return { ...CompiledAutocomplete.apply(steps, policy), steps };
+  }
+}

web/cli.ts

@@ -0,0 +1,41 @@
+// Run the browser runtime from the command line (Bun or Node >= 23.6):
+//   bun web/cli.ts "write a python fun" "how do i "
+//   bun web/cli.ts --json ...                 full suggestion chains (used by test_model.py)
+//   bun web/cli.ts --dist ...                 top-5 next-token log-probs after each text
+//   bun web/cli.ts --model path/to/model.bin  use another export (default web/model.bin)
+
+import { readFileSync } from "node:fs";
+import { join } from "node:path";
+import { preTokenize } from "../tokenizer.ts";
+import { CompiledAutocomplete } from "./autocomplete.ts";
+
+const dir = import.meta.dirname ?? new URL(".", import.meta.url).pathname;
+let args = process.argv.slice(2);
+let modelPath = join(dir, "model.bin");
+const mi = args.indexOf("--model");
+if (mi >= 0) {
+  modelPath = args[mi + 1];
+  args = [...args.slice(0, mi), ...args.slice(mi + 2)];
+}
+const mode = args[0] === "--json" || args[0] === "--dist" ? args.shift() : "";
+
+const buf = readFileSync(modelPath);
+const ac = new CompiledAutocomplete(buf.buffer.slice(buf.byteOffset, buf.byteOffset + buf.byteLength));
+
+if (mode === "--dist") {
+  const out = args.map((t) => {
+    const ids = [ac.eosId, ...preTokenize(t).map((p) => ac.vocab.indexOf(p)).map((i) => (i < 0 ? ac.unkId : i))];
+    const { probs } = ac.dist(ids);
+    const top = [...probs.keys()].sort((a, b) => probs[b] - probs[a]).slice(0, 5);
+    return { text: t, top: top.map((i) => [i, Math.log(probs[i])]) };
+  });
+  console.log(JSON.stringify(out));
+} else {
+  const results = args.map((t) => {
+    const start = performance.now();
+    const s = ac.suggest(t);
+    return { text: t, ...s, ms: performance.now() - start };
+  });
+  if (mode === "--json") console.log(JSON.stringify(results));
+  else for (const r of results) console.log(`${r.text}\x1b[90m${r.ghost}\x1b[0m   ${r.ghost ? Math.round(r.prob * 100) + "%" : "(no confident suggestion)"}  ${r.ms.toFixed(2)} ms`);
+}

web/demo.js

@@ -0,0 +1 @@
+function T(){let H=[];for(let q=33;q<=126;q++)H.push(q);for(let q=161;q<=172;q++)H.push(q);for(let q=174;q<=255;q++)H.push(q);let J=[...H],D=0;for(let q=0;q<256;q++)if(!H.includes(q))H.push(q),J.push(256+D++);return new Map(H.map((q,K)=>[q,String.fromCodePoint(J[K])]))}var $=T(),w=new Map([...$].map(([H,J])=>[J,H])),C=/'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+/gu,j=new TextEncoder,y=new TextDecoder;function v(H){let J=H.normalize("NFKC").toLowerCase();if(J&&!/^\s/.test(J))J=" "+J;return(J.match(C)??[]).map((D)=>[...j.encode(D)].map((q)=>$.get(q)).join(""))}function W(H){return y.decode(new Uint8Array([...H].map((J)=>w.get(J)??63)))}var g="TACB";class Y{vocab;nPrev;order;eosId;unkId;policy;bytes;contexts;text;pieceIndex;base;orders;oov;stopAfter;candCache=new Map;distCache=new Map;static async load(H,J){if(H instanceof ArrayBuffer)return new Y(H);let D=await fetch(H);if(!D.ok||!D.body)throw Error(`failed to fetch ${H}: ${D.status}`);let q=Number(D.headers.get("content-length"))||0,K=D.body.getReader(),O=[],R=0;for(;;){let{done:X,value:Z}=await K.read();if(X)break;if(O.push(Z),R+=Z.length,q)J?.(R/q)}let U=new Uint8Array(R),F=0;for(let X of O)U.set(X,F),F+=X.length;return J?.(1),new Y(U.buffer)}constructor(H){let J=new DataView(H);if(new TextDecoder().decode(new Uint8Array(H,0,4))!==g)throw Error("not a tiny-autocomplete model.bin");let q=J.getUint32(8,!0),K=JSON.parse(new TextDecoder().decode(new Uint8Array(H,12,q))),O=12+q+(8-(12+q)%8)%8,R=new Map;for(let V of K.sections)R.set(V.name,V);let U=K.id_bytes===2?Uint16Array:Uint32Array,F=(V)=>{let Q=R.get(V);return new Float32Array(H,O+Q.offset,Q.bytes/4)},X=(V)=>{let Q=R.get(V);return new Uint32Array(H,O+Q.offset,Q.bytes/4)},Z=(V)=>{let Q=R.get(V);return new U(H,O+Q.offset,Q.bytes/K.id_bytes)};this.vocab=K.vocab,this.text=this.vocab.map((V)=>W(V)),this.pieceIndex=new Map(this.vocab.map((V,Q)=>[V,Q])),this.nPrev=K.n_prev,this.order=this.nPrev+1,this.eosId=K.eos_id,this.unkId=K.unk_id,this.oov=K.oov_prefix_rates,this.policy=K.policy,this.stopAfter=new Set(K.stop_after),this.bytes=H.byteLength,this.base=F("base"),this.orders=K.orders.map(({k:V,contexts:Q})=>{let M=Z(`ctx${V}`),L=new Map;for(let B=0;B<Q;B++)L.set(M.subarray(B*V,(B+1)*V).join(","),B);return{k:V,ctx:M,support:F(`support${V}`),offsets:X(`offsets${V}`),dids:Z(`dids${V}`),dvals:F(`dvals${V}`),index:L}}),this.contexts=this.orders.reduce((V,Q)=>V+Q.index.size,0)}window(H){let J=H.length-1,D=[];for(let q=this.nPrev;q>=0;q--)D.push(H[Math.max(J-q,0)]);return D}dist(H){let J=this.window(H),D=J.join(","),q=this.distCache.get(D);if(q)return q;let K=Float64Array.from(this.base),O=0;for(let Z of this.orders){let V=Z.index.get(J.slice(-Z.k).join(","));if(V===void 0)continue;for(let Q=Z.offsets[V];Q<Z.offsets[V+1];Q++)K[Z.dids[Q]]+=Z.dvals[Q];O=Z.support[V]}let R=-1/0;for(let Z of K)if(Z>R)R=Z;let U=0,F=new Float64Array(K.length);for(let Z=0;Z<K.length;Z++)U+=F[Z]=Math.exp(K[Z]-R);for(let Z=0;Z<F.length;Z++)F[Z]/=U;if(this.distCache.size>2000)this.distCache.clear();let X={probs:F,support:O};return this.distCache.set(D,X),X}candidates(H){let J=this.candCache.get(H);if(!J){let D=[];for(let q=this.unkId+1;q<this.vocab.length;q++)if(this.vocab[q].startsWith(H))D.push(q);if(J=Int32Array.from(D),this.candCache.size>5000)this.candCache.clear();this.candCache.set(H,J)}return J}steps(H,J=this.policy.max_words){let D=v(H);if(!D.length)return[];let q=D[D.length-1],K=this.candidates(q);if(!K.length)return[];let O=[this.eosId,...D.slice(0,-1).map((Q)=>this.pieceIndex.get(Q)??this.unkId)],{probs:R,support:U}=this.dist(O),F=K[0],X=0;for(let Q of K)if(X+=R[Q],R[Q]>R[F])F=Q;let Z=(1-(this.oov[q]??0))*(R[F]/X),V=[{text:this.text[F].slice(W(q).length),prob:Z,support:U}];O.push(F);while(!this.stopAfter.has(this.vocab[O[O.length-1]])&&V.length<J){({probs:R,support:U}=this.dist(O));let Q=0;for(let M=1;M<R.length;M++)if(R[M]>R[Q])Q=M;if(Q===this.eosId)break;Z*=R[Q],V.push({text:this.text[Q],prob:Z,support:U}),O.push(Q)}return V}static apply(H,J){let D="",q=0;for(let K of H.slice(0,J.max_words)){if(K.prob<J.tau||K.support<J.min_support)break;D+=K.text,q=K.prob}return{ghost:D,prob:q}}suggest(H,J=this.policy){let D=this.steps(H,J.max_words);return{...Y.apply(D,J),steps:D}}}var G=(H)=>document.getElementById(H),N=G("prompt"),z=G("ghost"),f=G("status"),_=G("guess"),S=G("tau"),A=G("tau-label"),I=G("load"),k=["write a python fun","how do i ","Write a function to find the ","how to convert a ","can you explain ","I hope this email ","Once upon a ","what is the difference between "];function P(H){return H.replace(/[&<>]/g,(J)=>({"&":"&amp;","<":"&lt;",">":"&gt;"})[J])}async function E(){let H=performance.now(),J=await Y.load("model.bin",(R)=>I.textContent=`downloading model... ${Math.round(R*100)}%`),D=performance.now()-H,q={...J.policy};S.value=String(q.tau),A.textContent=`${Math.round(q.tau*100)}%`,I.textContent=`${(J.bytes/1e6).toFixed(1)} MB · ${J.contexts.toLocaleString()} compiled contexts · `+`${J.order}-token context · ${J.vocab.length.toLocaleString()} words · loaded in ${Math.round(D)} ms`,N.disabled=!1,N.focus();let K="",O=()=>{let R=N.value,U=performance.now(),F=J.suggest(R,q),X=performance.now()-U;K=F.ghost,z.innerHTML=P(R)+"<span>"+P(K)+"</span>",N.style.height="auto",N.style.height=Math.max(N.scrollHeight,z.scrollHeight)+"px",f.textContent=R?`${X.toFixed(2)} ms`:"";let Z=F.steps.findIndex((Q)=>Q.text),V=Z<0?"":F.steps.slice(0,Z+1).map((Q)=>Q.text).join("");if(!R)_.textContent="";else if(K)_.textContent=`suggesting with ${Math.round(F.prob*100)}% confidence`;else if(V)_.textContent=`staying quiet: best guess "${V}" is only ${Math.round(F.steps[Z].prob*100)}% likely`;else _.textContent="staying quiet: nothing known for this prefix"};N.addEventListener("input",O),N.addEventListener("keydown",(R)=>{if(R.key==="Tab"&&K)R.preventDefault(),N.value+=K,O()}),S.addEventListener("input",()=>{q.tau=Number(S.value),A.textContent=`${Math.round(q.tau*100)}%`,O()});for(let R of k){let U=document.createElement("button");U.className="ex",U.textContent=R.trim(),U.onclick=()=>{N.value=R,O(),N.focus()},G("examples").appendChild(U)}O()}E().catch((H)=>I.textContent=String(H));

web/demo.ts

@@ -0,0 +1,84 @@
+// Static demo for web/autocomplete.ts. Bundle with:
+//   bun build web/demo.ts --outfile web/demo.js --target browser --minify
+
+import { CompiledAutocomplete, type Policy } from "./autocomplete.ts";
+
+const $ = <T extends HTMLElement>(id: string) => document.getElementById(id) as T;
+const promptEl = $<HTMLTextAreaElement>("prompt");
+const ghostEl = $("ghost");
+const statusEl = $("status");
+const guessEl = $("guess");
+const tauEl = $<HTMLInputElement>("tau");
+const tauLabel = $("tau-label");
+const loadEl = $("load");
+
+const EXAMPLES = [
+  "write a python fun", "how do i ", "Write a function to find the ", "how to convert a ",
+  "can you explain ", "I hope this email ", "Once upon a ", "what is the difference between ",
+];
+
+function esc(s: string): string {
+  return s.replace(/[&<>]/g, (c) => ({ "&": "&amp;", "<": "&lt;", ">": "&gt;" })[c]!);
+}
+
+async function main(): Promise<void> {
+  const t0 = performance.now();
+  const ac = await CompiledAutocomplete.load("model.bin", (f) => (loadEl.textContent = `downloading model... ${Math.round(f * 100)}%`));
+  const loadMs = performance.now() - t0;
+  const policy: Policy = { ...ac.policy };
+  tauEl.value = String(policy.tau);
+  tauLabel.textContent = `${Math.round(policy.tau * 100)}%`;
+  loadEl.textContent =
+    `${(ac.bytes / 1e6).toFixed(1)} MB · ${ac.contexts.toLocaleString()} compiled contexts · ` +
+    `${ac.order}-token context · ${ac.vocab.length.toLocaleString()} words · loaded in ${Math.round(loadMs)} ms`;
+  promptEl.disabled = false;
+  promptEl.focus();
+
+  let ghost = "";
+  const render = () => {
+    const text = promptEl.value;
+    const start = performance.now();
+    const s = ac.suggest(text, policy);
+    const ms = performance.now() - start;
+    ghost = s.ghost;
+    ghostEl.innerHTML = esc(text) + "<span>" + esc(ghost) + "</span>";
+    promptEl.style.height = "auto";
+    promptEl.style.height = Math.max(promptEl.scrollHeight, ghostEl.scrollHeight) + "px";
+    statusEl.textContent = text ? `${ms.toFixed(2)} ms` : "";
+    // First step that adds text (the first step is empty when the typed word is already complete).
+    const i = s.steps.findIndex((st) => st.text);
+    const guess = i < 0 ? "" : s.steps.slice(0, i + 1).map((st) => st.text).join("");
+    if (!text) guessEl.textContent = "";
+    else if (ghost) guessEl.textContent = `suggesting with ${Math.round(s.prob * 100)}% confidence`;
+    else if (guess) guessEl.textContent = `staying quiet: best guess "${guess}" is only ${Math.round(s.steps[i].prob * 100)}% likely`;
+    else guessEl.textContent = "staying quiet: nothing known for this prefix";
+  };
+
+  promptEl.addEventListener("input", render);
+  promptEl.addEventListener("keydown", (e) => {
+    if (e.key === "Tab" && ghost) {
+      e.preventDefault();
+      promptEl.value += ghost;
+      render();
+    }
+  });
+  tauEl.addEventListener("input", () => {
+    policy.tau = Number(tauEl.value);
+    tauLabel.textContent = `${Math.round(policy.tau * 100)}%`;
+    render();
+  });
+  for (const ex of EXAMPLES) {
+    const b = document.createElement("button");
+    b.className = "ex";
+    b.textContent = ex.trim();
+    b.onclick = () => {
+      promptEl.value = ex;
+      render();
+      promptEl.focus();
+    };
+    $("examples").appendChild(b);
+  }
+  render();
+}
+
+main().catch((e) => (loadEl.textContent = String(e)));

web/index.html

@@ -0,0 +1,53 @@
+<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1">
+<title>tiny-autocomplete · in-browser</title>
+<style>
+  :root { --bg: #fafaf9; --fg: #1c1917; --muted: #78716c; --line: #e7e5e4; --accent: #2563eb; --ghost: #a8a29e; --card: #fff; }
+  @media (prefers-color-scheme: dark) {
+    :root { --bg: #0c0a09; --fg: #f5f5f4; --muted: #a8a29e; --line: #292524; --accent: #60a5fa; --ghost: #57534e; --card: #1c1917; }
+  }
+  * { box-sizing: border-box; }
+  body { margin: 0; font: 16px/1.5 system-ui, -apple-system, sans-serif; background: var(--bg); color: var(--fg); }
+  main { max-width: 720px; margin: 0 auto; padding: 48px 20px; }
+  h1 { font-size: 24px; margin: 0 0 4px; }
+  .sub { color: var(--muted); margin: 0 0 24px; font-size: 14px; }
+  .editor { position: relative; background: var(--card); border: 1px solid var(--line); border-radius: 12px; }
+  .editor:focus-within { border-color: var(--accent); }
+  .layer { font: 20px/1.5 system-ui, -apple-system, sans-serif; padding: 16px 18px; margin: 0; white-space: pre-wrap; word-wrap: break-word; }
+  #ghost { position: absolute; inset: 0; pointer-events: none; color: transparent; }
+  #ghost span { color: var(--ghost); }
+  #prompt { position: relative; width: 100%; min-height: 64px; border: 0; outline: 0; resize: none; background: transparent; color: var(--fg); overflow: hidden; }
+  .row { display: flex; justify-content: space-between; gap: 16px; color: var(--muted); font-size: 13px; margin: 8px 2px; }
+  #guess { min-height: 20px; }
+  kbd { border: 1px solid var(--line); border-radius: 4px; padding: 0 5px; font-size: 12px; background: var(--card); }
+  label { display: flex; align-items: center; gap: 10px; color: var(--muted); font-size: 14px; margin: 20px 2px 28px; }
+  input[type=range] { flex: 1; accent-color: var(--accent); }
+  h2 { font-size: 13px; text-transform: uppercase; letter-spacing: .06em; color: var(--muted); margin: 0 0 10px; font-weight: 600; }
+  .examples { display: flex; flex-wrap: wrap; gap: 8px; }
+  .ex { border: 1px solid var(--line); background: var(--card); color: inherit; border-radius: 999px; padding: 4px 12px; font: inherit; font-size: 14px; cursor: pointer; }
+  .ex:hover { border-color: var(--accent); }
+</style>
+</head>
+<body>
+<main>
+  <h1>tiny-autocomplete</h1>
+  <p class="sub" id="load">loading...</p>
+
+  <div class="editor">
+    <pre class="layer" id="ghost" aria-hidden="true"></pre>
+    <textarea class="layer" id="prompt" rows="1" placeholder="Type a coding prompt, e.g. write a python fun" spellcheck="false" disabled></textarea>
+  </div>
+  <div class="row"><span><kbd>Tab</kbd> accept suggestion</span><span id="status"></span></div>
+  <div class="row"><span id="guess"></span></div>
+
+  <label>Confidence threshold <input type="range" id="tau" min="0.3" max="0.99" step="0.01"> <span id="tau-label"></span></label>
+
+  <h2>Try</h2>
+  <div class="examples" id="examples"></div>
+</main>
+<script type="module" src="demo.js"></script>
+</body>
+</html>

web/model.bin

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