| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| """Train a small model that spots correctly spelled wrong words, like "don" typed for "do".""" |
|
|
| import argparse |
| import hashlib |
| import json |
| import os |
| import re |
| import time |
| import urllib.request |
| from collections import Counter, defaultdict |
| from pathlib import Path |
|
|
| import numpy as np |
| import onnxruntime as ort |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from datasets import load_dataset |
| from huggingface_hub import HfApi |
| from onnxruntime.quantization import QuantType, quantize_dynamic |
|
|
| LEFT_WINDOW = 16 |
| RIGHT_WINDOW = 4 |
| MAX_NEIGHBORS = 8 |
| CANDIDATES = MAX_NEIGHBORS + 1 |
| MIN_TYPED_LENGTH = 3 |
|
|
| PAD, UNK, NUMBER, START, NAME = 0, 1, 2, 3, 4 |
| SPECIAL_TOKENS = ["<pad>", "<unk>", "<num>", "<s>", "<name>"] |
|
|
| REAL_MISTAKE_SHARE = 0.6 |
| HOMOPHONE_WEIGHT = 20 |
| MIN_REAL_MISTAKE_COUNT = 2 |
| HOMOPHONES = [ |
| ("their", "there", "they're"), ("your", "you're"), ("its", "it's"), ("too", "two"), ("then", "than"), |
| ("were", "where", "we're"), ("whose", "who's"), ("hear", "here"), ("buy", "bye"), ("weather", "whether"), |
| ("peace", "piece"), ("affect", "effect"), ("lose", "loose"), ("accept", "except"), ("advice", "advise"), |
| ("quiet", "quite"), ("through", "though", "thorough"), ("passed", "past"), ("principal", "principle"), |
| ("complement", "compliment"), ("brake", "break"), ("waist", "waste"), ("right", "write"), ("sight", "site", "cite"), |
| ("wear", "where"), ("led", "lead"), ("meet", "meat"), ("week", "weak"), ("wood", "would"), ("hole", "whole"), |
| ("male", "mail"), ("sale", "sell"), ("role", "roll"), ("aloud", "allowed"), ("board", "bored"), ("fair", "fare"), |
| ] |
| TYPO_PAIRS_DATASET = "chirunder/github_typo_corrections" |
|
|
| 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 = {".", "!", "?", ":", '"'} |
| WORD = re.compile(r"^[a-z]+(?:'[a-z]+)?$") |
| DIGITS = re.compile(r"^\d") |
|
|
|
|
| def tokenize(text: str) -> list[str]: |
| """Lowercases words, except capitalized words mid-sentence, which become a single name token.""" |
| tokens = [] |
| for raw in TOKEN.findall(text.replace("’", "'")): |
| at_sentence_start = not tokens or tokens[-1] in SENTENCE_OPENERS |
| is_pronoun = raw == "I" or raw.startswith("I'") |
| tokens.append("<name>" if raw[0].isupper() and not at_sentence_start and not is_pronoun else raw.lower()) |
| return tokens |
|
|
|
|
| def fineweb_documents(): |
| stream = load_dataset("HuggingFaceFW/fineweb-edu", split="train", streaming=True).shuffle(seed=7, buffer_size=20000) |
| for row in stream: |
| yield row["text"] |
|
|
|
|
| def smoltalk_documents(): |
| stream = load_dataset("HuggingFaceTB/smoltalk", "all", split="train", streaming=True).shuffle(seed=7, buffer_size=20000) |
| for row in stream: |
| yield "\n".join(message["content"] for message in row["messages"] if message["role"] != "system") |
|
|
|
|
| def open_assistant_documents(): |
| for row in load_dataset("OpenAssistant/oasst1", split="train"): |
| if row["lang"] == "en": |
| yield row["text"] |
|
|
|
|
| def load_corpus(token_budgets: dict[str, int]) -> tuple[list[np.ndarray], dict[str, int]]: |
| sources = {"fineweb": fineweb_documents, "smoltalk": smoltalk_documents, "oasst": open_assistant_documents} |
| raw_ids = {"<s>": 0} |
| documents = [] |
| for name, budget in token_budgets.items(): |
| started, used = time.time(), 0 |
| for text in sources[name](): |
| tokens = tokenize(text) |
| if len(tokens) < 8: |
| continue |
| documents.append(np.array([0] + [raw_ids.setdefault(token, len(raw_ids)) for token in tokens], dtype=np.int32)) |
| used += len(tokens) |
| if used >= budget: |
| break |
| print(f"{name}: {used:,} tokens in {time.time() - started:.0f}s", flush=True) |
| return documents, raw_ids |
|
|
|
|
| def build_vocabulary(documents: list[np.ndarray], raw_ids: dict[str, int], size: 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)]] |
|
|
| token_to_id = {token: index for index, token in enumerate(vocabulary)} |
| lookup = np.full(len(raw_ids), UNK, dtype=np.int32) |
| for token, identifier in raw_ids.items(): |
| if token == "<s>": |
| lookup[identifier] = START |
| elif token == "<name>": |
| lookup[identifier] = NAME |
| elif DIGITS.match(token): |
| lookup[identifier] = NUMBER |
| elif token in token_to_id: |
| lookup[identifier] = token_to_id[token] |
| return vocabulary, lookup |
|
|
|
|
| def is_target_word(token: str) -> bool: |
| return len(token) >= 2 and token not in SPECIAL_TOKENS and WORD.match(token) is not None |
|
|
|
|
| def within_one_edit(first: str, second: str) -> bool: |
| if first == second or abs(len(first) - len(second)) > 1: |
| return False |
| if len(first) == len(second): |
| differences = [index for index, (a, b) in enumerate(zip(first, second)) if a != b] |
| if len(differences) == 1: |
| return True |
| return len(differences) == 2 and differences[1] == differences[0] + 1 and first[differences[0]] == second[differences[1]] and first[differences[1]] == second[differences[0]] |
| shorter, longer = sorted((first, second), key=len) |
| return any(longer[:index] + longer[index + 1 :] == shorter for index in range(len(longer))) |
|
|
|
|
| def strip_punctuation(token: str) -> str: |
| return token.strip("\"'()[]{}<>.,;:!?*_`#-=/\\") |
|
|
|
|
| def real_mistakes_for_training(vocabulary: set[str]) -> Counter: |
| """Counts real typed-for-expected word swaps, using only the half of the dev split that eval doesn't tune on.""" |
| counts = Counter() |
| for row in load_dataset(TYPO_PAIRS_DATASET, split="train"): |
| typed_tokens, fixed_tokens = (row["text"] or "").split(), (row["correction"] or "").split() |
| if len(typed_tokens) != len(fixed_tokens): |
| continue |
| different = [index for index, (first, second) in enumerate(zip(typed_tokens, fixed_tokens)) if first != second] |
| if len(different) != 1: |
| continue |
| typed, expected = strip_punctuation(typed_tokens[different[0]]), strip_punctuation(fixed_tokens[different[0]]) |
| digest = int(hashlib.sha1(f"{typed}|{expected}".lower().encode()).hexdigest(), 16) |
| if digest % 2 != 0 or (digest >> 1) % 2 != 0: |
| continue |
| typed, expected = typed.lower(), expected.lower() |
| if typed != expected and len(typed) >= MIN_TYPED_LENGTH and typed in vocabulary and expected in vocabulary and is_target_word(typed) and is_target_word(expected): |
| counts[(typed, expected)] += 1 |
| 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]]]: |
| """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 |
| confusions = defaultdict(list) |
| for (typed, expected), count in sorted(counts.items(), key=lambda item: (-item[1], item[0])): |
| confusions[typed].append((expected, count)) |
| return {typed: options[:MAX_NEIGHBORS] for typed, options in confusions.items()} |
|
|
|
|
| def build_real_mistake_tables(vocabulary: list[str], confusions: dict[str, list[tuple[str, int]]]) -> tuple[np.ndarray, np.ndarray]: |
| """For each intended word, lists the words typed by mistake and how often, for sampling.""" |
| token_to_id = {token: index for index, token in enumerate(vocabulary)} |
| mistakes = defaultdict(list) |
| for typed, options in confusions.items(): |
| for expected, count in options: |
| mistakes[token_to_id[expected]].append((token_to_id[typed], count)) |
| typed_table = np.zeros((len(vocabulary), MAX_NEIGHBORS), dtype=np.int32) |
| weights = np.zeros((len(vocabulary), MAX_NEIGHBORS), dtype=np.float32) |
| for expected, options in mistakes.items(): |
| options = sorted(options, key=lambda option: -option[1])[:MAX_NEIGHBORS] |
| typed_table[expected, : len(options)] = [typed for typed, _ in options] |
| weights[expected, : len(options)] = [count**0.5 for _, count in options] |
| return typed_table, weights |
|
|
|
|
| def build_neighbors(vocabulary: list[str], confusions: dict[str, list[tuple[str, int]]]) -> np.ndarray: |
| """Lists known mix-ups first, then the most common words one edit away, the same order the browser uses.""" |
| token_to_id = {token: index for index, token in enumerate(vocabulary)} |
| buckets = defaultdict(set) |
| for index, token in enumerate(vocabulary): |
| if not is_target_word(token): |
| continue |
| buckets[token].add(index) |
| for position in range(len(token)): |
| buckets[token[:position] + token[position + 1 :]].add(index) |
|
|
| table = np.zeros((len(vocabulary), MAX_NEIGHBORS), dtype=np.int32) |
| for index, token in enumerate(vocabulary): |
| if not is_target_word(token): |
| continue |
| nearby = set(buckets[token]) |
| for position in range(len(token)): |
| nearby |= buckets[token[:position] + token[position + 1 :]] |
| one_edit = sorted(other for other in nearby if within_one_edit(token, vocabulary[other])) |
| known_mixups = [token_to_id[expected] for expected, _ in confusions.get(token, [])] |
| found = list(dict.fromkeys(known_mixups + one_edit)) |
| table[index, : min(len(found), MAX_NEIGHBORS)] = found[:MAX_NEIGHBORS] |
| return table |
|
|
|
|
| class MatMulLinear(nn.Module): |
| """Linear layer that exports as MatMul so ONNX Runtime dynamic quantization applies.""" |
|
|
| def __init__(self, inputs: int, outputs: int): |
| super().__init__() |
| self.weight = nn.Parameter(nn.init.xavier_uniform_(torch.empty(inputs, outputs))) |
| self.bias = nn.Parameter(torch.zeros(outputs)) |
|
|
| def forward(self, values): |
| return values.matmul(self.weight) + self.bias |
|
|
|
|
| class Gru(nn.Module): |
| """GRU built from MatMul layers because ONNX Runtime cannot quantize its GRU operator.""" |
|
|
| def __init__(self, inputs: int, hidden: int): |
| super().__init__() |
| self.hidden = hidden |
| self.reset_gate = MatMulLinear(inputs + hidden, hidden) |
| self.update_gate = MatMulLinear(inputs + hidden, hidden) |
| self.proposal_input = MatMulLinear(inputs, hidden) |
| self.proposal_state = MatMulLinear(hidden, hidden) |
|
|
| def forward(self, sequence): |
| state = sequence[:, 0, :].sum(-1, keepdim=True).mul(0).repeat(1, self.hidden) |
| for step in range(sequence.shape[1]): |
| current = sequence[:, step, :] |
| joined = torch.cat([current, state], dim=-1) |
| reset = torch.sigmoid(self.reset_gate(joined)) |
| update = torch.sigmoid(self.update_gate(joined)) |
| proposal = torch.tanh(self.proposal_input(current) + reset * self.proposal_state(state)) |
| state = (1 - update) * proposal + update * state |
| return state |
|
|
|
|
| class WrongWordModel(nn.Module): |
| """Scores the typed word and its neighbors against the words on both sides.""" |
|
|
| def __init__(self, vocabulary_size: int, embedding: int, hidden: int): |
| super().__init__() |
| self.embedding = nn.Embedding(vocabulary_size, embedding, padding_idx=PAD) |
| self.word_bias = nn.Embedding(vocabulary_size, 1, padding_idx=PAD) |
| self.left = Gru(embedding, hidden) |
| self.right = Gru(embedding, hidden) |
| self.mix = MatMulLinear(hidden * 2, hidden) |
| self.project = MatMulLinear(hidden, embedding) |
| self.typed_bonus = nn.Parameter(torch.tensor(1.0)) |
| self.register_buffer("typed_position", F.one_hot(torch.tensor(0), CANDIDATES).float()) |
|
|
| def forward(self, left, right, candidates): |
| context = torch.cat([self.left(self.embedding(left)), self.right(self.embedding(right))], dim=-1) |
| summary = self.project(torch.tanh(self.mix(context))) |
| scores = (self.embedding(candidates) * summary.unsqueeze(1)).sum(-1) + self.word_bias(candidates).squeeze(-1) |
| return scores + self.typed_bonus * self.typed_position |
|
|
|
|
| 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): |
| 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.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) |
| self.is_target = torch.tensor([is_target_word(token) for token in vocabulary], device=device) |
|
|
| def batch(self, size: int, corrupt_rate: float): |
| positions = torch.randint(LEFT_WINDOW + 1, len(self.tokens) - RIGHT_WINDOW - 1, (size,), device=self.device) |
| target = self.tokens[positions] |
| counts = self.neighbor_counts[target] |
| choice = (torch.rand(size, device=self.device) * counts.clamp(min=1)).long() |
| 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) |
| 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) |
|
|
| positions, candidates, matches, is_corrupted = positions[valid], candidates[valid], matches[valid], is_corrupted[valid] |
| count = len(positions) |
| left_offsets = torch.arange(-LEFT_WINDOW, 0, device=self.device) |
| left = self.tokens[positions.unsqueeze(1) + left_offsets] |
| right = self.tokens[positions.unsqueeze(1) + torch.arange(1, RIGHT_WINDOW + 1, device=self.device)] |
|
|
| 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 |
| 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) |
| marks_start = is_short.unsqueeze(1) & (distance_from_end == kept + 1) |
| left = left.where(~hidden_left, torch.tensor(PAD, device=self.device)).where(~marks_start, torch.tensor(START, device=self.device)) |
| return left, right, candidates, matches.long().argmax(1), is_corrupted |
|
|
|
|
| def mask_empty(scores, candidates): |
| return scores.where(candidates != PAD, torch.tensor(-1e4, device=scores.device)) |
|
|
|
|
| @torch.no_grad() |
| def evaluate(model, sampler, batches: int, batch_size: int) -> dict: |
| model.eval() |
| clean_margins, corrupted_margins, corrupted_right = [], [], [] |
| for _ in range(batches): |
| left, right, candidates, label, is_corrupted = sampler.batch(batch_size, 0.5) |
| scores = mask_empty(model(left, right, candidates), candidates) |
| best = scores[:, 1:].max(1) |
| margin = best.values - scores[:, 0] |
| clean_margins.append(margin[~is_corrupted]) |
| corrupted_margins.append(margin[is_corrupted]) |
| corrupted_right.append(best.indices[is_corrupted] + 1 == label[is_corrupted]) |
| model.train() |
| clean = torch.cat(clean_margins) |
| margin = torch.cat(corrupted_margins) |
| right = torch.cat(corrupted_right) |
| results = {"clean_cases": len(clean), "corrupted_cases": len(margin)} |
| for false_rate in (0.005, 0.001, 0.0005): |
| threshold = torch.quantile(clean[:200000].float(), 1 - false_rate).item() |
| results[f"caught_at_false_{false_rate}"] = round(((margin >= threshold) & right).float().mean().item(), 4) |
| results[f"threshold_at_false_{false_rate}"] = round(threshold, 2) |
| return results |
|
|
|
|
| def export_onnx(model: WrongWordModel, output: Path, sampler: Sampler) -> dict: |
| model = model.cpu().eval() |
| float_path = output / "model.onnx" |
| example = (torch.ones(2, LEFT_WINDOW, dtype=torch.long), torch.ones(2, RIGHT_WINDOW, dtype=torch.long), torch.ones(2, CANDIDATES, dtype=torch.long)) |
| torch.onnx.export( |
| model, |
| example, |
| float_path, |
| input_names=["left", "right", "candidates"], |
| output_names=["scores"], |
| dynamic_axes={"left": {0: "batch"}, "right": {0: "batch"}, "candidates": {0: "batch"}, "scores": {0: "batch"}}, |
| opset_version=17, |
| dynamo=False, |
| ) |
| quantized = { |
| "model_quint8.onnx": ["MatMul", "Gather"], |
| "model_matmul_quint8.onnx": ["MatMul"], |
| } |
| 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") |
| 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(): |
| reference = model(left, right, candidates).numpy() |
|
|
| def run(path: Path) -> np.ndarray: |
| return ort.InferenceSession(str(path), providers=["CPUExecutionProvider"]).run(None, inputs)[0] |
|
|
| def margins(scores: np.ndarray) -> np.ndarray: |
| scores = np.where(inputs["candidates"] != PAD, scores, -1e4) |
| return scores[:, 1:].max(1) - scores[:, 0] |
|
|
| report = {"float_vs_torch_max_difference": float(np.abs(run(float_path) - reference).max())} |
| report["float_mb"] = round(float_path.stat().st_size / 1e6, 2) |
| for name in quantized: |
| scores = run(output / name) |
| report[name] = { |
| "mb": round((output / name).stat().st_size / 1e6, 2), |
| "pick_agreement": round(float((scores.argmax(1) == reference.argmax(1)).mean()), 4), |
| "mean_margin_difference": round(float(np.abs(margins(scores) - margins(reference)).mean()), 3), |
| } |
| return report |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--repo", default="cfahlgren1/tinyfix-realword") |
| parser.add_argument("--run-name", default="v4") |
| 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) |
| parser.add_argument("--vocabulary", type=int, default=20_000) |
| parser.add_argument("--embedding", type=int, default=64) |
| parser.add_argument("--hidden", type=int, default=256) |
| parser.add_argument("--steps", type=int, default=60_000) |
| parser.add_argument("--batch", type=int, default=4096) |
| 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") |
| args = parser.parse_args() |
|
|
| 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) |
| print(f"documents={len(documents):,} raw_types={len(raw_ids):,} vocabulary={len(vocabulary):,}", flush=True) |
|
|
| order = np.random.default_rng(7).permutation(len(documents)) |
| held_out = max(1, len(documents) // 20) |
| train_tokens = torch.from_numpy(lookup[np.concatenate([documents[index] for index in order[held_out:]])]).to(device).long() |
| validation_tokens = torch.from_numpy(lookup[np.concatenate([documents[index] for index in order[:held_out]])]).to(device).long() |
| 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) |
| (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) |
|
|
| model = WrongWordModel(len(vocabulary), args.embedding, args.hidden).to(device) |
| parameters = sum(parameter.numel() for parameter in model.parameters()) |
| print(f"parameters={parameters:,}", flush=True) |
| optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate, weight_decay=1e-6) |
| schedule = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=args.learning_rate, total_steps=args.steps, pct_start=0.1) |
|
|
| started = time.time() |
| 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) |
| loss = F.cross_entropy(mask_empty(model(left, right, candidates), candidates), label) |
| optimizer.zero_grad(set_to_none=True) |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| optimizer.step() |
| schedule.step() |
| running_loss += loss.item() |
| if step % report_every == 0: |
| print(f"step {step}/{args.steps} loss={running_loss / report_every:.4f} elapsed={time.time() - started:.0f}s", flush=True) |
| running_loss = 0.0 |
| if step % (args.steps // 5 or 1) == 0: |
| print(f"validation step {step}: {evaluate(model, validation_sampler, 60, args.batch)}", flush=True) |
|
|
| validation = evaluate(model, validation_sampler, 200, args.batch) |
| print(f"final validation: {validation}", flush=True) |
| export_report = export_onnx(model, output, validation_sampler) |
| print(f"export: {export_report}", flush=True) |
|
|
| metrics = { |
| "args": vars(args), |
| "parameters": parameters, |
| "validation": validation, |
| "export": export_report, |
| "train_tokens": len(train_tokens), |
| "train_minutes": round((time.time() - started) / 60, 1), |
| } |
| (output / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n") |
| print(json.dumps(metrics, indent=2), flush=True) |
|
|
| 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) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
