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 | """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() | |