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