LLM-course/chess-learning-v2

1"""2Custom Chess Tokenizer for the Chess Challenge.3 4This tokenizer treats each move as a single token using the extended UCI notation5from the Lichess dataset (e.g., WPe2e4, BNg8f6).6 7The dataset format uses:8- W/B prefix for White/Black9- Piece letter: P=Pawn, N=Knight, B=Bishop, R=Rook, Q=Queen, K=King10- Source and destination squares (e.g., e2e4)11- Special suffixes: (x)=capture, (+)=check, (+*)=checkmate, (o)/(O)=castling12"""13 14from __future__ import annotations15 16import json17import os18import re19from pathlib import Path20from typing import Dict, List, Optional, Tuple21 22from transformers import PreTrainedTokenizer23 24 25 26# Parse "WPe2e4(x+*)" etc.27_MOVE_RE = re.compile(28    r"^(?P<side>[WB])"29    r"(?P<piece>[PNBRQK])"30    r"(?P<src>[a-h][1-8])"31    r"(?P<dst>[a-h][1-8])"32    r"(?P<suffix>.*)$"33)34 35 36 37# Promotions like "=Q" or "=q"38_PROMO_RE = re.compile(r"=([QRBNqrbn])")39 40 41def _parse_suffix(suffix: str) -> Tuple[bool, bool, bool, Optional[str], Optional[str]]:42    """43    Returns:44      is_capture, is_check, is_mate, castle_kind, promo_piece45 46    castle_kind: "k" (kingside) or "q" (queenside) or None47    promo_piece: one of "q","r","b","n" or None48    """49    if not suffix:50        return False, False, False, None, None51 52    # Normalize 53    suf = suffix.strip()54 55    is_capture = "x" in suf56    is_check = "+" in suf57 58    # Mate indicator59    # We'll treat any "*" as mate.60    is_mate = "*" in suf61 62    # Castling: dataset uses (o)/(O) in the move string for king moves63    castle_kind = None64    if "(O)" in suf:65        castle_kind = "q"66    elif "(o)" in suf:67        castle_kind = "k"68 69    promo_piece = None70    m = _PROMO_RE.search(suf)71    if m:72        promo_piece = m.group(1).lower()73 74    return is_capture, is_check, is_mate, castle_kind, promo_piece75 76 77def _reindex_vocab(vocab: Dict[str, int]) -> Dict[str, int]:78    # sort by old id for stability79    items = sorted(vocab.items(), key=lambda kv: kv[1])80    return {tok: new_id for new_id, (tok, _) in enumerate(items)}81 82 83 84class ChessTokenizer(PreTrainedTokenizer):85    """86    A custom tokenizer for chess moves using extended UCI notation.87    88    This tokenizer maps each possible chess move to a unique token ID.89    The vocabulary is built from the training dataset to ensure all moves90    encountered during training have a corresponding token.91    92    Example:93        >>> tokenizer = ChessTokenizer()94        >>> tokenizer.encode("WPe2e4 BPe7e5")95        [1, 42, 87, 2]  # [BOS, e2e4, e7e5, EOS]96    """97    98    model_input_names = ["input_ids", "attention_mask"]99    vocab_files_names = {"vocab_file": "vocab.json"}100    101    # Special tokens102    PAD_TOKEN = "[PAD]"103    BOS_TOKEN = "[BOS]"104    EOS_TOKEN = "[EOS]"105    UNK_TOKEN = "[UNK]"106 107 108     # Component tokens109    SIDE_TOKENS = ("[W]", "[B]")110    PIECE_TOKENS = ("[P]", "[N]", "[B]", "[R]", "[Q]", "[K]")111    # flags112    FLAG_TOKENS = (113        "[x]",       # capture114        "[+]",       # check115        "[#]",       # mate116        "[O-O]",     # kingside castle marker (not required by evaluator)117        "[O-O-O]",   # queenside castle marker118        # promotions119        "[=q]", "[=r]", "[=b]", "[=n]",120    )   121    def __init__(122        self,123        vocab_file: Optional[str] = None,124        vocab: Optional[Dict[str, int]] = None,125        **kwargs,126    ):127        """128        Initialize the chess tokenizer.129        130        Args:131            vocab_file: Path to a JSON file containing the vocabulary mapping.132            vocab: Dictionary mapping tokens to IDs (alternative to vocab_file).133            **kwargs: Additional arguments passed to PreTrainedTokenizer.134        """135        # Initialize special tokens136        self._pad_token = self.PAD_TOKEN137        self._bos_token = self.BOS_TOKEN138        self._eos_token = self.EOS_TOKEN139        self._unk_token = self.UNK_TOKEN140 141        # Remove any duplicate special-token entries passed through kwargs142        # to avoid "multiple values for keyword" errors when loading from disk.143        kwargs.pop("pad_token", None)144        kwargs.pop("bos_token", None)145        kwargs.pop("eos_token", None)146        kwargs.pop("unk_token", None)147        148        # Load or create vocabulary149        if vocab is not None:150            self._vocab = vocab151        elif vocab_file is not None and os.path.exists(vocab_file):152            with open(vocab_file, "r", encoding="utf-8") as f:153                self._vocab = json.load(f)154        else:155            # Create a minimal vocabulary with just special tokens156            # The full vocabulary should be built from the dataset157            self._vocab = self._create_default_vocab()158        159        self._vocab = _reindex_vocab(self._vocab)160        161        # Create reverse mapping162        self._ids_to_tokens = {v: k for k, v in self._vocab.items()}163        164        # Call parent init AFTER setting up vocab165        super().__init__(166            pad_token=self._pad_token,167            bos_token=self._bos_token,168            eos_token=self._eos_token,169            unk_token=self._unk_token,170            **kwargs,171        )172    173    def _create_default_vocab(self) -> Dict[str, int]:174        """175        Create a minimal default vocabulary with just special tokens.176        177        For the full vocabulary, use `build_vocab_from_dataset()`.178        This minimal vocab is just a placeholder - you should build from data.179        """180        tokens: List[str] = [self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN]181        tokens += list(self.SIDE_TOKENS)182        tokens += list(self.PIECE_TOKENS)183 184        # Squares (64)185        for file in "abcdefgh":186            for rank in "12345678":187                tokens.append(f"[{file}{rank}]")188 189        tokens += list(self.FLAG_TOKENS)190 191        return {tok: idx for idx, tok in enumerate(tokens)}192    193    @classmethod194    def build_vocab_from_iterator(195        cls,196        iterator,197        min_frequency: int = 1,198    ) -> "ChessTokenizer":199        return cls()200    201    @classmethod202    def build_vocab_from_dataset(203        cls,204        dataset_name: str = "dlouapre/lichess_2025-01_1M",205        split: str = "train",206        column: str = "text",207        min_frequency: int = 500,208        max_samples: Optional[int] = 100000,209    ) -> "ChessTokenizer":210        return cls()211    212    @property213    def vocab_size(self) -> int:214        """Return the size of the vocabulary."""215        return len(self._vocab)216    217    def get_vocab(self) -> Dict[str, int]:218        """Return the vocabulary as a dictionary."""219        return dict(self._vocab)220    221    def _tokenize(self, text: str) -> List[str]:222        """223        Tokenize a string of moves into a list of tokens.224        225        Args:226            text: A string of space-separated moves.227        228        Returns:229            List of move tokens.230        """231        text = (text or "").strip()232        if not text:233            return []234 235        chunks = text.split()236        out: List[str] = []237 238        for chunk in chunks:239            # If chunk is pure uci like "e2e4" or "e7e8q"240            if re.fullmatch(r"[a-h][1-8][a-h][1-8][qrbn]?", chunk):241                src = chunk[0:2]242                dst = chunk[2:4]243                out.append(f"[{src}]")244                out.append(f"[{dst}]")245                if len(chunk) == 5 and chunk[4] in "qrbn":246                    out.append(f"[={chunk[4]}]")247                continue248 249            m = _MOVE_RE.match(chunk)250            if not m:251                out.append(self.UNK_TOKEN)252                continue253 254            side = "[W]" if m.group("side") == "W" else "[BL]"255            piece = m.group("piece")256            src = m.group("src")257            dst = m.group("dst")258            suffix = m.group("suffix") or ""259 260            out.append(side)261            out.append(f"[{piece}]")262            out.append(f"[{src}]")263            out.append(f"[{dst}]")264 265            is_cap, is_chk, is_mate, castle_kind, promo = _parse_suffix(suffix)266 267            # Castling markers (optional; evaluator doesn't need them)268            if castle_kind == "k":269                out.append("[O-O]")270            elif castle_kind == "q":271                out.append("[O-O-O]")272 273            if is_cap:274                out.append("[x]")275            if is_mate:276                out.append("[#]")277            elif is_chk:278                out.append("[+]")279 280            if promo in ("q", "r", "b", "n"):281                out.append(f"[={promo}]")282 283        return out284 285    286    def _convert_token_to_id(self, token: str) -> int:287        """Convert a token to its ID."""288        return self._vocab.get(token, self._vocab.get(self.UNK_TOKEN, 0))289    290    def _convert_id_to_token(self, index: int) -> str:291        """Convert an ID to its token."""292        return self._ids_to_tokens.get(index, self.UNK_TOKEN)293    294    def convert_tokens_to_string(self, tokens: List[str]) -> str:295        """Convert a list of tokens back to a string."""296        # Filter out special tokens for cleaner output297        special = {self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN}298        return " ".join(t for t in tokens if t not in special)299    300    def save_vocabulary(301        self,302        save_directory: str,303        filename_prefix: Optional[str] = None,304    ) -> tuple:305        """306        Save the vocabulary to a JSON file.307        308        Args:309            save_directory: Directory to save the vocabulary.310            filename_prefix: Optional prefix for the filename.311        312        Returns:313            Tuple containing the path to the saved vocabulary file.314        """315        if not os.path.isdir(save_directory):316            os.makedirs(save_directory, exist_ok=True)317        318        vocab_file = os.path.join(319            save_directory,320            (filename_prefix + "-" if filename_prefix else "") + "vocab.json",321        )322        323        with open(vocab_file, "w", encoding="utf-8") as f:324            json.dump(self._vocab, f, ensure_ascii=False, indent=2)325        326        return (vocab_file,)327 328 329def count_vocab_from_dataset(330    dataset_name: str = "dlouapre/lichess_2025-01_1M",331    split: str = "train",332    column: str = "text",333    max_samples: Optional[int] = 10000,334) -> Dict[str, int]:335    """336    Count token frequencies in a dataset (useful for vocabulary analysis).337    338    Args:339        dataset_name: Name of the dataset on Hugging Face Hub.340        split: Dataset split to use.341        column: Column containing the game strings.342        max_samples: Maximum number of samples to process.343    344    Returns:345        Dictionary mapping tokens to their frequencies.346    """347    from collections import Counter348    from datasets import load_dataset349 350    ds = load_dataset(dataset_name, split=split)351    if max_samples is not None:352        ds = ds.select(range(min(max_samples, len(ds))))353 354    tok = ChessTokenizer()355    counts = Counter()356    for ex in ds:357        counts.update(tok._tokenize(ex[column]))358    return dict(counts)359