Puzles y pares de preferencia: dos conjuntos que ya estaban hechos · Rukh

Qué vas a construir

Dos conjuntos de datos que el curso usa mucho y que casi no cuestan nada, porque alguien ya hizo el trabajo caro.

El primero es el banco de puzles: seis millones de posiciones tácticas de Lichess, filtradas por calidad, repartidas en tres tramos de dificultad y partidas en entrenamiento y test con una clave determinista. Su conjunto de test sostiene la tabla única de M2 a M6, así que no puede cambiar por el camino.

El segundo son los pares de preferencia para el DPO de M5: decenas de miles de tríos “posición, jugada buena, jugada mala” que salen de la columna pvs de la lección anterior sin ejecutar Stockfish. Lo único delicado es el signo. Son dos ficheros.

`puzzles.py`: un banco de test que no se mueve

"""Puzzle splits by difficulty band from ``Lichess/chess-puzzles`` (deterministic by id)."""

from __future__ import annotations

import hashlib

import polars as pl
from pydantic import Field

from rukh.config import BaseConfig
from rukh.data.db import DuckDbConfig, connect
from rukh.data.manifest import FileHash, Manifest
from rukh.data.uci import sha256_file
from rukh.paths import resolve

PUZZLES_FILE = "puzzles.parquet"
BANDS: list[tuple[str, int, int | None]] = [
    ("1000-1500", 1000, 1500),
    ("1500-2000", 1500, 2000),
    ("2000+", 2000, None),
]

src/rukh/data/puzzles.pylíneas 1-21 · p1

Los tramos viven en una constante y no en la configuración porque definen el banco de test: si se pudieran cambiar desde un YAML, dos ejecuciones darían dos conjuntos con la misma pinta y distinta composición, y las cifras de M2 y M6 dejarían de ser comparables sin que nada lo dijera.

class PuzzlesConfig(BaseConfig):
    """Remote dataset, quality filters and split sizes."""

    dataset: str = "Lichess/chess-puzzles"
    remote_glob: str = "hf://datasets/{dataset}/**/*.parquet"
    out_dir: str = "data/puzzles"
    max_rating_deviation: int = Field(default=100, ge=0)
    min_plays: int = Field(default=100, ge=0)
    n_test: int = Field(default=2_000, ge=1)
    n_train: int = Field(default=50_000, ge=1)
    seed: int = 42
    duckdb: DuckDbConfig = Field(default_factory=DuckDbConfig)


def band(rating: int) -> str | None:
    """Difficulty band for ``rating``; ``None`` below 1000."""
    for name, lo, hi in BANDS:
        if rating >= lo and (hi is None or rating < hi):
            return name
    return None

src/rukh/data/puzzles.pylíneas 24-43 · p1

Los filtros de calidad no son obvios. RatingDeviation es la incertidumbre de Lichess sobre la dificultad del puzle: con desviación 300, la dificultad declarada es poco más que una suposición. NbPlays >= 100 pide cien intentos para que la tasa de acierto signifique algo. Juntos tiran buena parte del catálogo, a cambio de que el tramo “2000+” tenga puzles que de verdad son de 2000+. Cualquier conjunto etiquetado por personas o por votos tiene la misma pregunta pendiente: cuántas opiniones hacen falta para fiarse de la etiqueta.

def _source(cfg: PuzzlesConfig) -> str:
    """Remote parquet glob (tests monkeypatch this to a local file)."""
    return cfg.remote_glob.format(dataset=cfg.dataset)


def load_filtered(cfg: PuzzlesConfig) -> pl.DataFrame:
    """Read the puzzles that pass the quality filters with DuckDB (predicate pushdown)."""
    con = connect(cfg.duckdb)
    try:
        table = con.execute(
            f"""
            SELECT PuzzleId AS puzzle_id, FEN AS fen, Moves AS moves, Rating AS rating,
                   Themes AS themes
            FROM read_parquet('{_source(cfg)}')
            WHERE RatingDeviation <= {cfg.max_rating_deviation}
              AND NbPlays >= {cfg.min_plays}
              AND Rating >= 1000
            """
        ).arrow()
    finally:
        con.close()
    return pl.from_arrow(table)  # type: ignore[return-value]

src/rukh/data/puzzles.pylíneas 46-67 · p1

Otra vez el predicado empujado al origen: cinco columnas de once y el filtro en el WHERE. El SELECT … AS pasa a snake_case en el mismo sitio donde se lee, así que nadie más tiene que acordarse de que Lichess escribe PuzzleId.

def split_key(puzzle_id: str, seed: int) -> int:
    """Stable 64-bit order key: ``blake2b(seed:puzzle_id)``."""
    digest = hashlib.blake2b(f"{seed}:{puzzle_id}".encode(), digest_size=8).digest()
    return int.from_bytes(digest, "little", signed=False)


def split_frame(frame: pl.DataFrame, cfg: PuzzlesConfig) -> pl.DataFrame:
    """Assign ``band`` and ``split``: per band, first ``n_test`` by key then up to ``n_train``."""
    keyed = frame.with_columns(
        pl.col("rating").map_elements(band, return_dtype=pl.String).alias("band"),
        pl.col("puzzle_id")
        .map_elements(lambda p: split_key(p, cfg.seed), return_dtype=pl.UInt64)
        .alias("key"),
        pl.col("themes").cast(pl.List(pl.String)),
    ).filter(pl.col("band").is_not_null())
    keyed = keyed.sort(["band", "key"]).with_columns(
        pl.int_range(pl.len()).over("band").alias("rank")
    )
    return (
        keyed.with_columns(
            pl.when(pl.col("rank") < cfg.n_test)
            .then(pl.lit("test"))
            .when(pl.col("rank") < cfg.n_test + cfg.n_train)
            .then(pl.lit("train"))
            .otherwise(None)
            .alias("split")
        )
        .filter(pl.col("split").is_not_null())
        .select("puzzle_id", "fen", "moves", "rating", "themes", "band", "split")
    )

src/rukh/data/puzzles.pylíneas 70-99 · p1

La partición se hace por una clave y no al azar. split_key es blake2b(f"{seed}:{puzzle_id}"), y eso da algo que un random.shuffle no da: es estable cuando el dataset crece. Si Lichess publica cien mil puzles nuevos y repites el paso, los que estaban en test siguen en test, porque su clave no depende de los demás. Con un barajado, cada ejecución reparte de nuevo y el banco de test de M6 no es el de M2. La semilla va dentro del hash: cambiarla da una partición alternativa entera, y concatenar texto no desborda como sumar.

// Teoría justaQué es polars, en SQL

polars es una librería de tablas en memoria (un dataframe), escrita en Rust: la usas como usarías un SELECT sobre una tabla local. Cada pl.col(...) es una expresión, no un valor, y split_frame se traduce casi palabra por palabra: with_columns(... .alias("band")) es SELECT *, … AS band; pl.int_range(pl.len()).over("band") es ROW_NUMBER() OVER (PARTITION BY band ORDER BY key) - 1, y pl.when().then().otherwise() es un CASE WHEN. map_elements es la excepción: llama a una función de Python fila a fila, lento pero inevitable para un hash. La versión perezosa, pl.scan_parquet(...)… .collect(), no lee nada hasta el collect(); antes planifica la consulta entera y solo trae las columnas y filas que hacen falta, como el predicado empujado de DuckDB.

Después se ordena por (band, key), se numera dentro de cada tramo y se reparte: los n_test primeros a test, los n_train siguientes a entrenamiento y el resto fuera. Salir del mismo orden garantiza que los dos conjuntos sean disjuntos, y descartar el sobrante deja el tamaño publicado acotado: 2 000 de test y hasta 50 000 de entrenamiento por tramo.

def run(cfg: PuzzlesConfig) -> Manifest:
    """Download, filter, split and write ``out_dir/puzzles.parquet`` plus the manifest."""
    out_dir = resolve(cfg.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    frame = split_frame(load_filtered(cfg), cfg)
    target = out_dir / PUZZLES_FILE
    frame.write_parquet(target.as_posix(), compression="zstd")
    counts = {
        f"{b}/{s}": int(n)
        for b, s, n in frame.group_by("band", "split").len().sort("band", "split").iter_rows()
    }
    manifest = Manifest(
        dataset=cfg.dataset,
        months=[],
        filters={
            "max_rating_deviation": cfg.max_rating_deviation,
            "min_plays": cfg.min_plays,
            "bands": [b for b, _, _ in BANDS],
            "n_test": cfg.n_test,
            "n_train": cfg.n_train,
            "seed": cfg.seed,
            "split_by": "blake2b(seed:PuzzleId)",
        },
        counts=counts,
        files=[
            FileHash(path=PUZZLES_FILE, sha256=sha256_file(target), bytes=target.stat().st_size)
        ],
    )
    (out_dir / "manifest.json").write_text(manifest.model_dump_json(indent=2) + "\n", "utf-8")
    return manifest

src/rukh/data/puzzles.pylíneas 102-131 · p1

Los conteos van por tramo/split ("2000+/test": 2000), y split_by guarda la fórmula de la clave como texto, igual que el tramo de Elo.

`pairs.py`: preferencias gratis

El DPO de M5 necesita tríos: una posición, una respuesta preferida y una rechazada. La columna pvs de la lección anterior ya los tiene: varias líneas por posición, de mejor a peor, con su valor. La primera jugada de la mejor es la elegida; la de una suficientemente peor, la rechazada.

"""DPO pairs from the multi-PV evaluations: best first move versus a clearly worse one.

The lines are ranked with ``rukh.data.scoring``, the same ordering and mate convention the
``evals`` consolidation uses, so ``chosen`` is always the consolidated ``best_move``. The
margin is measured from the side to move (for Black a lower ``cp`` is better). Both moves are
checked for legality with python-chess. One pair per FEN, balanced by phase.
"""

from __future__ import annotations

import logging
from pathlib import Path

import polars as pl
import pyarrow.parquet as pq
from pydantic import Field

from rukh.config import BaseConfig
from rukh.data.manifest import FileHash, Manifest
from rukh.data.scoring import (
    BEST_LINE_DOC,
    MATE_SCORE,
    line_rank_key,
    score_mover,
    score_white,
)
from rukh.data.uci import sha256_file
from rukh.paths import resolve

LOGGER = logging.getLogger(__name__)
PAIRS_FILE = "pairs.parquet"
PHASES = ("opening", "middlegame", "endgame")
BATCH_ROWS = 50_000

src/rukh/data/pairs.pylíneas 1-33 · p1

class PairsConfig(BaseConfig):
    """Input evaluations, margin and the per-phase cap."""

    positions_eval: str = "data/evals/positions-eval.parquet"
    out_dir: str = "data/pairs"
    min_delta_cp: int = Field(default=100, ge=1)
    max_per_phase: int = Field(default=200_000, ge=1)
    seed: int = 42

src/rukh/data/pairs.pylíneas 36-43 · p1

min_delta_cp: 100 es un peón entero. Con un margen menor saldrían más pares y peores: veinte centipeones a profundidad 20 son ruido del motor, y preferir una línea sobre otra por eso es enseñar ruido.

def rank_lines(fen: str, pvs: list[dict[str, object]]) -> list[tuple[str, int, int]]:
    """``(move, score_white, score_mover)`` for every scored line, best line first."""
    turn = fen.split()[1]
    scored: list[tuple[tuple[bool, int, int, str], str, int, int]] = []
    for pv in pvs:
        move = pv.get("move")
        cp, mate, depth = pv.get("cp"), pv.get("mate"), pv.get("depth")
        white = score_white(cp, mate)  # type: ignore[arg-type]
        if not move or white is None:
            continue
        key = line_rank_key(cp, mate, depth, str(move), turn)  # type: ignore[arg-type]
        scored.append((key, str(move), white, score_mover(cp, mate, turn)))  # type: ignore[arg-type]
    scored.sort(key=lambda item: item[0])
    return [(move, white, mover) for _, move, white, mover in scored]

src/rukh/data/pairs.pylíneas 46-59 · p1

rank_lines ordena con line_rank_key, la misma regla que la consolidación, así que la elegida es siempre el best_move del paso anterior. Devuelve los dos valores: desde las blancas, para guardar el cp con el convenio de Lichess, y desde el que mueve, para medir el margen.

def make_pair(
    fen: str, pvs: list[dict[str, object]], min_delta_cp: int
) -> dict[str, object] | None:
    """``chosen``/``rejected`` for one FEN or ``None`` when no line is bad enough or illegal."""
    import chess

    scored = rank_lines(fen, pvs)
    if len(scored) < 2:
        return None
    chosen, cp_chosen, best_mover = scored[0]
    rejected: tuple[str, int, int] | None = None
    for move, white, mover in scored[1:]:
        if move != chosen and best_mover - mover >= min_delta_cp:
            rejected = (move, white, mover)
            break
    if rejected is None:
        return None
    try:
        board = chess.Board(f"{fen} 0 1")
    except ValueError:
        return None
    legal = {m.uci() for m in board.legal_moves}
    if chosen not in legal or rejected[0] not in legal:
        return None
    return {
        "fen": fen,
        "chosen": chosen,
        "rejected": rejected[0],
        "cp_chosen": cp_chosen,
        "cp_rejected": rejected[1],
    }

src/rukh/data/pairs.pylíneas 62-92 · p1

La primera decisión: el margen se mide desde el que mueve (best_mover - mover). Para las blancas coincide con el convenio de Lichess; para las negras está negado.

// Error que enseñaEl signo que construye un dataset perfectamente al revés

Con el margen medido desde las blancas (cp_chosen - cp_rejected >= 100), el dataset saldría bien en las posiciones de las blancas e invertido en las de las negras, porque para ellas un cp más bajo es mejor.

No da ningún síntoma: mismas filas, jugadas legales, margen medio esperado, y la mitad de los pares correctos. El DPO aprende tan contento a mejorar con blancas y a empeorar con negras, y tres módulos después ves un modelo que no mejora sin que nada apunte a este >=. Lo atrapa en un minuto mirar diez pares con turno de negras: cp_chosen tiene que ser menor que cp_rejected. El mismo tropiezo espera en cualquier dato de preferencias que mezcle una métrica donde menos es mejor (una latencia, un error) con otra donde más es mejor.

La segunda: las dos jugadas se verifican legales aunque vengan de un motor, porque cuesta microsegundos (el " 0 1" repone los contadores que python-chess necesita para montar el tablero).

La tercera: un solo par por posición, y la rechazada es la menos mala de las que están suficientemente mal. La peor de todas daría pares demasiado fáciles (distinguir la mejor jugada de una catástrofe enseña poco), y varios pares por posición la ponderarían sin querer.

def build_pairs(evals: Path, min_delta_cp: int) -> pl.DataFrame:
    """One pair per FEN, reading the evaluations parquet batch by batch."""
    reader = pq.ParquetFile(evals)
    rows: list[dict[str, object]] = []
    for batch in reader.iter_batches(batch_size=BATCH_ROWS, columns=["fen", "phase", "pvs"]):
        records = zip(
            batch.column("fen").to_pylist(),
            batch.column("phase").to_pylist(),
            batch.column("pvs").to_pylist(),
            strict=True,
        )
        for fen, ph, pvs in records:
            pair = make_pair(fen, pvs, min_delta_cp)
            if pair is not None:
                pair["phase"] = ph
                rows.append(pair)
    schema = {
        "fen": pl.String,
        "chosen": pl.String,
        "rejected": pl.String,
        "cp_chosen": pl.Int32,
        "cp_rejected": pl.Int32,
        "phase": pl.String,
    }
    return pl.DataFrame(rows, schema=schema)

src/rukh/data/pairs.pylíneas 95-119 · p1

El schema explícito del final hace falta: si no sale ningún par (una cobertura mala en el join), polars no tiene de dónde inferir los tipos, y el error llegaría más abajo hablando de columnas inexistentes en vez de decir que no hay pares.

def phase_counts(frame: pl.DataFrame) -> dict[str, int]:
    """Rows per phase, always with the three keys."""
    return {ph: int(frame.filter(pl.col("phase") == ph).height) for ph in PHASES}


def balance(frame: pl.DataFrame, max_per_phase: int, seed: int) -> pl.DataFrame:
    """Same number of pairs per phase: the smallest phase count, capped at ``max_per_phase``."""
    counts = phase_counts(frame)
    n = min(min(counts.values()), max_per_phase)
    if n == 0:
        return frame.clear()
    parts = [
        frame.filter(pl.col("phase") == ph).sample(n=n, seed=seed, shuffle=True) for ph in PHASES
    ]
    return pl.concat(parts)

src/rukh/data/pairs.pylíneas 122-136 · p1

phase_counts devuelve siempre las tres claves, aunque valgan cero, para que nadie tenga que acordarse de un .get(). balance iguala las fases al tamaño de la más pequeña, los finales (las partidas de 1800 a ritmo rápido se deciden antes). Sin equilibrar, el conjunto sería casi todo aperturas y el modelo alineado con él no habría visto un final. Se muestrea con semilla en vez de coger las primeras, que serían las de las primeras partidas.

def run(cfg: PairsConfig) -> Manifest:
    """Build, balance and write ``out_dir/pairs.parquet`` plus the manifest."""
    evals = resolve(cfg.positions_eval)
    out_dir = resolve(cfg.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    raw = build_pairs(evals, cfg.min_delta_cp)
    candidates = phase_counts(raw)
    empty = [ph for ph, n in candidates.items() if n == 0]
    if empty:
        LOGGER.warning(
            "no DPO candidates for %s: the balanced output is empty (candidates %s). "
            "Check the coverage of %s.",
            ", ".join(empty),
            candidates,
            evals,
        )
    balanced = balance(raw, cfg.max_per_phase, cfg.seed)
    target = out_dir / PAIRS_FILE
    balanced.write_parquet(target.as_posix(), compression="zstd")
    counts = phase_counts(balanced)
    counts["candidates"] = int(raw.height)
    manifest = Manifest(
        dataset="Lichess/chess-position-evaluations",
        months=[],
        filters={
            "min_delta_cp": cfg.min_delta_cp,
            "max_per_phase": cfg.max_per_phase,
            "seed": cfg.seed,
            "mate_score": MATE_SCORE,
            "cp_point_of_view": "white",
            "best_line": BEST_LINE_DOC,
            "legality": "python-chess",
            "candidates_by_phase": candidates,
            "empty_phases": empty,
        },
        counts=counts,
        files=[FileHash(path=PAIRS_FILE, sha256=sha256_file(target), bytes=target.stat().st_size)],
    )
    (out_dir / "manifest.json").write_text(manifest.model_dump_json(indent=2) + "\n", "utf-8")
    return manifest

src/rukh/data/pairs.pylíneas 139-178 · p1

El LOGGER.warning es la única traza del pipeline, y está donde hace falta: si una fase se queda a cero, balance devuelve un dataframe vacío y el paso terminaría en verde con un parquet sin filas. El aviso, y el manifiesto por si se pierde, dicen qué fase faltó. Los counts llevan además candidates, el total antes de equilibrar: si es 69 246 y las fases quedan en 4 629, es que los finales eran 4 629.

Un par está definido por un FEN, que un modelo de secuencias de jugadas no sabe leer. Reconstruir el prefijo de jugadas que llevó a cada posición se hace en M5, donde se necesita.

Los tests de los puzles

Los tests fijan el contrato del banco de puzles; léelos al teclearlos, no hacen falta para seguir la lección. El «Ejecútalo» de los dos pasos los pasa, con test_pairs.py, plegado después: si no los tecleas, tráelos antes con git checkout p1 -- tests/unit/test_puzzles.py tests/unit/test_pairs.py.

// Teststests/unit/test_puzzles.py: filtros, reparto estable y la misma semilla, el mismo banco

"""Tests for rukh.data.puzzles: filters, bands and the deterministic disjoint split."""

from __future__ import annotations

from pathlib import Path

import polars as pl
import pytest

from rukh.data import puzzles as puzzles_module
from rukh.data.puzzles import PuzzlesConfig, band, run, split_frame, split_key

pytestmark = pytest.mark.unit


def test_band_edges() -> None:
    assert band(999) is None
    assert band(1000) == "1000-1500"
    assert band(1499) == "1000-1500"
    assert band(1500) == "1500-2000"
    assert band(2000) == "2000+"
    assert band(3100) == "2000+"


def test_split_key_is_stable() -> None:
    assert split_key("abc12", 42) == split_key("abc12", 42)
    assert split_key("abc12", 42) != split_key("abc12", 43)
    assert split_key("abc12", 42) != split_key("abc13", 42)

tests/unit/test_puzzles.pylíneas 1-28 · p1

test_band_edges prueba los bordes por los dos lados, porque los errores de una función de tramos están siempre en el >= contra el >.

def _remote(path: Path) -> None:
    rows = []
    for i in range(60):
        rows.append(
            {
                "PuzzleId": f"p{i:05d}",
                "GameId": f"g{i}",
                "FEN": "r1bqkbnr/pppp1ppp/2n5/4p3/4P3/5N2/PPPP1PPP/RNBQKB1R w KQkq - 2 3",
                "Moves": "f1b5 a7a6",
                "Rating": 900 + i * 40,
                "RatingDeviation": 150 if i % 10 == 0 else 80,
                "Popularity": 90,
                "NbPlays": 50 if i % 7 == 0 else 500,
                "Themes": ["fork", "short"],
                "OpeningTags": ["Ruy_Lopez"],
            }
        )
    pl.DataFrame(rows).write_parquet(path.as_posix())


def test_run_filters_bands_and_splits(rukh_home: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    remote = rukh_home / "remote.parquet"
    _remote(remote)
    monkeypatch.setattr(puzzles_module, "_source", lambda cfg: remote.as_posix())
    cfg = PuzzlesConfig(n_test=3, n_train=5)
    manifest = run(cfg)
    frame = pl.read_parquet((rukh_home / "data" / "puzzles" / "puzzles.parquet").as_posix())
    assert frame.columns == ["puzzle_id", "fen", "moves", "rating", "themes", "band", "split"]
    assert frame["rating"].min() >= 1000
    assert frame["puzzle_id"].n_unique() == frame.height
    for b in ("1000-1500", "1500-2000", "2000+"):
        test = frame.filter((pl.col("band") == b) & (pl.col("split") == "test"))
        train = frame.filter((pl.col("band") == b) & (pl.col("split") == "train"))
        assert test.height == 3 and 0 < train.height <= 5
        assert not set(test["puzzle_id"]) & set(train["puzzle_id"])
    assert manifest.counts["2000+/test"] == 3
    # Rows dropped by quality filters never appear.
    ids = set(frame["puzzle_id"].to_list())
    assert "p00000" not in ids and "p00007" not in ids and "p00010" not in ids

tests/unit/test_puzzles.pylíneas 31-69 · p1

_remote fabrica sesenta puzles y mete a propósito los que hay que tirar (desviación alta, pocas partidas), y el test comprueba por id que no aparecen. La aserción que vale por todo el fichero es not set(test) & set(train): que no hay fuga entre las dos mitades, comprobado y no supuesto.

def test_split_is_deterministic_and_seed_sensitive(rukh_home: Path) -> None:
    remote = rukh_home / "remote.parquet"
    _remote(remote)
    frame = (
        pl.read_parquet(remote.as_posix())
        .rename(
            {
                "PuzzleId": "puzzle_id",
                "FEN": "fen",
                "Moves": "moves",
                "Rating": "rating",
                "Themes": "themes",
            }
        )
        .select("puzzle_id", "fen", "moves", "rating", "themes")
    )
    cfg = PuzzlesConfig(n_test=4, n_train=6, seed=42)
    a = split_frame(frame, cfg)
    b = split_frame(frame, cfg)
    assert a.equals(b)
    c = split_frame(frame, cfg.model_copy(update={"seed": 7}))
    assert set(a.filter(pl.col("split") == "test")["puzzle_id"]) != set(
        c.filter(pl.col("split") == "test")["puzzle_id"]
    )

tests/unit/test_puzzles.pylíneas 72-95 · p1

Este comprueba las dos caras del determinismo: con la misma semilla, particiones idénticas; con otra, un test distinto. La segunda mitad es la que no se suele escribir, y la única que detectaría una split_key que ignorase la semilla.

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

"""Tests for rukh.data.pairs: mover-relative margins, mates, legality and phase balance."""

from __future__ import annotations

from pathlib import Path

import chess
import polars as pl
import pytest

from rukh.data import evals as evals_module
from rukh.data.evals import EvalsConfig
from rukh.data.pairs import PairsConfig, balance, make_pair, run, score_mover

pytestmark = pytest.mark.unit

WHITE = "rnbqkbnr/pppp1ppp/8/4p3/4P3/8/PPPP1PPP/RNBQKBNR w KQkq -"
BLACK = "rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR b KQkq -"


def _pv(
    move: str, cp: int | None = None, mate: int | None = None, depth: int = 30
) -> dict[str, object]:
    return {"move": move, "cp": cp, "mate": mate, "depth": depth}


def test_score_mover_sign_and_mate() -> None:
    assert score_mover(150, None, "w") == 150
    assert score_mover(150, None, "b") == -150
    assert score_mover(None, 2, "w") == 9_998
    assert score_mover(None, 2, "b") == -9_998
    assert score_mover(None, -1, "b") == 9_999
    assert score_mover(None, None, "w") is None


def test_make_pair_white_and_black_perspective() -> None:
    pair = make_pair(WHITE, [_pv("g1f3", 40), _pv("b1c3", 30), _pv("f2f4", -70)], 100)
    assert pair is not None
    assert (pair["chosen"], pair["rejected"]) == ("g1f3", "f2f4")
    assert (pair["cp_chosen"], pair["cp_rejected"]) == (40, -70)
    # For Black a lower cp is better: -20 beats 90 by 110.
    pair = make_pair(BLACK, [_pv("e7e5", -20), _pv("c7c5", -10), _pv("f7f5", 90)], 100)
    assert pair is not None
    assert (pair["chosen"], pair["rejected"]) == ("e7e5", "f7f5")
    assert (pair["cp_chosen"], pair["cp_rejected"]) == (-20, 90)


def test_make_pair_margin_and_mates() -> None:
    assert make_pair(WHITE, [_pv("g1f3", 40), _pv("b1c3", -50)], 100) is None
    pair = make_pair(WHITE, [_pv("g1f3", None, 2), _pv("b1c3", 300)], 100)
    assert pair is not None and pair["chosen"] == "g1f3" and pair["cp_chosen"] == 9_998
    pair = make_pair(BLACK, [_pv("e7e5", None, -3), _pv("c7c5", -400)], 100)
    assert pair is not None and pair["chosen"] == "e7e5" and pair["cp_chosen"] == -9_997
    assert pair["rejected"] == "c7c5"
    # A mate in 1 outranks a mate in 5, but the margin between them is only 4 centipawns.
    assert make_pair(WHITE, [_pv("g1f3", None, 5), _pv("b1c3", None, 1)], 100) is None
    # Depth outranks the score: the deepest line is chosen, a shallow better one is no pair.
    assert make_pair(WHITE, [_pv("g1f3", 900, depth=12), _pv("b1c3", 40, depth=30)], 100) is None
    pair = make_pair(WHITE, [_pv("g1f3", 40, depth=30), _pv("b1c3", -400, depth=12)], 100)
    assert pair is not None and pair["chosen"] == "g1f3" and pair["rejected"] == "b1c3"
    assert make_pair(WHITE, [_pv("g1f3", 40)], 100) is None


def test_make_pair_rejects_illegal_moves() -> None:
    assert make_pair(WHITE, [_pv("e1e5", 40), _pv("b1c3", -200)], 100) is None
    assert make_pair(WHITE, [_pv("g1f3", 40), _pv("e4e6", -200)], 100) is None
    assert make_pair("not a fen", [_pv("g1f3", 40), _pv("b1c3", -200)], 100) is None


def test_balance_uses_smallest_phase() -> None:
    frame = pl.DataFrame(
        {
            "fen": [f"f{i}" for i in range(9)],
            "chosen": ["a"] * 9,
            "rejected": ["b"] * 9,
            "cp_chosen": [0] * 9,
            "cp_rejected": [0] * 9,
            "phase": ["opening"] * 4 + ["middlegame"] * 3 + ["endgame"] * 2,
        }
    )
    out = balance(frame, max_per_phase=10, seed=1)
    assert out.group_by("phase").len()["len"].to_list() == [2, 2, 2]
    assert balance(frame, max_per_phase=1, seed=1).height == 3
    assert balance(frame.filter(pl.col("phase") != "endgame"), 10, 1).height == 0


def _evals_parquet(path: Path) -> None:
    board = chess.Board()
    fens, phases, pvs = [], [], []
    moves = "e2e4 e7e5 g1f3 b8c6 f1b5 a7a6 b5a4 g8f6 e1g1 f8e7 f1e1 b7b5 a4b3 d7d6 c2c3 e8g8"
    for ply, move in enumerate(moves.split(), start=1):
        board.push_uci(move)
        legal = [m.uci() for m in list(board.legal_moves)[:3]]
        sign = 1 if board.turn == chess.WHITE else -1
        fens.append(" ".join(board.fen().split()[:4]))
        phases.append("opening" if ply <= 10 else "middlegame" if ply <= 13 else "endgame")
        pvs.append([_pv(legal[0], sign * 60), _pv(legal[1], sign * 20), _pv(legal[2], -sign * 90)])
    pl.DataFrame({"fen": fens, "phase": phases, "pvs": pvs}).write_parquet(path.as_posix())


def test_run_writes_balanced_pairs(rukh_home: Path) -> None:
    evals = rukh_home / "data" / "evals" / "positions-eval.parquet"
    evals.parent.mkdir(parents=True)
    _evals_parquet(evals)
    manifest = run(PairsConfig(max_per_phase=2))
    frame = pl.read_parquet((rukh_home / "data" / "pairs" / "pairs.parquet").as_posix())
    assert frame.columns == ["fen", "chosen", "rejected", "cp_chosen", "cp_rejected", "phase"]
    assert manifest.counts["candidates"] == 16
    assert manifest.counts["opening"] == 2
    assert manifest.counts["middlegame"] == 2 and manifest.counts["endgame"] == 2
    assert frame.height == 6
    for fen, chosen, rejected, cp_c, cp_r, _ in frame.iter_rows():
        board = chess.Board(fen + " 0 1")
        legal = {m.uci() for m in board.legal_moves}
        assert chosen in legal and rejected in legal and chosen != rejected
        mover = 1 if board.turn == chess.WHITE else -1
        assert mover * (cp_c - cp_r) >= 100


def test_chosen_is_the_consolidated_best_move(
    rukh_home: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """``pairs`` and the ``evals`` consolidation must agree, mates included."""
    board = chess.Board()
    board.push_uci("f2f3")
    board.push_uci("e7e5")
    board.push_uci("g2g4")  # 2...Qh4# is mate in 1 for Black
    mate_fen = " ".join(board.fen().split()[:4])
    quiet = " ".join(chess.Board().fen().split()[:4])
    quiet_moves = [m.uci() for m in list(chess.Board().legal_moves)[:3]]
    rows = [
        # Black to move: the mate in 1 must win over the deeper, better-looking cp line.
        (mate_fen, "d8h4", 12, None, -1),
        (mate_fen, "d7d5", 40, -300, None),
        (mate_fen, "b8c6", 40, 200, None),
        # White to move: deepest line first, then the best cp.
        (quiet, quiet_moves[0], 40, 30, None),
        (quiet, quiet_moves[1], 40, -200, None),
        (quiet, quiet_moves[2], 12, 900, None),
    ]
    remote = rukh_home / "remote.parquet"
    pl.DataFrame(
        [
            {
                "fen": fen,
                "line": f"{move} a7a6",
                "depth": depth,
                "knodes": 10,
                "cp": cp,
                "mate": mate,
            }
            for fen, move, depth, cp, mate in rows
        ]
    ).write_parquet(remote.as_posix())
    positions = rukh_home / "data" / "positions" / "positions.parquet"
    positions.parent.mkdir(parents=True)
    pl.DataFrame(
        {
            "fen4": [mate_fen, quiet],
            "game_id": [1, 2],
            "ply": [3, 0],
            "last_move": ["g2g4", "e2e4"],
            "result": ["0-1", "1-0"],
            "phase": ["opening", "opening"],
            "n_seen": [1, 1],
        }
    ).write_parquet(positions.as_posix())
    monkeypatch.setattr(evals_module, "_source", lambda cfg, index: remote.as_posix())
    evals_module.run(EvalsConfig(n_files=1), files="0")

    consolidated = pl.read_parquet(
        (rukh_home / "data" / "evals" / "positions-eval.parquet").as_posix()
    )
    assert consolidated.height == 2
    for row in consolidated.to_dicts():
        pair = make_pair(row["fen"], row["pvs"], 100)
        assert pair is not None, row["fen"]
        assert pair["chosen"] == row["best_move"]
    best_by_fen = {r["fen"]: r["best_move"] for r in consolidated.to_dicts()}
    assert best_by_fen[mate_fen] == "d8h4"

tests/unit/test_pairs.pylíneas 1-180 · p1

// EjecútaloDescomenta los dos pasos y pasa sus tests

Con puzzles.py y pairs.py escritos, descomenta en pipeline.py y pipeline.yaml lo de estos dos pasos: las importaciones de PairsConfig y PuzzlesConfig, sus dos campos y las secciones puzzles: y pairs:. Borra la DUCKDB_STEPS provisional y descomenta la original, que ya tiene sus cuatro pasos. Solo queda comentado publish, hasta «Publicar el dataset».

Pasa los tests de los puzles y teclea también tests/unit/test_pairs.py, plegado justo encima: es el catálogo de casos de la caja del signo —el margen de las negras escrito dos veces, con un mate y con centipeones, el mate que gana a cualquier ventaja material, la jugada ilegal que descarta el par, el equilibrado por fase— y un último test que comprueba que la elegida es la misma mejor línea que consolida evals.py. Si algún día tocas make_pair, es lo primero que hay que pasar:

git diff p1 -- src/rukh/data/pipeline.py configs/data/pipeline.yaml
uv run pytest -m unit -q tests/unit/test_puzzles.py tests/unit/test_pairs.py tests/unit/test_db.py

En el git diff ya solo tiene que salir publish: su importación, su campo y su sección, cada línea con - la de p1 y con + tu versión comentada.

Quince tests en verde (cuatro de puzzles, siete de pairs y los cuatro de test_db.py, ya sin -k), sin red: los dos run se prueban con un parquet fabricado en el propio test, y test_shipped_pipeline_yaml_declares_duckdb confirma que tu pipeline.yaml descomentado valida y reparte su duckdb: a puzzles.

// EjecútaloLos dos conjuntos, de verdad

Igual que en las dos lecciones anteriores, con la CLI y la configuración del pipeline.yaml:

uv run rukh data puzzles
uv run rukh data pairs

El primero necesita red (unos minutos) y deja data/puzzles/puzzles.parquet: en los conteos, <tramo>/test tiene que valer 2 000 en los tres tramos. El segundo lee data/evals/positions-eval.parquet y tarda segundos: las tres fases salen con el mismo número y candidates dice de cuántos se partió. Con una sola parte de las evaluaciones los números son pequeños; los definitivos salen cuando el join de la lección anterior tenga sus veinte partes.

// Ejercicio 01Comprueba el signo de los pares de las negras

Con data/pairs/pairs.parquet generado, escribe una consulta que se quede con los pares en los que le toca mover a las negras y compruebe la relación entre cp_chosen y cp_rejected. ¿Cuál tiene que ser? Después escribe la misma consulta para las blancas. ¿Qué consulta única comprobaría las dos a la vez? (Se ejecutan como dice la caja «Cómo ejecutar una consulta SQL» de la lección anterior: en consulta.sql, con uv run python -c "import duckdb, pathlib; print(duckdb.sql(pathlib.Path('consulta.sql').read_text()))".)

// SoluciónVer la solución

El turno está en el segundo campo del FEN: split_part(fen, ' ', 2) = 'b'. Para las negras, cp_chosen < cp_rejected en todos los pares, porque los cp están desde el punto de vista de las blancas y para las negras menos es mejor. Para las blancas, cp_chosen > cp_rejected.

La consulta única usa el signo del que mueve, que es lo que hace mover_sign:

SELECT count(*) AS mal
FROM read_parquet('data/pairs/pairs.parquet')
WHERE (CASE WHEN split_part(fen, ' ', 2) = 'w' THEN 1 ELSE -1 END)
      * (cp_chosen - cp_rejected) < 100

Tiene que dar cero: el margen, medido desde el que mueve, es siempre de al menos cien centipeones. Guárdate la consulta: es la que ejecutarías el día que el DPO de M5 no mejore nada.

Qué has aprendido

Has construido un banco de test que no cambia cuando el catálogo de origen crece, y un conjunto de preferencias sacado de lo que otro paso iba a tirar, midiendo la diferencia desde el lado correcto.

Cómo se mide: uv run rukh data puzzles deja data/puzzles/puzzles.parquet con 2 000 puzles de test por tramo y hasta 50 000 de entrenamiento, y un manifiesto con el conteo por tramo/split. uv run rukh data pairs deja data/pairs/pairs.parquet con las tres fases igualadas y candidates diciendo de cuántos se partió. uv run pytest -m unit -q tests/unit/test_puzzles.py tests/unit/test_pairs.py pasa en verde.

Con esto el pipeline de datos está completo. La otra mitad del módulo convierte las partidas en números: tres vocabularios y una tabla medida que decide cuál se usa.