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, pad_to: int | None = None
) -> dict[str, torch.Tensor]:
    """Passages as the model's input: the four per-byte sequences, CLS first, padded to the
    longest (or to pad_to, for the fixed shapes torch.compile wants), 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 = max(1 + int(nb.max(initial=0)), pad_to or 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, enabled: bool = True):
        self.run = None
        self.wandb = None
        if not enabled or 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, batch_size: int, rank: int = 0, pad_to: int | None = None):
        self.texts, self.label_ids = texts, label_ids
        self.args, self.cfg = args, cfg
        self.batch_size, self.rank, self.pad_to = batch_size, rank, pad_to

    def __iter__(self):
        info = torch.utils.data.get_worker_info()
        # every worker of every GPU's process draws its own passages
        seed = self.args.seed * 1000 + self.rank * 100 + (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.batch_size)
            batch = encode_batch(texts, self.cfg.max_len, self.cfg.ngram_buckets, rng, pad_to=self.pad_to)
            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 _free_port() -> int:
    import socket

    with socket.socket() as s:
        s.bind(("127.0.0.1", 0))
        return s.getsockname()[1]


def train(args: TrainArgs) -> None:
    """Train on every GPU the machine has (or args.gpus of them): one process per GPU, the batch
    split between them and the gradients averaged (DistributedDataParallel). Without CUDA,
    args.gpus > 1 starts that many CPU processes, which is how the distributed path is tested."""
    if torch.cuda.is_available():
        n = torch.cuda.device_count()
        world = n if args.gpus <= 0 else min(args.gpus, n)
    else:
        world = max(1, args.gpus)
    if world == 1:
        _train(0, 1, 0, args)
        return
    if args.batch_size % world:
        sys.exit(f"batch size {args.batch_size} does not split over {world} processes")
    load_dataset_dir(args.dataset)  # download once, before the processes start
    import torch.multiprocessing as mp

    mp.spawn(_train, args=(world, _free_port(), args), nprocs=world, join=True)


def _train(rank: int, world: int, port: int, args: TrainArgs) -> None:
    import pandas as pd
    from safetensors.torch import save_file

    cuda = torch.cuda.is_available()
    distributed = world > 1
    if distributed:
        import torch.distributed as dist

        os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"] = "127.0.0.1", str(port)
        if cuda:
            torch.cuda.set_device(rank)
        dist.init_process_group("nccl" if cuda else "gloo", rank=rank, world_size=world)
    main = rank == 0
    say = print if main else (lambda *a, **k: None)
    torch.manual_seed(args.seed)
    if args.threads:
        torch.set_num_threads(args.threads)
    device = f"cuda:{rank}" if cuda else "cpu"
    device_type = "cuda" if cuda 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)
        say(f"extra training passages: {len(extra):,}", flush=True)
    trained = set(tr["label"])
    label_ids = tr["label"].map(labels.index).to_numpy()
    say(f"{device_type} x{world}; 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)
    model = net
    if distributed:
        from torch.nn.parallel import DistributedDataParallel

        model = DistributedDataParallel(net, device_ids=[rank] if cuda else None)
    step_net = torch.compile(model) if args.compile and cuda else model
    params = sum(p.numel() for p in net.parameters())
    say(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)), "gpus": world,
         "trained_labels": len(trained), "vocabulary": len(labels), "validation_passages": int(len(va))},
        enabled=main,
    )  # 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=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)))

    workers = args.workers
    if workers < 0:
        workers = 0 if os.name == "nt" else max(0, min(12, ((os.cpu_count() or 2) - 2) // world))
    script_table()  # once, before the workers fork, so they share it
    say(f"{workers} data workers per process, {args.batch_size // world} passages per process per step", flush=True)
    pad_to = cfg.max_len if args.compile and cuda else None  # fixed shapes: compiled once
    loader = torch.utils.data.DataLoader(
        TrainStream(tr["text"].tolist(), label_ids, args, cfg, args.batch_size // world, rank, pad_to),
        batch_size=None,
        num_workers=workers,
        pin_memory=cuda,
        prefetch_factor=4 if workers else None,
        persistent_workers=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") if main else None
    best, best_step, started, seen = -1.0, 0, time.monotonic(), 0
    use_bf16 = 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_type, 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 and main:
            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:
            # The first process evaluates and saves; the others wait for its verdict.
            stop = False
            if main:
                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)
                    stop = True
                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,
                )
            if distributed:
                flag = torch.tensor([int(stop)], device=device)
                dist.broadcast(flag, 0)
                stop = bool(flag.item())
            if stop:
                break
    if distributed:
        dist.destroy_process_group()
    if not main:
        return
    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,
        "gpus": world,
        "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", {})
    compare = [ext for ext, on in (("glotlid", args.compare_glotlid), ("fasttext", args.compare_fasttext)) if on]
    results, boards = evaluate_all(model, root, out / "reports", compare, 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 against GlotLID and fastText ---------------------------------------------
#
# Both are fastText models: GlotLID (cis-lmu, 2,000+ labels) and Meta's NLLB language identifier
# (218 labels), published as facebook/fasttext-language-identification. Both write labels as
# ISO 639-3 + script; their answers are mapped onto MillID's spelling of the same language
# before scoring, so that a right answer spelled differently is not counted wrong.

#: name: (Hub repo, file, environment variable for a copy on disk, display name, GB in memory)
EXTERNAL = {
    "glotlid": ("cis-lmu/glotlid", "model.bin", "GLOTLID_PATH", "GlotLID", 1.8),
    "fasttext": ("facebook/fasttext-language-identification", "model.bin", "FASTTEXT_LID_PATH", "fastText (Meta, NLLB)", 1.3),
}
#: Spellings MillID's aliases do not cover: a macrolanguage for its one member in MillID
#: (Estonian, Akan as Twi, Guarani as Paraguayan Guarani, Kikongo as Koongo).
EXTERNAL_ALIASES = {"est_Latn": "ekk_Latn", "aka_Latn": "twi_Latn", "grn_Latn": "gug_Latn", "kon_Latn": "kng_Latn"}

_lid = None
_lid_aliases: dict[str, str] = {}


def to_millid_label(label: str, aliases: dict[str, str]) -> str:
    """An external model's label in MillID's spelling: Simplified and Traditional Han are both
    Hani, then MillID's aliases, then EXTERNAL_ALIASES."""
    code, _, script = label.partition("_")
    if script in ("Hans", "Hant"):
        label = f"{code}_Hani"
    label = aliases.get(label, label)
    return EXTERNAL_ALIASES.get(label, label)


def _lid_init(path: str, aliases: dict[str, str]) -> None:
    global _lid, _lid_aliases
    import fasttext

    _lid, _lid_aliases = fasttext.load_model(path), aliases


def _lid_chunk(texts: list[str]) -> list[list[tuple[str, float]]]:
    labels, probs = _lid.f.multilinePredict([normalize(t) for t in texts], 50, 0.0, "strict")
    out = []
    for ls, ps in zip(labels, probs, strict=True):
        seen, pairs = set(), []
        for lab, p in zip(ls, ps, strict=True):
            lab = to_millid_label(lab.removeprefix("__label__"), _lid_aliases)
            if lab not in seen:  # two of its labels can be one of MillID's: keep the better
                seen.add(lab)
                pairs.append((lab, float(p)))
        out.append(pairs)
    return out


def _lid_labels() -> list[str]:
    return sorted({to_millid_label(lab.removeprefix("__label__"), _lid_aliases) for lab in _lid.get_labels()})


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 FastTextLID:
    """A fastText language identifier's 50 best answers per passage, in MillID's label
    spelling, computed in parallel processes (fastText holds the GIL, and each process loads the
    model once). Both views come from the same answers: as it is, and restricted to MillID's
    labels."""

    def __init__(self, name: str = "glotlid", aliases: dict[str, str] | None = None, workers: int | None = None):
        from huggingface_hub import hf_hub_download

        repo, filename, env, self.display, gb = EXTERNAL[name]
        self.name, self.aliases = name, aliases or {}
        # A copy already on disk, rather than a gigabyte or two from the Hub again.
        self.path = os.environ.get(env) or hf_hub_download(repo, filename)
        self.workers = workers or _cpu_workers(gb_each=gb)
        self._labels: set[str] | None = None

    def labels(self) -> set[str]:
        """Every label it can answer, in MillID's spelling."""
        if self._labels is None:
            _lid_init(self.path, self.aliases)
            self._labels = set(_lid_labels())
        return self._labels

    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:
            _lid_init(self.path, self.aliases)
            return [r for part in parts for r in _lid_chunk(part)]
        from concurrent.futures import ProcessPoolExecutor

        with ProcessPoolExecutor(self.workers, initializer=_lid_init, initargs=(self.path, self.aliases)) as pool:
            return [r for result in pool.map(_lid_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


GlotLID = FastTextLID  # the name earlier code and reports use


def leaderboard(tables: dict, names: dict[str, str]):
    """One row per language: MillID against each external model, as it is and restricted to
    MillID's labels. `delta` is MillID's F1 minus restricted GlotLID's, `delta_fasttext` minus
    restricted fastText's: the like-for-like comparisons."""
    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)
    for ext, suffix, winner in (("glotlid", "", "GlotLID"), ("fasttext", "_fasttext", "fastText")):
        if f"{ext}@millid" in tables:
            d = out["millid_f1"] - out[f"{ext}@millid_f1"]
            out[f"delta{suffix}"] = d
            out[f"winner{suffix}"] = np.where(d > 0.01, "MillID", np.where(d < -0.01, winner, "tie"))
    out.index.name = "label"
    return out.sort_values("support", ascending=False)


def _score(name: str, frame, rankings: dict, out: Path, names: dict) -> tuple[dict, object]:
    d = out / name
    d.mkdir(parents=True, exist_ok=True)
    results, 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[mname], tables[mname] = sm, table
        print(f"{name:24} {mname:16} top-1 {sm['top1_accuracy']:.4f} top-3 {sm['top3_accuracy']:.4f} macro F1 {sm['macro_f1']:.4f}", flush=True)
    board = None
    if len(tables) > 1:
        board = leaderboard(tables, names)
        board.round(4).to_csv(d / "leaderboard.csv")
    return results, board


def evaluate_all(
    model: MillID, root: Path, out: Path, compare=False, which=("validation", "test", "udhr"), names=None,
    lids: dict[str, FastTextLID] | None = None, cache: dict | None = None,
) -> tuple[dict, dict]:  # fmt: skip
    """MillID on validation, test and UDHR, and the external models in `compare` (True means
    GlotLID alone) on test and UDHR, each 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.

    fastText knows 218 labels, so on sets of 1,600 it is mostly scored on languages it cannot
    answer. `<set>-common` scores every model on the passages whose language all of them know:
    the head-to-head on equal terms.

    An external model's ranking of a set does not depend on the MillID it is compared with:
    pass the same `lids` and `cache` to evaluate several models and each set is read once."""
    names = names or {}
    compare = ("glotlid",) if compare is True else tuple(compare or ())
    lids = {} if lids is None else lids
    cache = {} if cache is None else cache
    for ext in compare:
        if ext not in lids:
            lids[ext] = FastTextLID(ext, model.labels.aliases)
    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
    common = set(model.trained)
    if len(compare) > 1:
        for ext in compare:
            common &= lids[ext].labels()

    results, boards = {}, {}
    for name, frame in frames.items():
        texts = frame["text"].tolist()
        rankings = {"millid": model.predict(texts, k=5, batch_size=1024)}
        if name != "validation":
            for ext in compare:
                if name in ("udhr", "udhr-all"):
                    if (ext, "udhr") not in cache:
                        cache[ext, "udhr"] = lids[ext].rank(frames["udhr-all"]["text"].tolist())
                    full = cache[ext, "udhr"]
                    ranked = [full[i] for i in frame.index]
                else:
                    if (ext, name) not in cache:
                        cache[ext, name] = lids[ext].rank(texts)
                    ranked = cache[ext, name]
                rankings[ext] = FastTextLID.view(ranked)
                rankings[f"{ext}@millid"] = FastTextLID.view(ranked, model.trained)
        results[name], board = _score(name, frame, rankings, out, names)
        if board is not None:
            boards[name] = board
        if len(compare) > 1 and name in ("test", "udhr"):
            keep = frame["label"].isin(common).to_numpy()
            sub = {m: [r for r, k in zip(ranked, keep, strict=True) if k] for m, ranked in rankings.items()}
            results[f"{name}-common"], board = _score(f"{name}-common", frame[keep], sub, out, names)
            boards[f"{name}-common"] = board
    (out / "summary.json").write_text(json.dumps(results, indent=2), encoding="utf-8")
    return results, boards


def _board_md(board, set_name: str) -> str:
    ext = [(name, col, w) for name, col, w in (("GlotLID", "glotlid", "winner"), ("fastText", "fasttext", "winner_fasttext")) if w in board]
    if not ext:
        return ""
    counts = "; ".join(
        f"against {name} restricted to MillID's labels, MillID is better on **{(board[w] == 'MillID').sum()}** languages, "
        f"worse on **{(board[w] == name).sum()}** and within 0.01 on {(board[w] == 'tie').sum()}"
        for name, _, w in ext
    )
    key = "delta" if "delta" in board else "delta_fasttext"

    def table(frame) -> str:
        head = "| Label | Language | Passages | MillID F1 | " + " | ".join(f"{n} F1 | {n} F1, restricted" for n, _, _ in ext) + " |"
        lines = [head, "|---|---|---|---|" + "---|---|" * len(ext)]
        for lab, r in frame.iterrows():
            cells = " | ".join(f"{r[f'{c}_f1']:.3f} | {r[f'{c}@millid_f1']:.3f}" for _, c, _ in ext)
            lines.append(f"| `{lab}` | {r['language']} | {int(r['support'])} | {r['millid_f1']:.3f} | {cells} |")
        return "\n".join(lines)

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

MillID's F1 against the others', language by language: {counts}. 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 over {ext[0][0]}:

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

Where MillID loses most to {ext[0][0]}:

{table(board.sort_values(key).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 and fastText

`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. `fasttext` is
Meta's [fastText language identifier](https://huggingface.co/facebook/fasttext-language-identification) (NLLB,
218 labels), its answers mapped onto MillID's spelling (Simplified and Traditional Han as `Hani`, `pes`/`prs` as
`fas`, `est` as `ekk`). It cannot answer most of the 1,600 languages here, so `test-common` and `udhr-common`
score every model on the passages whose language all of them know: the head-to-head on equal terms.

{chr(10).join(rows)}

Top-1 accuracy by passage length:

{chr(10).join(size_rows)}

{chr(10).join(_board_md(boards[s], s) for s in ('test', 'udhr', 'test-common') if s in boards)}
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"""
    compare_fasttext: bool = False
    """score Meta's fastText language identifier (NLLB, 218 labels) too"""
    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 per training process; -1 picks from the cores available, 0 encodes in the training process"""
    compile: bool = False
    """torch.compile the model for training (GPU only); batches are then padded to max_len, so it compiles once"""
    gpus: int = 0
    """GPUs to train on, one process each, the batch split between them; 0 uses every GPU the machine has"""
    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 trained models on validation, test and UDHR, for instance again after the
    dataset's evaluation splits were corrected."""

    model: list[str]
    """model directories or Hub repos; GlotLID and fastText read each set once for all of them"""
    dataset: str = "baobabtech/millid-dataset"
    """Hub dataset id or local build directory"""
    compare_glotlid: bool = False
    """score GlotLID too"""
    compare_fasttext: bool = False
    """score Meta's fastText language identifier (NLLB, 218 labels) too"""
    output_dir: str = "millid-eval"
    """where the reports are written, one directory per model"""
    push: bool = False
    """push each Hub model's new reports and model card back to its repo"""


@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 evaluate_models(args: EvaluateArgs) -> None:
    from types import SimpleNamespace

    from huggingface_hub import HfApi, snapshot_download

    root = load_dataset_dir(args.dataset)
    names = json.loads((root / "labels.json").read_text(encoding="utf-8")).get("names", {})
    compare = [ext for ext, on in (("glotlid", args.compare_glotlid), ("fasttext", args.compare_fasttext)) if on]
    lids: dict = {}
    cache: dict = {}
    for ref in args.model:
        local = Path(ref).exists()
        src = Path(ref) if local else Path(snapshot_download(ref, allow_patterns=["*.json", "*.safetensors"]))
        out = Path(args.output_dir) / ref.replace("/", "__")
        print(f"== {ref}", flush=True)
        model = MillID.from_pretrained(str(src))
        results, boards = evaluate_all(model, root, out / "reports", compare, names=names, lids=lids, cache=cache)
        config = json.loads((src / "config.json").read_text(encoding="utf-8"))
        training = SimpleNamespace(**config["training"])
        write_card(out, config, results, training, boards)
        if args.push and not local:
            HfApi().upload_folder(repo_id=ref, folder_path=str(out), commit_message=f"Evaluated again on {args.dataset}")
            print(f"pushed https://huggingface.co/{ref}", flush=True)


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_models(args)
    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:])