# /// script
# requires-python = ">=3.11,<3.14"
# dependencies = [
#   "torch>=2.6,<2.8",
#   "datasets>=3.0",
#   "numpy<3",
#   "onnx>=1.17",
#   "onnxruntime>=1.20",
#   "huggingface_hub>=0.30",
# ]
# ///
"""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()