tokenization_chess_policy.py

"""Chess-board preprocessing for the Hugging Face model.

This is deliberately a board encoder rather than a text tokenizer.  Each
board becomes 64 mover-relative piece IDs, and each legal move is represented
by the tokenized board after that move.  Successors remain in the original
position's perspective.
"""

from __future__ import annotations

from collections.abc import Sequence
from typing import Any

import chess
import torch


EMPTY = 0
MY_PAWN = 1
MY_KNIGHT = 2
MY_BISHOP = 3
MY_ROOK = 4
MY_QUEEN = 5
MY_KING = 6
OPPONENT_OFFSET = 6
NUM_PIECE_STATES = 13


def canonical_square(square: chess.Square, perspective: chess.Color) -> int:
    """Map a square into mover-relative canonical order."""

    return square if perspective == chess.WHITE else square ^ 56


def tokenize_board(
    board: chess.Board, perspective: chess.Color | None = None
) -> torch.Tensor:
    """Return a board as a ``[64]`` uint8 tensor of piece-role IDs."""

    if perspective is None:
        perspective = board.turn
    tokens = torch.zeros(64, dtype=torch.uint8)
    for square, piece in board.piece_map().items():
        role = piece.piece_type
        if piece.color != perspective:
            role += OPPONENT_OFFSET
        tokens[canonical_square(square, perspective)] = role
    return tokens


def build_model_inputs(
    boards: Sequence[chess.Board],
    *,
    include_candidate_uci: bool = False,
) -> dict[str, Any]:
    """Build the tensors expected by ``ChessTransitionPolicy``.

    Candidate order is the order returned by ``python-chess``.  If
    ``include_candidate_uci`` is true, the returned dictionary also contains
    ``candidate_uci`` for mapping output scores back to moves.
    """

    if not boards:
        raise ValueError("At least one board is required")

    current_boards: list[torch.Tensor] = []
    successor_boards: list[torch.Tensor] = []
    candidate_uci: list[tuple[str, ...]] = []
    candidate_counts: list[int] = []

    for board in boards:
        perspective = board.turn
        legal_moves = list(board.legal_moves)
        if not legal_moves:
            raise ValueError("Every board must have at least one legal move")
        current_boards.append(tokenize_board(board, perspective))
        candidate_counts.append(len(legal_moves))
        if include_candidate_uci:
            candidate_uci.append(tuple(move.uci() for move in legal_moves))
        for move in legal_moves:
            board.push(move)
            successor_boards.append(tokenize_board(board, perspective))
            board.pop()

    counts = torch.tensor(candidate_counts, dtype=torch.int32)
    offsets = torch.zeros(len(boards) + 1, dtype=torch.int64)
    offsets[1:] = torch.cumsum(counts.to(torch.int64), dim=0)
    owners = torch.repeat_interleave(
        torch.arange(len(boards), dtype=torch.int64), counts.to(torch.int64)
    )
    max_candidates = max(candidate_counts)
    candidate_mask = (
        torch.arange(max_candidates).unsqueeze(0)
        < counts.unsqueeze(1)
    )

    result: dict[str, Any] = {
        "current_boards": torch.stack(current_boards),
        "successor_boards": torch.stack(successor_boards),
        "candidate_owner": owners,
        "candidate_offsets": offsets,
        "candidate_mask": candidate_mask,
    }
    if include_candidate_uci:
        result["candidate_uci"] = tuple(candidate_uci)
    return result