retrain on more mix-up groups: fixes about 40% of wrong words on the new test bench, up from 24%

README.md

@@ -51,23 +51,30 @@ tinyfix reads the words on both sides of a correctly spelled word and picks the
 <table>
   <tr>
     <td align="center" width="33%"><h2>94%</h2>of its fixes are right</td>
-    <td align="center" width="33%"><h2>31%</h2>of real-word typos fixed<br><sub>in Eager mode</sub></td>
+    <td align="center" width="33%"><h2>40%</h2>of wrong words fixed<br><sub>24% before this version</sub></td>
     <td align="center" width="33%"><h2>~3 ms</h2>per check in Chrome</td>
   </tr>
 </table>
 
-Measured on 2,460 real-word typos from GitHub commits and 45,000 correct words of email and documentation that it never trained on.
+Measured by typing test text into the code [offline.page](https://github.com/cfahlgren1/offline-page) runs, on text the model never trained on and that nothing was tuned on. The 40% is the average of five test sets:
 
-| Setting | Real-word typos fixed | Fixes that were right | Correct email words changed | Correct GitHub words changed |
-|---|---|---|---|---|
-| Careful *(default)* | **19%** | **94%** | about 1 in 2,500 | about 1 in 1,700 |
-| Eager *(demo and offline.page)* | **31%** | **94%** | about 1 in 910 | about 1 in 430 |
+| Test set | Wrong words fixed | Fixes that were right |
+|---|---|---|
+| Classic mix-ups put into email, chat and encyclopedia text (their/there, to/too, its/it's…) | **69%** | 99% |
+| Hand-written sentences | **58%** | 100% |
+| Finger slips that make a real word | **29%** | 87% |
+| Real errors from children's writing | **28%** | 91% |
+| Real typos from GitHub commits | **18%** | 96% |
+
+On about 110,000 words of correct email, chat, encyclopedia, news, documentation and fiction it wrongly changed about 1 word in 1,800, the same as the version before it.
 
-It also helps choose the right fix for misspelled words, raising that fix rate from 56% to 61.5% while 96.6% of fixes stay right.
+The first two sets are built from the mix-up pairs the model was trained to watch for, so they flatter it. The last three are the fairer picture of everyday typing.
 
 ### What it misses
 
-- It misses more real-word typos than it catches, including “don” typed for “do”.
+- It still misses more wrong words than it catches, and about four in five real typos from GitHub commits.
+- When both words make sense it stays quiet: “writing this form a plane” is left alone, because “writing this form” is something people say.
+- Two-letter words are hit and miss. It fixes “I am to tired” but not every “to” that should be “too”.
 - It can only judge the 20,000 words it knows, so a rarer word like “ware” typed for “were” is left alone.
 - It needs the words that come after, so these fixes land a couple of words later.
 - English only.
@@ -97,14 +104,16 @@ IDs are line numbers in `wrong-word-vocab.txt`, starting from 0. Words are lower
 |---|---|
 | `wrong-word.onnx` | The quantized model |
 | `wrong-word-vocab.txt` | The 20,000 words it knows, one per line |
-| `wrong-word-confusions.txt` | Mix-ups like “their” and “there” that are far apart on the keyboard |
+| `wrong-word-confusions.txt` | 234 words people mix up, like “their” and “there”, each with the words they are mistaken for |
 | `train_realword.py` | The training script |
 
+The model only scores words. The rules that decide when to make a swap (which words to check, the extra confidence needed at the start of a sentence, and so on) are in offline.page's [`src/tinyfix`](https://github.com/cfahlgren1/offline-page/tree/main/src/tinyfix).
+
 </details>
 
 ## Training
 
-A word-level GRU that reads 16 words before and up to 4 words after, trained for 25 minutes on an A10G. It learned from 118M tokens of [FineWeb-Edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu), [SmolTalk](https://huggingface.co/datasets/HuggingFaceTB/smoltalk) and [OpenAssistant](https://huggingface.co/datasets/OpenAssistant/oasst1), by swapping real words for nearby real words and learning to swap them back.
+A word-level GRU that reads 16 words before and up to 4 words after, trained for 120,000 steps on one GPU. It learned from [FineWeb-Edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu), [SmolTalk](https://huggingface.co/datasets/HuggingFaceTB/smoltalk) and [OpenAssistant](https://huggingface.co/datasets/OpenAssistant/oasst1) text, by swapping real words for nearby real words and learning to swap them back. About one example in four is centered on a known mix-up, half of them used correctly, so it also learns when to leave a word alone.
 
 <br>
 

train_realword.py

@@ -53,6 +53,32 @@ HOMOPHONES = [
 ]
 TYPO_PAIRS_DATASET = "chirunder/github_typo_corrections"
 
+# Sound-alikes and classic mix-ups. Every member can be typed for every other member.
+EXTENDED_GROUPS = HOMOPHONES + [
+    ("to", "too", "two"), ("of", "off"), ("no", "know", "now"), ("by", "buy", "bye"), ("for", "four"), ("one", "won"),
+    ("new", "knew"), ("wait", "weight"), ("way", "weigh"), ("plane", "plain"), ("made", "maid"), ("pair", "pear"),
+    ("tale", "tail"), ("flour", "flower"), ("threw", "through"), ("son", "sun"), ("bought", "brought"), ("wear", "were"),
+    ("bare", "bear"), ("dear", "deer"), ("hour", "our"), ("knight", "night"), ("knot", "not"), ("sea", "see"),
+    ("steal", "steel"), ("stair", "stare"), ("tide", "tied"), ("which", "witch"), ("course", "coarse"), ("sent", "scent", "cent"),
+    ("die", "dye"), ("flew", "flu"), ("great", "grate"), ("heal", "heel"), ("heard", "herd"), ("lessen", "lesson"),
+    ("main", "mane"), ("medal", "meddle"), ("missed", "mist"), ("morning", "mourning"), ("pain", "pane"), ("pause", "paws"),
+    ("peak", "peek"), ("pole", "poll"), ("poor", "pour", "pore"), ("pray", "prey"), ("rain", "reign", "rein"), ("read", "red"),
+    ("real", "reel"), ("road", "rode"), ("root", "route"), ("rose", "rows"), ("scene", "seen"), ("seam", "seem"),
+    ("sole", "soul"), ("some", "sum"), ("stationary", "stationery"), ("throne", "thrown"), ("vain", "vein"), ("waive", "wave"),
+    ("ate", "eight"), ("blew", "blue"), ("cell", "sell"), ("cereal", "serial"), ("council", "counsel"), ("desert", "dessert"),
+    ("dual", "duel"), ("fourth", "forth"), ("guessed", "guest"), ("hall", "haul"), ("higher", "hire"), ("loan", "lone"),
+    ("hours", "ours"), ("presence", "presents"), ("sail", "sale"), ("stake", "steak"), ("breath", "breathe"), ("chose", "choose"),
+    ("clothes", "cloths"), ("dose", "does"), ("later", "latter"), ("personal", "personnel"), ("precede", "proceed"),
+    ("quite", "quit"), ("woman", "women"), ("form", "from"), ("lose", "loss"), ("whit", "with"), 
+    ("grown", "groan"), ("idle", "idol"), ("sweet", "suite"), 
+    ("ensure", "insure"), ("emigrate", "immigrate"),
+    ("we", "wee"), ("so", "sew"),
+]
+# Mix-ups that only make sense one way ("could of" for "could have").
+ONE_WAY = [("of", "have"), ("then", "than"), ("loose", "lose"), ("defiantly", "definitely")]
+GROUP_SETS = {"shipped": HOMOPHONES, "extended": EXTENDED_GROUPS}
+SLIP_WORDS = ["fro", "nit", "grate", "manger", "bye", "whit", "wee", "thee", "lest", "sew"]
+
 WORD_LIST_URL = "https://raw.githubusercontent.com/wolfgarbe/SymSpell/master/SymSpell.FrequencyDictionary/en-80k.txt"
 TOKEN = re.compile(r"[A-Za-z]+(?:'[A-Za-z]+)?|\d+(?:[.,]\d+)*|[^\s\w]")
 SENTENCE_OPENERS = {".", "!", "?", ":", '"'}
@@ -106,14 +132,70 @@ def load_corpus(token_budgets: dict[str, int]) -> tuple[list[np.ndarray], dict[s
     return documents, raw_ids
 
 
-def build_vocabulary(documents: list[np.ndarray], raw_ids: dict[str, int], size: int) -> tuple[list[str], np.ndarray]:
+CACHE_REPO = "cfahlgren1/tinyfix-lab"
+
+
+def load_corpus_cached(token_budgets: dict[str, int], api: HfApi) -> tuple[list[np.ndarray], dict[str, int]]:
+    """Tokenizing 120M tokens takes a while, so the first run saves the result to the private lab repo."""
+    folder = f"data/corpus-{token_budgets['fineweb']}-{token_budgets['smoltalk']}"
+    try:
+        paths = {name: api.hf_hub_download(CACHE_REPO, f"{folder}/{name}") for name in ("tokens.npy", "lengths.npy", "raw_tokens.json")}
+        tokens, lengths = np.load(paths["tokens.npy"]), np.load(paths["lengths.npy"])
+        raw_tokens = json.loads(Path(paths["raw_tokens.json"]).read_text())
+        print(f"loaded cached corpus {folder}: {len(tokens):,} tokens", flush=True)
+        return np.split(tokens, np.cumsum(lengths)[:-1]), {token: index for index, token in enumerate(raw_tokens)}
+    except Exception as error:
+        print(f"no cached corpus ({type(error).__name__}); building it", flush=True)
+    documents, raw_ids = load_corpus(token_budgets)
+    cache = Path("/tmp/corpus-cache")
+    cache.mkdir(exist_ok=True)
+    np.save(cache / "tokens.npy", np.concatenate(documents))
+    np.save(cache / "lengths.npy", np.array([len(document) for document in documents], dtype=np.int64))
+    (cache / "raw_tokens.json").write_text(json.dumps(list(raw_ids)))
+    api.create_repo(CACHE_REPO, repo_type="model", private=True, exist_ok=True)
+    api.upload_folder(folder_path=str(cache), repo_id=CACHE_REPO, path_in_repo=folder, repo_type="model")
+    print(f"cached corpus to {CACHE_REPO}/{folder}", flush=True)
+    return documents, raw_ids
+
+
+def slip_words(ranked: list[str], counts_by_token: dict[str, int], start: int, how_many: int, min_count: int) -> list[str]:
+    """Words just outside the vocabulary that are one edit from a common word, so a slip that lands on them can be judged."""
+    common = ranked[:3000]
+    buckets = defaultdict(list)
+    for word in common:
+        if len(word) < 3:
+            continue
+        for position in range(len(word)):
+            buckets[word[:position] + word[position + 1 :]].append(word)
+        buckets[word].append(word)
+    found = []
+    for word in ranked[start:]:
+        if len(found) >= how_many or counts_by_token[word] < min_count:
+            break
+        if len(word) < 3 or "'" in word:
+            continue
+        keys = {word} | {word[:position] + word[position + 1 :] for position in range(len(word))}
+        if any(within_one_edit(word, other) for key in keys for other in buckets.get(key, [])):
+            found.append(word)
+    return found
+
+
+def build_vocabulary(documents: list[np.ndarray], raw_ids: dict[str, int], size: int, forced: list[str], slips: int) -> tuple[list[str], np.ndarray]:
     with urllib.request.urlopen(WORD_LIST_URL) as response:
         dictionary = {line.split()[0] for line in response.read().decode().splitlines() if line.strip()}
     counts = np.bincount(np.concatenate(documents), minlength=len(raw_ids))
     allowed = [(token, identifier) for token, identifier in raw_ids.items() if (WORD.match(token) and token in dictionary) or (len(token) == 1 and not token.isalnum())]
     allowed = [pair for pair in allowed if pair[0] != "<s>"]
     allowed.sort(key=lambda pair: (-counts[pair[1]], pair[0]))
-    vocabulary = SPECIAL_TOKENS + [token for token, _ in allowed[: size - len(SPECIAL_TOKENS)]]
+    ranked = [token for token, _ in allowed]
+    room = size - len(SPECIAL_TOKENS)
+    extra = [word for word in forced if word in raw_ids and word not in ranked[:room]]
+    if slips:
+        counts_by_token = {token: int(counts[identifier]) for token, identifier in allowed}
+        extra += [word for word in slip_words(ranked, counts_by_token, room, slips, 20) if word not in extra]
+    missing = [word for word in forced if word not in raw_ids]
+    print(f"vocabulary: forcing {len(extra)} extra words, e.g. {extra[:40]}; not in corpus: {missing}", flush=True)
+    vocabulary = SPECIAL_TOKENS + ranked[: room - len(extra)] + extra
 
     token_to_id = {token: index for index, token in enumerate(vocabulary)}
     lookup = np.full(len(raw_ids), UNK, dtype=np.int32)
@@ -169,16 +251,33 @@ def real_mistakes_for_training(vocabulary: set[str]) -> Counter:
     return Counter({pair: count for pair, count in counts.items() if count >= MIN_REAL_MISTAKE_COUNT})
 
 
-def build_confusions(vocabulary: list[str], use_real_mistakes: bool) -> dict[str, list[tuple[str, int]]]:
+PAIRS_FILE: list[tuple[str, str]] = []
+
+
+def load_pairs_file(api: HfApi, repo: str, path: str) -> list[tuple[str, str]]:
+    """Reads a mix-up list ("typed meant1 meant2 ...") from the lab repo, so training and the engine use the same pairs."""
+    text = Path(api.hf_hub_download(repo, path)).read_text()
+    pairs = [(parts[0], meant) for parts in (line.split() for line in text.splitlines()) if parts for meant in parts[1:]]
+    return list(dict.fromkeys(pairs))
+
+
+def group_pairs(groups: str) -> list[tuple[str, str]]:
+    if PAIRS_FILE:
+        return PAIRS_FILE
+    pairs = [(typed, expected) for group in GROUP_SETS[groups] for typed in group for expected in group if typed != expected]
+    if groups != "shipped":
+        pairs += ONE_WAY
+    return list(dict.fromkeys(pairs))
+
+
+def build_confusions(vocabulary: list[str], use_real_mistakes: bool, groups: str, min_typed: int) -> dict[str, list[tuple[str, int]]]:
     """Maps a typed word to words people really meant by it, most likely first."""
     known = set(vocabulary)
     counts = real_mistakes_for_training(known) if use_real_mistakes else Counter()
     print(f"real mistake pairs={len(counts):,}", flush=True)
-    for group in HOMOPHONES:
-        for typed in group:
-            for expected in group:
-                if typed != expected and typed in known and expected in known and len(typed) >= MIN_TYPED_LENGTH:
-                    counts[(typed, expected)] += HOMOPHONE_WEIGHT
+    for typed, expected in group_pairs(groups):
+        if typed in known and expected in known and len(typed) >= min_typed:
+            counts[(typed, expected)] += HOMOPHONE_WEIGHT
     confusions = defaultdict(list)
     for (typed, expected), count in sorted(counts.items(), key=lambda item: (-item[1], item[0])):
         confusions[typed].append((expected, count))
@@ -282,41 +381,98 @@ class WrongWordModel(nn.Module):
         return scores + self.typed_bonus * self.typed_position
 
 
+INFLECTION_ENDINGS = {"", "s", "es", "d", "ed", "ing", "r", "er", "ly"}
+
+
+def is_inflection(first: str, second: str) -> bool:
+    """True for pairs like friendship/friendships and continued/continues that only differ in an ending."""
+    shared = 0
+    while shared < min(len(first), len(second)) and first[shared] == second[shared]:
+        shared += 1
+    return shared >= 3 and first[shared:] in INFLECTION_ENDINGS and second[shared:] in INFLECTION_ENDINGS
+
+
 class Sampler:
     """Draws training cases by swapping real words for nearby real words."""
 
-    def __init__(self, tokens: torch.Tensor, vocabulary: list[str], neighbors: np.ndarray, real_typed: np.ndarray, real_weights: np.ndarray, device: str):
+    def __init__(self, tokens: torch.Tensor, vocabulary: list[str], neighbors: np.ndarray, real_typed: np.ndarray, real_weights: np.ndarray, device: str, options: dict):
+        self.init_args = (vocabulary, neighbors, real_typed, real_weights, options)
         self.tokens = tokens
         self.device = device
         self.vocabulary = vocabulary
-        self.neighbor_table = neighbors
-        self.real_typed_table = real_typed
-        self.real_weights_table = real_weights
+        self.options = options
         self.real_typed = torch.from_numpy(real_typed).to(device).long()
         self.real_weights = torch.from_numpy(real_weights).to(device)
         self.has_real = self.real_weights.sum(1) > 0
         self.neighbors = torch.from_numpy(neighbors).to(device).long()
         self.neighbor_counts = (self.neighbors != PAD).sum(1)
-        self.lengths = torch.tensor([len(token) for token in vocabulary], device=device)
+        neighbor_weights = (neighbors != PAD).astype(np.float32)
+        if options["inflection_weight"] != 1:
+            for row in np.nonzero(neighbors[:, 0] != PAD)[0]:
+                for column in range(MAX_NEIGHBORS):
+                    other = neighbors[row, column]
+                    if other != PAD and is_inflection(vocabulary[row], vocabulary[other]):
+                        neighbor_weights[row, column] = options["inflection_weight"]
+        neighbor_weights[neighbor_weights.sum(1) == 0] = 1
+        self.neighbor_weights = torch.from_numpy(neighbor_weights).to(device)
+        lengths = torch.tensor([len(token) for token in vocabulary], device=device)
+        typed_ok = lengths >= options["min_typed"]
+        if options["short_listed"]:
+            typed_ok |= (lengths >= 2) & self.has_listed(real_typed, len(vocabulary), device)
+        self.typed_ok = typed_ok
         self.is_target = torch.tensor([is_target_word(token) for token in vocabulary], device=device)
 
-    def batch(self, size: int, corrupt_rate: float):
+        self.group_share = options["group_share"]
+        if self.group_share > 0:
+            group_ids = torch.tensor(sorted(options["group_ids"]), device=device)
+            inside = torch.zeros(len(tokens), dtype=torch.bool, device=device)
+            inside[LEFT_WINDOW + 1 : len(tokens) - RIGHT_WINDOW - 1] = True
+            where = torch.nonzero(torch.isin(tokens, group_ids) & inside).squeeze(1)
+            order = torch.argsort(tokens[where], stable=True)
+            self.group_positions = where[order]
+            type_counts = torch.bincount(tokens[self.group_positions], minlength=len(vocabulary))
+            self.group_starts = torch.cumsum(type_counts, 0) - type_counts
+            self.group_counts = type_counts
+            self.group_type_weights = type_counts.float().pow(options["group_power"])
+            print(f"group positions={len(where):,} types={(type_counts > 0).sum().item()}", flush=True)
+
+    @staticmethod
+    def has_listed(real_typed: np.ndarray, size: int, device: str) -> torch.Tensor:
+        listed = torch.zeros(size, dtype=torch.bool, device=device)
+        typed = torch.from_numpy(real_typed[real_typed != PAD]).to(device).long()
+        listed[typed] = True
+        return listed
+
+    def to(self, device: str) -> "Sampler":
+        return Sampler(self.tokens.to(device), *self.init_args[:4], device, self.init_args[4])
+
+    def positions(self, size: int) -> torch.Tensor:
         positions = torch.randint(LEFT_WINDOW + 1, len(self.tokens) - RIGHT_WINDOW - 1, (size,), device=self.device)
+        if self.group_share <= 0:
+            return positions
+        from_groups = torch.rand(size, device=self.device) < self.group_share
+        types = torch.multinomial(self.group_type_weights, size, replacement=True)
+        offsets = (torch.rand(size, device=self.device) * self.group_counts[types]).long()
+        chosen = self.group_positions[(self.group_starts[types] + offsets).clamp(max=len(self.group_positions) - 1)]
+        return torch.where(from_groups, chosen, positions)
+
+    def batch(self, size: int, corrupt_rate: float):
+        positions = self.positions(size)
         target = self.tokens[positions]
         counts = self.neighbor_counts[target]
-        choice = (torch.rand(size, device=self.device) * counts.clamp(min=1)).long()
+        choice = torch.multinomial(self.neighbor_weights[target], 1).squeeze(1)
         corrupted = self.neighbors[target, choice]
         weights = self.real_weights[target]
         weights = weights.where(self.has_real[target].unsqueeze(1), torch.ones_like(weights))
         real_choice = torch.multinomial(weights, 1).squeeze(1)
-        uses_real = self.has_real[target] & (torch.rand(size, device=self.device) < REAL_MISTAKE_SHARE)
+        uses_real = self.has_real[target] & (torch.rand(size, device=self.device) < self.options["real_share"])
         corrupted = torch.where(uses_real, self.real_typed[target, real_choice], corrupted)
         counts = counts.where(~uses_real, torch.ones_like(counts))
         is_corrupted = (torch.rand(size, device=self.device) < corrupt_rate) & (counts > 0)
         typed = torch.where(is_corrupted, corrupted, target)
         candidates = torch.cat([typed.unsqueeze(1), self.neighbors[typed]], dim=1)
         matches = candidates == target.unsqueeze(1)
-        valid = self.is_target[target] & (self.lengths[typed] >= MIN_TYPED_LENGTH) & (self.neighbor_counts[typed] > 0) & matches.any(1)
+        valid = self.is_target[target] & self.typed_ok[typed] & (self.neighbor_counts[typed] > 0) & matches.any(1)
 
         positions, candidates, matches, is_corrupted = positions[valid], candidates[valid], matches[valid], is_corrupted[valid]
         count = len(positions)
@@ -327,7 +483,7 @@ class Sampler:
         visible_right = torch.randint(0, RIGHT_WINDOW + 1, (count,), device=self.device).unsqueeze(1)
         right = right.where(torch.arange(RIGHT_WINDOW, device=self.device).unsqueeze(0) < visible_right, torch.tensor(PAD, device=self.device)).flip(1)
 
-        is_short = torch.rand(count, device=self.device) < 0.2
+        is_short = torch.rand(count, device=self.device) < self.options["short_share"]
         kept = torch.randint(0, LEFT_WINDOW, (count,), device=self.device).unsqueeze(1)
         distance_from_end = (-left_offsets).unsqueeze(0)
         hidden_left = is_short.unsqueeze(1) & (distance_from_end > kept + 1)
@@ -385,7 +541,7 @@ def export_onnx(model: WrongWordModel, output: Path, sampler: Sampler) -> dict:
     for name, operators in quantized.items():
         quantize_dynamic(float_path, output / name, weight_type=QuantType.QUInt8, op_types_to_quantize=operators)
 
-    sampler_cpu = Sampler(sampler.tokens.cpu(), sampler.vocabulary, sampler.neighbor_table, sampler.real_typed_table, sampler.real_weights_table, "cpu")
+    sampler_cpu = sampler.to("cpu")
     left, right, candidates, label, _ = sampler_cpu.batch(4096, 0.5)
     inputs = {"left": left.numpy().astype(np.int64), "right": right.numpy().astype(np.int64), "candidates": candidates.numpy().astype(np.int64)}
     with torch.no_grad():
@@ -412,8 +568,19 @@ def export_onnx(model: WrongWordModel, output: Path, sampler: Sampler) -> dict:
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--repo", default="cfahlgren1/tinyfix-realword")
-    parser.add_argument("--run-name", default="v4")
+    parser.add_argument("--repo", default="cfahlgren1/tinyfix-lab")
+    parser.add_argument("--run-name", required=True)
+    parser.add_argument("--groups", choices=sorted(GROUP_SETS), default="shipped")
+    parser.add_argument("--min-typed", type=int, default=MIN_TYPED_LENGTH)
+    parser.add_argument("--short-listed", action="store_true", help="Also train two-letter typed words that have listed mix-ups.")
+    parser.add_argument("--real-share", type=float, default=REAL_MISTAKE_SHARE)
+    parser.add_argument("--group-share", type=float, default=0.0, help="Share of cases centered on a word from a mix-up group.")
+    parser.add_argument("--group-power", type=float, default=0.5)
+    parser.add_argument("--inflection-weight", type=float, default=1.0)
+    parser.add_argument("--short-share", type=float, default=0.2)
+    parser.add_argument("--corrupt-rate", type=float, default=0.5)
+    parser.add_argument("--force-group-words", action="store_true")
+    parser.add_argument("--slip-words", type=int, default=0)
     parser.add_argument("--real-mistakes", action="store_true", help="Also train on GitHub word swaps. Scored worse than homophones alone in v3.")
     parser.add_argument("--fineweb-tokens", type=int, default=70_000_000)
     parser.add_argument("--smoltalk-tokens", type=int, default=50_000_000)
@@ -425,15 +592,26 @@ def main():
     parser.add_argument("--learning-rate", type=float, default=2e-3)
     parser.add_argument("--output", default="/tmp/tinyfix-realword")
     parser.add_argument("--no-upload", action="store_true")
+    parser.add_argument("--pairs-file", default="", help="Path in --repo of a mix-up list to train on instead of --groups.")
+    parser.add_argument("--seed", type=int, default=0)
     args = parser.parse_args()
+    torch.manual_seed(args.seed)
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
     output = Path(args.output)
     output.mkdir(parents=True, exist_ok=True)
     print(f"device={device} args={vars(args)}", flush=True)
 
-    documents, raw_ids = load_corpus({"fineweb": args.fineweb_tokens, "smoltalk": args.smoltalk_tokens, "oasst": 10**12})
-    vocabulary, lookup = build_vocabulary(documents, raw_ids, args.vocabulary)
+    api = HfApi(token=os.environ["HF_TOKEN"])
+    if args.pairs_file:
+        PAIRS_FILE.extend(load_pairs_file(api, args.repo, args.pairs_file))
+        print(f"pairs from {args.pairs_file}: {len(PAIRS_FILE)}", flush=True)
+    documents, raw_ids = load_corpus_cached({"fineweb": args.fineweb_tokens, "smoltalk": args.smoltalk_tokens, "oasst": 10**12}, api)
+    pairs = group_pairs(args.groups)
+    forced = list(dict.fromkeys(word for pair in pairs for word in pair)) if args.force_group_words else []
+    if args.slip_words:
+        forced += [word for word in SLIP_WORDS if word not in forced]
+    vocabulary, lookup = build_vocabulary(documents, raw_ids, args.vocabulary, forced, args.slip_words)
     print(f"documents={len(documents):,} raw_types={len(raw_ids):,} vocabulary={len(vocabulary):,}", flush=True)
 
     order = np.random.default_rng(7).permutation(len(documents))
@@ -443,15 +621,22 @@ def main():
     print(f"train={len(train_tokens):,} validation={len(validation_tokens):,} unknown={(train_tokens == UNK).float().mean().item():.3f}", flush=True)
     del documents
 
-    confusions = build_confusions(vocabulary, args.real_mistakes)
+    min_typed = 2 if args.short_listed else args.min_typed
+    confusions = build_confusions(vocabulary, args.real_mistakes, args.groups, min_typed)
     (output / "confusions.txt").write_text("".join(f"{typed} {' '.join(expected for expected, _ in options)}\n" for typed, options in confusions.items()))
     neighbors = build_neighbors(vocabulary, confusions)
     real_typed, real_weights = build_real_mistake_tables(vocabulary, confusions)
     print(f"words with neighbors={(neighbors[:, 0] != PAD).sum():,}", flush=True)
     (output / "vocab.txt").write_text("\n".join(vocabulary) + "\n")
 
-    train_sampler = Sampler(train_tokens, vocabulary, neighbors, real_typed, real_weights, device)
-    validation_sampler = Sampler(validation_tokens, vocabulary, neighbors, real_typed, real_weights, device)
+    token_to_id = {token: index for index, token in enumerate(vocabulary)}
+    options = {
+        "min_typed": args.min_typed, "short_listed": args.short_listed, "inflection_weight": args.inflection_weight,
+        "group_share": args.group_share, "group_power": args.group_power, "real_share": args.real_share, "short_share": args.short_share,
+        "group_ids": {token_to_id[word] for pair in pairs for word in pair if word in token_to_id},
+    }
+    train_sampler = Sampler(train_tokens, vocabulary, neighbors, real_typed, real_weights, device, options)
+    validation_sampler = Sampler(validation_tokens, vocabulary, neighbors, real_typed, real_weights, device, {**options, "group_share": 0.0})
 
     model = WrongWordModel(len(vocabulary), args.embedding, args.hidden).to(device)
     parameters = sum(parameter.numel() for parameter in model.parameters())
@@ -463,7 +648,7 @@ def main():
     running_loss = 0.0
     report_every = max(args.steps // 40, 1)
     for step in range(1, args.steps + 1):
-        left, right, candidates, label, _ = train_sampler.batch(args.batch, 0.5)
+        left, right, candidates, label, _ = train_sampler.batch(args.batch, args.corrupt_rate)
         loss = F.cross_entropy(mask_empty(model(left, right, candidates), candidates), label)
         optimizer.zero_grad(set_to_none=True)
         loss.backward()
@@ -479,6 +664,7 @@ def main():
 
     validation = evaluate(model, validation_sampler, 200, args.batch)
     print(f"final validation: {validation}", flush=True)
+    torch.save(model.state_dict(), output / "model.pt")
     export_report = export_onnx(model, output, validation_sampler)
     print(f"export: {export_report}", flush=True)
 
@@ -496,10 +682,9 @@ def main():
     if args.no_upload:
         print(f"Skipping upload; artifacts in {output}", flush=True)
         return
-    api = HfApi(token=os.environ["HF_TOKEN"])
     api.create_repo(args.repo, repo_type="model", private=True, exist_ok=True)
-    api.upload_folder(folder_path=str(output), repo_id=args.repo, path_in_repo=args.run_name, repo_type="model")
-    print(f"Uploaded to {args.repo}/{args.run_name}", flush=True)
+    api.upload_folder(folder_path=str(output), repo_id=args.repo, path_in_repo=f"combine/{args.run_name}", repo_type="model")
+    print(f"Uploaded to {args.repo}/combine/{args.run_name}", flush=True)
 
 
 if __name__ == "__main__":

wrong-word-confusions.txt

@@ -5,13 +5,15 @@ your you're
 you're your
 its it's
 it's its
-too two
-two too
+too two to
+two too to
+to too two
 then than
-than then
+than then that
 were where we're
-where were we're
-we're were where
+where were we're wear
+we're were where wear
+wear were where we're
 whose who's
 who's whose
 hear here
@@ -22,17 +24,18 @@ peace piece
 piece peace
 affect effect
 effect affect
-lose loose
-loose lose
+lose loose loss
+loose lose loss
 accept except
 except accept
 advice advise
 advise advice
 quiet quite
 quite quiet
-through though thorough
-though through thorough
-thorough through though
+through though thorough threw
+though through thorough threw
+thorough through though threw
+threw through though thorough
 passed past
 past passed
 principal principle
@@ -43,13 +46,12 @@ brake break
 break brake
 waist waste
 waste waist
-right write
-write right
+right write rite
+write right rite
+rite right write
 sight site cite
 site sight cite
 cite sight site
-wear where
-where wear
 led lead
 lead led
 meet meat
@@ -62,8 +64,8 @@ hole whole
 whole hole
 male mail
 mail male
-sale sell
-sell sale
+sale sell sail
+sell sale cell
 role roll
 roll role
 aloud allowed
@@ -72,3 +74,162 @@ board bored
 bored board
 fair fare
 fare fair
+of off
+off of
+no know
+know no now
+now no know
+by buy bye
+buy by bye
+bye by buy
+for four
+four for
+one won
+won one
+son sun
+sun son
+our hour are
+hour our are
+are our hour
+way weigh
+weigh way
+wait weight
+weight wait
+bear bare
+bare bear
+plane plain
+plain plane
+made maid
+maid made
+pair pear
+pear pair
+tale tail
+tail tale
+flour flower
+flower flour
+knew new
+new knew
+die dye
+dye die
+bee be
+see sea
+sea see
+heard herd
+herd heard
+red read
+steal steel
+steel steal
+feet feat
+feat feet
+night knight
+knight night
+sum some
+some sum
+do due dew
+due do dew
+dew do due
+form from
+from form
+desert dessert
+dessert desert
+breath breathe
+breathe breath
+diary dairy
+dairy diary
+trial trail
+trail trial
+let's lets
+cell sell
+cent scent sent
+scent cent sent
+council counsel
+counsel council
+dear deer
+deer dear
+hair hare
+hare hair
+heal heel
+heel heal
+idle idol
+idol idle
+knot not
+knows nose
+nose knows
+lessen lesson
+lesson lessen
+loan lone
+lone loan
+main mane
+mane main
+miner minor
+minor miner
+morning mourning
+mourning morning
+naval navel
+navel naval
+pail pale
+pale pail
+pain pane
+pane pain
+peak peek
+peek peak
+pole poll
+poll pole
+pray prey
+prey pray
+rain reign rein
+reign rain rein
+rein rain reign
+real reel
+reel real
+road rode rowed
+rode road rowed
+rowed road rode
+root route
+route root
+sail sale
+scene seen
+seen scene
+seam seem
+seem seam
+sole soul
+soul sole
+stair stare
+stare stair
+stationary stationery
+stationery stationary
+suite sweet
+sweet suite
+tide tied
+tied tide
+vain vein
+vein vain
+warn worn
+worn warn
+yoke yolk
+yolk yoke
+throne thrown
+thrown throne
+shone shown
+shown shone
+farther further
+loss lose loose
+chose choose
+choose chose
+that than
+altar alter
+base bass
+bass base
+berth birth
+birth berth
+creak creek
+creek creak
+flea flee
+flee flea
+mist missed
+nun none
+pause paws
+paws pause
+raise rays
+rays raise
+rows rose

wrong-word-vocab.txt

@@ -19976,25 +19976,25 @@ grating
 icing
 immunology
 inoculation
-instantaneously
-intelligible
-laura
-malignancy
-mocking
-momentarily
-objected
-oft
-parenthood
-proclaiming
-proprietor
-racket
-rang
-reconciled
-reflexive
-restrain
-sanitized
-sanitizing
-sharpening
-shrunk
-stout
-supercomputer
+bye
+mane
+navel
+pail
+rein
+rowed
+stair
+stationery
+yoke
+berth
+creak
+nun
+cloths
+defiantly
+emigrate
+immigrate
+grate
+groan
+meddle
+waive
+wee
+whit

wrong-word.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:43e3660ebc01749470e8c66bf0c94bcb1c838e81b7154d8040be6f65d853407c
-size 2092693
+oid sha256:35d76c0f2cc1fbb5450188a749f4c10b0d569d8806bb939bd66fa926561b47e9
+size 2092692