""" Custom Chess Tokenizer for the Chess Challenge. This tokenizer treats chess moves using a 'Square-Aware' Character strategy. Instead of full moves (e.g., WPe2e4), it splits them into meaningful atomic parts: - Pieces/Colors: W, B, P, N, B, R, Q, K - Full Squares: e2, e4, h8 (keeps coordinates together for geometric understanding) - Separators: Space " " Example: "WPe2e4" -> ["W", "P", "e2", "e4"] """ from __future__ import annotations import re import json import os from typing import Dict, List, Optional from transformers import PreTrainedTokenizer class ChessTokenizer(PreTrainedTokenizer): """ A custom tokenizer for chess moves using 'Square-Aware' tokenization. It maps atomic chess components (squares like 'e4', pieces like 'P') to IDs. This creates a small, dense vocabulary (~80 tokens) allowing deeper models. """ model_input_names = ["input_ids", "attention_mask"] vocab_files_names = {"vocab_file": "vocab.json"} # Special tokens PAD_TOKEN = "[PAD]" BOS_TOKEN = "[BOS]" EOS_TOKEN = "[EOS]" UNK_TOKEN = "[UNK]" def __init__( self, vocab_file: Optional[str] = None, vocab: Optional[Dict[str, int]] = None, **kwargs, ): # Initialize special tokens self._pad_token = self.PAD_TOKEN self._bos_token = self.BOS_TOKEN self._eos_token = self.EOS_TOKEN self._unk_token = self.UNK_TOKEN # Clean kwargs to avoid conflicts kwargs.pop("pad_token", None) kwargs.pop("bos_token", None) kwargs.pop("eos_token", None) kwargs.pop("unk_token", None) # Load or create vocabulary if vocab is not None: self._vocab = vocab elif vocab_file is not None and os.path.exists(vocab_file): with open(vocab_file, "r", encoding="utf-8") as f: self._vocab = json.load(f) else: # Minimal default vocab (placeholder) self._vocab = self._create_default_vocab() # Create reverse mapping (ID -> Token) self._ids_to_tokens = {v: k for k, v in self._vocab.items()} super().__init__( pad_token=self._pad_token, bos_token=self._bos_token, eos_token=self._eos_token, unk_token=self._unk_token, **kwargs, ) def _create_default_vocab(self) -> Dict[str, int]: special_tokens = [self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN] return {token: idx for idx, token in enumerate(special_tokens)} @classmethod def build_vocab_from_iterator( cls, iterator, min_frequency: int = 1, ) -> "ChessTokenizer": """ Build vocabulary by scanning the dataset. Splits text into pieces (W, P) and full squares (e2, e4). """ from collections import Counter token_counts = Counter() for game in iterator: # 1. Nettoyage : on enlève les suffixes (x), (+) game = re.sub(r'\(.*?\)', '', game) # 2. Découpage par coups pour gérer les espaces correctement moves = game.strip().split() for i, move in enumerate(moves): # Regex : Capture soit une case [a-h][1-8], soit n'importe quel autre char (.) tokens = re.findall(r'[a-h][1-8]|.', move) token_counts.update(tokens) # Ajout explicite de l'espace entre les coups (sauf après le dernier) if i < len(moves) - 1: token_counts.update([" "]) # Filter and sort tokens tokens = [t for t, count in token_counts.items() if count >= min_frequency] tokens = sorted(tokens) # Build final vocabulary dict special_tokens = [cls.PAD_TOKEN, cls.BOS_TOKEN, cls.EOS_TOKEN, cls.UNK_TOKEN] vocab = {token: idx for idx, token in enumerate(special_tokens + tokens)} return cls(vocab=vocab) @classmethod def build_vocab_from_dataset( cls, dataset_name: str = "dlouapre/lichess_2025-01_1M", split: str = "train", column: str = "text", min_frequency: int = 1, # Keep at 1 to catch all squares/pieces max_samples: Optional[int] = 50000, ) -> "ChessTokenizer": from datasets import load_dataset dataset = load_dataset(dataset_name, split=split) if max_samples is not None: dataset = dataset.select(range(min(max_samples, len(dataset)))) def game_iterator(): for example in dataset: yield example[column] return cls.build_vocab_from_iterator(game_iterator(), min_frequency=min_frequency) @property def vocab_size(self) -> int: return len(self._vocab) def get_vocab(self) -> Dict[str, int]: return dict(self._vocab) def _tokenize(self, text: str) -> List[str]: """ Tokenize input text using the Square-Aware logic. "WPe2e4" -> ["W", "P", "e2", "e4"] """ # 1. Remove suffixes text = re.sub(r'\(.*?\)', '', text) # 2. Split into moves to manage spaces moves = text.strip().split() all_tokens = [] for i, move in enumerate(moves): # Regex match: squares OR single chars tokens = re.findall(r'[a-h][1-8]|.', move) all_tokens.extend(tokens) # Re-insert space token between moves if i < len(moves) - 1: all_tokens.append(" ") return all_tokens def _convert_token_to_id(self, token: str) -> int: return self._vocab.get(token, self._vocab.get(self.UNK_TOKEN, 0)) def _convert_id_to_token(self, index: int) -> str: return self._ids_to_tokens.get(index, self.UNK_TOKEN) def convert_tokens_to_string(self, tokens: List[str]) -> str: """ Convert tokens back to string. IMPORTANT: Join with empty string "" because space " " is already a token. """ special = {self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN} filtered_tokens = [t for t in tokens if t not in special] # Join with "" because the space character is treated as a token in our vocab return "".join(filtered_tokens) def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> tuple: if not os.path.isdir(save_directory): os.makedirs(save_directory, exist_ok=True) vocab_file = os.path.join( save_directory, (filename_prefix + "-" if filename_prefix else "") + "vocab.json", ) with open(vocab_file, "w", encoding="utf-8") as f: json.dump(self._vocab, f, ensure_ascii=False, indent=2) return (vocab_file,) def count_vocab_from_dataset( dataset_name: str = "dlouapre/lichess_2025-01_1M", split: str = "train", column: str = "text", max_samples: Optional[int] = 10000, ) -> Dict[str, int]: # Utility function remains similar but should use the new regex logic if needed for analysis # For simple counting, split() is often enough approximation, but let's be precise: from collections import Counter from datasets import load_dataset import re dataset = load_dataset(dataset_name, split=split) if max_samples is not None: dataset = dataset.select(range(min(max_samples, len(dataset)))) token_counts = Counter() for example in dataset: text = re.sub(r'\(.*?\)', '', example[column]) moves = text.strip().split() for move in moves: tokens = re.findall(r'[a-h][1-8]|.', move) token_counts.update(tokens) token_counts.update([" "]) return dict(token_counts)