millid.py

# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "torch>=2.4",
#     "numpy",
#     "pandas",
#     "pyarrow",
#     "huggingface_hub>=0.30",
#     "safetensors",
#     "fasttext-wheel",
#     "wandb",
#     "tyro>=0.8",
# ]
# ///
"""MillID: MultiLingual Language IDentifier.

A byte-level, encoder-only transformer that names the language of a passage of text, over
the label set of baobabtech/millid-dataset (1,000+ language-script labels).

One file, three uses:

    # train on Hugging Face Jobs, evaluate against GlotLID, push to the Hub
    hf jobs uv run --flavor h200 --timeout 4h --secrets HF_TOKEN \\
        millid/millid.py train --dataset baobabtech/millid-dataset --push-to baobabtech/millid

    # evaluate a trained model
    python millid.py evaluate --model baobabtech/millid --dataset baobabtech/millid-dataset

    # tune hyperparameters with a Weights & Biases sweep (agents run as Jobs too)
    wandb sweep millid/sweep.yaml          # prints ENTITY/PROJECT/SWEEP_ID
    hf jobs uv run --flavor h200 --secrets HF_TOKEN --secrets WANDB_API_KEY \\
        millid/millid.py sweep-agent --sweep-id ENTITY/PROJECT/SWEEP_ID --count 6

    # use it
    from millid import MillID
    model = MillID.from_pretrained("baobabtech/millid")
    model.predict(["Lave men ou ak savon."], k=3)  # [[("hat_Latn", 0.98), ...]]

One universal encoding
----------------------
Text: Unicode NFKC, whitespace collapsed, then UTF-8 bytes. Every script in every language is the
same 256 symbols, so there is no tokenizer to train, no vocabulary to outgrow and no unknown
token; a script the model has never seen is still bytes it can read.

Labels: ISO 639-3 plus ISO 15924 script (`eng_Latn`, `amh_Ethi`), the dataset's labels.json
fixing label i to labels[i]. A model trained with extra sources keeps those ids and appends its
own labels after them.

Architecture (ARCHITECTURE.md has the full design)
------------
    UTF-8 bytes, each embedded as byte + script + byte-in-character + hashed char 2/3/4-grams
                -> depthwise convolution stem -> strided convolution (2 bytes per token)
                -> N pre-norm blocks: RMSNorm, attention with RoPE, RMSNorm, SwiGLU
                -> masked mean pooling -> language head (label + language + script vectors)

Presets: tiny ~2M and small ~12M for trials, medium ~114M (default), large ~190M.

The model is two parts, and the split is deliberate. `MillIDTextEncoder` turns bytes into one
d_model vector; `LanguageHead` turns a vector into a language. MillID v2 adds images: an image
encoder producing vectors in the same d_model space, fused with the text vector (or used alone
for a scan with no text layer) before the same head. Checkpoint keys are `text.*` and `head.*`
from v1 on, so v1 weights load into v2 unchanged.

Configuration is CleanRL-style: each command is a typed dataclass (`TrainArgs`, ...) whose
fields are the flags and whose field docstrings are the help, parsed by tyro. `python millid.py
train --help` lists every setting; the resolved settings are saved with each run, in
config.json and in W&B.

Tracking: with WANDB_API_KEY set (or --wandb-project), every run logs its hyperparameters,
training loss, learning rate, gradient norm and throughput, the validation curve, and the
final per-set metrics and per-language tables to Weights & Biases.
"""

from __future__ import annotations

import json
import math
import os
import random
import shutil
import sys
import time
import unicodedata
import zlib
from dataclasses import asdict, dataclass, field
from functools import lru_cache
from pathlib import Path
from typing import Literal

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

# --- the universal text encoding ---------------------------------------------------------
#
# A passage becomes four aligned sequences, one entry per UTF-8 byte:
#
#   ids      the byte, shifted by OFFSET; CLS first
#   script   the Unicode script of the character the byte belongs to
#   charpos  1 a one-byte character, 2 the first byte of a longer one, 3 a continuation byte
#   ngrams   hashed character n-grams of 2, 3 and 4 characters ending at the byte's character
#
# Bytes make every language the same 256 symbols. The other three tell the model what bytes
# alone hide: which writing system it is reading, where characters begin, and the character
# n-grams that fastText and GlotLID identify languages by.

PAD, CLS, EOS = 0, 1, 2
OFFSET = 3  # byte b is token b + OFFSET
VOCAB_SIZE = 256 + OFFSET
SCRIPT_BUCKETS = 512
SCRIPT_COMMON = 1  # digits, punctuation, symbols, spaces: shared by every script
NGRAM_ORDERS = (2, 3, 4)


def normalize(text: str) -> str:
    """NFKC folds what PDF text layers and keyboards differ on (ligatures, full-width forms,
    Arabic presentation forms) without changing the language of anything. Runs of whitespace
    become one space."""
    return " ".join(unicodedata.normalize("NFKC", str(text)).split())


@lru_cache(maxsize=1 << 16)
def char_script(ch: str) -> int:
    """The script of a letter or mark, as an id: the first word of its Unicode name (LATIN,
    ETHIOPIC, DEVANAGARI, CJK, HANGUL, ADLAM...), hashed. Python has no Unicode script property;
    the name's first word is the same grouping for letters and marks, and hashing keeps the ids
    stable when a newer Unicode adds scripts."""
    if unicodedata.category(ch)[0] not in "LM":
        return SCRIPT_COMMON
    word = unicodedata.name(ch, "").split(" ", 1)[0]
    if not word:
        return SCRIPT_COMMON
    return 2 + zlib.crc32(word.encode()) % (SCRIPT_BUCKETS - 2)


_SCRIPT_TABLE_SIZE = 0x40000  # planes 0 to 3, every script in use; beyond, char_script per character
_script_table: np.ndarray | None = None


def script_table() -> np.ndarray:
    """char_script for every code point below 0x40000, built once per process (about a second).
    Built before the DataLoader forks, the workers share it."""
    global _script_table
    if _script_table is None:
        uncached = char_script.__wrapped__
        _script_table = np.fromiter((uncached(chr(c)) for c in range(_SCRIPT_TABLE_SIZE)), dtype=np.int64, count=_SCRIPT_TABLE_SIZE)
    return _script_table


# The n-gram hash: a polynomial over lowercased code points, read from the character back,
# then splitmix64, so neighbouring n-grams land in unrelated buckets. Before the passage every
# position holds _START, which no code point equals, so a word at the start of a passage has
# n-grams of its own. numpy's uint64 arithmetic wraps, which is the modulus.
_START = 0x110000
_PRIME = np.uint64(0x100000001B3)
_MIX1, _MIX2 = np.uint64(0xBF58476D1CE4E5B9), np.uint64(0x94D049BB133111EB)
_ORDER_SEED = {n: np.uint64((0x9E3779B97F4A7C15 * n) & 0xFFFFFFFFFFFFFFFF) for n in NGRAM_ORDERS}


def _splitmix(x: np.ndarray) -> np.ndarray:
    x = (x ^ (x >> np.uint64(30))) * _MIX1
    x = (x ^ (x >> np.uint64(27))) * _MIX2
    return x ^ (x >> np.uint64(31))


def _codepoints(text: str) -> np.ndarray:
    # "replace": a lone surrogate (broken PDF text) becomes "?", one code point for one
    return np.frombuffer(text.encode("utf-32-le", "replace"), dtype=np.uint32)


_UTF8_LEAD = np.array([0, 0x00, 0xC0, 0xE0, 0xF0], dtype=np.int64)  # by sequence length


def encode_batch(texts: list[str], max_len: int = 256, ngram_buckets: int = 32768, rng: random.Random | None = None) -> dict[str, torch.Tensor]:
    """Passages as the model's input: the four per-byte sequences, CLS first, padded to the
    longest, and the mask. A passage longer than max_len bytes is cut to a window of whole
    characters: its start at inference, a random place in training.

    The whole batch is one string, and every step one numpy operation over it, because the
    training data workers run this for every batch: per passage it would be dozens of small
    array operations, whose overhead is most of the cost."""
    texts = [normalize(t) for t in texts]
    b, budget = len(texts), max_len - 1
    joined = "".join(texts)
    cp = _codepoints(joined)
    low = joined.lower()  # lower() may lengthen a character; then keep the first of its lowercase
    low = _codepoints(low) if len(low) == len(joined) else _codepoints("".join(ch.lower()[0] for ch in joined))
    lens = np.fromiter(map(len, texts), dtype=np.int64, count=b)
    first = np.cumsum(lens) - lens  # each passage's first character in cp
    after = first + lens
    sizes = (1 + (cp >= 0x80) + (cp >= 0x800) + (cp >= 0x10000)).astype(np.int64)  # UTF-8 bytes per character
    ends = np.concatenate([[0], np.cumsum(sizes)])  # ends[k]: the bytes before character k

    # Each passage's window, characters lo..hi-1 of cp.
    lo = first.copy()
    if rng is not None:
        long = np.flatnonzero(ends[after] - ends[first] > budget)
        if len(long):
            last = np.searchsorted(ends, ends[after[long]] - budget, side="left")  # the first start that fits the rest
            u = np.array([rng.random() for _ in long])
            lo[long] += (u * (last - first[long])).astype(np.int64)  # every start before `last` fills the window
    hi = np.minimum(np.searchsorted(ends, ends[lo] + budget, side="right") - 1, after)
    w = hi - lo
    at = np.repeat(lo - (np.cumsum(w) - w), w) + np.arange(w.sum())  # every window character, in order
    c = cp[at].astype(np.int64)

    inside = c < _SCRIPT_TABLE_SIZE
    script = script_table()[np.where(inside, c, 0)]
    for k in np.flatnonzero(~inside):
        script[k] = char_script(chr(c[k]))

    # n-grams: the characters from this one back, _START before the start of its passage
    into = at - np.repeat(first, w)
    grams = np.empty((len(at), len(NGRAM_ORDERS)), dtype=np.int64)
    h = np.zeros(len(at), dtype=np.uint64)
    for n in range(1, max(NGRAM_ORDERS) + 1):
        h = h * _PRIME + np.where(into >= n - 1, low[np.maximum(at - (n - 1), 0)], _START).astype(np.uint64)
        if n in _ORDER_SEED:
            grams[:, NGRAM_ORDERS.index(n)] = (_splitmix(h ^ _ORDER_SEED[n]) % np.uint64(ngram_buckets - 1)).astype(np.int64) + 1

    # UTF-8, from the code points: byte k of an s-byte character
    s = sizes[at]
    cb, sb = np.repeat(c, s), np.repeat(s, s)
    k = np.arange(len(cb)) - np.repeat(np.cumsum(s) - s, s)
    shifted = cb >> (6 * (sb - 1 - k))
    raw = np.where(k == 0, shifted | _UTF8_LEAD[sb], (shifted & 0x3F) | 0x80)
    charpos = np.where(sb == 1, 1, np.where(k == 0, 2, 3))

    nb = ends[hi] - ends[lo]  # bytes per passage
    row = np.repeat(np.arange(b), nb)
    col = 1 + np.arange(len(cb)) - np.repeat(np.cumsum(nb) - nb, nb)
    width = 1 + int(nb.max(initial=0))
    out = {name: np.zeros((b, width), dtype=np.int64) for name in ("ids", "script", "charpos")}
    out["ngrams"] = np.zeros((b, width, len(NGRAM_ORDERS)), dtype=np.int64)
    out["ids"][:, 0], out["script"][:, 0], out["charpos"][:, 0] = CLS, SCRIPT_COMMON, 1
    out["ids"][row, col] = raw + OFFSET
    out["script"][row, col] = np.repeat(script, s)
    out["charpos"][row, col] = charpos
    out["ngrams"][row, col] = np.repeat(grams, s, axis=0)
    out = {name: torch.from_numpy(v) for name, v in out.items()}
    out["mask"] = out["ids"] != PAD
    return out


@dataclass
class Encoded:
    """One passage as four aligned arrays, one entry per byte, CLS first."""

    ids: np.ndarray  # [m]
    script: np.ndarray  # [m]
    charpos: np.ndarray  # [m]
    ngrams: np.ndarray  # [m, len(NGRAM_ORDERS)]


def encode(text: str, max_len: int = 256, ngram_buckets: int = 32768, rng: random.Random | None = None) -> Encoded:
    """One passage, unpadded: encode_batch of one."""
    x = encode_batch([text], max_len, ngram_buckets, rng)
    m = int(x["mask"][0].sum())
    return Encoded(*(x[name][0, :m].numpy() for name in ("ids", "script", "charpos", "ngrams")))


def collate(items: list[Encoded]) -> dict[str, torch.Tensor]:
    n, g = max(len(e.ids) for e in items), len(NGRAM_ORDERS)
    ids = np.zeros((len(items), n), dtype=np.int64)
    script = np.zeros_like(ids)
    charpos = np.zeros_like(ids)
    ngrams = np.zeros((len(items), n, g), dtype=np.int64)
    for i, e in enumerate(items):
        m = len(e.ids)
        ids[i, :m], script[i, :m], charpos[i, :m], ngrams[i, :m] = e.ids, e.script, e.charpos, e.ngrams
    out = {k: torch.from_numpy(v) for k, v in (("ids", ids), ("script", script), ("charpos", charpos), ("ngrams", ngrams))}
    out["mask"] = out["ids"] != PAD
    return out


# --- the universal label encoding --------------------------------------------------------


@dataclass
class Labels:
    labels: list[str]
    aliases: dict[str, str] = field(default_factory=dict)
    dropped: list[str] = field(default_factory=list)

    def __post_init__(self) -> None:
        self.index = {lab: i for i, lab in enumerate(self.labels)}

    def __len__(self) -> int:
        return len(self.labels)

    def canonical(self, label: str) -> str | None:
        if label in self.dropped:
            return None
        return self.aliases.get(label, label)

    def extend(self, more) -> Labels:
        """The same ids, with new labels appended after them."""
        new = sorted({lab for lab in more if lab not in self.index})
        return Labels(self.labels + new, self.aliases, self.dropped)

    @classmethod
    def load(cls, path: Path) -> Labels:
        d = json.loads(Path(path).read_text(encoding="utf-8"))
        return cls(d["labels"], d.get("aliases", {}), d.get("dropped", []))

    def save(self, path: Path, trained: set[str] | None = None) -> None:
        d = {
            "scheme": "ISO 639-3 language + '_' + ISO 15924 script, e.g. eng_Latn, amh_Ethi",
            "text_encoding": "Unicode NFKC, whitespace collapsed, then UTF-8 bytes",
            "labels": self.labels,
            "aliases": self.aliases,
            "dropped": self.dropped,
        }
        if trained is not None:
            d["trained"] = sorted(trained)
        Path(path).write_text(json.dumps(d, indent=1, ensure_ascii=False), encoding="utf-8")


# --- the model ---------------------------------------------------------------------------


@dataclass
class MillIDConfig:
    num_labels: int = 0
    max_len: int = 256
    d_model: int = 768
    n_layers: int = 12
    n_heads: int = 12
    dropout: float = 0.1
    conv_kernel: int = 5
    downsample: int = 2
    ngram_buckets: int = 32768
    script_buckets: int = SCRIPT_BUCKETS
    rope_theta: float = 10_000.0
    vocab_size: int = VOCAB_SIZE
    modalities: tuple[str, ...] = ("text",)
    model_type: str = "millid-byte-encoder"

    @property
    def d_ff(self) -> int:
        """SwiGLU width: 8/3 d, rounded to 64, the same parameters as a 4d GELU layer."""
        return max(64, round(8 * self.d_model / 3 / 64) * 64)


PRESETS = {
    "tiny": dict(d_model=192, n_layers=3, n_heads=3, downsample=1, ngram_buckets=4096),
    "small": dict(d_model=384, n_layers=6, n_heads=6, downsample=2, ngram_buckets=16384),
    "medium": dict(d_model=768, n_layers=12, n_heads=12, downsample=2, ngram_buckets=32768),
    "large": dict(d_model=1024, n_layers=12, n_heads=16, downsample=2, ngram_buckets=32768),
}


def rope_tables(n: int, head_dim: int, theta: float) -> tuple[torch.Tensor, torch.Tensor]:
    inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
    angles = torch.outer(torch.arange(n).float(), inv)
    return angles.cos(), angles.sin()


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """Rotate each pair of features by its position's angle (x: batch, heads, length, head_dim)."""
    c = cos[: x.size(2)].to(x.dtype)[None, None]
    s = sin[: x.size(2)].to(x.dtype)[None, None]
    x1, x2 = x[..., 0::2], x[..., 1::2]
    return torch.stack((x1 * c - x2 * s, x1 * s + x2 * c), dim=-1).flatten(-2)


class Attention(nn.Module):
    def __init__(self, cfg: MillIDConfig):
        super().__init__()
        self.h = cfg.n_heads
        self.hd = cfg.d_model // cfg.n_heads
        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)
        self.out = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
        self.p = cfg.dropout

    def forward(self, x, cos, sin, key_mask):
        b, n, _ = x.shape
        q, k, v = self.qkv(x).view(b, n, 3, self.h, self.hd).permute(2, 0, 3, 1, 4)
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        y = F.scaled_dot_product_attention(
            q, k, v, attn_mask=key_mask[:, None, None, :], dropout_p=self.p if self.training else 0.0
        )
        return self.out(y.transpose(1, 2).reshape(b, n, -1))


class SwiGLU(nn.Module):
    def __init__(self, cfg: MillIDConfig):
        super().__init__()
        self.w1 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
        self.w3 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
        self.w2 = nn.Linear(cfg.d_ff, cfg.d_model, bias=False)
        self.drop = nn.Dropout(cfg.dropout)

    def forward(self, x):
        return self.w2(self.drop(F.silu(self.w1(x)) * self.w3(x)))


class Block(nn.Module):
    """Pre-norm: RMSNorm, attention with RoPE, residual; RMSNorm, SwiGLU, residual."""

    def __init__(self, cfg: MillIDConfig):
        super().__init__()
        self.norm1 = nn.RMSNorm(cfg.d_model)
        self.attn = Attention(cfg)
        self.norm2 = nn.RMSNorm(cfg.d_model)
        self.ffn = SwiGLU(cfg)
        self.drop = nn.Dropout(cfg.dropout)

    def forward(self, x, cos, sin, key_mask):
        x = x + self.drop(self.attn(self.norm1(x), cos, sin, key_mask))
        return x + self.drop(self.ffn(self.norm2(x)))


class MillIDTextEncoder(nn.Module):
    """A passage's four byte sequences to one d_model vector: the representation a future image
    branch is fused with."""

    def __init__(self, cfg: MillIDConfig):
        super().__init__()
        d = cfg.d_model
        self.byte_emb = nn.Embedding(cfg.vocab_size, d, padding_idx=PAD)
        self.script_emb = nn.Embedding(cfg.script_buckets, d, padding_idx=PAD)
        self.charpos_emb = nn.Embedding(4, d, padding_idx=PAD)
        self.ngram_emb = nn.Embedding(cfg.ngram_buckets, d, padding_idx=PAD)
        # Depthwise over each feature, then mixed: n-grams of conv_kernel bytes, per position.
        self.stem = nn.Sequential(
            nn.Conv1d(d, d, cfg.conv_kernel, padding=cfg.conv_kernel // 2, groups=d),
            nn.GELU(),
            nn.Conv1d(d, d, 1),
        )
        # Then every `downsample` bytes become one token: the transformer reads half as many.
        self.r = cfg.downsample
        self.down = nn.Conv1d(d, d, self.r, stride=self.r) if self.r > 1 else None
        self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layers))
        self.norm = nn.RMSNorm(d)
        cos, sin = rope_tables(cfg.max_len // self.r + 2, d // cfg.n_heads, cfg.rope_theta)
        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)

    def tokens(self, batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
        """Every token's vector after the last layer, and which tokens are real."""
        mask = batch["mask"]
        x = (
            self.byte_emb(batch["ids"])
            + self.script_emb(batch["script"])
            + self.charpos_emb(batch["charpos"])
            + self.ngram_emb(batch["ngrams"]).sum(2)
        )
        m = mask.unsqueeze(-1).to(x.dtype)
        x = (x + self.stem((x * m).transpose(1, 2)).transpose(1, 2)) * m
        if self.down is not None:
            pad = (-x.size(1)) % self.r
            if pad:
                x, mask = F.pad(x, (0, 0, 0, pad)), F.pad(mask, (0, pad))
            x = self.down(x.transpose(1, 2)).transpose(1, 2)
            mask = mask.view(mask.size(0), -1, self.r).any(-1)
        for block in self.blocks:
            x = block(x, self.cos, self.sin, mask)
        return self.norm(x), mask

    def forward(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
        x, mask = self.tokens(batch)
        m = mask.unsqueeze(-1).to(x.dtype)
        return (x * m).sum(1) / m.sum(1).clamp(min=1.0)


class LanguageHead(nn.Module):
    """A vector to a score per language, shared by every modality.

    Each language's vector is built from its parts: its own, its language's (shared by
    `srp_Cyrl` and `srp_Latn`) and its script's (shared by every Ethiopic-script language). A
    language with little data borrows from what it shares, and `language_vectors()` are the
    learned language embeddings themselves.
    """

    def __init__(self, cfg: MillIDConfig, labels: list[str]):
        super().__init__()
        codes = sorted({lab.split("_", 1)[0] for lab in labels})
        scripts = sorted({lab.split("_", 1)[1] if "_" in lab else "" for lab in labels})
        code_id = {c: i for i, c in enumerate(codes)}
        script_id = {sc: i for i, sc in enumerate(scripts)}
        self.register_buffer("code_of", torch.tensor([code_id[lab.split("_", 1)[0]] for lab in labels]))
        self.register_buffer(
            "script_of", torch.tensor([script_id[lab.split("_", 1)[1] if "_" in lab else ""] for lab in labels])
        )
        d = cfg.d_model
        self.label_emb = nn.Embedding(len(labels), d)
        self.code_emb = nn.Embedding(len(codes), d)
        self.script_emb = nn.Embedding(len(scripts), d)
        self.bias = nn.Parameter(torch.zeros(len(labels)))
        self.drop = nn.Dropout(cfg.dropout)

    def language_vectors(self) -> torch.Tensor:
        return self.label_emb.weight + self.code_emb(self.code_of) + self.script_emb(self.script_of)

    def forward(self, z: torch.Tensor) -> torch.Tensor:
        return F.linear(self.drop(z), self.language_vectors(), self.bias)


class MillIDEncoder(nn.Module):
    def __init__(self, cfg: MillIDConfig, labels: list[str]):
        super().__init__()
        self.cfg = cfg
        self.text = MillIDTextEncoder(cfg)
        self.head = LanguageHead(cfg, labels)
        self.apply(self._init)
        # Each residual branch's last projection starts small, so a deep stack starts close to
        # the identity (GPT-2's scaling).
        for block in self.text.blocks:
            for w in (block.attn.out.weight, block.ffn.w2.weight):
                nn.init.normal_(w, std=0.02 / math.sqrt(2 * cfg.n_layers))

    @staticmethod
    def _init(m: nn.Module) -> None:
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, std=0.02)
            if m.padding_idx is not None:
                with torch.no_grad():
                    m.weight[m.padding_idx].zero_()

    def embed(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
        """The passage as one vector, before the head: for similarity, clustering, fusion."""
        return self.text(batch)

    def forward(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
        return self.head(self.text(batch))


class MillID:
    """A trained model, ready to answer: ranked (label, probability) per text."""

    def __init__(self, model: MillIDEncoder, labels: Labels, trained: set[str] | None = None, device: str | None = None):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.model = model.to(self.device).eval()
        self.labels = labels
        # Logits of labels the model never saw a passage of are not answers.
        self.trained = trained or set(labels.labels)
        keep = torch.tensor([lab in self.trained for lab in labels.labels], device=self.device)
        self.bias = torch.where(keep, 0.0, float("-inf"))

    @classmethod
    def from_pretrained(cls, path_or_repo: str, device: str | None = None) -> MillID:
        from safetensors.torch import load_file

        path = Path(path_or_repo)
        if not path.exists():
            from huggingface_hub import snapshot_download

            path = Path(snapshot_download(path_or_repo, allow_patterns=["config.json", "labels.json", "model.safetensors"]))
        cfg = MillIDConfig(**json.loads((path / "config.json").read_text())["model"])
        meta = json.loads((path / "labels.json").read_text(encoding="utf-8"))
        labels = Labels(meta["labels"], meta.get("aliases", {}), meta.get("dropped", []))
        model = MillIDEncoder(cfg, labels.labels)
        model.load_state_dict(load_file(str(path / "model.safetensors")))
        return cls(model, labels, set(meta.get("trained", meta["labels"])), device)

    def batch(self, texts: list[str]) -> dict[str, torch.Tensor]:
        cfg = self.model.cfg
        b = encode_batch(texts, cfg.max_len, cfg.ngram_buckets)
        return {k: v.to(self.device) for k, v in b.items()}

    @torch.no_grad()
    def predict(self, texts: list[str], k: int = 5, batch_size: int = 512) -> list[list[tuple[str, float]]]:
        out = []
        for i in range(0, len(texts), batch_size):
            with torch.autocast(self.device, dtype=torch.bfloat16, enabled=self.device == "cuda"):
                logits = self.model(self.batch(texts[i : i + batch_size])).float()
            probs = F.softmax(logits + self.bias, dim=-1)
            p, idx = probs.topk(k, dim=-1)
            for pr, ix in zip(p.tolist(), idx.tolist(), strict=True):
                out.append([(self.labels.labels[j], v) for j, v in zip(ix, pr, strict=True)])
        return out

    @torch.no_grad()
    def embed(self, texts: list[str]) -> torch.Tensor:
        """Passage vectors, B x d_model."""
        return self.model.embed(self.batch(texts)).float().cpu()


# --- metrics (the same definitions as scripts/metrics.py, tested there) -----------------


def per_language(true: list[str], pred: list[str]):
    import pandas as pd

    t, p = pd.Series(list(true)), pd.Series(list(pred))
    n = len(t)
    out = pd.DataFrame({"support": t.value_counts()})
    out["tp"] = t[t.values == p.values].value_counts().reindex(out.index, fill_value=0)
    out["fp"] = p.value_counts().reindex(out.index, fill_value=0) - out["tp"]
    out["fn"] = out["support"] - out["tp"]
    out["tn"] = n - out["tp"] - out["fp"] - out["fn"]
    with np.errstate(divide="ignore", invalid="ignore"):
        out["precision"] = np.where(out["tp"] + out["fp"] > 0, out["tp"] / (out["tp"] + out["fp"]), 0.0)
        out["recall"] = out["tp"] / out["support"]
        pr = out["precision"] + out["recall"]
        out["f1"] = np.where(pr > 0, 2 * out["precision"] * out["recall"] / pr, 0.0)
        out["specificity"] = out["tn"] / (out["tn"] + out["fp"])
        out["fp_per_million"] = out["fp"] / (n - out["support"]) * 1e6
    out["sensitivity"] = out["recall"]
    out.index.name = "label"
    return out.sort_index()


def summarize(table, true: list[str], ranked: list[list[str]]) -> dict:
    n = max(1, len(true))
    s = {f"top{k}_accuracy": sum(t in r[:k] for t, r in zip(true, ranked, strict=True)) / n for k in (1, 3, 5)}
    w = table["support"] / table["support"].sum()
    s.update(
        passages=int(table["support"].sum()),
        languages=int(len(table)),
        macro_precision=float(table["precision"].mean()),
        macro_recall=float(table["recall"].mean()),
        macro_f1=float(table["f1"].mean()),
        macro_specificity=float(table["specificity"].mean()),
        weighted_f1=float((table["f1"] * w).sum()),
        median_f1=float(table["f1"].median()),
        languages_f1_at_least_0_9=int((table["f1"] >= 0.9).sum()),
        languages_f1_below_0_5=int((table["f1"] < 0.5).sum()),
        mean_fp_per_million=float(table["fp_per_million"].mean()),
    )
    s["micro_f1"] = s["top1_accuracy"]
    return s


# --- data --------------------------------------------------------------------------------


def load_dataset_dir(dataset: str) -> Path:
    """A local build directory, or the dataset repo downloaded (parquet and labels.json only)."""
    if Path(dataset).exists():
        return Path(dataset)
    from huggingface_hub import snapshot_download

    return Path(snapshot_download(dataset, repo_type="dataset", allow_patterns=["data/*", "udhr/*", "labels.json", "stats.json"]))


def read_split(root: Path, name: str):
    import pandas as pd

    if name == "udhr":
        return pd.read_parquet(root / "udhr" / "test.parquet")
    return pd.read_parquet(root / "data" / f"{name}.parquet")


class Sampler:
    """Training passages, drawn so that small languages are seen often enough to learn.

    A language is drawn with probability proportional to n^alpha, n being its passage count:
    alpha 1 is the data as it is, alpha 0 every language equally; 0.3 in between. Then a
    passage of it is drawn, and sometimes cut short, because a short snippet is where
    language identification is hardest and what a page holds in headers and table cells.
    """

    def __init__(self, texts: list[str], label_ids: np.ndarray, alpha: float, short_p: float, seed: int):
        self.texts = texts
        self.rng = random.Random(seed)
        self.np_rng = np.random.default_rng(seed)
        order = np.argsort(label_ids, kind="stable")
        self.by_label = np.split(order, np.flatnonzero(np.diff(label_ids[order])) + 1)
        self.label_of = np.array([label_ids[g[0]] for g in self.by_label])
        sizes = np.array([len(g) for g in self.by_label], dtype=float)
        p = sizes**alpha
        self.p = p / p.sum()
        self.short_p = short_p

    def draw(self, n: int) -> tuple[list[str], np.ndarray]:
        groups = self.np_rng.choice(len(self.by_label), size=n, p=self.p)
        texts, labels = [], []
        for g in groups:
            idx = self.by_label[g][self.rng.randrange(len(self.by_label[g]))]
            texts.append(self._maybe_short(self.texts[idx]))
            labels.append(self.label_of[g])
        return texts, np.array(labels)

    def _maybe_short(self, text: str) -> str:
        if self.rng.random() >= self.short_p or len(text) <= 40:
            return text
        n = self.rng.randint(20, 80)
        start = self.rng.randrange(0, max(1, len(text) - n))
        if " " in text:  # start and end on word boundaries where the script has them
            left = text.rfind(" ", 0, start) + 1
            right = text.find(" ", left + n)
            return text[left : right if right > 0 else len(text)]
        return text[start : start + n]


# --- tracking ----------------------------------------------------------------------------


class Tracker:
    """Weights & Biases when there is a key or a project, and nothing otherwise, so the same run
    works locally without an account. Inside a sweep agent, the agent's run is used."""

    def __init__(self, args, config: dict):
        self.run = None
        self.wandb = None
        if args.no_wandb or not (args.wandb_project or os.environ.get("WANDB_API_KEY")):
            return
        import wandb

        self.wandb = wandb
        clean = asdict(args)
        if wandb.run is not None:
            self.run = wandb.run
            self.run.config.update({**config}, allow_val_change=True)
            self.owned = False
        else:
            self.run = wandb.init(
                project=args.wandb_project or "millid",
                entity=args.wandb_entity,
                name=args.wandb_run_name,
                tags=args.wandb_tags or None,
                config={**clean, **config},
                job_type="train",
            )
            self.owned = True
        self.run.define_metric("train/*", step_metric="step")
        self.run.define_metric("val/*", step_metric="step")

    def log(self, data: dict, step: int) -> None:
        if self.run:
            self.run.log({"step": step, **data})

    def summary(self, data: dict) -> None:
        if self.run:
            self.run.summary.update(data)

    def table(self, key: str, frame) -> None:
        if self.run:
            self.run.log({key: self.wandb.Table(dataframe=frame.reset_index())})

    def finish(self) -> None:
        if self.run and self.owned:
            self.run.finish()


# --- training ----------------------------------------------------------------------------


class TrainStream(torch.utils.data.IterableDataset):
    """Endless training batches, encoded in DataLoader workers so the GPU never waits for
    Python. Each worker draws with its own seed."""

    def __init__(self, texts, label_ids, args, cfg: MillIDConfig):
        self.texts, self.label_ids = texts, label_ids
        self.args, self.cfg = args, cfg

    def __iter__(self):
        info = torch.utils.data.get_worker_info()
        seed = self.args.seed * 1000 + (info.id if info else 0)
        sampler = Sampler(self.texts, self.label_ids, self.args.alpha, self.args.short_p, seed)
        rng = random.Random(seed)
        while True:
            texts, y = sampler.draw(self.args.batch_size)
            batch = encode_batch(texts, self.cfg.max_len, self.cfg.ngram_buckets, rng)
            batch["y"] = torch.from_numpy(y)
            yield batch




def stratified(frame, per_label: int, seed: int = 7):
    return frame.sample(frac=1, random_state=seed).groupby("label").head(per_label)


def evaluate_frame(model: MillID, frame, batch_size: int = 1024) -> tuple[dict, object, list]:
    ranked = model.predict(frame["text"].tolist(), k=5, batch_size=batch_size)
    pred = [r[0][0] for r in ranked]
    table = per_language(frame["label"].tolist(), pred)
    return summarize(table, frame["label"].tolist(), [[lab for lab, _ in r] for r in ranked]), table, ranked


def train(args: TrainArgs) -> None:
    import pandas as pd
    from safetensors.torch import save_file

    torch.manual_seed(args.seed)
    if args.threads:
        torch.set_num_threads(args.threads)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    root = load_dataset_dir(args.dataset)
    labels = Labels.load(root / "labels.json")
    tr = read_split(root, "train")
    va = read_split(root, "validation")
    if args.commercial_only:
        tr = tr[~tr["license"].str.contains("NC", na=False)]
        va = va[~va["license"].str.contains("NC", na=False)]
    if args.extra:
        extra = pd.concat([pd.read_parquet(p, columns=["text", "label"]) for p in args.extra], ignore_index=True)
        extra["label"] = extra["label"].map(labels.canonical)
        extra = extra.dropna(subset=["label"])
        evals = set(va["text"].map(normalize)) | set(read_split(root, "test")["text"].map(normalize)) | set(read_split(root, "udhr")["text"].map(normalize))
        extra = extra[~extra["text"].map(normalize).isin(evals)]
        labels = labels.extend(extra["label"].unique())
        tr = pd.concat([tr[["text", "label"]], extra], ignore_index=True)
        print(f"extra training passages: {len(extra):,}", flush=True)
    trained = set(tr["label"])
    label_ids = tr["label"].map(labels.index).to_numpy()
    print(f"{device}; train {len(tr):,} passages, {len(trained)} labels; validation {len(va):,}; vocabulary {len(labels)}", flush=True)

    cfg = MillIDConfig(num_labels=len(labels), max_len=args.max_len, dropout=args.dropout, **PRESETS[args.preset])
    net = MillIDEncoder(cfg, labels.labels).to(device)
    step_net = torch.compile(net) if args.compile and device == "cuda" else net
    params = sum(p.numel() for p in net.parameters())
    print(f"MillID {args.preset}: {params / 1e6:.1f}M parameters, {cfg}", flush=True)
    tracker = Tracker(
        args,
        {"model": asdict(cfg), "parameters": params, "train_passages": int(len(tr)),
         "trained_labels": len(trained), "vocabulary": len(labels), "validation_passages": int(len(va))},
    )  # fmt: skip

    decay = [p for n, p in net.named_parameters() if p.dim() >= 2 and "emb" not in n]
    no_decay = [p for n, p in net.named_parameters() if p.dim() < 2 or "emb" in n]
    opt = torch.optim.AdamW(
        [{"params": decay, "weight_decay": args.weight_decay}, {"params": no_decay, "weight_decay": 0.0}],
        lr=args.lr, betas=(0.9, 0.98), eps=1e-8, fused=device == "cuda",
    )  # fmt: skip

    def lr_at(step: int) -> float:
        if step < args.warmup:
            return args.lr * step / args.warmup
        progress = (step - args.warmup) / max(1, args.steps - args.warmup)
        return args.lr * (0.05 + 0.95 * 0.5 * (1 + math.cos(math.pi * progress)))

    if args.workers < 0:
        args.workers = 0 if os.name == "nt" else max(0, min(12, (os.cpu_count() or 2) - 2))
    script_table()  # once, before the workers fork, so they share it
    print(f"{args.workers} data workers", flush=True)
    loader = torch.utils.data.DataLoader(
        TrainStream(tr["text"].tolist(), label_ids, args, cfg),
        batch_size=None,
        num_workers=args.workers,
        pin_memory=device == "cuda",
        prefetch_factor=4 if args.workers else None,
        persistent_workers=args.workers > 0,
    )
    stream = iter(loader)
    va_small = stratified(va, args.eval_per_label)
    out = Path(args.output_dir)
    out.mkdir(parents=True, exist_ok=True)
    log = open(out / "training_log.jsonl", "w", encoding="utf-8")
    best, best_step, started, seen = -1.0, 0, time.monotonic(), 0
    use_bf16 = device == "cuda" and torch.cuda.is_bf16_supported()

    for step in range(1, args.steps + 1):
        net.train()
        batch = next(stream)
        y = batch.pop("y").to(device, non_blocking=True)
        batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}
        for g in opt.param_groups:
            g["lr"] = lr_at(step)
        with torch.autocast(device, dtype=torch.bfloat16, enabled=use_bf16):
            logits = step_net(batch)
        loss = F.cross_entropy(logits.float(), y, label_smoothing=args.label_smoothing)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        grad_norm = torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0)
        opt.step()
        seen += args.batch_size

        if step % args.log_every == 0:
            rate = seen / (time.monotonic() - started)
            print(f"step {step}/{args.steps} loss {loss.item():.4f} lr {lr_at(step):.2e} {rate:,.0f} passages/s", flush=True)
            tracker.log(
                {"train/loss": loss.item(), "train/lr": lr_at(step), "train/grad_norm": float(grad_norm),
                 "train/passages_per_s": rate, "train/epochs": seen / len(tr)},
                step,
            )  # fmt: skip
        if step % args.eval_every == 0 or step == args.steps:
            s, _, _ = evaluate_frame(MillID(net, labels, trained, device), va_small)
            entry = {"step": step, "loss": round(loss.item(), 5), **{k: round(s[k], 5) for k in ("macro_f1", "top1_accuracy", "top3_accuracy")}}
            log.write(json.dumps(entry) + "\n")
            log.flush()
            print(f"  validation @ {step}: macro F1 {s['macro_f1']:.4f} top-1 {s['top1_accuracy']:.4f}", flush=True)
            if s["macro_f1"] > best:
                best, best_step = s["macro_f1"], step
                save_file({k: v.detach().cpu().contiguous() for k, v in net.state_dict().items()}, str(out / "model.safetensors"))
            elif step - best_step >= args.patience:
                print(f"no gain on validation for {args.patience} steps: stopping at {step}", flush=True)
                break
            tracker.log(
                {f"val/{k}": s[k] for k in ("macro_f1", "top1_accuracy", "top3_accuracy", "macro_precision", "macro_recall")}
                | {"val/best_macro_f1": best},
                step,
            )
    log.close()
    tracker.summary({"best_step": best_step, "val/best_macro_f1": best})

    labels.save(out / "labels.json", trained)
    config = {
        "model": asdict(cfg),
        "training": asdict(args),
        "best_step": best_step,
        "parameters": params,
        "train_passages": int(len(tr)),
        "trained_labels": len(trained),
        "seconds": round(time.monotonic() - started),
    }
    (out / "config.json").write_text(json.dumps(config, indent=2, default=str), encoding="utf-8")
    shutil.copy(__file__, out / "millid.py")
    doc = Path(__file__).with_name("ARCHITECTURE.md")
    if doc.exists():  # beside the script locally; a Job ships only the script
        shutil.copy(doc, out / "ARCHITECTURE.md")
    print(f"best validation macro F1 {best:.4f} at step {best_step}", flush=True)

    if args.final_eval == "none":
        tracker.finish()
        return
    model = MillID.from_pretrained(str(out))
    sets = ("validation",) if args.final_eval == "validation" else ("validation", "test", "udhr")
    names = json.loads((root / "labels.json").read_text(encoding="utf-8")).get("names", {})
    results, boards = evaluate_all(model, root, out / "reports", args.compare_glotlid, sets, names)
    tracker.summary(
        {f"eval/{set_name}/{m}/{k}": v for set_name, by in results.items() for m, sm in by.items()
         for k, v in sm.items() if isinstance(v, (int, float))}
    )  # fmt: skip
    for set_name, by in results.items():
        rows = [{"model": m, **{k: v for k, v in sm.items() if isinstance(v, (int, float))}} for m, sm in by.items()]
        tracker.table(f"leaderboard/{set_name}/models", pd.DataFrame(rows).set_index("model"))
    for set_name, board in boards.items():
        tracker.table(f"leaderboard/{set_name}/languages", board.round(4))
        if "winner" in board:
            tracker.summary({f"eval/{set_name}/languages_{w.lower()}": int((board["winner"] == w).sum()) for w in ("MillID", "GlotLID", "tie")})
    write_card(out, config, results, args, boards)
    tracker.finish()
    if args.push_to:
        from huggingface_hub import HfApi

        api = HfApi()
        api.create_repo(args.push_to, private=not args.public, exist_ok=True)
        api.upload_folder(repo_id=args.push_to, folder_path=str(out), commit_message=f"MillID {args.preset}, step {best_step}")
        print(f"pushed https://huggingface.co/{args.push_to}", flush=True)


# --- evaluation and the GlotLID comparison ----------------------------------------------

_glot = None


def _glot_init(path: str) -> None:
    global _glot
    import fasttext

    _glot = fasttext.load_model(path)


def _glot_chunk(texts: list[str]) -> list[list[tuple[str, float]]]:
    labels, probs = _glot.f.multilinePredict([normalize(t) for t in texts], 50, 0.0, "strict")
    return [
        [(lab.removeprefix("__label__"), float(p)) for lab, p in zip(ls, ps, strict=True)]
        for ls, ps in zip(labels, probs, strict=True)
    ]


def _cpu_workers(gb_each: float) -> int:
    """Processes the machine has cores and memory for, at most 8."""
    cpus = os.cpu_count() or 2
    try:
        ram = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") / 2**30
    except (ValueError, OSError, AttributeError):
        ram = 16.0
    return max(1, min(8, cpus - 2, int(ram // (gb_each * 2))))


class GlotLID:
    """GlotLID's 50 best answers per passage, computed once in parallel processes (fastText
    holds the GIL, and each process loads the 1.7 GB model once). Both views come from the
    same answers: as it is, and restricted to MillID's labels."""

    def __init__(self, workers: int | None = None):
        from huggingface_hub import hf_hub_download

        # GLOTLID_PATH: a copy already on disk, rather than 1.7 GB from the Hub again.
        self.path = os.environ.get("GLOTLID_PATH") or hf_hub_download("cis-lmu/glotlid", "model.bin")
        self.workers = workers or _cpu_workers(gb_each=1.8)

    def rank(self, texts: list[str], chunk: int = 10_000) -> list[list[tuple[str, float]]]:
        parts = [texts[i : i + chunk] for i in range(0, len(texts), chunk)]
        if self.workers <= 1 or len(parts) == 1:
            _glot_init(self.path)
            return [r for part in parts for r in _glot_chunk(part)]
        from concurrent.futures import ProcessPoolExecutor

        with ProcessPoolExecutor(self.workers, initializer=_glot_init, initargs=(self.path,)) as pool:
            return [r for result in pool.map(_glot_chunk, parts) for r in result]

    @staticmethod
    def view(ranked, allowed: set[str] | None = None, k: int = 5) -> list[list[tuple[str, float]]]:
        out = []
        for pairs in ranked:
            if allowed is not None:
                pairs = [(lab, p) for lab, p in pairs if lab in allowed]
            out.append(pairs[:k] or [("unk", 0.0)])
        return out


def leaderboard(tables: dict, names: dict[str, str]):
    """One row per language: MillID against GlotLID, as it is and restricted to MillID's
    labels. `delta` is MillID's F1 minus restricted GlotLID's: the like-for-like comparison."""
    import pandas as pd

    m = tables["millid"]
    out = pd.DataFrame({"language": [names.get(lab, lab) for lab in m.index], "support": m["support"]}, index=m.index)
    for model, t in tables.items():
        for col in ("f1", "precision", "recall"):
            out[f"{model}_{col}"] = t[col].reindex(m.index)
    if "glotlid@millid" in tables:
        out["delta"] = out["millid_f1"] - out["glotlid@millid_f1"]
        out["winner"] = np.where(out["delta"] > 0.01, "MillID", np.where(out["delta"] < -0.01, "GlotLID", "tie"))
    out.index.name = "label"
    return out.sort_values("support", ascending=False)


def evaluate_all(model: MillID, root: Path, out: Path, compare: bool, which=("validation", "test", "udhr"), names=None) -> tuple[dict, dict]:
    """MillID on validation, test and UDHR; with compare, GlotLID on test and UDHR, as it is and
    restricted to MillID's labels. Writes per-language tables, summaries and a leaderboard per
    set under out/<set>/. Returns the summaries and the leaderboards."""
    names = names or {}
    frames = {}
    if "validation" in which:
        frames["validation"] = read_split(root, "validation")
    if "test" in which:
        frames["test"] = read_split(root, "test")
    if "udhr" in which:
        u = read_split(root, "udhr").reset_index(drop=True)
        frames["udhr"] = u[u["label"].isin(model.trained)]
        frames["udhr-all"] = u
    glot = GlotLID() if compare else None
    glot_cache: dict[str, list] = {}
    results, boards = {}, {}
    for name, frame in frames.items():
        texts = frame["text"].tolist()
        rankings = {"millid": model.predict(texts, k=5, batch_size=1024)}
        if glot is not None and name != "validation":
            if name in ("udhr", "udhr-all"):
                if "udhr" not in glot_cache:
                    glot_cache["udhr"] = glot.rank(frames["udhr-all"]["text"].tolist())
                full = glot_cache["udhr"]
                ranked = [full[i] for i in frame.index]
            else:
                ranked = glot.rank(texts)
            rankings["glotlid"] = GlotLID.view(ranked)
            rankings["glotlid@millid"] = GlotLID.view(ranked, model.trained)
        d = out / name
        d.mkdir(parents=True, exist_ok=True)
        results[name], tables = {}, {}
        for mname, ranked in rankings.items():
            pred = [r[0][0] for r in ranked]
            table = per_language(frame["label"].tolist(), pred)
            sm = summarize(table, frame["label"].tolist(), [[lab for lab, _ in r] for r in ranked])
            if "size" in frame:
                f = frame.assign(pred=pred)
                sm["top1_by_size"] = {k: float((g["label"] == g["pred"]).mean()) for k, g in f.groupby("size")}
            table.round(6).to_csv(d / f"{mname}_per_language.csv")
            (d / f"{mname}_summary.json").write_text(json.dumps(sm, indent=2), encoding="utf-8")
            results[name][mname], tables[mname] = sm, table
            print(f"{name:10} {mname:15} top-1 {sm['top1_accuracy']:.4f} top-3 {sm['top3_accuracy']:.4f} macro F1 {sm['macro_f1']:.4f}", flush=True)
        if len(tables) > 1:
            boards[name] = leaderboard(tables, names)
            boards[name].round(4).to_csv(d / "leaderboard.csv")
    (out / "summary.json").write_text(json.dumps(results, indent=2), encoding="utf-8")
    return results, boards


def _board_md(board, set_name: str) -> str:
    if "delta" not in board:
        return ""
    wins = (board["winner"] == "MillID").sum()
    losses = (board["winner"] == "GlotLID").sum()
    ties = (board["winner"] == "tie").sum()

    def table(frame) -> str:
        lines = ["| Label | Language | Passages | MillID F1 | GlotLID F1 | GlotLID F1, restricted |", "|---|---|---|---|---|---|"]
        for lab, r in frame.iterrows():
            lines.append(
                f"| `{lab}` | {r['language']} | {int(r['support'])} | {r['millid_f1']:.3f} | {r['glotlid_f1']:.3f} | {r['glotlid@millid_f1']:.3f} |"
            )
        return "\n".join(lines)

    return f"""### Leaderboard: {set_name}

MillID's F1 against GlotLID's, language by language. Against GlotLID restricted to MillID's labels, MillID is
better on **{wins}** languages, worse on **{losses}**, and within 0.01 on {ties}. The full table, with precision and
recall, is `reports/{set_name}/leaderboard.csv`.

The 20 languages with the most passages:

{table(board.head(20))}

Where MillID gains most:

{table(board.sort_values("delta", ascending=False).head(15))}

Where MillID loses most:

{table(board.sort_values("delta").head(15))}
"""


def write_card(out: Path, config: dict, results: dict, args, boards: dict | None = None) -> None:
    boards = boards or {}
    rows = ["| Set | Model | Top-1 | Top-3 | Top-5 | Macro P | Macro R | Macro F1 | Median F1 | Languages |", "|---|---|---|---|---|---|---|---|---|---|"]
    for set_name, by_model in results.items():
        for mname, sm in by_model.items():
            rows.append(
                f"| {set_name} | {mname} | {sm['top1_accuracy']:.4f} | {sm['top3_accuracy']:.4f} | {sm['top5_accuracy']:.4f} | "
                f"{sm['macro_precision']:.4f} | {sm['macro_recall']:.4f} | {sm['macro_f1']:.4f} | {sm['median_f1']:.4f} | {sm['languages']} |"
            )
    size_rows = ["| Set | Model | short | sentence | paragraph |", "|---|---|---|---|---|"]
    for set_name, by_model in results.items():
        for mname, sm in by_model.items():
            b = sm.get("top1_by_size") or {}
            if b:
                size_rows.append(f"| {set_name} | {mname} | " + " | ".join(f"{b.get(k, float('nan')):.4f}" for k in ("short", "sentence", "paragraph")) + " |")
    m = config["model"]
    repo = args.push_to or "baobabtech/millid"
    card = f"""---
library_name: pytorch
pipeline_tag: text-classification
tags:
- language-identification
datasets:
- {args.dataset}
---

# MillID {args.preset}: MultiLingual Language IDentifier

A byte-level, encoder-only transformer that names the language of a passage: **{config['trained_labels']} language-script
labels** (ISO 639-3 + script, e.g. `eng_Latn`, `amh_Ethi`). {config['parameters'] / 1e6:.1f}M parameters, trained from
scratch on [{args.dataset}](https://huggingface.co/datasets/{args.dataset}).

```python
from millid import MillID   # millid.py, in this repository

model = MillID.from_pretrained("{repo}")
model.predict(["Lave men ou ak savon anvan ou manje."], k=3)
```

## Architecture

UTF-8 bytes (after Unicode NFKC), each embedded as byte + Unicode script + byte-in-character + hashed character
2-, 3- and 4-grams ({m['ngram_buckets']:,} buckets) -> depthwise convolution stem (kernel {m['conv_kernel']}) -> strided
convolution, {m['downsample']} bytes per token -> {m['n_layers']} pre-norm blocks (RMSNorm, attention with RoPE,
{m['n_heads']} heads, SwiGLU; d={m['d_model']}) -> masked mean pooling -> language head, each language's vector the sum of
its own, its language's and its script's. Up to {m['max_len']} bytes per passage. See `ARCHITECTURE.md`.

## Evaluation against GlotLID

`validation` chose the checkpoint; `test` is split by document from training; `udhr` is the Universal Declaration
of Human Rights, labelled by its translators, over the languages MillID knows, and `udhr-all` over all of them.
`glotlid` is [GlotLID](https://huggingface.co/cis-lmu/glotlid) as it is; `glotlid@millid` is GlotLID restricted to
MillID's labels, the like-for-like comparison. GlotLID's own labels came from training on data that includes UDHR
and labelled FinePDFs and FineWeb-2, so its scores on `udhr` and on those sources are optimistic.

{chr(10).join(rows)}

Top-1 accuracy by passage length:

{chr(10).join(size_rows)}

{_board_md(boards['test'], 'test') if 'test' in boards else ''}
{_board_md(boards['udhr'], 'udhr') if 'udhr' in boards else ''}
Per-language precision, recall (sensitivity), F1 and specificity for every model are in `reports/<set>/`.

## Training

{config['train_passages']:,} passages, best validation macro F1 at step {config['best_step']}, {config['seconds'] / 3600:.1f} hours.
Languages sampled in proportion to n^{args.alpha}; {int(args.short_p * 100)}% of passages cut to a 20 to 80 character
snippet; batch {args.batch_size}, AdamW at {args.lr}, cosine schedule, label smoothing {args.label_smoothing}.
"""
    (out / "README.md").write_text(card, encoding="utf-8")


# --- command line ------------------------------------------------------------------------


@dataclass
class TrainArgs:
    """Train MillID, evaluate it against GlotLID, and optionally push it to the Hub."""

    dataset: str = "baobabtech/millid-dataset"
    """Hub dataset id, or a local build directory with data/, udhr/ and labels.json"""
    preset: Literal["tiny", "small", "medium", "large"] = "medium"
    """model size: tiny ~2M and small ~12M for trials, medium ~114M, large ~190M parameters"""
    max_len: int = 256
    """bytes per passage, [CLS] included; a longer passage is cut to a window"""
    steps: int = 60_000
    """optimiser steps"""
    batch_size: int = 512
    """passages per step"""
    lr: float = 3e-4
    """peak learning rate, after warm-up"""
    warmup: int = 2_000
    """steps of linear warm-up, then cosine decay to 5% of the peak"""
    weight_decay: float = 0.01
    """AdamW weight decay, on weight matrices only"""
    dropout: float = 0.1
    """dropout in the encoder and before the head"""
    label_smoothing: float = 0.1
    """cross-entropy label smoothing"""
    alpha: float = 0.3
    """language sampling: a language is drawn in proportion to n^alpha (1: as the data is, 0: uniform)"""
    short_p: float = 0.25
    """share of training passages cut to a 20 to 80 character snippet"""
    eval_every: int = 2_000
    """steps between validation passes, which choose the checkpoint"""
    eval_per_label: int = 20
    """validation passages per language at each pass"""
    patience: int = 12_000
    """stop after this many steps without a better validation macro F1"""
    commercial_only: bool = False
    """leave out rows licensed for non-commercial use only"""
    extra: list[str] = field(default_factory=list)
    """more training parquet files (text, label): trained on, never evaluated on"""
    compare_glotlid: bool = False
    """score GlotLID too, as it is and restricted to MillID's labels"""
    final_eval: Literal["all", "validation", "none"] = "all"
    """which sets to score after training"""
    output_dir: str = "millid-out"
    """where the model, its card and reports are written"""
    push_to: str | None = None
    """Hub model repo to push to, e.g. baobabtech/millid"""
    public: bool = False
    """make the pushed repo public; private otherwise"""
    seed: int = 7
    """random seed"""
    log_every: int = 100
    """steps between training-loss lines and W&B points"""
    threads: int = 0
    """CPU threads for PyTorch; 0 leaves PyTorch's default"""
    workers: int = -1
    """DataLoader workers encoding batches; -1 picks from the cores available, 0 encodes in the training process"""
    compile: bool = False
    """torch.compile the model for training (GPU only)"""
    wandb_project: str | None = os.environ.get("WANDB_PROJECT")
    """Weights & Biases project; with WANDB_API_KEY set, 'millid' when not given"""
    wandb_entity: str | None = os.environ.get("WANDB_ENTITY")
    """Weights & Biases team or user"""
    wandb_run_name: str | None = None
    """run name in Weights & Biases"""
    wandb_tags: list[str] = field(default_factory=list)
    """run tags in Weights & Biases"""
    no_wandb: bool = False
    """never log to Weights & Biases, even with a key"""


@dataclass
class EvaluateArgs:
    """Evaluate a trained model on validation, test and UDHR."""

    model: str
    """model directory or Hub repo"""
    dataset: str = "baobabtech/millid-dataset"
    """Hub dataset id or local build directory"""
    compare_glotlid: bool = False
    """score GlotLID too"""
    output_dir: str = "millid-eval"
    """where the reports are written"""


@dataclass
class SweepArgs:
    """Run Weights & Biases sweep trials in this process."""

    sweep_id: str
    """ENTITY/PROJECT/SWEEP_ID, from `wandb sweep millid/sweep.yaml`"""
    count: int = 4
    """trials this agent runs"""
    dataset: str = "baobabtech/millid-dataset"
    """used unless the sweep sets it"""


@dataclass
class PredictArgs:
    """The language of each text."""

    model: str
    """model directory or Hub repo"""
    texts: list[str] = field(default_factory=list)
    """texts to identify"""
    k: int = 3
    """answers per text"""


def sweep_agent(args: SweepArgs) -> None:
    """Each trial takes its settings from the sweep and the TrainArgs defaults for the rest,
    and reports val/best_macro_f1. No pushing, no GlotLID, validation only: a sweep chooses
    settings, and the chosen ones are then trained in full."""
    import wandb

    def trial() -> None:
        run = wandb.init()
        a = TrainArgs(dataset=args.dataset)
        for key, value in dict(run.config).items():
            name = key.replace("-", "_")
            if hasattr(a, name):
                setattr(a, name, value)
        a.push_to, a.compare_glotlid, a.final_eval = None, False, "validation"
        a.output_dir = f"sweep/{run.id}"
        train(a)
        run.finish()

    wandb.agent(args.sweep_id, function=trial, count=args.count)


def main(argv: list[str] | None = None) -> None:
    import tyro

    args = tyro.extras.subcommand_cli_from_dict(
        {"train": TrainArgs, "evaluate": EvaluateArgs, "sweep-agent": SweepArgs, "predict": PredictArgs},
        args=argv,
        prog="millid",
    )
    if isinstance(args, TrainArgs):
        train(args)
    elif isinstance(args, EvaluateArgs):
        evaluate_all(MillID.from_pretrained(args.model), load_dataset_dir(args.dataset), Path(args.output_dir), args.compare_glotlid)
    elif isinstance(args, SweepArgs):
        sweep_agent(args)
    else:
        for text, answers in zip(args.texts, MillID.from_pretrained(args.model).predict(args.texts, args.k), strict=True):
            print(text, "->", answers)


if __name__ == "__main__":
    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
    main(sys.argv[1:])