Empaquetar y servir lotes · Rukh

Qué vas a construir

El puente entre los vocabularios de la lección anterior y la GPU de M2, en tres piezas:

Un empaquetador que escribe cada mes (unos tres millones de partidas codificadas) en un único array de uint16 de 480 MB, con un índice aparte de dónde empieza cada una, sin tenerlas nunca todas en memoria.
Un dataset que sirve ventanas de 200 tokens desde ese array con memmap, alineadas al <bos> de una partida y con el objetivo y desplazado una posición: el contrato que consume el decoder de M2.
Un orquestador, el de rukh data tokenize, con la fixture de paridad que el navegador tendrá que reproducir id a id.

Escribir en flujo, indexar aparte y leer con memmap sirve para cualquier corpus que no cabe en memoria, sea de partidas, de texto o de fragmentos para un RAG.

`pack.py`: tres codificadores con la misma forma

"""Pack encoded games into one ``uint16`` stream (``tokens.npy``) plus a game index.

The stream is the concatenation of every encoded game (``<bos> ... <eos>``), untruncated;
``starts.npy`` holds the offset of each ``<bos>`` so a loader can cut ``block``-sized windows
aligned to game starts (D-019). ``meta.json`` records counts, the scheme and a hash of the
vocabulary so a stale pack is never mixed with a different tokenizer.
"""

from __future__ import annotations

import hashlib
import json
from pathlib import Path
from typing import TYPE_CHECKING, Protocol

import numpy as np
import pyarrow.parquet as pq
from pydantic import BaseModel, ConfigDict

from rukh.data.pipeline import TokenizeConfig
from rukh.paths import resolve
from rukh.tokenize.san_chars import SanCharTokenizer, san_text, uci_to_san
from rukh.tokenize.uci_vocab import RESULT_TOKENS, UciTokenizer

if TYPE_CHECKING:
    from tokenizers import Tokenizer

TOKENS_FILE = "tokens.npy"
STARTS_FILE = "starts.npy"
META_FILE = "meta.json"
CHUNK_GAMES = 20_000
COPY_TOKENS = (64 << 20) // 2  # 64 MiB of uint16 per copy into the memmap

src/rukh/tokenize/pack.pylíneas 1-32 · p1

El flujo va en uint16: dos bytes por token dan para 65 536 ids, de sobra para los tres vocabularios. Con int32 el fichero pesaría el doble sin guardar nada más.

Los codificadores son adaptadores con la misma forma sobre los vocabularios de la lección anterior; puedes saltarlos en la primera lectura y volver al teclearlos. El del BPE, justo debajo, sí tiene una decisión que importa al entrenamiento. Si no vas a teclear lo plegado de pack.py, trae el fichero entero antes de empezarlo, con git checkout p1 -- src/rukh/tokenize/pack.py, y lee sobre tu copia lo que queda a la vista; si ya tecleaste una parte, no lo traigas (pisaría lo tuyo): git diff p1 -- src/rukh/tokenize/pack.py te marca con - las líneas que te faltan.

// Fontaneríapack.py: el Protocol GameEncoder y los codificadores UCI y SAN

class GameEncoder(Protocol):
    """What ``pack_month`` needs from a tokenizer."""

    scheme: str
    pad_id: int
    bos_id: int
    eos_id: int

    def vocab_size(self) -> int: ...

    def vocab_hash(self) -> str: ...

    def encode_game(
        self, uci: str, white_elo: int, black_elo: int, result: str, max_len: int
    ) -> list[int]: ...

src/rukh/tokenize/pack.pylíneas 35-49 · p1

GameEncoder es un Protocol: los tres codificadores encajan sin heredar de nada porque tienen los métodos. Así pack_month no importa ninguno, y empaquetar con UCI no arrastra tokenizers.

class UciGameEncoder:
    """The fixed UCI vocabulary as a ``GameEncoder``."""

    scheme = "uci"

    def __init__(self, tokenizer: UciTokenizer | None = None) -> None:
        self.tokenizer = tokenizer or UciTokenizer()
        self.pad_id = self.tokenizer.pad_id
        self.bos_id = self.tokenizer.bos_id
        self.eos_id = self.tokenizer.eos_id

    def vocab_size(self) -> int:
        return len(self.tokenizer)

    def vocab_hash(self) -> str:
        return _sha256_text(json.dumps(self.tokenizer.ids))

    def encode_game(
        self, uci: str, white_elo: int, black_elo: int, result: str, max_len: int
    ) -> list[int]:
        return self.tokenizer.encode_game(uci, white_elo, black_elo, result, max_len=max_len)

src/rukh/tokenize/pack.pylíneas 52-72 · p1

class SanGameEncoder:
    """Char-level SAN: ``[<bos>, chars of '1.e4 e5 ... 1-0', <eos>]`` (no Elo tokens)."""

    scheme = "san"

    def __init__(self, tokenizer: SanCharTokenizer | None = None) -> None:
        self.tokenizer = tokenizer or SanCharTokenizer()
        self.pad_id = self.tokenizer.pad_id
        self.bos_id = self.tokenizer.bos_id
        self.eos_id = self.tokenizer.eos_id

    def vocab_size(self) -> int:
        return len(self.tokenizer)

    def vocab_hash(self) -> str:
        return _sha256_text(json.dumps(self.tokenizer.ids))

    def encode_game(
        self, uci: str, white_elo: int, black_elo: int, result: str, max_len: int
    ) -> list[int]:
        return self.tokenizer.encode(san_text(uci_to_san(uci), result))[:max_len]

src/rukh/tokenize/pack.pylíneas 75-95 · p1

El encode_game de SanGameEncoder recibe los dos Elo y los ignora: es la limitación del alfabeto sin tokens de Elo, visible en la firma. Reconstruir el SAN pide un tablero por partida, y por eso este esquema tarda hora y media en empaquetarse y los otros dos, minutos.

class BpeGameEncoder:
    """BPE over spaceless UCI: ``[<bos>, bpe ids..., <result>, <eos>]`` (no Elo tokens)."""

    scheme = "bpe"

    def __init__(self, bpe: Tokenizer) -> None:
        self.bpe = bpe
        self.pad_id = UciTokenizer.pad_id
        self.bos_id = UciTokenizer.bos_id
        self.eos_id = UciTokenizer.eos_id
        for token in UciTokenizer.specials:
            if bpe.token_to_id(token) != UciTokenizer().vocab[token]:
                raise ValueError(f"BPE special token {token} does not share the UCI id")

    def vocab_size(self) -> int:
        return self.bpe.get_vocab_size()

    def vocab_hash(self) -> str:
        return _sha256_text(self.bpe.to_str())

    def encode_game(
        self, uci: str, white_elo: int, black_elo: int, result: str, max_len: int
    ) -> list[int]:
        from rukh.tokenize.bpe import bpe_text

        if result not in RESULT_TOKENS:
            raise ValueError(f"unknown result {result!r}")
        ids = [self.bos_id, *self.bpe.encode(bpe_text(uci)).ids]
        ids.append(self.bpe.token_to_id(RESULT_TOKENS[result]))
        ids.append(self.eos_id)
        return ids[:max_len]

src/rukh/tokenize/pack.pylíneas 98-128 · p1

El constructor de BpeGameEncoder exige que los ocho especiales tengan los mismos ids que en el UCI. Con un BPE entrenado sin special_tokens, <pad> no sería el 0 y ignore_index=0 estaría ignorando la jugada que ocupe ese id; con la comprobación, el pack falla al construirse. La secuencia BPE lleva tres tokens de encuadre en vez de cinco, porque tampoco hay Elo.

class PackInfo(BaseModel):
    """Contents of ``meta.json``."""

    model_config = ConfigDict(extra="forbid")

    n_games: int
    n_tokens: int
    scheme: str
    vocab_size: int
    vocab_hash: str
    source: str


def _sha256_text(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def make_encoder(scheme: str, bpe: Tokenizer | None = None) -> GameEncoder:
    if scheme == "uci":
        return UciGameEncoder()
    if scheme == "san":
        return SanGameEncoder()
    if scheme == "bpe":
        if bpe is None:
            raise ValueError("scheme 'bpe' needs a trained tokenizer")
        return BpeGameEncoder(bpe)
    raise ValueError(f"unknown scheme {scheme!r}")

src/rukh/tokenize/pack.pylíneas 131-157 · p1

El campo que importa de PackInfo es vocab_hash: permite que el entrenamiento de M2 rechace un pack cuyo vocabulario no es el del modelo. Sin él, entrenar sobre un pack de otra enumeración funcionaría, con ids que significan otra cosa, y la pérdida bajaría igual. En make_encoder, solo la rama del BPE pide algo: el tokenizador entrenado.

def pack_month(games: Path, encoder: GameEncoder, out: Path) -> PackInfo:
    """Encode every game of ``games`` (a UCI parquet) into ``out/tokens.npy`` + ``starts.npy``.

    The parquet is read row group by row group and the token stream is appended to a raw file;
    once its length is known the raw file is copied into the ``.npy`` memmap in 64 MiB chunks,
    so neither the games nor the whole stream are ever held in memory at once.
    """
    if encoder.vocab_size() > np.iinfo(np.uint16).max + 1:
        raise ValueError("vocabulary does not fit in uint16")
    out = Path(out)
    out.mkdir(parents=True, exist_ok=True)
    raw_path = out / "tokens.bin"
    starts: list[int] = []
    n_tokens = 0
    reader = pq.ParquetFile(Path(games))
    with raw_path.open("wb") as raw:
        for batch in reader.iter_batches(
            batch_size=CHUNK_GAMES, columns=["uci", "white_elo", "black_elo", "result"]
        ):
            buffer: list[int] = []
            rows = zip(
                batch.column("uci").to_pylist(),
                batch.column("white_elo").to_pylist(),
                batch.column("black_elo").to_pylist(),
                batch.column("result").to_pylist(),
                strict=True,
            )
            for uci, w_elo, b_elo, result in rows:
                ids = encoder.encode_game(uci, int(w_elo), int(b_elo), result, max_len=1 << 30)
                starts.append(n_tokens)
                n_tokens += len(ids)
                buffer.extend(ids)
            raw.write(np.asarray(buffer, dtype=np.uint16).tobytes())
    tokens = np.lib.format.open_memmap(
        out / TOKENS_FILE, mode="w+", dtype=np.uint16, shape=(n_tokens,)
    )
    _copy_raw(raw_path, tokens, n_tokens)
    tokens.flush()
    del tokens
    raw_path.unlink()
    np.save(out / STARTS_FILE, np.asarray(starts, dtype=np.int64))
    info = PackInfo(
        n_games=len(starts),
        n_tokens=n_tokens,
        scheme=encoder.scheme,
        vocab_size=encoder.vocab_size(),
        vocab_hash=encoder.vocab_hash(),
        source=Path(games).name,
    )
    (out / META_FILE).write_text(info.model_dump_json(indent=2) + "\n", encoding="utf-8")
    return info

src/rukh/tokenize/pack.pylíneas 160-210 · p1

pack_month es la lección del fichero, y se lee por la memoria. open_memmap necesita la forma final del array, y la longitud del flujo no se conoce hasta codificar la última partida. Acumularlo todo serían varios gigas y recorrer el parquet dos veces, codificar seis millones de partidas dos veces. La salida es escribir los ids a un fichero crudo mientras se cuentan; al acabar, n_tokens es exacto, se crea el memmap y se copia.

El max_len=1 << 30 hace que el pack guarde las partidas enteras: el corte a 200 lo hace el loader. Así, cambiar la ventana de contexto en M2 no obliga a volver a tokenizar. Lo único que crece en memoria es starts, un entero por partida (24 MB para tres millones).

El resto del fichero es copia y rutas, y se puede saltar en la primera lectura, pero no al teclear: pack_month llama a _copy_raw y la CLI a pack_scheme. Si lo saltas, git diff p1 -- src/rukh/tokenize/pack.py te marca con - las líneas que faltan, o trae el fichero como se dice arriba.

// Fontaneríapack.py, el final: la copia por trozos, leer el meta.json y empaquetar los dos meses

def _copy_raw(raw_path: Path, tokens: np.ndarray, n_tokens: int) -> None:
    """Copy the raw token file into the memmap in ``COPY_TOKENS``-sized chunks."""
    offset = 0
    with raw_path.open("rb") as raw:
        while offset < n_tokens:
            chunk = np.fromfile(raw, dtype=np.uint16, count=min(COPY_TOKENS, n_tokens - offset))
            if not len(chunk):
                raise OSError(f"{raw_path}: expected {n_tokens} tokens, got {offset}")
            tokens[offset : offset + len(chunk)] = chunk
            offset += len(chunk)

src/rukh/tokenize/pack.pylíneas 213-222 · p1

def read_pack_info(directory: Path) -> PackInfo:
    return PackInfo.model_validate_json((Path(directory) / META_FILE).read_text("utf-8"))


def month_parquet(uci_dir: Path, month: str) -> Path:
    year, mm = month.split("-")
    return Path(uci_dir) / f"year={year}" / f"month={mm}" / "games.parquet"


def pack_scheme(cfg: TokenizeConfig, scheme: str, bpe: Tokenizer | None = None) -> list[Path]:
    """Pack ``train_month`` and ``val_month`` under ``out_dir/<scheme>/{train,val}``."""
    encoder = make_encoder(scheme, bpe)
    uci_dir = resolve(cfg.uci_dir)
    written: list[Path] = []
    for split, month in (("train", cfg.train_month), ("val", cfg.val_month)):
        out = resolve(cfg.out_dir) / scheme / split
        pack_month(month_parquet(uci_dir, month), encoder, out)
        written.append(out / META_FILE)
    return written

src/rukh/tokenize/pack.pylíneas 225-243 · p1

pack_scheme convierte la partición temporal de la primera lección en dos directorios: enero en train/ y febrero en val/.

`loader.py`: el contrato con el decoder

// Teoría justaTensores en cinco minutos

Aquí entra torch por primera vez. Un tensor es un array de NumPy con forma (shape) y tipo (dtype) fijos que además puede vivir en la GPU; torch.from_numpy lo crea sin copiar. Un Dataset es cualquier clase con __len__ y __getitem__, como un array que se calcula al pedirlo, y un DataLoader lo recorre en lotes barajados: un iterador de páginas. De aquí salen x e y de forma (64, 200) y tipo int64, porque la entropía cruzada exige enteros de 64 bits. El resto de PyTorch (formas, autograd, el bucle) está en PyTorch mínimo, la lección que abre M2.

"""``PackedDataset`` and ``make_loader``: ``<bos>``-aligned windows over a packed token stream."""

from __future__ import annotations

from pathlib import Path

import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset

from rukh.tokenize.pack import STARTS_FILE, TOKENS_FILE, PackInfo, read_pack_info

IGNORE_INDEX = 0  # ``<pad>`` in every scheme; the loss must use ``ignore_index=0``.

src/rukh/tokenize/loader.pylíneas 1-13 · p1

IGNORE_INDEX = 0 vale en los tres esquemas gracias a la comprobación del BpeGameEncoder, y es de donde sale el ignore_index de la pérdida de M2.

class PackedDataset(Dataset[tuple[torch.Tensor, torch.Tensor]]):
    """Windows of ``block`` tokens over ``tokens.npy``.

    With ``start_at_game`` every item starts at a ``<bos>`` (item ``i`` is game ``i``);
    otherwise the stream is cut into consecutive blocks. A window that runs past the end of
    the stream is padded with ``<pad>`` (0). Each item is ``(x, y)`` with ``y = x[1:] + [<pad>]``.
    """

    def __init__(self, directory: Path, block: int = 200, start_at_game: bool = True) -> None:
        if block < 2:
            raise ValueError("block must be at least 2")
        self.directory = Path(directory)
        self.block = block
        self.start_at_game = start_at_game
        self.info: PackInfo = read_pack_info(self.directory)
        self.tokens = np.load(self.directory / TOKENS_FILE, mmap_mode="r")
        self.starts = np.load(self.directory / STARTS_FILE)
        self.pad_id = 0

    def __len__(self) -> int:
        if self.start_at_game:
            return int(len(self.starts))
        return max(1, int(len(self.tokens)) // self.block) if len(self.tokens) else 0

    def window(self, index: int) -> np.ndarray:
        """``block + 1`` tokens from the window origin, padded with ``<pad>`` past the end."""
        if index < 0 or index >= len(self):
            raise IndexError(index)
        origin = int(self.starts[index]) if self.start_at_game else index * self.block
        chunk = np.asarray(self.tokens[origin : origin + self.block + 1], dtype=np.int64)
        if len(chunk) < self.block + 1:
            chunk = np.concatenate(
                [chunk, np.full(self.block + 1 - len(chunk), self.pad_id, dtype=np.int64)]
            )
        return chunk

    def __getitem__(self, index: int) -> tuple[torch.Tensor, torch.Tensor]:
        chunk = self.window(index)
        x = torch.from_numpy(chunk[: self.block].copy())
        y = torch.from_numpy(chunk[1 : self.block + 1].copy())
        y[-1] = self.pad_id
        return x, y

src/rukh/tokenize/loader.pylíneas 16-57 · p1

PackedDataset es el contrato entre los datos y el modelo. Con mmap_mode="r", np.load no lee el fichero: lo mapea, y el sistema operativo trae las páginas cuando se tocan. Los 480 MB no ocupan memoria del proceso y varios workers comparten las mismas páginas, el libro de consulta de la biblioteca de la primera lección.

__len__ cuenta partidas, así que barajar es barajar partidas. Con start_at_game=False el mismo dataset corta bloques sin alinear, la variante de nanoGPT, para poder medirla en M2 en vez de discutirla.

window pide block + 1 tokens porque y[t] es el token que viene después de x[t]. Solo la última partida del flujo necesita relleno. y[-1] va a <pad> porque su objetivo real está fuera de la ventana, y es más honesto no entrenarlo que inventarlo. El paso a int64 es lo que esperan la entropía cruzada y nn.Embedding, la primera capa del modelo: una tabla con un vector aprendido por id, el embedding, que se consulta con el id como índice. El .copy() es obligatorio porque un trozo de memmap es de solo lectura.

def make_loader(
    dataset: PackedDataset,
    batch_size: int,
    seed: int = 0,
    workers: int = 0,
    shuffle: bool = True,
    drop_last: bool = True,
) -> DataLoader[tuple[torch.Tensor, torch.Tensor]]:
    """A ``DataLoader`` whose shuffling is fully determined by ``seed``."""
    generator = torch.Generator()
    generator.manual_seed(seed)
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        drop_last=drop_last,
        num_workers=workers,
        generator=generator,
        persistent_workers=workers > 0,
    )

src/rukh/tokenize/loader.pylíneas 60-79 · p1

workers llega al DataLoader como num_workers: cuántos procesos aparte preparan lotes mientras la GPU calcula el anterior. Con 0, el valor por defecto, los prepara el propio proceso, que es lo que quieres en un test; en M2 van cuatro.

drop_last=True hace que todos los lotes tengan la misma forma. M2 acelera el modelo con la función torch.compile, que genera código de GPU especializado en las formas que ve, y un lote corto al final de cada época dispararía una recompilación. El torch.Generator sembrado es lo único que decide el orden, así que la misma semilla da los mismos lotes (una semilla fija el azar). Y persistent_workers evita arrancar intérpretes nuevos en cada época, que con spawn cuesta segundos de GPU parada.

// EjecútaloUn lote, de cerca

Con pack.py y loader.py escritos ya puedes ver un lote de verdad, sin esperar a los packs grandes: empaqueta las veinte partidas de tests/fixtures/games.parquet (la fixture de la lección del pipeline), saca un lote, pásalo por un embedding y calcula la pérdida que tendría un modelo que no sabe nada. Guarda esto como lote.py dentro de .tmp/, que está en el .gitignore:

import torch

from rukh.tokenize.loader import PackedDataset, make_loader
from rukh.tokenize.pack import UciGameEncoder, pack_month
from rukh.tokenize.uci_vocab import UciTokenizer

pack_month("tests/fixtures/games.parquet", UciGameEncoder(), ".tmp/pack-prueba")
xb, yb = next(iter(make_loader(PackedDataset(".tmp/pack-prueba"), batch_size=4)))
print(xb.shape, xb.dtype)
print(UciTokenizer().decode(xb[0, :8].tolist()))
print(xb[0, :8].tolist(), yb[0, :8].tolist())

emb = torch.nn.Embedding(2030, 8)
print(emb(xb).shape)

torch.manual_seed(0)
logits = torch.randn(4, 200, 2030)
loss = torch.nn.CrossEntropyLoss(ignore_index=0)(logits.view(-1, 2030), yb.view(-1))
print(round(loss.item(), 2))

uv run python .tmp/lote.py

La primera línea es torch.Size([4, 200]) torch.int64: cuatro ventanas de 200 ids. La segunda decodifica los ocho primeros de la primera ventana: <bos>, los dos tramos de Elo y cinco jugadas UCI. En la tercera, y es x desplazada una posición: cada objetivo es el id que viene después. emb(xb) tiene forma (4, 200, 8), un vector de ocho números por id; en M2 son 512. La pérdida ronda 8,1: un modelo que reparte la probabilidad por igual entre 2 030 ids paga ln 2030 ≈ 7,6 por token, y el ruido de los logits al azar la sube un poco. De ahí arranca la curva de pérdida de M2. Lo que hace cada una de esas llamadas por dentro está en PyTorch mínimo.

`run.py`: lo que hace `rukh data tokenize`

run.py es el orquestador de rukh data tokenize. Casi todo es fontanería que puedes saltar en la primera lectura; lo que enseña es compute_stats, la tabla de lo que cuesta cada esquema. Sin los dos pliegues el comando no existe, así que si no los tecleas, trae run.py entero antes de empezarlo, con git checkout p1 -- src/rukh/tokenize/run.py, y lee compute_stats sobre tu copia.

// Fontaneríarun.py, primera parte: el informe, el BPE que se carga o se reentrena y los artefactos para la web

"""Orchestration behind ``rukh data tokenize``: artifacts, parity fixture, statistics, packing."""

from __future__ import annotations

import json
from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np
from pydantic import BaseModel, ConfigDict

from rukh.data.pipeline import TokenizeConfig
from rukh.paths import resolve
from rukh.tokenize.fixture import build_fixture, write_fixture
from rukh.tokenize.san_chars import SanCharTokenizer, san_text, uci_to_san
from rukh.tokenize.uci_vocab import UciTokenizer

if TYPE_CHECKING:
    from tokenizers import Tokenizer

SCHEMES = ("uci", "san", "bpe")


class TokenizeReport(BaseModel):
    """What ``run`` wrote, as absolute POSIX paths."""

    model_config = ConfigDict(extra="forbid")

    scheme: str
    vocab_size: int
    written: list[str]

src/rukh/tokenize/run.pylíneas 1-31 · p1

def _bpe_path(cfg: TokenizeConfig) -> Path:
    return resolve(cfg.artifacts_dir) / "bpe.json"


def _ensure_bpe(cfg: TokenizeConfig, games: Path, retrain: bool) -> Tokenizer:
    """Load ``bpe.json``; train it from ``games`` when missing or ``retrain``."""
    from rukh.tokenize.bpe import load_bpe, train_bpe_from_parquet

    path = _bpe_path(cfg)
    if retrain or not path.is_file():
        return train_bpe_from_parquet(games, path, cfg.bpe_vocab_size, cfg.bpe_train_games)
    return load_bpe(path)


def export_artifacts(cfg: TokenizeConfig, bpe: Tokenizer | None) -> list[Path]:
    """Write ``vocab.json`` and ``fixtures/games.json`` (with ``bpe_ids`` when ``bpe`` given)."""
    artifacts = resolve(cfg.artifacts_dir)
    tokenizer = UciTokenizer()
    vocab_path = artifacts / "vocab.json"
    tokenizer.export(vocab_path)
    entries = build_fixture(
        resolve(cfg.fixture_pgn), tokenizer, SanCharTokenizer(), bpe, max_len=cfg.max_len
    )
    fixture_path = artifacts / "fixtures" / "games.json"
    write_fixture(entries, fixture_path)
    return [vocab_path, fixture_path]

src/rukh/tokenize/run.pylíneas 34-59 · p1

_ensure_bpe: con --scheme bpe se reentrena; si solo hace falta uno (estadísticas, bpe_ids de la fixture) se carga el que haya. De ahí el orden de los comandos: primero --scheme bpe y al final el uci con --export-fixture. export_artifacts escribe los dos ficheros que copian rukh-lab y rukh-web.

def _summary(lengths: list[int], max_len: int) -> dict[str, float]:
    arr = np.asarray(lengths, dtype=np.int64)
    return {
        "mean": round(float(arr.mean()), 2),
        "p50": round(float(np.percentile(arr, 50)), 2),
        "p95": round(float(np.percentile(arr, 95)), 2),
        "pct_le_max_len": round(float((arr <= max_len).mean() * 100), 2),
    }


def compute_stats(cfg: TokenizeConfig, games: Path, bpe: Tokenizer) -> dict[str, object]:
    """Tokens per game for the three schemes over the first ``stats_n_games`` of ``games``.

    Every count is what the scheme's encoder really emits (``rukh.tokenize.pack``): the UCI
    scheme frames a game with five tokens (``<bos>``, two Elo tokens, result, ``<eos>``), BPE
    with three (``<bos>``, result, ``<eos>``) and char-level SAN with two, because there the
    result is part of the text.
    """
    import polars as pl

    from rukh.tokenize.bpe import bpe_text, longest_tokens

    frame = (
        pl.scan_parquet(Path(games).as_posix())
        .select("uci", "white_elo", "black_elo", "result")
        .head(cfg.stats_n_games)
        .collect()
    )
    uci_tok = UciTokenizer()
    san_tok = SanCharTokenizer()
    uci_len: list[int] = []
    san_len: list[int] = []
    bpe_len: list[int] = []
    for uci, w_elo, b_elo, result in frame.iter_rows():
        uci_len.append(len(uci_tok.encode_game(uci, w_elo, b_elo, result, max_len=1 << 30)))
        san_len.append(len(san_tok.encode(san_text(uci_to_san(uci), result))))
        bpe_len.append(len(bpe.encode(bpe_text(uci)).ids) + 3)
    return {
        "source": {"games": Path(games).name, "n_games": frame.height, "max_len": cfg.max_len},
        "schemes": {
            "uci": {"vocab_size": len(uci_tok), **_summary(uci_len, cfg.max_len)},
            "san": {"vocab_size": len(san_tok), **_summary(san_len, cfg.max_len)},
            "bpe": {"vocab_size": bpe.get_vocab_size(), **_summary(bpe_len, cfg.max_len)},
        },
        "bpe_longest_tokens": longest_tokens(bpe, 10),
    }

src/rukh/tokenize/run.pylíneas 62-107 · p1

Aquí sale la tabla de estadísticas del módulo. Cada esquema cuenta los tokens que su codificador emite de verdad, encuadre incluido (el + 3 del BPE), porque eso es lo que va a pagar el entrenamiento. La media es el coste de una época, el p95 la ventana que necesitas para no truncar casi nada y el porcentaje, lo que cubre el bloque del loader.

Lo que salió medido sobre las primeras 20 000 partidas de enero:

Esquema	Vocabulario	Tokens/partida (media)	p50	p95	% ≤ 200 tokens
UCI (vocabulario fijo)	2030	80,64	76	139	99,8
SAN (carácter)	35	422,44	393	754	7,77
BPE	4096	68,05	62	129	99,84

// tokens BPE más largos

e2e4c7c5g1f3d7d6d2d4c5d4f3d4g8f6b1c3a7a6 · 10 jugadas
e2e4e7e6d2d4d7d5e4e5c7c5c2c3b8c6g1f3d8b6 · 10 jugadas
e2e4c7c5g1f3d7d6d2d4c5d4f3d4g8f6b1c3 · 9 jugadas
e2e4e7e6d2d4d7d5e4e5c7c5c2c3b8c6g1f3 · 9 jugadas
e2e4c7c5g1f3b8c6d2d4c5d4f3d4g8f6b1c3 · 9 jugadas
e2e4c7c5g1f3b8c6d2d4c5d4f3d4e7e5d4b5 · 9 jugadas
e2e4e7e5g1f3b8c6f1c4g8f6f3g5d7d5e4d5 · 9 jugadas
d2d4g8f6c2c4g7g6b1c3f8g7e2e4d7d6 · 8 jugadas
e2e4e7e6g1f3d7d5e4d5e6d5d2d4g8f6 · 8 jugadas
e2e4c7c6g1f3d7d5e4d5c6d5d2d4b8c6 · 8 jugadas

Fuente: rukh data tokenize --stats · 20.000 partidas

UCI fijo: media 80,64 tokens por partida, p50 76, p95 139 y el 99,8 % por debajo de 200.
SAN por carácter: media 422,44, p50 393, p95 754 y solo el 7,77 % cabe en 200.
BPE de 4 096: media 68,05, p50 62, p95 129 y el 99,84 % por debajo de 200.

El 5,2× entre UCI y SAN se repite sobre el corpus entero (240 068 954 tokens contra 1 258 749 609 para las mismas 2 949 514 partidas de enero), y se paga en tiempo de GPU: con el presupuesto fijo de M2, 20 000 pasos × 51 200 tokens, UCI hace algo más de cuatro épocas sobre enero y SAN no completa una. Como la atención crece con el cuadrado de la longitud, igualar las épocas costaría unas 27 veces más, no 5,2.

Escribir el JSON, ordenar los pasos y exponerlos en la CLI no añade ninguna idea: plegado. El «Ejecútalo» de los tokens lo necesita entero; si no lo tecleas, run.py llega con el git checkout de arriba y el subcomando tokenize de cli.py se copia del pliegue, porque un checkout de cli.py pisaría lo que llevas escrito.

// Fontaneríarun.py, el resto, y el subcomando rukh data tokenize de cli.py

def write_stats(cfg: TokenizeConfig, games: Path, bpe: Tokenizer) -> Path:
    path = resolve(cfg.web_artifacts_dir) / "tokenizer-stats.json"
    path.parent.mkdir(parents=True, exist_ok=True)
    stats = compute_stats(cfg, games, bpe)
    path.write_text(json.dumps(stats, indent=2) + "\n", encoding="utf-8", newline="\n")
    return path

src/rukh/tokenize/run.pylíneas 110-115 · p1

Otra vez newline="\n": este JSON lo copia pnpm sync:data al sitio del curso, y con saltos de línea distintos por sistema habría un diff en cada máquina.

def run(
    cfg: TokenizeConfig,
    scheme: str = "uci",
    export_fixture: bool = False,
    stats: bool = False,
    games: Path | None = None,
    pack: bool = False,
) -> TokenizeReport:
    """Run the requested tokenizer step(s) for ``scheme``.

    ``--scheme bpe`` (re)trains ``bpe.json`` from ``games``; ``--stats`` needs a BPE, so it
    trains one from ``games`` when none exists. ``--export-fixture`` adds ``bpe_ids`` to the
    fixture whenever ``bpe.json`` is available. ``--pack`` writes ``train_month`` and
    ``val_month`` as memmap token streams under ``out_dir/<scheme>/``.
    """
    if scheme not in SCHEMES:
        raise ValueError(f"scheme must be one of {SCHEMES}, got {scheme!r}")
    games_path = games or resolve(cfg.stats_games)
    written: list[Path] = []
    bpe: Tokenizer | None = None

    if scheme == "bpe" or stats:
        bpe = _ensure_bpe(cfg, games_path, retrain=scheme == "bpe")
        if scheme == "bpe":
            written.append(_bpe_path(cfg))
    elif (export_fixture or pack) and _bpe_path(cfg).is_file():
        from rukh.tokenize.bpe import load_bpe

        bpe = load_bpe(_bpe_path(cfg))

    if stats and bpe is not None:
        written.append(write_stats(cfg, games_path, bpe))
    if pack:
        from rukh.tokenize.pack import pack_scheme

        written.extend(pack_scheme(cfg, scheme, bpe))
    if export_fixture:
        written.extend(export_artifacts(cfg, bpe))

    vocab_size = len(UciTokenizer())
    if scheme == "san":
        vocab_size = len(SanCharTokenizer())
    elif scheme == "bpe" and bpe is not None:
        vocab_size = bpe.get_vocab_size()
    return TokenizeReport(
        scheme=scheme, vocab_size=vocab_size, written=[p.as_posix() for p in written]
    )

src/rukh/tokenize/run.pylíneas 118-164 · p1

La primera mitad de run decide si hace falta un BPE y de dónde sale; la segunda ejecuta lo pedido en orden (estadísticas, empaquetado, fixture). El elif carga el BPE si existe aunque no se pida, y por eso la fixture lleva bpe_ids y el orden de los comandos importa.

@data_app.command("tokenize")
def data_tokenize(
    config: PipelineOption = None,
    scheme: Annotated[
        str, typer.Option("--scheme", help="Tokenization scheme: uci, san or bpe.")
    ] = "uci",
    export_fixture: Annotated[
        bool,
        typer.Option(
            "--export-fixture",
            help="Write artifacts/tokenizer/vocab.json and fixtures/games.json.",
        ),
    ] = False,
    stats: Annotated[
        bool,
        typer.Option("--stats", help="Write artifacts/web/tokenizer-stats.json for all schemes."),
    ] = False,
    games: Annotated[
        Path | None,
        typer.Option(
            "--games",
            exists=True,
            dir_okay=False,
            help="UCI games parquet for BPE training and statistics (default from config).",
        ),
    ] = None,
    pack: Annotated[
        bool,
        typer.Option("--pack", help="Pack train/val months into data/tokens/<scheme>/."),
    ] = False,
) -> None:
    """Build tokenizer artifacts, the parity fixture, statistics and packed token streams."""
    from rukh.data.pipeline import load_pipeline
    from rukh.tokenize.run import SCHEMES, run

    if scheme not in SCHEMES:
        typer.echo(f"error: --scheme must be one of {', '.join(SCHEMES)}", err=True)
        raise typer.Exit(code=2)
    cfg = load_pipeline(config).tokenize
    report = run(
        cfg, scheme=scheme, export_fixture=export_fixture, stats=stats, games=games, pack=pack
    )
    typer.echo(f"scheme:     {report.scheme}")
    typer.echo(f"vocab_size: {report.vocab_size}")
    _echo_written(report.written)

src/rukh/cli.pylíneas 141-185 · p1

Es el único subcomando del grupo data que no imprime un manifiesto, y comprueba el esquema antes de cargar la configuración para salir con el código 2 de typer en vez de una traza.

`fixture.py`: veinte partidas que el navegador tendrá que reproducir

La fixture de paridad es un contrato entre repositorios, no de modelado; puedes saltarla en la primera lectura, pero la necesita el «Ejecútalo» de debajo: si no la tecleas, tráela antes con git checkout p1 -- src/rukh/tokenize/fixture.py.

// Fontaneríafixture.py: veinte partidas codificadas con los tres esquemas para comparar con TypeScript

"""Parity fixture: 20 PGN games encoded with every scheme, consumed by the TypeScript twin."""

from __future__ import annotations

import json
from pathlib import Path
from typing import TYPE_CHECKING

import chess.pgn

from rukh.tokenize.bpe import bpe_text
from rukh.tokenize.uci_vocab import UciTokenizer

if TYPE_CHECKING:
    from tokenizers import Tokenizer

    from rukh.tokenize.san_chars import SanCharTokenizer

src/rukh/tokenize/fixture.pylíneas 1-17 · p1

def read_pgn_games(path: Path) -> list[dict[str, object]]:
    """Read every game of a PGN file as ``{id, white_elo, black_elo, result, uci, san}``.

    ``san`` is the numbered movetext without comments or result (``1.e4 e5 2.Nf3 ...``) and
    ``uci`` the space-separated UCI moves. Every move is replayed with python-chess, so an
    illegal game raises instead of producing a corrupt fixture.
    """
    games: list[dict[str, object]] = []
    with Path(path).open(encoding="utf-8") as fh:
        while True:
            game = chess.pgn.read_game(fh)
            if game is None:
                break
            if game.errors:
                raise ValueError(f"{path}: {game.headers.get('Site')}: {game.errors[0]}")
            board = game.board()
            uci: list[str] = []
            san: list[str] = []
            for move in game.mainline_moves():
                if move not in board.legal_moves:
                    raise ValueError(f"{path}: illegal move {move.uci()} in {board.fen()}")
                if board.turn == chess.WHITE:
                    san.append(f"{board.fullmove_number}.{board.san(move)}")
                else:
                    san.append(board.san(move))
                uci.append(move.uci())
                board.push(move)
            games.append(
                {
                    "id": game.headers.get("Site", f"game-{len(games) + 1}"),
                    "white_elo": int(game.headers["WhiteElo"]),
                    "black_elo": int(game.headers["BlackElo"]),
                    "result": game.headers["Result"],
                    "uci": " ".join(uci),
                    "san": " ".join(san),
                }
            )
    return games

src/rukh/tokenize/fixture.pylíneas 20-57 · p1

read_pgn_games lanza ante un error o una jugada ilegal, al revés que san_to_uci. Allí una partida mala entre seis millones no podía tumbar el lote; aquí son veinte partidas versionadas, y una mala significa una fixture corrupta, que es peor que ninguna: las dos implementaciones coincidirían en reproducir la corrupción. El id es la clave con la que TypeScript empareja sus resultados.

def build_fixture(
    pgn_path: Path,
    uci_tokenizer: UciTokenizer,
    san_tokenizer: SanCharTokenizer | None = None,
    bpe: Tokenizer | None = None,
    max_len: int = 200,
) -> list[dict[str, object]]:
    """Encode the PGN games with the fixed vocabulary and, when given, char-level SAN and BPE."""
    entries = read_pgn_games(pgn_path)
    for entry in entries:
        entry["uci_ids"] = uci_tokenizer.encode_game(
            str(entry["uci"]),
            int(entry["white_elo"]),  # type: ignore[arg-type]
            int(entry["black_elo"]),  # type: ignore[arg-type]
            str(entry["result"]),
            max_len=max_len,
        )
        if san_tokenizer is not None:
            entry["san_ids"] = san_tokenizer.encode(f"{entry['san']} {entry['result']}")
        if bpe is not None:
            entry["bpe_ids"] = bpe.encode(bpe_text(str(entry["uci"]))).ids
    return entries


def write_fixture(entries: list[dict[str, object]], path: Path) -> None:
    """Write the fixture deterministically (``indent=0``, ASCII, trailing newline)."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    text = json.dumps(entries, indent=0, ensure_ascii=True) + "\n"
    path.write_text(text, encoding="utf-8", newline="\n")

src/rukh/tokenize/fixture.pylíneas 60-89 · p1

san_ids y bpe_ids son opcionales, para poder generar la fixture sin tokenizers instalado. write_fixture repite el determinismo de UciTokenizer.export, porque el fichero se compara byte a byte entre tres repositorios.

// FontaneríaCódigo completo: tests/unit/test_bpe.py

"""Tests for rukh.tokenize.bpe: training, exact round trip, saved file and statistics."""

from __future__ import annotations

import json
from pathlib import Path

import pytest
from tokenizers import Tokenizer
from typer.testing import CliRunner

from rukh.cli import app
from rukh.tokenize.bpe import (
    bpe_text,
    decode,
    encode,
    load_bpe,
    longest_tokens,
    save,
    split_moves,
    train_bpe,
)
from rukh.tokenize.fixture import read_pgn_games
from rukh.tokenize.uci_vocab import UciTokenizer

pytestmark = pytest.mark.unit


@pytest.fixture(scope="module")
def games(repo_root: Path) -> list[str]:
    return [str(g["uci"]) for g in read_pgn_games(repo_root / "tests" / "fixtures" / "games.pgn")]


@pytest.fixture(scope="module")
def trained(games: list[str]) -> Tokenizer:
    return train_bpe((bpe_text(g) for g in games), vocab_size=600)


@pytest.fixture(scope="module")
def repo_root() -> Path:
    return Path(__file__).resolve().parents[2]


def test_specials_and_config(trained: Tokenizer) -> None:
    assert trained.get_vocab_size() == 600
    for i, token in enumerate(UciTokenizer.specials):
        assert trained.token_to_id(token) == i
    payload = json.loads(trained.to_str())
    assert payload["model"]["type"] == "BPE"
    assert payload["model"]["unk_token"] == "<unk>"
    assert payload["model"]["continuing_subword_prefix"] is None
    assert payload["model"]["end_of_word_suffix"] is None
    assert payload["pre_tokenizer"]["type"] == "WhitespaceSplit"
    assert payload["normalizer"] is None
    assert payload["decoder"] is None


def test_round_trip_is_exact(trained: Tokenizer, games: list[str]) -> None:
    for uci in games:
        ids = encode(trained, bpe_text(uci))
        assert decode(trained, ids) == uci
        assert UciTokenizer.unk_id not in ids


def test_learns_multi_move_tokens(trained: Tokenizer) -> None:
    longest = longest_tokens(trained, 10)
    assert len(longest) == 10
    assert longest[0]["moves"] >= 2
    assert longest[0]["length"] == len(longest[0]["token"])


def test_save_and_load(tmp_path: Path, trained: Tokenizer, games: list[str]) -> None:
    path = tmp_path / "bpe.json"
    save(trained, path)
    again = Tokenizer.from_file(str(path))
    assert again.encode(bpe_text(games[3])).ids == trained.encode(bpe_text(games[3])).ids
    assert load_bpe(path).get_vocab_size() == 600


def test_split_moves_handles_promotion_ambiguity() -> None:
    assert split_moves("e2e4e7e5") == ["e2e4", "e7e5"]
    # 4...Bxc1 is a bishop move from rank 2 to rank 1 followed by 5.b4: the 'b' is a file.
    bishop = "d2d3 e7e6 e1d2 f8b4 d2e3 b4d2 e3f3 d2c1 b2b4"
    assert split_moves(bishop.replace(" ", "")) == bishop.split()
    # 5.gxh8=B then 5...Nc6: the first 'b' is the promotion piece, the second a file.
    promo = "h2h4 g7g5 h4g5 h7h6 g5g6 f7f6 g6g7 f6f5 g7h8b b8c6"
    assert split_moves(promo.replace(" ", "")) == promo.split()
    # After an illegal move the split falls back to the rank rule.
    assert split_moves("e2e5e7e8qa1a2") == ["e2e5", "e7e8q", "a1a2"]


def test_committed_bpe_matches_fixture(repo_root: Path) -> None:
    bpe = load_bpe(repo_root / "artifacts" / "tokenizer" / "bpe.json")
    fixture = json.loads(
        (repo_root / "artifacts" / "tokenizer" / "fixtures" / "games.json").read_text("utf-8")
    )
    for entry in fixture:
        assert entry["bpe_ids"] == encode(bpe, bpe_text(entry["uci"])), entry["id"]
        assert decode(bpe, entry["bpe_ids"]) == entry["uci"]


def test_stats_lengths_match_the_encoders(repo_root: Path, trained: Tokenizer) -> None:
    """The statistics count exactly what each encoder emits (BPE frames with three tokens)."""
    import polars as pl

    from rukh.data.pipeline import TokenizeConfig
    from rukh.tokenize.pack import BpeGameEncoder, SanGameEncoder, UciGameEncoder
    from rukh.tokenize.run import compute_stats

    games = repo_root / "tests" / "fixtures" / "games.parquet"
    stats = compute_stats(TokenizeConfig(), games, trained)
    frame = pl.read_parquet(games.as_posix()).select("uci", "white_elo", "black_elo", "result")
    encoders = {
        "uci": UciGameEncoder(),
        "san": SanGameEncoder(),
        "bpe": BpeGameEncoder(trained),
    }
    for scheme, encoder in encoders.items():
        lengths = [
            len(encoder.encode_game(uci, int(w), int(b), result, max_len=1 << 30))
            for uci, w, b, result in frame.iter_rows()
        ]
        assert stats["schemes"][scheme]["mean"] == round(sum(lengths) / len(lengths), 2), scheme


def test_cli_stats_and_bpe_training(rukh_home: Path, repo_root: Path) -> None:
    games = (repo_root / "tests" / "fixtures" / "games.parquet").as_posix()
    pgn = (repo_root / "tests" / "fixtures" / "games.pgn").as_posix()
    cfg = rukh_home / "pipeline.yaml"
    cfg.write_text(
        f"tokenize:\n  fixture_pgn: {pgn}\n  bpe_vocab_size: 300\n  bpe_train_games: 20\n",
        encoding="utf-8",
    )
    result = CliRunner().invoke(
        app,
        [
            "data",
            "tokenize",
            "--config",
            str(cfg),
            "--scheme",
            "bpe",
            "--stats",
            "--export-fixture",
            "--games",
            games,
        ],
    )
    assert result.exit_code == 0, result.output
    assert "vocab_size: 300" in result.output
    bpe_path = rukh_home / "artifacts" / "tokenizer" / "bpe.json"
    assert Tokenizer.from_file(str(bpe_path)).get_vocab_size() == 300
    stats = json.loads((rukh_home / "artifacts" / "web" / "tokenizer-stats.json").read_text())
    assert set(stats["schemes"]) == {"uci", "san", "bpe"}
    assert stats["schemes"]["uci"]["vocab_size"] == 2030
    assert stats["schemes"]["san"]["vocab_size"] == 35
    assert stats["schemes"]["bpe"]["vocab_size"] == 300
    assert stats["schemes"]["san"]["mean"] > stats["schemes"]["uci"]["mean"]
    assert stats["source"]["n_games"] == 20
    assert len(stats["bpe_longest_tokens"]) == 10
    fixture = json.loads(
        (rukh_home / "artifacts" / "tokenizer" / "fixtures" / "games.json").read_text("utf-8")
    )
    assert all("bpe_ids" in e and "san_ids" in e for e in fixture)

tests/unit/test_bpe.pylíneas 1-164 · p1

// EjecútaloLa enumeración y el BPE, fijados por sus tests

Con fixture.py escrito ya se pueden pasar los dos ficheros de test de la lección anterior. Los dos leen tests/fixtures/games.pgn, las veinte partidas de la fixture de paridad, elegidas para cubrir los tres resultados, los enroques de los dos bandos y una promoción a caballo; tráelo de p1 con git checkout. Teclea también tests/unit/test_bpe.py, plegado justo encima (o tráelo con git checkout p1 -- tests/unit/test_bpe.py): de sus ocho tests, test_split_moves_handles_promotion_ambiguity es el caso del callout de split_moves, con un alfil que acaba en la columna b y una promoción a alfil seguida de una jugada desde la b.

git checkout p1 -- tests/fixtures/games.pgn
uv run pytest -m unit -q tests/unit/test_uci_vocab.py tests/unit/test_bpe.py -k "not committed"

Treinta y un tests en verde (veinticuatro de uci_vocab, contando los diez casos de elo_bin, y siete de bpe) y tres deseleccionados. Entran los que pasan por rukh data tokenize y los de estadísticas, porque run.py y pack.py solo necesitan el TokenizeConfig de pipeline.py, que carga con publish todavía comentado. Los tres que se quedan fuera comparan contra los artefactos versionados (vocab.json, bpe.json, fixtures/games.json), que tú generas con la CLI al final de la lección.

Los tests del loader

Los tests fijan el contrato del loader; léelos al teclearlos, no hacen falta para seguir la lección. El «Ejecútalo» del final los pasa: si no los tecleas, tráelos con git checkout p1 -- tests/unit/test_loader.py.

// Teststests/unit/test_loader.py: ventanas alineadas, relleno y reproducibilidad

"""Tests for rukh.tokenize.loader: <bos>-aligned windows, padding, determinism."""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest
import torch

from rukh.tokenize.loader import IGNORE_INDEX, PackedDataset, make_loader
from rukh.tokenize.pack import UciGameEncoder, pack_month
from rukh.tokenize.uci_vocab import UciTokenizer

pytestmark = pytest.mark.unit


@pytest.fixture(scope="module")
def pack_dir(tmp_path_factory: pytest.TempPathFactory) -> Path:
    root = Path(__file__).resolve().parents[2]
    out = tmp_path_factory.mktemp("pack")
    pack_month(root / "tests" / "fixtures" / "games.parquet", UciGameEncoder(), out)
    return out

tests/unit/test_loader.pylíneas 1-23 · p1

La fixture empaqueta de verdad, con pack_month, así que los tests del loader prueban las dos piezas juntas sobre un pack real.

def test_items_are_bos_aligned_windows(pack_dir: Path) -> None:
    ds = PackedDataset(pack_dir, block=200)
    assert len(ds) == 20
    for i in range(len(ds)):
        x, y = ds[i]
        assert x.shape == (200,) and y.shape == (200,)
        assert x.dtype == torch.int64
        assert x[0] == UciTokenizer.bos_id
        assert torch.equal(y[:-1], x[1:])
        assert y[-1] == IGNORE_INDEX


def test_last_window_is_padded(pack_dir: Path) -> None:
    ds = PackedDataset(pack_dir, block=400)
    x, y = ds[len(ds) - 1]
    tokens = np.load(pack_dir / "tokens.npy")
    tail = len(tokens) - int(ds.starts[-1])
    assert tail < 400
    assert x[tail - 1] == UciTokenizer.eos_id
    assert torch.all(x[tail:] == IGNORE_INDEX)
    assert torch.all(y[tail - 1 :] == IGNORE_INDEX)


def test_contiguous_blocks(pack_dir: Path) -> None:
    ds = PackedDataset(pack_dir, block=64, start_at_game=False)
    tokens = np.load(pack_dir / "tokens.npy")
    assert len(ds) == len(tokens) // 64
    x, _ = ds[1]
    assert torch.equal(x, torch.from_numpy(tokens[64:128].astype(np.int64)))
    with pytest.raises(IndexError):
        ds[len(ds)]

tests/unit/test_loader.pylíneas 26-56 · p1

El primer test comprueba en cada partida las cinco promesas que el entrenamiento de M2 da por hechas: forma, tipo, <bos> al principio, y[:-1] == x[1:] y relleno al final. test_last_window_is_padded usa un bloque mayor que lo que queda, calcula cuántos tokens reales hay desde el último <bos> y comprueba que después del <eos> todo es relleno, sin escribir el número a mano.

def test_same_seed_same_batches(pack_dir: Path) -> None:
    ds = PackedDataset(pack_dir, block=200)
    a = make_loader(ds, batch_size=4, seed=7)
    b = make_loader(ds, batch_size=4, seed=7)
    c = make_loader(ds, batch_size=4, seed=8)
    batches_a = [x for x, _ in a]
    batches_b = [x for x, _ in b]
    batches_c = [x for x, _ in c]
    assert len(batches_a) == 5
    assert all(torch.equal(p, q) for p, q in zip(batches_a, batches_b, strict=True))
    assert any(not torch.equal(p, q) for p, q in zip(batches_a, batches_c, strict=True))
    assert batches_a[0].shape == (4, 200)


def test_thousand_samples_in_range(pack_dir: Path) -> None:
    ds = PackedDataset(pack_dir, block=200)
    rng = np.random.default_rng(0)
    for i in rng.integers(0, len(ds), size=1000):
        x, y = ds[int(i)]
        assert int(x.max()) < 2030 and int(y.max()) < 2030
        assert int(x.min()) >= 0 and int(y.min()) >= 0

tests/unit/test_loader.pylíneas 59-79 · p1

La reproducibilidad va otra vez con sus dos mitades, porque la segunda (semilla distinta, algún lote distinto) es la que detectaría un generador ignorado. test_thousand_samples_in_range es un test de propiedad: no prueba un caso concreto, sino que ninguna de mil ventanas al azar se sale del vocabulario.

// FontaneríaCódigo completo: tests/unit/test_pack.py

"""Tests for rukh.tokenize.pack: stream layout, index, metadata and the three encoders."""

from __future__ import annotations

import hashlib
import json
from pathlib import Path

import numpy as np
import pytest

from rukh.tokenize.bpe import load_bpe
from rukh.tokenize.pack import (
    STARTS_FILE,
    TOKENS_FILE,
    BpeGameEncoder,
    SanGameEncoder,
    UciGameEncoder,
    make_encoder,
    pack_month,
    read_pack_info,
)
from rukh.tokenize.uci_vocab import UciTokenizer

pytestmark = pytest.mark.unit


@pytest.fixture
def games_parquet(repo_root: Path) -> Path:
    return repo_root / "tests" / "fixtures" / "games.parquet"


def test_pack_month_layout(tmp_path: Path, games_parquet: Path) -> None:
    info = pack_month(games_parquet, UciGameEncoder(), tmp_path / "pack")
    tokens = np.load(tmp_path / "pack" / "tokens.npy", mmap_mode="r")
    starts = np.load(tmp_path / "pack" / "starts.npy")
    assert tokens.dtype == np.uint16 and starts.dtype == np.int64
    assert info.n_games == 20 and len(starts) == 20
    assert info.n_tokens == len(tokens)
    assert tokens[-1] == UciTokenizer.eos_id
    assert all(tokens[s] == UciTokenizer.bos_id for s in starts)
    assert starts[0] == 0 and np.all(np.diff(starts) > 0)
    # Every game is stored whole: <bos>, two Elo tokens, plies, result, <eos>.
    import polars as pl

    plies = pl.read_parquet(games_parquet.as_posix())["n_plies"].to_list()
    lengths = np.diff(np.append(starts, len(tokens)))
    assert lengths.tolist() == [p + 5 for p in plies]
    assert not (tmp_path / "pack" / "tokens.bin").exists()
    meta = json.loads((tmp_path / "pack" / "meta.json").read_text("utf-8"))
    assert meta["scheme"] == "uci" and meta["vocab_size"] == 2030
    assert len(meta["vocab_hash"]) == 64
    assert read_pack_info(tmp_path / "pack") == info


def test_pack_month_bytes_are_stable(tmp_path: Path, games_parquet: Path, repo_root: Path) -> None:
    """The packed bytes are a published contract: streaming must not change them."""
    expected = {
        "uci": (
            "a9df9aa63875adee99b42d4d1ef8ed2db7ba4dd087c5b0533dc88705bb82eb30",
            "91ffa8971d5552917893774c773245cfd980c0c7db1fe5c103140360cd08b1da",
        ),
        "bpe": (
            "028a4281a0cade6ccbb55abbe68a393354da394d74554caea12626c3afda5a18",
            "cc8eda0b46c069c982b6e9ad1a4ea1ac39dcf37f303a220a292e33919c72d60f",
        ),
    }
    encoders = {
        "uci": UciGameEncoder(),
        "bpe": BpeGameEncoder(load_bpe(repo_root / "artifacts" / "tokenizer" / "bpe.json")),
    }
    for scheme, encoder in encoders.items():
        out = tmp_path / scheme
        pack_month(games_parquet, encoder, out)
        digests = tuple(
            hashlib.sha256((out / name).read_bytes()).hexdigest()
            for name in (TOKENS_FILE, STARTS_FILE)
        )
        assert digests == expected[scheme], scheme


def test_san_and_bpe_encoders(tmp_path: Path, games_parquet: Path, repo_root: Path) -> None:
    san = SanGameEncoder()
    ids = san.encode_game("e2e4 e7e5", 1800, 1800, "1-0", max_len=1000)
    assert ids[0] == 1 and ids[-1] == 2
    assert san.tokenizer.decode(ids) == "1.e4 e5 1-0"
    info = pack_month(games_parquet, san, tmp_path / "san")
    assert info.scheme == "san" and info.vocab_size == 35

    bpe = BpeGameEncoder(load_bpe(repo_root / "artifacts" / "tokenizer" / "bpe.json"))
    ids = bpe.encode_game("e2e4 e7e5", 1800, 1800, "0-1", max_len=1000)
    assert ids[0] == 1 and ids[-2] == 6 and ids[-1] == 2
    info = pack_month(games_parquet, bpe, tmp_path / "bpe")
    assert info.scheme == "bpe" and info.vocab_size == 600
    tokens = np.load(tmp_path / "bpe" / "tokens.npy")
    assert tokens[-1] == 2 and tokens[0] == 1


def test_make_encoder() -> None:
    assert make_encoder("uci").scheme == "uci"
    assert make_encoder("san").scheme == "san"
    with pytest.raises(ValueError):
        make_encoder("bpe")
    with pytest.raises(ValueError):
        make_encoder("nope")


def test_cli_pack_writes_train_and_val(rukh_home: Path, games_parquet: Path) -> None:
    import shutil

    from typer.testing import CliRunner

    from rukh.cli import app

    for month in ("01", "02"):
        target = rukh_home / "data" / "uci" / "year=2025" / f"month={month}" / "games.parquet"
        target.parent.mkdir(parents=True)
        shutil.copy(games_parquet, target)
    result = CliRunner().invoke(app, ["data", "tokenize", "--scheme", "uci", "--pack"])
    assert result.exit_code == 0, result.output
    for split in ("train", "val"):
        info = read_pack_info(rukh_home / "data" / "tokens" / "uci" / split)
        assert info.n_games == 20

tests/unit/test_pack.pylíneas 1-123 · p1

// EjecútaloLos tokens, empaquetados

El comando rukh data tokenize ya funciona entero. Lanza los tres comandos en su orden (las estadísticas y los bpe_ids de la fixture que escribe el tercero salen del bpe.json del primero) y con ellos pasan los tests que comparan contra los artefactos que acabas de generar. Teclea también tests/unit/test_pack.py, plegado justo encima (o tráelo con git checkout p1 -- tests/unit/test_pack.py): fija los bytes de un pack con un hash y comprueba que meta.json es coherente con los arrays (n_tokens == len(tokens), n_games == len(starts)), que es lo primero que hay que mirar si un pack sale raro.

uv run rukh data tokenize --scheme bpe --pack
uv run rukh data tokenize --scheme san --pack   # opcional: hora y media
uv run rukh data tokenize --scheme uci --stats --export-fixture --pack
uv run pytest -m unit -q tests/unit/test_uci_vocab.py tests/unit/test_bpe.py tests/unit/test_san_chars.py tests/unit/test_pack.py tests/unit/test_loader.py --deselect tests/unit/test_pack.py::test_pack_month_bytes_are_stable --deselect tests/unit/test_pack.py::test_san_and_bpe_encoders

Los tres necesitan data/uci/ de la conversión. El BPE tarda unos minutos y el UCI entre cinco y diez; el último imprime vocab_size: 2030 y deja artifacts/tokenizer/vocab.json, fixtures/games.json, artifacts/web/tokenizer-stats.json y data/tokens/uci/. Los tests son 47 en verde y dos deseleccionados, ya sin -k: test_committed_vocab_matches comprueba que tu vocab.json tiene el sha256 del de referencia.

Los dos deseleccionados no prueban tu código sino un fichero que no tienes. En p1, test_pack.py carga artifacts/tokenizer/bpe.json y da por hecho que es el BPE versionado en el repositorio de referencia, uno de 600 entradas entrenado para los tests: exige vocab_size == 600 y fija el hash del pack que produce. El tuyo es el de 4 096 que acaba de entrenar --scheme bpe, así que esos dos fallan aunque todo esté bien; sin --deselect verás 2 failed, 47 passed. Las versiones posteriores de rukh lo arreglaron entrenando dentro del test su propio BPE de 600. No traigas el bpe.json de p1 para ponerlos en verde: pisarías el de 4 096 con el que están hechos tus packs y el que lee el lab 5.

// Ejercicio 01Mide el precio del padding

Escribe una variante PaddedDataset que devuelva una partida por ejemplo, rellenada con <pad> hasta 200 y truncada si es más larga. Con un lote de 64, calcula la fracción de objetivos que son <pad> y compárala con la de PackedDataset. ¿Cuántos tokens “útiles” por lote sirve cada uno?

// SoluciónVer la solución

Con la media medida arriba (80,64 tokens por partida), una partida por fila deja unos 119 <pad> por ejemplo: el 60 % del lote es relleno (lo dibuja la figura 03 de la primera lección). ignore_index lo descarta de la pérdida, pero la atención ya lo ha calculado.

PackedDataset mide 0,005 —medio por ciento— en un lote de 64 × 200 sobre el pack de enero, porque el relleno solo aparece en la última ventana del flujo. Son unas 2,5 veces más posiciones contribuyendo a la pérdida con el mismo presupuesto de GPU.

Para medirlo tú: (yb == 0).float().mean().item() sobre el primer lote de cada uno.

Qué has aprendido

Has escrito seis millones de partidas en dos ficheros de 480 MB, uno por mes, sin tenerlas en memoria, y sirves desde él lotes (x, y) que no gastan la mitad del cálculo en relleno. El hash del vocabulario en meta.json impide que un pack viejo se mezcle con un tokenizador nuevo.

Cómo se mide: los tres comandos, en este orden, dejan los tres packs y los artefactos:

uv run rukh data tokenize --scheme bpe --pack
uv run rukh data tokenize --scheme san --pack
uv run rukh data tokenize --scheme uci --stats --export-fixture --pack

El último imprime vocab_size: 2030 y la lista de lo escrito. El pack de enero con el esquema UCI tiene 2 949 514 partidas y 240 068 954 tokens; el de febrero, 2 946 874 y 238 657 571. uv run pytest -m unit -q tests/unit/test_pack.py tests/unit/test_loader.py pasa en verde salvo los dos tests de test_pack.py que esperan el BPE de 600 de la referencia (los del --deselect de la última caja).

Lo siguiente es sacar todo esto del disco: las cards y el comando que sube seis datasets y un tokenizador a Hugging Face.

Qué vas a construir

pack.py: tres codificadores con la misma forma

loader.py: el contrato con el decoder

run.py: lo que hace rukh data tokenize

fixture.py: veinte partidas que el navegador tendrá que reproducir