LLM-course/amine-final

1"""2Custom Chess Tokenizer for the Chess Challenge.3 4This tokenizer uses a DECOMPOSED format compatible with the evaluator:5    "WPe2e4" -> ["WP", "e2_f", "e4_t"]6 7The decomposed format uses:8- Piece token: "WP", "BN", etc. (color + piece)9- Source square with _f suffix: "e2_f", "g1_f", etc.10- Destination square with _t suffix: "e4_t", "f3_t", etc.11- Optional suffix for annotations: "(x)", "(+)", "(+*)", "(o)", "(O)"12 13The dataset format uses:14- W/B prefix for White/Black15- Piece letter: P=Pawn, N=Knight, B=Bishop, R=Rook, Q=Queen, K=King16- Source and destination squares (e.g., e2e4)17- Special suffixes: (x)=capture, (+)=check, (+*)=checkmate, (o)/(O)=castling18"""19 20from __future__ import annotations21 22import json23import os24from pathlib import Path25from typing import Dict, List, Optional26 27from transformers import PreTrainedTokenizer28 29 30class ChessTokenizer(PreTrainedTokenizer):31    """32    A custom tokenizer for chess moves using DECOMPOSED format.33 34    This tokenizer decomposes each move into sub-tokens:35    - Piece: "WP", "BN", etc.36    - Source square with _f suffix: "e2_f", "g1_f", etc.37    - Destination square with _t suffix: "e4_t", "f3_t", etc.38    - Optional suffix: "(x)", "(+)", etc.39 40    This format is compatible with the evaluator's 'decomposed' detection.41 42    Example:43        >>> tokenizer = ChessTokenizer.build_vocab_from_dataset()44        >>> tokenizer.tokenize("WPe2e4 BPe7e5")45        ['WP', 'e2_f', 'e4_t', 'BP', 'e7_f', 'e5_t']46    """47    48    model_input_names = ["input_ids", "attention_mask"]49    vocab_files_names = {"vocab_file": "vocab.json"}50    51    # Special tokens52    PAD_TOKEN = "[PAD]"53    BOS_TOKEN = "[BOS]"54    EOS_TOKEN = "[EOS]"55    UNK_TOKEN = "[UNK]"56    57    def __init__(58        self,59        vocab_file: Optional[str] = None,60        vocab: Optional[Dict[str, int]] = None,61        **kwargs,62    ):63        """64        Initialize the chess tokenizer.65        66        Args:67            vocab_file: Path to a JSON file containing the vocabulary mapping.68            vocab: Dictionary mapping tokens to IDs (alternative to vocab_file).69            **kwargs: Additional arguments passed to PreTrainedTokenizer.70        """71        # Initialize special tokens72        self._pad_token = self.PAD_TOKEN73        self._bos_token = self.BOS_TOKEN74        self._eos_token = self.EOS_TOKEN75        self._unk_token = self.UNK_TOKEN76 77        # Remove any duplicate special-token entries passed through kwargs78        # to avoid "multiple values for keyword" errors when loading from disk.79        kwargs.pop("pad_token", None)80        kwargs.pop("bos_token", None)81        kwargs.pop("eos_token", None)82        kwargs.pop("unk_token", None)83        84        # Load or create vocabulary85        if vocab is not None:86            self._vocab = vocab87        elif vocab_file is not None and os.path.exists(vocab_file):88            with open(vocab_file, "r", encoding="utf-8") as f:89                self._vocab = json.load(f)90        else:91            # Create a minimal vocabulary with just special tokens92            # The full vocabulary should be built from the dataset93            self._vocab = self._create_default_vocab()94        95        # Create reverse mapping96        self._ids_to_tokens = {v: k for k, v in self._vocab.items()}97        98        # Call parent init AFTER setting up vocab99        super().__init__(100            pad_token=self._pad_token,101            bos_token=self._bos_token,102            eos_token=self._eos_token,103            unk_token=self._unk_token,104            **kwargs,105        )106    107    def _create_default_vocab(self) -> Dict[str, int]:108        """109        Create a minimal default vocabulary with just special tokens.110        111        For the full vocabulary, use `build_vocab_from_dataset()`.112        This minimal vocab is just a placeholder - you should build from data.113        """114        special_tokens = [self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN]115        vocab = {token: idx for idx, token in enumerate(special_tokens)}116        return vocab117    118    @classmethod119    def build_vocab_from_iterator(120        cls,121        iterator,122        min_frequency: int = 1,123    ) -> "ChessTokenizer":124        """125        Build a tokenizer vocabulary from an iterator of game strings.126 127        Decomposes each move into tokens: piece, source_f, dest_t, and optional suffix.128 129        Args:130            iterator: An iterator yielding game strings (space-separated moves).131            min_frequency: Minimum frequency for a token to be included.132 133        Returns:134            A ChessTokenizer with the built vocabulary.135        """136        from collections import Counter137 138        token_counts = Counter()139 140        for game in iterator:141            moves = game.strip().split()142            for move in moves:143                if len(move) < 6:144                    token_counts[move] += 1145                    continue146 147                # Decompose move into tokens148                piece = move[:2]           # e.g., "WP", "BN"149                source = move[2:4] + "_f"  # e.g., "e2_f"150                dest = move[4:6] + "_t"    # e.g., "e4_t"151                suffix = move[6:] if len(move) > 6 else None152 153                token_counts[piece] += 1154                token_counts[source] += 1155                token_counts[dest] += 1156                if suffix:157                    token_counts[suffix] += 1158 159        # Filter by frequency160        tokens = [161            token for token, count in token_counts.items()162            if count >= min_frequency163        ]164 165        # Sort for reproducibility166        tokens = sorted(tokens)167 168        # Build vocabulary169        special_tokens = [cls.PAD_TOKEN, cls.BOS_TOKEN, cls.EOS_TOKEN, cls.UNK_TOKEN]170        vocab = {token: idx for idx, token in enumerate(special_tokens + tokens)}171 172        return cls(vocab=vocab)173    174    @classmethod175    def build_vocab_from_dataset(176        cls,177        dataset_name: str = "dlouapre/lichess_2025-01_1M",178        split: str = "train",179        column: str = "text",180        min_frequency: int = 500,181        max_samples: Optional[int] = 100000,182    ) -> "ChessTokenizer":183        """184        Build a tokenizer vocabulary from a Hugging Face dataset.185        186        Args:187            dataset_name: Name of the dataset on Hugging Face Hub.188            split: Dataset split to use.189            column: Column containing the game strings.190            min_frequency: Minimum frequency for a token to be included (default: 500).191            max_samples: Maximum number of samples to process (default: 100k).192        193        Returns:194            A ChessTokenizer with the built vocabulary.195        """196        from datasets import load_dataset197        198        dataset = load_dataset(dataset_name, split=split)199        200        if max_samples is not None:201            dataset = dataset.select(range(min(max_samples, len(dataset))))202        203        def game_iterator():204            for example in dataset:205                yield example[column]206        207        return cls.build_vocab_from_iterator(game_iterator(), min_frequency=min_frequency)208    209    @property210    def vocab_size(self) -> int:211        """Return the size of the vocabulary."""212        return len(self._vocab)213    214    def get_vocab(self) -> Dict[str, int]:215        """Return the vocabulary as a dictionary."""216        return dict(self._vocab)217    218    def _tokenize(self, text: str) -> List[str]:219        """220        Tokenize a string of moves into decomposed tokens.221 222        Each move like "WPe2e4" becomes ["WP", "e2_f", "e4_t"].223        Moves with suffixes like "WPe2e4(x)" become ["WP", "e2_f", "e4_t", "(x)"].224 225        Args:226            text: A string of space-separated moves.227 228        Returns:229            List of decomposed tokens.230        """231        moves = text.strip().split()232        tokens = []233 234        for move in moves:235            if len(move) < 6:236                # Invalid move format, add as unknown237                tokens.append(move)238                continue239 240            # Split move into components241            piece = move[:2]           # e.g., "WP", "BN"242            source = move[2:4] + "_f"  # e.g., "e2_f", "g1_f"243            dest = move[4:6] + "_t"    # e.g., "e4_t", "f3_t"244            suffix = move[6:] if len(move) > 6 else None  # e.g., "(x)", "(+)"245 246            tokens.extend([piece, source, dest])247            if suffix:248                tokens.append(suffix)249 250        return tokens251    252    def _convert_token_to_id(self, token: str) -> int:253        """Convert a token to its ID."""254        return self._vocab.get(token, self._vocab.get(self.UNK_TOKEN, 0))255    256    def _convert_id_to_token(self, index: int) -> str:257        """Convert an ID to its token."""258        return self._ids_to_tokens.get(index, self.UNK_TOKEN)259    260    def convert_tokens_to_string(self, tokens: List[str]) -> str:261        """262        Convert decomposed tokens back to a string of moves.263 264        Reconstructs moves from [piece, source_f, dest_t, optional_suffix] format.265        E.g., ["WP", "e2_f", "e4_t"] -> "WP e2_f e4_t"266 267        For the evaluator's decomposed format, we keep the tokens space-separated.268        """269        # Filter out special tokens270        special = {self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN}271        filtered = [t for t in tokens if t not in special]272        return " ".join(filtered)273    274    def save_vocabulary(275        self,276        save_directory: str,277        filename_prefix: Optional[str] = None,278    ) -> tuple:279        """280        Save the vocabulary to a JSON file.281        282        Args:283            save_directory: Directory to save the vocabulary.284            filename_prefix: Optional prefix for the filename.285        286        Returns:287            Tuple containing the path to the saved vocabulary file.288        """289        if not os.path.isdir(save_directory):290            os.makedirs(save_directory, exist_ok=True)291        292        vocab_file = os.path.join(293            save_directory,294            (filename_prefix + "-" if filename_prefix else "") + "vocab.json",295        )296        297        with open(vocab_file, "w", encoding="utf-8") as f:298            json.dump(self._vocab, f, ensure_ascii=False, indent=2)299        300        return (vocab_file,)301 302 303def count_vocab_from_dataset(304    dataset_name: str = "dlouapre/lichess_2025-01_1M",305    split: str = "train",306    column: str = "text",307    max_samples: Optional[int] = 10000,308) -> Dict[str, int]:309    """310    Count decomposed token frequencies in a dataset (useful for vocabulary analysis).311 312    Args:313        dataset_name: Name of the dataset on Hugging Face Hub.314        split: Dataset split to use.315        column: Column containing the game strings.316        max_samples: Maximum number of samples to process.317 318    Returns:319        Dictionary mapping decomposed tokens to their frequencies.320    """321    from collections import Counter322    from datasets import load_dataset323 324    dataset = load_dataset(dataset_name, split=split)325 326    if max_samples is not None:327        dataset = dataset.select(range(min(max_samples, len(dataset))))328 329    token_counts = Counter()330 331    for example in dataset:332        moves = example[column].strip().split()333        for move in moves:334            if len(move) < 6:335                token_counts[move] += 1336                continue337 338            # Decompose move339            piece = move[:2]340            source = move[2:4] + "_f"341            dest = move[4:6] + "_t"342            suffix = move[6:] if len(move) > 6 else None343 344            token_counts[piece] += 1345            token_counts[source] += 1346            token_counts[dest] += 1347            if suffix:348                token_counts[suffix] += 1349 350    return dict(token_counts)351