Los tokens de Elo y el corpus que faltaba

The months are read **shard by shard, downloaded first**, not streamed over ``hf://``. Streaming
was the original design and it is the obvious one: DuckDB reads the remote parquet directly and
only the rows that pass the filter ever land on disk. It stopped working. A month is 72 files of
about a gigabyte, and reading them remotely means thousands of small range requests against one
host, which earns an ``HTTP 429`` partway through and then a retry schedule that doubles until
the throughput is measured in kilobytes. Measured on 2026-09-20: streaming stalled at 0.02 MB/s
after fifty minutes with nothing written, while downloading the same shard whole ran at 14.8 MB/s.

So each shard is fetched once, with ``huggingface_hub`` (which caches, resumes and backs off
properly), filtered locally into a part file and deleted again. Peak disk is one shard, the work
resumes where it stopped, and ``limit`` now stops the download early instead of only trimming the
result -- for a slice of the rating range that is the difference between 8 shards and 72.

Finally, it signs the requests to ``hf://``. Reading a month of Lichess means a few hundred range
requests against the same host, and anonymous ones get ``HTTP 429 Too Many Requests`` partway
through -- which arrives as a failed ``COPY`` after minutes of work, with no partial file to
resume from. A token raises the allowance and identifies the traffic, and the retry budget is
raised from three attempts to eight because a 429 is a "wait", not a "no".

class DuckDbConfig(BaseConfig):
    """Settings applied to every DuckDB connection (``duckdb:`` in ``pipeline.yaml``)."""

    memory_limit: str = "32GB"
    threads: int = Field(default=0, ge=0)
    """``0`` = ``cpu_count()``."""
    preserve_insertion_order: bool = False
    """Keep the input order in the output. Off: it buffers the whole result of a ``COPY``."""
    http_retries: int = Field(default=8, ge=0)
    """Attempts per HTTP request. The default of 3 is not enough for a rate-limited host."""
    http_retry_wait_ms: int = Field(default=500, ge=0)
    """First backoff, doubled by ``http_retry_backoff`` on each attempt."""
    http_timeout: int = Field(default=120, ge=1)
    """Seconds before a single request is given up on; the default 30 is tight over ``hf://``."""
    authenticate: bool = True
    """Sign ``hf://`` requests with the locally stored Hugging Face token when there is one."""

def hf_token() -> str | None:
    """The Hugging Face token this machine is logged in with, or None.

    Read through ``huggingface_hub`` rather than from the environment so it picks up the token
    that ``huggingface-cli login`` wrote, which is where it usually lives on a workstation.
    """
    try:
        from huggingface_hub import get_token
    except ImportError:  # pragma: no cover - the hub is a hard dependency
        return None
    token = get_token()
    return token or None


def authenticate(con: duckdb.DuckDBPyConnection, token: str | None = None) -> bool:
    """Give the connection a Hugging Face secret; return whether one was installed."""
    secret = token if token is not None else hf_token()
    if not secret:
        log.info("no Hugging Face token found: hf:// reads go out anonymous and rate-limited")
        return False
    con.execute("INSTALL httpfs")
    con.execute("LOAD httpfs")
    con.execute("CREATE OR REPLACE SECRET hf (TYPE huggingface, TOKEN ?)", [secret])
    return True

def connect(cfg: DuckDbConfig | None = None) -> duckdb.DuckDBPyConnection:
    """An in-memory DuckDB connection with the spill directory, memory limit and threads set."""
    import duckdb

    settings = cfg or DuckDbConfig()
    con = duckdb.connect()
    con.execute("SET temp_directory = ?", [temp_directory().as_posix()])
    con.execute("SET memory_limit = ?", [settings.memory_limit])
    con.execute("SET threads = ?", [settings.threads or (os.cpu_count() or 4)])
    con.execute("SET preserve_insertion_order = ?", [settings.preserve_insertion_order])
    con.execute("SET http_retries = ?", [settings.http_retries])
    con.execute("SET http_retry_wait_ms = ?", [settings.http_retry_wait_ms])
    con.execute("SET http_timeout = ?", [settings.http_timeout])
    if settings.authenticate:
        authenticate(con)
    return con

class FetchConfig(BaseConfig):
    """What to fetch and how to filter it."""

    dataset: str = "Lichess/standard-chess-games"
    months: list[str] = Field(min_length=1)
    min_elo: int = Field(default=1800, ge=0)
    max_elo: int | None = Field(default=None, ge=0)
    """Upper bound on both ratings, so a second corpus can cover the band the first one left out.

    The project's whole corpus was fetched with ``min_elo: 1800`` and no ceiling, which means the
    Elo header tokens below ``<w1800>`` were never trained: conditioning the model on 1500 shows it
    a vector that is still at its initialisation. Filling that band needs a fetch that is bounded
    above as well, and a bound that is applied to *both* players for the same reason ``min_elo``
    is -- a 1200 against a 2400 is not a 1200-rated game, it is a mismatch."""

    min_base_seconds: int = Field(default=180, ge=0)
    terminations: list[str] = Field(default=["Normal", "Time forfeit"], min_length=1)
    min_plies: int = Field(default=20, ge=0)
    exclude_variants: bool = True
    out_dir: str = "data/raw"
    limit: int | None = Field(default=None, ge=1)
    """Rows to keep per month. It also stops the download: once the parts hold this many games
    the remaining shards are never fetched."""
    cache_dir: str = "data/hf-cache"
    """Where downloaded shards are parked before being filtered, and deleted from afterwards."""
    keep_shards: bool = False
    """Keep the downloaded shards instead of deleting each one once it has been filtered.
    72 GB per month, so it is off; useful when the same month is going to be re-filtered."""
    duckdb: DuckDbConfig = Field(default_factory=DuckDbConfig)

    @field_validator("months")
    @classmethod
    def _check_months(cls, months: list[str]) -> list[str]:
        for month in months:
            if not MONTH_RE.match(month):
                raise ValueError(f"month {month!r} must look like YYYY-MM")
        return months

    @model_validator(mode="after")
    def _check_elo_band(self) -> FetchConfig:
        if self.max_elo is not None and self.max_elo < self.min_elo:
            raise ValueError(f"max_elo {self.max_elo} is below min_elo {self.min_elo}")
        return self

def where_clause(cfg: FetchConfig) -> str:
    """The filter itself, one condition per line: the only thing that decides what is kept.

    ``TRY_CAST`` drops rows whose base time is not an integer (correspondence games export
    ``TimeControl = '-'``). It is built once and used twice -- in the query the dry run prints
    and against every downloaded shard -- so what is documented is what actually runs.
    """
    conditions = [
        f"WhiteElo >= {cfg.min_elo} AND BlackElo >= {cfg.min_elo}",
        f"TRY_CAST(split_part(TimeControl, '+', 1) AS INTEGER) >= {cfg.min_base_seconds}",
        "Termination IN (" + ", ".join(_sql_string(t) for t in cfg.terminations) + ")",
    ]
    if cfg.max_elo is not None:
        conditions.insert(1, f"WhiteElo <= {cfg.max_elo} AND BlackElo <= {cfg.max_elo}")
    if cfg.exclude_variants:
        conditions.append("Event NOT ILIKE '%variant%'")
    return "\n    AND ".join(conditions)

def month_shards(cfg: FetchConfig, month: str) -> list[str]:
    """Repository paths of the parquet shards of one month, in order."""
    from huggingface_hub import HfApi

    year, mm = _split_month(month)
    prefix = f"data/year={year}/month={mm}/"
    api = HfApi()
    files = api.list_repo_files(cfg.dataset, repo_type="dataset")
    return sorted(name for name in files if name.startswith(prefix) and name.endswith(".parquet"))


def _download(cfg: FetchConfig, shard: str) -> Path:
    """One shard as a real file on disk, resumed if a previous run was cut off partway.

    ``local_dir`` and not ``cache_dir``. The Hub cache stores the bytes under ``blobs/`` and hands
    back a link into ``snapshots/``, so deleting what this returns frees nothing: the first run
    that way left 5.7 GB of shards behind after claiming to have cleaned up. With ``local_dir``
    the path *is* the file, and ``unlink`` means what it says.
    """
    from huggingface_hub import hf_hub_download

    return Path(
        hf_hub_download(
            cfg.dataset,
            shard,
            repo_type="dataset",
            local_dir=paths.resolve(cfg.cache_dir).as_posix(),
        )
    )

def _filter_shard(
    con: Any, cfg: FetchConfig, month: str, local: Path, target: Path, remaining: int | None
) -> int:
    """Write the rows of ``local`` that pass the filter into ``target``; return how many."""
    limit = "" if remaining is None else f"\nLIMIT {remaining}"
    con.execute(
        f"""
        COPY (
          SELECT *, {_sql_string(month)} AS month
          FROM read_parquet({_sql_string(local.as_posix())}, hive_partitioning = false)
          WHERE {where_clause(cfg)}{limit}
        ) TO {_sql_string(target.as_posix())} (FORMAT PARQUET)
        """
    )
    row = con.execute(
        f"SELECT count(*) FROM read_parquet({_sql_string(target.as_posix())})"
    ).fetchone()
    return int(row[0]) if row else 0

def _fetch_month(con: Any, cfg: FetchConfig, month: str, target: Path) -> int:
    """Download, filter and delete the month's shards until ``limit`` is reached.

    Parts are kept next to the target until the month is done, so an interrupted month resumes
    from the shard it stopped at instead of starting over.
    """
    parts_dir = target.parent / "parts"
    parts_dir.mkdir(parents=True, exist_ok=True)
    shards = month_shards(cfg, month)
    if not shards:
        raise FileNotFoundError(f"{cfg.dataset} has no parquet shards for {month}")
    kept = 0
    for index, shard in enumerate(shards):
        part = parts_dir / f"part-{index:05d}.parquet"
        if part.is_file() and part.stat().st_size:
            kept += _count(con, part)
            log.info("part %s/%s already filtered (%s games so far)", index + 1, len(shards), kept)
        else:
            remaining = None if cfg.limit is None else cfg.limit - kept
            if remaining is not None and remaining <= 0:
                break
            local = _download(cfg, shard)
            kept += _filter_shard(con, cfg, month, local, part, remaining)
            if not cfg.keep_shards:
                local.unlink(missing_ok=True)
            log.info(
                "shard %s/%s of %s filtered: %s games so far", index + 1, len(shards), month, kept
            )
        if cfg.limit is not None and kept >= cfg.limit:
            log.info("%s reached the limit of %s games", month, cfg.limit)
            break
    parts = sorted(path for path in parts_dir.glob("part-*.parquet") if path.stat().st_size)
    # Counted in rows, not in bytes: a parquet with no rows still has a header and a footer, so
    # an empty month would sail through a size check and land downstream looking like a corpus.
    if not parts or not kept:
        raise RuntimeError(f"no games of {month} passed the filter")
    files = ", ".join(_sql_string(path.as_posix()) for path in parts)
    # `hive_partitioning = false` again: the parts sit under `year=/month=/parts/`, so without it
    # DuckDB would read the directory names as columns and collide with the explicit `month`.
    con.execute(
        f"COPY (SELECT * FROM read_parquet([{files}], hive_partitioning = false, "
        f"union_by_name = true)) TO {_sql_string(target.as_posix())} (FORMAT PARQUET)"
    )
    total = _count(con, target)
    shutil.rmtree(parts_dir, ignore_errors=True)
    return total

def _is_empty(path: Path) -> bool:
    """Whether the month still has to be fetched: missing, or the stub a failed COPY leaves."""
    return not path.is_file() or path.stat().st_size == 0

    fetch_plan = plan(cfg)
    if dry_run:
        return fetch_plan

    out_dir = Path(fetch_plan.out_dir)
    counts: dict[str, int] = {}
    files: list[FileHash] = []
    con = connect(cfg.duckdb)
    try:
        for month, out_path in zip(cfg.months, fetch_plan.out_paths, strict=True):
            target = out_dir / out_path
            target.parent.mkdir(parents=True, exist_ok=True)
            if not (overwrite or _is_empty(target)):
                log.info("%s is already there, keeping it (pass overwrite=True to refetch)", target)
                counts[month] = _count(con, target)
            else:
                log.info("fetching %s into %s", month, target)
                counts[month] = _fetch_month(con, cfg, month, target)
            files.append(
                FileHash(path=out_path, sha256=_sha256(target), bytes=target.stat().st_size)
            )

# The P4 corpus: the 1000-1799 band that `pipeline-v4.yaml` never had, converted and balanced.
#
# Only the steps P4 needs are configured. `uci` reads `data/raw-low` into its own `data/uci-low`
# so `data/uci` keeps meaning "what the published models trained on"; `elo_bins` is the only step
# that joins the two, and it is the step whose whole job is to teach the Elo header.
duckdb:
  memory_limit: 32GB
  threads: 0
uci:
  raw_dir: data/raw-low
  out_dir: data/uci-low
  min_plies: 20
  max_plies: 300
  workers: 0
elo_bins:
  # Both halves of the axis, in bands of 200 Elo. 200 is a compromise: 100 would put the four
  # club bands at the same resolution as the master ones and then starve them of games, and 400
  # would blur "1500" and "1800" into one token's worth of meaning.
  sources: [data/uci-low, data/uci]
  out_dir: data/elo-bins-v2
  bin_width: 200
  min_bin: 1000
  max_bin: 2600
  # Flat by construction. The real corpus is a pyramid and that is exactly what teaches the model
  # that `<w1500>` is rare rather than what it means; the manifest publishes what each band got,
  # because the top ones cannot fill their quota and the lesson says so.
  n_per_bin: 150000
  seed: 42
tokenize:
  uci_dir: data/uci
  out_dir: data/tokens-elo
  artifacts_dir: artifacts/tokenizer
  web_artifacts_dir: artifacts/web
  fixture_pgn: tests/fixtures/games.pgn
  # No months: the training split is the balanced sample and nothing else, so the fine-tune sees
  # the flat distribution and not the pyramid it came from.
  train_months: []
  extra_train_parquets: ['data/elo-bins-v2/games.parquet']
  val_month: '2025-02'
  val_games: 100000
  # The balanced sample is the whole training split: the rest of the validation month must not
  # join it, or 2.85 M games of 1800+ play would flatten the very shape this corpus exists for.
  val_remainder_trains: false
  stats_games: data/uci/year=2025/month=01/games.parquet
  bpe_vocab_size: 4096
  bpe_train_games: 200000
  stats_n_games: 20000
  max_len: 200

# Token pack of the Elite Database alone, for the masters fine-tune of M4.
#
# The elite games are already part of `tokens-v4`, mixed in with the 1800+ corpus; this pack has
# nothing else in it. That is the whole experiment: `docs/spec/02` predicts that a last fine-tune
# on masters raises the Elo and lowers the diversity, and neither half of that sentence can be
# read off a corpus where masters are one ingredient among several.
#
# Same validation split as every other run -- the first 100 000 games of 2025-02 -- and its
# remainder stays out of training, so what the model is fine-tuned on is masters and only masters.
tokenize:
  uci_dir: data/uci
  out_dir: data/tokens-masters
  train_months: []
  extra_train_parquets: ['data/elite/games.parquet']
  val_month: '2025-02'
  val_games: 100000
  val_remainder_trains: false
  stats_games: data/uci/year=2025/month=01/games.parquet
  max_len: 200

# Lever 2: months beyond the two P1 fetched. Same filters, a higher cap -- the 3,000,000 of
# `lichess-2025-01-02.yaml` was our own ceiling, not the dataset's (7,136,017,970 games).
# `threads` is held below cpu_count() on purpose: this runs while the GPU trains and the
# dataloader needs cores.
dataset: Lichess/standard-chess-games
months: ["2025-03"]
min_elo: 1800
min_base_seconds: 180
terminations: ["Normal", "Time forfeit"]
min_plies: 20
exclude_variants: true
out_dir: data/raw
limit: 6000000
duckdb:
  memory_limit: 16GB
  threads: 3   # 8 tripped HTTP 429 three times running; fewer parallel range requests

@dataclass(frozen=True)
class GameSlice:
    """A contiguous run of games inside one UCI parquet.

    ``skip`` games are dropped from the front and at most ``take`` are kept (``None`` = to the
    end). It exists so one month can feed both splits: the validation slice is the first
    ``val_games`` rows and training gets the remainder, instead of a whole month sitting idle
    as validation when validation only ever reads a few hundred thousand tokens per eval.
    """

    path: Path
    skip: int = 0
    take: int | None = None

    def label(self) -> str:
        """``month=02/games.parquet[100000:]``: what ``meta.json`` records as the source."""
        name = f"{self.path.parent.name}/{self.path.name}"
        if not self.skip and self.take is None:
            return name
        end = "" if self.take is None else str(self.skip + self.take)
        return f"{name}[{self.skip}:{end}]"

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


def pack_games(sources: Sequence[GameSlice], encoder: GameEncoder, out: Path) -> PackInfo:
    """Encode every game of ``sources``, in order, into ``out/tokens.npy`` + ``starts.npy``.

    Each 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. Several sources
    concatenate into a single stream, which is why ``starts.npy`` is written last: its offsets
    are absolute over the concatenation.
    """
    if encoder.vocab_size() > np.iinfo(np.uint16).max + 1:
        raise ValueError("vocabulary does not fit in uint16")
    if not sources:
        raise ValueError("pack_games needs at least one source")

    with raw_path.open("wb") as raw:
        for source in sources:
            reader = pq.ParquetFile(source.path)
            seen = 0
            taken = 0
            for batch in reader.iter_batches(
                batch_size=CHUNK_GAMES, columns=["uci", "white_elo", "black_elo", "result"]
            ):
                rows_here = batch.num_rows
                lo = min(max(source.skip - seen, 0), rows_here)
                hi = rows_here
                if source.take is not None:
                    hi = min(hi, lo + source.take - taken)
                seen += rows_here
                if hi <= lo:
                    if source.take is not None and taken >= source.take:
                        break
                    continue
                batch = batch.slice(lo, hi - lo)
                taken += hi - lo
                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())
                if source.take is not None and taken >= source.take:
                    break

    train_months: list[str] = Field(default_factory=lambda: ["2025-01"])
    val_month: str = "2025-02"
    val_games: int = Field(default=0, ge=0)
    extra_train_parquets: list[str] = Field(default_factory=list)
    """UCI parquets appended to the training split, outside the ``year=/month=`` layout."""
    val_remainder_trains: bool = True
    """Whether the part of ``val_month`` that is not held out joins the training split.

    Off for a corpus whose shape is the point: the Elo-balanced sample of M4 has the same number
    of games per rating band on purpose, and 2.85 M games of rated-1800+ play poured on top would
    undo it."""
    """Games of ``val_month`` held out; the rest of it trains. 0 = the whole month validates."""

def pack_scheme(cfg: TokenizeConfig, scheme: str, bpe: Tokenizer | None = None) -> list[Path]:
    """Pack the ``train`` and ``val`` splits under ``out_dir/<scheme>/{train,val}``."""
    encoder = make_encoder(scheme, bpe)
    written: list[Path] = []
    for split, sources in split_sources(cfg).items():
        out = resolve(cfg.out_dir) / scheme / split
        pack_games(sources, encoder, out)
        written.append(out / META_FILE)
    return written

    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.pad_id = 0
        self._open()

    def _open(self) -> None:
        """Map the two arrays from disk. Called on construction and in every worker process."""
        self.tokens = np.load(self.directory / TOKENS_FILE, mmap_mode="r")
        self.starts = np.load(self.directory / STARTS_FILE, mmap_mode="r")

    def __getstate__(self) -> dict[str, object]:
        """Ship the paths to a worker, never the arrays.

        ``DataLoader`` workers are spawned on Windows, so the dataset is pickled once per worker.
        ``starts`` is one int64 per game: 95 MB at 12.5 M games, 145 MB at 18.9 M, and at 145 MB
        the spawn died with ``UnpicklingError: pickle data was truncated`` before a single step
        ran. Both arrays live in files the worker can map itself (D-067).
        """
        return {k: v for k, v in self.__dict__.items() if k not in ("tokens", "starts")}

    def __setstate__(self, state: dict[str, object]) -> None:
        self.__dict__.update(state)
        self._open()