# /// script
# requires-python = ">=3.11,<3.14"
# dependencies = [
#   "torch>=2.6,<2.8",
#   "datasets>=3.0",
#   "numpy<3",
#   "pyarrow>=15",
#   "tokenizers>=0.20",
#   "transformers>=4.46",
#   "onnx>=1.17",
#   "onnxruntime>=1.20",
#   "safetensors>=0.4",
#   "huggingface_hub>=0.30",
#   "nvidia-ml-py>=12",
#   "trackio>=0.2",
# ]
# ///
"""tinynext v2: a small word-piece transformer for in-browser autocomplete that can also fill in the middle.

One script does everything: it trains a tokenizer, builds a mix of general text, conversation and email,
trains the model, scores it with eval.py (the same harness as tinynext v1 and SmolLM2), exports ONNX with a
key/value cache for the browser, and uploads the lot to the model repo under v2/. Training curves, throughput
and GPU use go to a Trackio dashboard.

  hf jobs uv run --namespace huggingface --flavor a100-large --secrets HF_TOKEN train.py
  hf jobs uv run --namespace huggingface --flavor a10g-small --secrets HF_TOKEN train.py --smoke --no-push

The data is prepared on every CPU core before training starts, and the whole token array then sits in GPU
memory, so the GPU never waits on Python between steps.
"""

import argparse
import importlib.util
import json
import math
import os
import random
import shutil
import sys
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

HERE = Path(__file__).parent
SPECIAL = ["<|endoftext|>", "<|pre|>", "<|suf|>", "<|mid|>"]
EOT, PRE, SUF, MID = range(4)

FINEWEB = "hf://datasets/HuggingFaceFW/fineweb-edu/sample/10BT/{:03d}_00000.parquet"
SMOLTALK = "hf://datasets/HuggingFaceTB/smoltalk/data/all/train-{:05d}-of-00009.parquet"
ENRON = "hf://datasets/corbt/enron-emails/data/train-{:05d}-of-00003.parquet"


def cpus() -> int:
    return len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count() or 1


# ---- Shared eval code ------------------------------------------------------------------------------------------


def fetch_shared(repo: str, folder: str, *names: str) -> None:
    """Shared files (eval.py, eval_set.json, compress.py) come from next to this script or from the repo."""
    from huggingface_hub import hf_hub_download

    for name in names:
        if not (HERE / name).exists():
            shutil.copy(hf_hub_download(repo, f"{folder}/{name}"), HERE / name)


def load_compress_module(repo: str, folder: str):
    fetch_shared(repo, folder, "compress.py")
    spec = importlib.util.spec_from_file_location("tinynext_compress", HERE / "compress.py")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def load_eval_module(repo: str, folder: str):
    """Uses eval.py and eval_set.json next to this script, or fetches them from the repo, so every model is scored
    by one implementation on the same text. Registered as a module so worker processes can use it too."""
    fetch_shared(repo, folder, "eval.py", "eval_set.json")
    path = HERE / "eval.py"
    spec = importlib.util.spec_from_file_location("tinynext_eval", path)
    module = importlib.util.module_from_spec(spec)
    sys.modules["tinynext_eval"] = module
    spec.loader.exec_module(module)
    return module


# ---- Data ------------------------------------------------------------------------------------------------------


def to_ascii(text: str) -> str:
    """The browser model only sees keyboard characters, so it trains on the same."""
    import tinynext_eval as ev

    return text.translate(ev.SMART_CHARS).encode("ascii", "ignore").decode().strip()


def smoltalk_text(batch):
    return {"text": [to_ascii("\n\n".join(m["content"] for m in turns if m["role"] in ("user", "assistant"))) for turns in batch["messages"]]}


def enron_text(batch):
    """Training senders only: the eval's held-out senders and enron-autocomplete's validation senders are skipped."""
    import hashlib

    import tinynext_eval as ev

    texts = []
    for sender, body in zip(batch["from"], batch["body"]):
        sender = (sender or "").lower()
        keep = "@enron.com" in sender and not ev.BULK_SENDER.search(sender)
        keep = keep and int(hashlib.sha1(sender.encode()).hexdigest(), 16) % 100 >= 10
        texts.append(to_ascii("\n\n".join(ev.clean_email(body or ""))) if keep else "")
    return {"text": texts}


# Byte-level BPE with an 8K vocabulary averages a little over 4 characters per token on English. Reading 40% more
# text than that covers the variance, and anything short gets repeated (and logged) in flatten().
CHARS_PER_TOKEN = 4.0
TEXT_MARGIN = 1.4


def enough_rows(ds, budget_tokens: int, lengths: np.ndarray, seed: int):
    """A random subset of rows whose text adds up to roughly the token budget, so only that much gets cleaned."""
    order = np.random.default_rng(seed).permutation(len(ds))
    needed = np.searchsorted(np.cumsum(lengths[order]), budget_tokens * CHARS_PER_TOKEN * TEXT_MARGIN) + 1
    return ds.select(np.sort(order[:needed]))


def load_sources(args, budgets: dict) -> dict:
    """Downloads parquet files, keeps only as many documents as each source's budget needs, and cleans them on
    every core."""
    import pyarrow.compute as pc
    from datasets import load_dataset

    procs = cpus()
    fineweb = load_dataset("parquet", data_files=[FINEWEB.format(i) for i in range(args.fineweb_files)], split="train")
    fineweb = fineweb.select_columns(["text"])
    fineweb = enough_rows(fineweb, budgets["fineweb"], pc.utf8_length(fineweb.data.column("text")).to_numpy(), seed=0)
    fineweb = fineweb.map(lambda b: {"text": [to_ascii(t) for t in b["text"]]}, batched=True, num_proc=procs)
    smoltalk = load_dataset("parquet", data_files=[SMOLTALK.format(i) for i in range(args.smoltalk_files)], split="train")
    sample = smoltalk.shuffle(seed=0).select(range(min(2000, len(smoltalk))))
    average = np.mean([sum(len(m["content"]) for m in row["messages"]) for row in sample])
    smoltalk = enough_rows(smoltalk, budgets["smoltalk"], np.full(len(smoltalk), average), seed=1)
    smoltalk = smoltalk.map(smoltalk_text, batched=True, num_proc=procs, remove_columns=smoltalk.column_names)
    enron = load_dataset("parquet", data_files=[ENRON.format(i) for i in range(3)], split="train")
    enron = enron.map(enron_text, batched=True, num_proc=procs, remove_columns=enron.column_names)
    sources = {"fineweb": fineweb, "smoltalk": smoltalk, "enron": enron}
    return {name: ds.filter(lambda b: [len(t) >= 40 for t in b["text"]], batched=True, num_proc=procs) for name, ds in sources.items()}


def train_tokenizer(sources: dict, weights: dict, vocab: int, sample_docs: int):
    from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers

    def sample():
        for name, weight in weights.items():
            ds = sources[name].shuffle(seed=0)
            for row in ds.select(range(min(len(ds), int(sample_docs * weight)))):
                yield row["text"]

    tokenizer = Tokenizer(models.BPE())
    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tokenizer.decoder = decoders.ByteLevel()
    trainer = trainers.BpeTrainer(vocab_size=vocab, special_tokens=SPECIAL, min_frequency=2,
                                  initial_alphabet=pre_tokenizers.ByteLevel.alphabet())
    tokenizer.train_from_iterator(sample(), trainer=trainer)
    return tokenizer


def fim(ids: list[int], rng: random.Random, max_middle: int, max_suffix: int, max_prefix: int) -> list[int]:
    """Suffix first, then prefix, then the middle, so typing at the cursor only appends to the prompt."""
    if len(ids) < 8:
        return ids + [EOT]
    start = rng.randint(1, len(ids) - 2)
    end = min(len(ids), start + rng.randint(1, max_middle))
    prefix, middle, suffix = ids[max(0, start - max_prefix):start], ids[start:end], ids[end:end + max_suffix]
    if rng.random() < 0.25:
        suffix = []  # also learn the plain "nothing after the cursor" case in this format
    return [SUF, *suffix, PRE, *prefix, MID, *middle, EOT]


def make_encoder(tokenizer_json: str, seq: int, fim_rate: float):
    """Tokenizes a batch and turns each document into training pieces: plain chunks, or fill-in-the-middle examples."""

    def encode(batch, indices):
        from tokenizers import Tokenizer

        tokenizer = Tokenizer.from_str(tokenizer_json)
        rng = random.Random(indices[0])
        out = []
        for encoding in tokenizer.encode_batch(batch["text"]):
            ids, pieces = encoding.ids, []
            for at in range(0, len(ids), seq - 8):
                chunk = ids[at: at + seq - 8]
                pieces += fim(chunk, rng, 24, 64, seq - 100) if rng.random() < fim_rate else chunk + [EOT]
            out.append(pieces)
        return {"tokens": out}

    return encode


def flatten(ds, budget: int) -> np.ndarray:
    """All of a source's tokens as one array, cut to its share of the budget (repeated if the source is small)."""
    import pyarrow.compute as pc

    column = ds.data.column("tokens")
    tokens = pc.list_flatten(column).to_numpy().astype(np.uint16)
    if len(tokens) < budget:
        print(f"    only {len(tokens) / 1e6:.0f}M tokens, repeating to reach {budget / 1e6:.0f}M", flush=True)
        tokens = np.tile(tokens, math.ceil(budget / len(tokens)))
    return tokens[:budget]


def build_rows(sources: dict, tokenizer, weights: dict, args) -> tuple[np.ndarray, np.ndarray, dict]:
    encode = make_encoder(tokenizer.to_str(), args.seq, args.fim_rate)
    train_parts, val_parts, counts = [], [], {}
    for name, weight in weights.items():
        started = time.time()
        split = sources[name].train_test_split(test_size=0.002, seed=0)
        encoded = {part: split[part].map(encode, batched=True, batch_size=512, with_indices=True, num_proc=cpus(),
                                         remove_columns=["text"]) for part in ("train", "test")}
        budget = int(args.tokens * weight)
        train_parts.append(flatten(encoded["train"], budget))
        val_parts.append(flatten(encoded["test"], int(args.val_tokens * weight)))
        counts[name] = {"docs": len(split["train"]), "train_tokens": int(len(train_parts[-1]))}
        print(f"  {name}: {len(split['train']):,} docs → {len(train_parts[-1]) / 1e6:.0f}M tokens in {time.time() - started:.0f}s", flush=True)

    def rows(parts):
        tokens = np.concatenate(parts)
        count = len(tokens) // (args.seq + 1)
        return tokens[: count * (args.seq + 1)].reshape(count, args.seq + 1)

    return rows(train_parts), rows(val_parts), counts


# ---- Model -----------------------------------------------------------------------------------------------------


class RMSNorm(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * self.weight


def rotate(x, cos, sin):
    half = x.shape[-1] // 2
    x1, x2 = x[..., :half], x[..., half:]
    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)


class Block(nn.Module):
    def __init__(self, dim: int, heads: int, ffn: int):
        super().__init__()
        self.heads, self.head_dim = heads, dim // heads
        self.norm1, self.norm2 = RMSNorm(dim), RMSNorm(dim)
        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
        self.out = nn.Linear(dim, dim, bias=False)
        self.gate_up = nn.Linear(dim, 2 * ffn, bias=False)
        self.down = nn.Linear(ffn, dim, bias=False)

    def forward(self, x, cos, sin, past_k, past_v, mask):
        batch, length, _ = x.shape
        q, k, v = self.qkv(self.norm1(x)).view(batch, length, 3, self.heads, self.head_dim).unbind(2)
        q, k, v = rotate(q.transpose(1, 2), cos, sin), rotate(k.transpose(1, 2), cos, sin), v.transpose(1, 2)
        if past_k is not None:
            k, v = torch.cat([past_k, k], dim=2), torch.cat([past_v, v], dim=2)
        if mask is None:
            attended = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        else:
            attended = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        x = x + self.out(attended.transpose(1, 2).reshape(batch, length, -1))
        gate, up = self.gate_up(self.norm2(x)).chunk(2, dim=-1)
        return x + self.down(F.silu(gate) * up), k, v


class TinyNext(nn.Module):
    def __init__(self, vocab: int, dim: int, layers: int, heads: int, ffn: int):
        super().__init__()
        self.config = {"vocab": vocab, "dim": dim, "layers": layers, "heads": heads, "ffn": ffn}
        self.embed = nn.Embedding(vocab, dim)
        self.blocks = nn.ModuleList(Block(dim, heads, ffn) for _ in range(layers))
        self.norm = RMSNorm(dim)
        self.head = None  # set only for export: an untied copy of the output weights, so int8 quantization covers it
        self.head_dim = dim // heads
        inverse = 1.0 / (10000 ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim))
        self.register_buffer("inverse_frequency", inverse, persistent=False)
        nn.init.normal_(self.embed.weight, std=0.02)

    def forward(self, ids, past_k=None, past_v=None):
        """ids: [batch, length]. past_k/past_v: [layers, batch, heads, past, head_dim] or None.
        Returns logits for every position and the updated cache."""
        past = 0 if past_k is None else past_k.shape[3]
        length = ids.shape[1]
        positions = torch.arange(length, device=ids.device) + past
        angles = positions[:, None].float() * self.inverse_frequency[None, :]
        cos, sin = angles.cos()[None, None], angles.sin()[None, None]
        mask = None
        if past_k is not None:
            keys = torch.arange(past + length, device=ids.device)
            mask = keys[None, :] <= positions[:, None]
        x = self.embed(ids)
        new_k, new_v = [], []
        for index, block in enumerate(self.blocks):
            x, k, v = block(x, cos, sin, None if past_k is None else past_k[index], None if past_v is None else past_v[index], mask)
            new_k.append(k)
            new_v.append(v)
        logits = self.norm(x) @ (self.embed.weight.T if self.head is None else self.head)
        return logits, torch.stack(new_k), torch.stack(new_v)


class BrowserStep(nn.Module):
    """The ONNX graph: always takes a cache (possibly empty) and returns only the last position's logits."""

    def __init__(self, model: TinyNext):
        super().__init__()
        self.model = model

    def forward(self, ids, past_k, past_v):
        logits, k, v = self.model(ids, past_k, past_v)
        return logits[:, -1], k, v


# ---- Training --------------------------------------------------------------------------------------------------


class GpuMonitor:
    """Reads GPU utilization and memory through NVML, so the logs show whether the GPU is actually busy."""

    def __init__(self):
        try:
            import pynvml

            pynvml.nvmlInit()
            self.nvml, self.handle = pynvml, pynvml.nvmlDeviceGetHandleByIndex(0)
        except Exception as error:  # no NVIDIA GPU (a laptop smoke run): report nothing rather than fail
            print(f"GPU monitoring off: {error}", flush=True)
            self.nvml = None

    def read(self) -> dict:
        if not self.nvml:
            return {}
        rates = self.nvml.nvmlDeviceGetUtilizationRates(self.handle)
        memory = self.nvml.nvmlDeviceGetMemoryInfo(self.handle)
        return {"gpu_util_percent": rates.gpu, "gpu_memory_gb": round(memory.used / 1e9, 1)}


def evaluate_loss(model, rows: torch.Tensor, device, batch: int) -> float:
    model.eval()
    losses = []
    with torch.no_grad(), torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"):
        for at in range(0, len(rows), batch):
            chunk = rows[at: at + batch].long()
            logits, _, _ = model(chunk[:, :-1])
            losses.append(F.cross_entropy(logits.float().flatten(0, 1), chunk[:, 1:].flatten()).item())
    model.train()
    return float(np.mean(losses))


def train(model, rows: torch.Tensor, val_rows: torch.Tensor, args, device, tracker) -> list[dict]:
    steps = int(args.epochs * len(rows) // args.batch)
    decay = [p for p in model.parameters() if p.dim() >= 2]
    no_decay = [p for p in model.parameters() if p.dim() < 2]
    optimizer = torch.optim.AdamW([{"params": decay, "weight_decay": 0.1}, {"params": no_decay, "weight_decay": 0.0}],
                                  lr=args.lr, betas=(0.9, 0.95), fused=device.type == "cuda")
    warmup = max(1, steps // 50)
    schedule = lambda step: min(1, step / warmup) * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * min(1, step / steps))))
    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, schedule)
    compiled = torch.compile(model) if device.type == "cuda" else model
    parameters = sum(p.numel() for p in model.parameters())
    gpu, history = GpuMonitor(), []
    generator = torch.Generator(device=device).manual_seed(0)
    window_started, window_tokens, started = time.time(), 0, time.time()
    print(f"training {steps:,} steps of {args.batch * args.seq:,} tokens", flush=True)
    for step in range(steps):
        index = torch.randint(0, len(rows), (args.batch,), device=device, generator=generator)
        chunk = rows[index].long()
        with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"):
            logits, _, _ = compiled(chunk[:, :-1])
            loss = F.cross_entropy(logits.float().flatten(0, 1), chunk[:, 1:].flatten())
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()
        window_tokens += args.batch * args.seq
        if step % args.log_every == 0 or step == steps - 1:
            loss_value = loss.item()  # syncs the GPU, so throughput below is real
            seconds = time.time() - window_started
            tokens_per_second = window_tokens / max(seconds, 1e-9)
            entry = {
                "step": step, "train_loss": round(loss_value, 4), "lr": round(scheduler.get_last_lr()[0], 6),
                "tokens_per_second": round(tokens_per_second), "minutes": round((time.time() - started) / 60, 2),
                # Training compute is about 6 FLOPs per parameter per token; an A100 does 312 TFLOP/s in bf16.
                "mfu_percent": round(100 * 6 * parameters * tokens_per_second / 312e12, 1), **gpu.read(),
            }
            if step % (args.log_every * 10) == 0 or step == steps - 1:
                entry["val_loss"] = round(evaluate_loss(model, val_rows, device, args.batch), 4)
            history.append(entry)
            tracker.log(entry, step=step)
            print(json.dumps(entry), flush=True)
            window_started, window_tokens = time.time(), 0
    return history


# ---- Eval, export --------------------------------------------------------------------------------------------


def as_token_lm(ev, name: str, model: TinyNext, tokenizer, device, fim: bool = True):
    texts = [tokenizer.decode([i], skip_special_tokens=False) if i >= len(SPECIAL) else "" for i in range(tokenizer.get_vocab_size())]

    @torch.no_grad()
    def next_logits(ids, cache):
        past_k, past_v = cache if cache else (None, None)
        logits, k, v = model(torch.tensor([ids], device=device), past_k, past_v)
        return logits[0, -1].float().cpu().numpy(), (k, v)

    encode = lambda text: tokenizer.encode(text).ids or [EOT]
    # A model trained without fill-in-the-middle never saw the special tokens, so it only gets the text before.
    prompt = (lambda prefix, suffix: prefix if not suffix else [SUF, *suffix, PRE, *prefix, MID]) if fim else None
    return ev.TokenLM(name, next_logits, encode, texts, set(range(len(SPECIAL))), prompt=prompt)


def parity_prompts(ev, tokenizer, count: int = 300) -> list[list[int]]:
    """Real text from the eval set, cut at random points, as token prompts."""
    rng = random.Random(0)
    texts = [p for domain in ev.load_eval_set().values() for doc in domain for p in doc]
    prompts = []
    for text in rng.sample(texts, min(count, len(texts))):
        ids = tokenizer.encode(text[: rng.randint(40, len(text))]).ids[-300:]
        if len(ids) >= 4:
            prompts.append(ids)
    return prompts


def greedy_tokens(step, ids: list[int], empty, steps: int) -> list[int]:
    """Prompt, then `steps` cached greedy steps; returns the top token at each point."""
    logits, k, v = step(np.array([ids]), empty, empty)
    tops = [int(logits.argmax())]
    for _ in range(steps):
        logits, k, v = step(np.array([[tops[-1]]]), k, v)
        tops.append(int(logits.argmax()))
    return tops


def export(model: TinyNext, output: Path, device, prompts: list[list[int]], compressor) -> dict:
    """Float ONNX, and the browser format: the same graph with int8 weights (one scale per 64 values) that the page
    expands back to float32. Both are checked against PyTorch on real prompts over several cached steps.

    onnxruntime-web's int8 kernels were 8-11x slower than its float ones for this model and int8 activations changed
    the top word too often, so only the download is compressed."""
    import onnxruntime as ort

    model = model.float().eval().to("cpu")
    cfg = model.config
    head_dim = cfg["dim"] // cfg["heads"]
    empty = np.zeros((cfg["layers"], 1, cfg["heads"], 0, head_dim), np.float32)
    # Trace with a non-empty cache so no shape gets fixed at zero; the checks below start from an empty one.
    with torch.no_grad():
        _, trace_k, trace_v = model(torch.tensor([[1, 2]]))
    float_path, web = output / "model.onnx", output / "web"
    torch.onnx.export(
        BrowserStep(model), (torch.tensor([[5, 6, 7]]), trace_k, trace_v), float_path, opset_version=17,
        input_names=["input_ids", "past_k", "past_v"], output_names=["logits", "present_k", "present_v"],
        dynamic_axes={"input_ids": {1: "length"}, "past_k": {3: "past"}, "past_v": {3: "past"},
                      "present_k": {3: "total"}, "present_v": {3: "total"}},
        dynamo=False,
    )
    sizes = compressor.compress(float_path, web)
    expanded = compressor.expand(web)
    shutil.copy(output / "tokenizer.json", web / "tokenizer.json")

    @torch.no_grad()
    def torch_step(ids, k, v):
        past = None if k.shape[3] == 0 else (torch.from_numpy(k), torch.from_numpy(v))
        logits, nk, nv = model(torch.from_numpy(ids), *(past or (None, None)))
        return logits[0, -1].numpy(), nk.numpy(), nv.numpy()

    expected = [greedy_tokens(torch_step, ids, empty, 4) for ids in prompts]
    report = {"parity_prompts": len(prompts), "parity_steps_per_prompt": 5, "web_sizes": sizes}
    for name, path in (("float", float_path), ("web_q8", web / "graph.onnx")):
        session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
        onnx_step = lambda ids, k, v: (lambda out: (out[0][0], out[1], out[2]))(
            session.run(None, {"input_ids": ids.astype(np.int64), "past_k": k, "past_v": v}))
        got = [greedy_tokens(onnx_step, ids, empty, 4) for ids in prompts]
        report[name] = {
            "megabytes": round(sizes["weights_q8_mb"] if name == "web_q8" else path.stat().st_size / 1e6, 2),
            "same_first_token": round(float(np.mean([g[0] == e[0] for g, e in zip(got, expected)])), 4),
            "same_five_greedy_tokens": round(float(np.mean([g == e for g, e in zip(got, expected)])), 4),
        }
    expanded.unlink()  # the browser rebuilds this; only weights.q8 is published
    model.to(device)
    return report


def headline(scored: dict) -> dict:
    """The few eval numbers worth a chart: typing saved and 2+ word suggestions at the app's cutoff."""
    table = scored["cutoffs"]["0.5"]
    out = {}
    for key, m in table.items():
        slug = key.replace("/", "_")
        out[f"eval/{slug}/saved_with_tab"] = m["saved_with_tab"]
        out[f"eval/{slug}/right"] = m["right"]
        out[f"eval/{slug}/multiword_shown"] = m["multiword_shown"]
    return out


# ---- Prepared data ---------------------------------------------------------------------------------------------


def data_key(args) -> str:
    """Prepared data depends only on these settings, so a training run can find the matching files."""
    return f"{args.name}-v{args.vocab}-s{args.seq}-t{int(args.tokens / 1e6)}M-fim{args.fim_rate}-{args.mix}".replace(",", "_").replace("=", "")


def prepare(args, weights: dict, output: Path):
    budgets = {name: int((args.tokens + args.val_tokens) * weight) for name, weight in weights.items()}
    started = time.time()
    sources = load_sources(args, budgets)
    print(f"sources ready in {time.time() - started:.0f}s: " + ", ".join(f"{n} {len(d):,} docs" for n, d in sources.items()), flush=True)
    tokenizer = train_tokenizer(sources, weights, args.vocab, args.tokenizer_docs)
    tokenizer.save(str(output / "tokenizer.json"))
    rows, val_rows, counts = build_rows(sources, tokenizer, weights, args)
    return tokenizer, rows, val_rows, counts


def upload_prepared(args, output: Path, rows: np.ndarray, val_rows: np.ndarray, counts: dict, minutes: float):
    from huggingface_hub import HfApi

    folder = output / "data"
    folder.mkdir(exist_ok=True)
    rows.astype(np.uint16).tofile(folder / "train.bin")
    val_rows.astype(np.uint16).tofile(folder / "val.bin")
    shutil.copy(output / "tokenizer.json", folder / "tokenizer.json")
    meta = {"seq": args.seq, "train_rows": len(rows), "val_rows": len(val_rows), "counts": counts, "minutes": round(minutes, 1), "args": vars(args)}
    (folder / "meta.json").write_text(json.dumps(meta, indent=1))
    api = HfApi(token=os.environ.get("HF_TOKEN"))
    api.create_repo(args.data_repo, repo_type="dataset", private=True, exist_ok=True)
    api.upload_folder(folder_path=str(folder), repo_id=args.data_repo, repo_type="dataset", path_in_repo=data_key(args),
                      commit_message=f"prepared data: {data_key(args)}")
    print(f"Prepared data at https://huggingface.co/datasets/{args.data_repo}/tree/main/{data_key(args)}", flush=True)


def download_prepared(args, output: Path):
    from huggingface_hub import snapshot_download
    from tokenizers import Tokenizer

    local = Path(snapshot_download(args.data_repo, repo_type="dataset", allow_patterns=[f"{data_key(args)}/*"], max_workers=16))
    folder = local / data_key(args)
    meta = json.loads((folder / "meta.json").read_text())
    shutil.copy(folder / "tokenizer.json", output / "tokenizer.json")
    rows = np.fromfile(folder / "train.bin", dtype=np.uint16).reshape(meta["train_rows"], meta["seq"] + 1)
    val_rows = np.fromfile(folder / "val.bin", dtype=np.uint16).reshape(meta["val_rows"], meta["seq"] + 1)
    return Tokenizer.from_file(str(output / "tokenizer.json")), rows, val_rows, meta["counts"]


# ---- Main ------------------------------------------------------------------------------------------------------


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--vocab", type=int, default=8192)
    parser.add_argument("--dim", type=int, default=320)
    parser.add_argument("--layers", type=int, default=8)
    parser.add_argument("--heads", type=int, default=5)
    parser.add_argument("--ffn", type=int, default=864)
    parser.add_argument("--seq", type=int, default=512)
    parser.add_argument("--tokens", type=float, default=1.5e9, help="Training tokens")
    parser.add_argument("--val-tokens", type=float, default=4e6)
    parser.add_argument("--epochs", type=float, default=1.0)
    parser.add_argument("--batch", type=int, default=512)
    parser.add_argument("--lr", type=float, default=3e-3)
    parser.add_argument("--fim-rate", type=float, default=0.5)
    parser.add_argument("--mix", default="fineweb=0.6,smoltalk=0.3,enron=0.1")
    parser.add_argument("--fineweb-files", type=int, default=2, help="Each is about 2 GB of text, roughly 0.9B tokens")
    parser.add_argument("--smoltalk-files", type=int, default=9)
    parser.add_argument("--tokenizer-docs", type=int, default=200_000)
    parser.add_argument("--eval-pauses", type=int, default=1000)
    parser.add_argument("--log-every", type=int, default=50)
    parser.add_argument("--repo", default="cfahlgren1/tinynext")
    parser.add_argument("--folder", default="v2", help="Where this run's outputs are uploaded")
    parser.add_argument("--eval-folder", default="v2", help="Where the shared eval.py and eval_set.json live")
    parser.add_argument("--name", default="tinynext-v2")
    parser.add_argument("--trackio-space", default="cfahlgren1/tinynext-trackio")
    parser.add_argument("--stage", choices=("all", "prepare", "train"), default="all",
                        help="prepare: build the tokenizer and token arrays on a CPU job and save them to --data-repo. "
                             "train: download them and train on the GPU. all: both in one job.")
    parser.add_argument("--data-repo", default="cfahlgren1/tinynext-data")
    parser.add_argument("--smoke", action="store_true", help="Tiny settings to check the pipeline end to end")
    parser.add_argument("--no-push", action="store_true")
    args = parser.parse_args()
    if args.smoke:
        args.vocab, args.dim, args.layers, args.heads, args.ffn, args.seq = 2048, 64, 2, 2, 128, 256
        args.tokens, args.val_tokens, args.batch, args.tokenizer_docs = 5e6, 2e5, 64, 5000
        args.fineweb_files, args.smoltalk_files, args.eval_pauses, args.log_every = 1, 1, 20, 10
        args.name = "tinynext-v2-smoke"
    return args


class NoTracker:
    def log(self, *_, **__):
        pass

    def finish(self):
        pass


def start_tracker(args):
    if args.no_push and not os.environ.get("HF_TOKEN"):
        return NoTracker()
    import trackio

    trackio.init(project="tinynext", name=args.name, space_id=args.trackio_space, config=vars(args))
    return trackio


def main():
    args = parse_args()
    device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
    torch.manual_seed(0)
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    output = HERE / "output"
    output.mkdir(exist_ok=True)
    ev = load_eval_module(args.repo, args.eval_folder)
    weights = {k: float(v) for k, v in (part.split("=") for part in args.mix.split(","))}
    print(f"device {device}, {cpus()} CPUs, mix {weights}", flush=True)
    tracker = start_tracker(args)

    started = time.time()
    if args.stage == "train":
        tokenizer, rows, val_rows, counts = download_prepared(args, output)
    else:
        tokenizer, rows, val_rows, counts = prepare(args, weights, output)
    data_minutes = (time.time() - started) / 60
    print(f"data ready in {data_minutes:.1f} min: {rows.size / 1e6:.0f}M train tokens, {val_rows.size / 1e6:.1f}M val", flush=True)
    if args.stage == "prepare":
        upload_prepared(args, output, rows, val_rows, counts, data_minutes)
        return

    # The whole token array lives on the GPU; each step just indexes into it.
    rows_gpu = torch.from_numpy(rows.astype(np.int32)).to(device)
    val_gpu = torch.from_numpy(val_rows.astype(np.int32)).to(device)
    model = TinyNext(tokenizer.get_vocab_size(), args.dim, args.layers, args.heads, args.ffn).to(device)
    parameters = sum(p.numel() for p in model.parameters())
    print(f"{parameters / 1e6:.2f}M parameters", flush=True)
    started = time.time()
    history = train(model, rows_gpu, val_gpu, args, device, tracker)
    train_minutes = (time.time() - started) / 60

    model.eval()
    from safetensors.torch import save_file

    save_file({k: v.contiguous() for k, v in model.state_dict().items()}, output / "model.safetensors")
    scored = ev.score(as_token_lm(ev, args.name, model, tokenizer, device, fim=args.fim_rate > 0), ev.pauses(ev.load_eval_set(), args.eval_pauses))
    (output / "metrics").mkdir(exist_ok=True)
    (output / "metrics" / f"{args.name}.json").write_text(json.dumps(scored, indent=1))
    tracker.log(headline(scored))
    exported = export(model, output, device, parity_prompts(ev, tokenizer), load_compress_module(args.repo, args.eval_folder))

    config = {  # also written next to the browser files, which need it

        **model.config, "parameters": parameters, "seq": args.seq, "special_tokens": dict(zip(SPECIAL, range(len(SPECIAL)))),
        "prompt": "No text after the cursor: the text before it. Otherwise: <|suf|> after <|pre|> before <|mid|>.",
    }
    (output / "config.json").write_text(json.dumps(config, indent=1))
    shutil.copy(output / "config.json", output / "web" / "config.json")
    report = {"args": vars(args), "parameters": parameters, "data": counts, "data_minutes": round(data_minutes, 1),
              "train_minutes": round(train_minutes, 1), "history": history, "exports": exported}
    (output / "report.json").write_text(json.dumps(report, indent=1))
    print(json.dumps(exported, indent=1))
    print(ev.table([output / "metrics" / f"{args.name}.json"]))
    tracker.finish()

    for name in ("train.py", "eval.py", "eval_set.json", "compress.py"):
        shutil.copy(HERE / name, output / name)
    if (HERE / "metrics").exists():
        for path in (HERE / "metrics").glob("*.json"):
            if not (output / "metrics" / path.name).exists():
                shutil.copy(path, output / "metrics" / path.name)
    if args.no_push:
        print(f"Skipping upload; everything is in {output}")
        return
    from huggingface_hub import HfApi

    HfApi(token=os.environ.get("HF_TOKEN")).upload_folder(
        folder_path=str(output), repo_id=args.repo, path_in_repo=args.folder, repo_type="model",
        commit_message=f"{args.name}: {parameters / 1e6:.1f}M params, {train_minutes:.0f} min on GPU",
    )
    print(f"Uploaded to https://huggingface.co/{args.repo}/tree/main/{args.folder}")


if __name__ == "__main__":
    main()