Confidence-gated ghost text for coding prompts: Kneser-Ney compile, 4.5 MB browser export + zero-dep runtime, held-out precision/keystrokes-saved benchmarks
4113449 verified | """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() | |