LLM-course/basic_tokenizer

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 os18from pathlib import Path19from typing import Dict, List, Optional20 21from transformers import PreTrainedTokenizer22 23import re24 25 26 27class ChessTokenizer(PreTrainedTokenizer):28    """29    A custom tokenizer for chess moves using extended UCI notation.30    31    This tokenizer maps each possible chess move to a unique token ID.32    The vocabulary is built from the training dataset to ensure all moves33    encountered during training have a corresponding token.34    35    Example:36        >>> tokenizer = ChessTokenizer()37        >>> tokenizer.encode("WPe2e4 BPe7e5")38        [1, 42, 87, 2]  # [BOS, e2e4, e7e5, EOS]39    """40    41    model_input_names = ["input_ids", "attention_mask"]42    vocab_files_names = {"vocab_file": "vocab.json"}43    44    # Special tokens45    PAD_TOKEN = "[PAD]"46    BOS_TOKEN = "[BOS]"47    EOS_TOKEN = "[EOS]"48    UNK_TOKEN = "[UNK]"49    50    _MOVE_RE = re.compile(51        r'^(?P<color>[WB])(?P<piece>[PNBRQK])(?P<from>[a-h][1-8])(?P<to>[a-h][1-8])(?P<rest>.*)$'52    )53 54    _SUFFIX_MAP = {55        "(x)": "cap",56        "(+)": "check",57        "(+*)": "mate",58        "(o)": "castle_k",59        "(O)": "castle_q",60    }61 62    _PROMO_RE = re.compile(r'=?([QRBNqrbn])')  # accepts "=Q" or "q" style63 64 65    def __init__(66        self,67        vocab_file: Optional[str] = None,68        vocab: Optional[Dict[str, int]] = None,69        **kwargs,70    ):71        """72        Initialize the chess tokenizer.73        74        Args:75            vocab_file: Path to a JSON file containing the vocabulary mapping.76            vocab: Dictionary mapping tokens to IDs (alternative to vocab_file).77            **kwargs: Additional arguments passed to PreTrainedTokenizer.78        """79        # Initialize special tokens80        self._pad_token = self.PAD_TOKEN81        self._bos_token = self.BOS_TOKEN82        self._eos_token = self.EOS_TOKEN83        self._unk_token = self.UNK_TOKEN84 85        # Remove any duplicate special-token entries passed through kwargs86        # to avoid "multiple values for keyword" errors when loading from disk.87        kwargs.pop("pad_token", None)88        kwargs.pop("bos_token", None)89        kwargs.pop("eos_token", None)90        kwargs.pop("unk_token", None)91        92        # Load or create vocabulary93        if vocab is not None:94            self._vocab = vocab95        elif vocab_file is not None and os.path.exists(vocab_file):96            with open(vocab_file, "r", encoding="utf-8") as f:97                self._vocab = json.load(f)98        else:99            # Create a minimal vocabulary with just special tokens100            # The full vocabulary should be built from the dataset101            self._vocab = self._create_default_vocab()102        103        # Create reverse mapping104        self._ids_to_tokens = {v: k for k, v in self._vocab.items()}105        106        # Call parent init AFTER setting up vocab107        super().__init__(108            pad_token=self._pad_token,109            bos_token=self._bos_token,110            eos_token=self._eos_token,111            unk_token=self._unk_token,112            **kwargs,113        )114    115    def _decompose_move(self, tok: str) -> List[str]:116        """117        Convert e.g. 'WPe2e4(x)' -> ['WP', 'e2_f', 'e4_t', 'cap']118        """119        m = self._MOVE_RE.match(tok)120        if not m:121            return [self.UNK_TOKEN]122 123        color = m.group("color")124        piece = m.group("piece")125        from_sq = m.group("from")126        to_sq = m.group("to")127        rest = m.group("rest") or ""128 129        out = [f"{color}{piece}", f"{from_sq}_f", f"{to_sq}_t"]130 131        for raw, mapped in self._SUFFIX_MAP.items():132            if raw in rest:133                out.append(mapped)134 135        # Promotion (rare depending on dataset formatting, but safe)136        pm = self._PROMO_RE.search(rest)137        if pm:138            p = pm.group(1).lower()139            if p in ("q", "r", "b", "n"):140                out.append(f"promo_{p}")141 142        return out143    def _create_default_vocab(self) -> Dict[str, int]:144        """145        Create a minimal default vocabulary with just special tokens.146        147        For the full vocabulary, use `build_vocab_from_dataset()`.148        This minimal vocab is just a placeholder - you should build from data.149        """150        special_tokens = [self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN]151        vocab = {token: idx for idx, token in enumerate(special_tokens)}152        return vocab153    154    @classmethod155    def build_vocab_from_iterator(156        cls,157        iterator,158        min_frequency: int = 1,159    ) -> "ChessTokenizer":160        """161        Build a tokenizer vocabulary from an iterator of game strings.162        163        Args:164            iterator: An iterator yielding game strings (space-separated moves).165            min_frequency: Minimum frequency for a token to be included.166        167        Returns:168            A ChessTokenizer with the built vocabulary.169        """170        from collections import Counter171        172        token_counts = Counter()173        174        for game in iterator:175            moves = game.strip().split()176            token_counts.update(moves)177        178        # Filter by frequency179        tokens = [180            token for token, count in token_counts.items()181            if count >= min_frequency182        ]183        184        # Sort for reproducibility185        tokens = sorted(tokens)186        187        # Build vocabulary188        special_tokens = [cls.PAD_TOKEN, cls.BOS_TOKEN, cls.EOS_TOKEN, cls.UNK_TOKEN]189        vocab = {token: idx for idx, token in enumerate(special_tokens + tokens)}190        191        return cls(vocab=vocab)192    193    @classmethod194    def build_vocab_from_dataset(195        cls,196        dataset_name: str = "dlouapre/lichess_2025-01_1M",197        split: str = "train",198        column: str = "text",199        min_frequency: int = 500,200        max_samples: Optional[int] = 100000,201    ) -> "ChessTokenizer":202        """203        Build a tokenizer vocabulary from a Hugging Face dataset.204        205        Args:206            dataset_name: Name of the dataset on Hugging Face Hub.207            split: Dataset split to use.208            column: Column containing the game strings.209            min_frequency: Minimum frequency for a token to be included (default: 500).210            max_samples: Maximum number of samples to process (default: 100k).211        212        Returns:213            A ChessTokenizer with the built vocabulary.214        """215        from datasets import load_dataset216        217        dataset = load_dataset(dataset_name, split=split)218        219        if max_samples is not None:220            dataset = dataset.select(range(min(max_samples, len(dataset))))221        222        def game_iterator():223            for example in dataset:224                yield example[column]225        226        return cls.build_vocab_from_iterator(game_iterator(), min_frequency=min_frequency)227    228    @property229    def vocab_size(self) -> int:230        """Return the size of the vocabulary."""231        return len(self._vocab)232    233    def get_vocab(self) -> Dict[str, int]:234        """Return the vocabulary as a dictionary."""235        return dict(self._vocab)236    237    # def _tokenize(self, text: str) -> List[str]:238    #     """239    #     Tokenize a string of moves into a list of tokens.240        241    #     Args:242    #         text: A string of space-separated moves.243        244    #     Returns:245    #         List of move tokens.246    #     """247    #     return text.strip().split()248    def _tokenize(self, text: str) -> List[str]:249        parts = text.strip().split()250        out: List[str] = []251        special = {self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN}252 253        for p in parts:254            if p in special:255                out.append(p)256            else:257                out.extend(self._decompose_move(p))258        return out259 260 261    @classmethod262    def build_structured_vocab(cls) -> "ChessTokenizer":263        special = [cls.PAD_TOKEN, cls.BOS_TOKEN, cls.EOS_TOKEN, cls.UNK_TOKEN]264 265        # 12 color+piece tokens266        cp = [f"{c}{p}" for c in ("W", "B") for p in ("P", "N", "B", "R", "Q", "K")]267 268        files = "abcdefgh"269        ranks = "12345678"270 271        from_tokens = [f"{f}{r}_f" for f in files for r in ranks]  # 64272        to_tokens = [f"{f}{r}_t" for f in files for r in ranks]    # 64273 274        suffix = ["cap", "check", "mate", "castle_k", "castle_q"]275        promo = [f"promo_{p}" for p in ("q", "r", "b", "n")]276 277        tokens = special + cp + from_tokens + to_tokens + suffix + promo278        vocab = {t: i for i, t in enumerate(tokens)}279        return cls(vocab=vocab)280 281    def _convert_token_to_id(self, token: str) -> int:282        """Convert a token to its ID."""283        return self._vocab.get(token, self._vocab.get(self.UNK_TOKEN, 0))284    285    def _convert_id_to_token(self, index: int) -> str:286        """Convert an ID to its token."""287        return self._ids_to_tokens.get(index, self.UNK_TOKEN)288    289    def convert_tokens_to_string(self, tokens: List[str]) -> str:290        """Convert a list of tokens back to a string."""291        # Filter out special tokens for cleaner output292        special = {self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN}293        return " ".join(t for t in tokens if t not in special)294    295    def save_vocabulary(296        self,297        save_directory: str,298        filename_prefix: Optional[str] = None,299    ) -> tuple:300        """301        Save the vocabulary to a JSON file.302        303        Args:304            save_directory: Directory to save the vocabulary.305            filename_prefix: Optional prefix for the filename.306        307        Returns:308            Tuple containing the path to the saved vocabulary file.309        """310        if not os.path.isdir(save_directory):311            os.makedirs(save_directory, exist_ok=True)312        313        vocab_file = os.path.join(314            save_directory,315            (filename_prefix + "-" if filename_prefix else "") + "vocab.json",316        )317        318        with open(vocab_file, "w", encoding="utf-8") as f:319            json.dump(self._vocab, f, ensure_ascii=False, indent=2)320        321        return (vocab_file,)322 323 324def count_vocab_from_dataset(325    dataset_name: str = "dlouapre/lichess_2025-01_1M",326    split: str = "train",327    column: str = "text",328    max_samples: Optional[int] = 10000,329) -> Dict[str, int]:330    """331    Count token frequencies in a dataset (useful for vocabulary analysis).332    333    Args:334        dataset_name: Name of the dataset on Hugging Face Hub.335        split: Dataset split to use.336        column: Column containing the game strings.337        max_samples: Maximum number of samples to process.338    339    Returns:340        Dictionary mapping tokens to their frequencies.341    """342    from collections import Counter343    from datasets import load_dataset344    345    dataset = load_dataset(dataset_name, split=split)346    347    if max_samples is not None:348        dataset = dataset.select(range(min(max_samples, len(dataset))))349    350    token_counts = Counter()351    352    for example in dataset:353        moves = example[column].strip().split()354        token_counts.update(moves)355    356    return dict(token_counts)357