tokenizer.py

"""
Custom Chess Tokenizer V3 for the Chess Challenge.

Enhanced version with additional chess-specific tokens for:
- Castling moves (O-O, O-O-O)
- Check/checkmate indicators (+, #)
- Capture indicator (x)
- Turn indicators ([WHITE], [BLACK])

This provides richer context while keeping vocabulary minimal (81 tokens total).
"""

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Dict, List, Optional
import re

from transformers import PreTrainedTokenizer


class ChessTokenizer(PreTrainedTokenizer):
    """
    Enhanced chess tokenizer with special chess notation tokens.
    
    Vocabulary (79 tokens):
    - 4 special tokens: [PAD], [BOS], [EOS], [UNK]
    - 64 square tokens: a1-h8
    - 4 promotion tokens: q, r, b, n
    - 2 castling tokens: O-O, O-O-O
    - 3 modifier tokens: +, #, x (check, checkmate, capture)
    - 2 turn tokens: [WHITE], [BLACK]
    """
    
    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]"
    WHITE_TOKEN = "[WHITE]"
    BLACK_TOKEN = "[BLACK]"
    
    def __init__(
        self,
        vocab_file: Optional[str] = None,
        vocab: Optional[Dict[str, int]] = None,
        **kwargs,
    ):
        self._pad_token = self.PAD_TOKEN
        self._bos_token = self.BOS_TOKEN
        self._eos_token = self.EOS_TOKEN
        self._unk_token = self.UNK_TOKEN

        kwargs.pop("pad_token", None)
        kwargs.pop("bos_token", None)
        kwargs.pop("eos_token", None)
        kwargs.pop("unk_token", None)
        
        # Enhanced regex pattern for chess notation
        # Matches: squares, promotions, castling, modifiers, turn indicators
        self.token_pattern = re.compile(
            r'O-O-O|O-O|'  # Castling (match O-O-O first!)
            r'\[WHITE\]|\[BLACK\]|'  # Turn indicators
            r'[a-h][1-8]|'  # Squares
            r'[qrbn]|'  # Promotions
            r'[+#x]'  # Check, checkmate, capture
        )

        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:
            self._vocab = self._create_default_vocab()
        
        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]:
        """
        Create the complete vocabulary with all chess-specific tokens.
        
        Total: 79 tokens
        """
        vocab = {}
        idx = 0
        
        # Special tokens (0-3)
        for token in [self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN]:
            vocab[token] = idx
            idx += 1
        
        # Squares (4-67)
        for f in 'abcdefgh':
            for r in '12345678':
                vocab[f"{f}{r}"] = idx
                idx += 1
        
        # Promotions (68-71)
        for p in ['q', 'r', 'b', 'n']:
            vocab[p] = idx
            idx += 1
        
        # Castling (72-73)
        vocab["O-O"] = idx
        idx += 1
        vocab["O-O-O"] = idx
        idx += 1
        
        # Modifiers (74-76)
        vocab["+"] = idx  # Check
        idx += 1
        vocab["#"] = idx  # Checkmate
        idx += 1
        vocab["x"] = idx  # Capture
        idx += 1
        
        # Turn indicators (77-78)
        vocab[self.WHITE_TOKEN] = idx
        idx += 1
        vocab[self.BLACK_TOKEN] = idx
        idx += 1
        
        return vocab
    
    def _tokenize(self, text: str) -> List[str]:
        """
        Enhanced tokenization with preprocessing for common chess notation variants.
        
        Handles:
        - Lichess format: (Q) → q, (x) → x, (+) → +, (#) → #
        - Standard notation: keeps O-O, O-O-O, +, #, x as-is
        - Extracts squares, promotions, castling, and modifiers
        """
        # Normalize Lichess-style parentheses notation
        text = (text.replace("(Q)", "q")
                    .replace("(R)", "r")
                    .replace("(B)", "b")
                    .replace("(N)", "n")
                    .replace("(x)", "x")
                    .replace("(+)", "+")
                    .replace("(#)", "#")
                    .replace("(+*)", "#")  # Checkmate variant
                    .replace("(o)", "O-O")  # Kingside castling
                    .replace("(O)", "O-O-O"))  # Queenside castling
        
        # Extract all chess tokens
        return self.token_pattern.findall(text)
    
    def _convert_token_to_id(self, token: str) -> int:
        """Convert a token to its ID."""
        return self._vocab.get(token, self._vocab.get(self.UNK_TOKEN, 0))
    
    def _convert_id_to_token(self, index: int) -> str:
        """Convert an ID to its token."""
        return self._ids_to_tokens.get(index, self.UNK_TOKEN)
    
    def convert_tokens_to_string(self, tokens: List[str]) -> str:
        """
        Reconstructs chess moves in standard UCI format with modifiers.
        
        Intelligently groups tokens:
        - Combines squares into moves: e2, e4 → e2e4
        - Attaches promotions: a7, a8, q → a7a8q
        - Keeps modifiers separate: e2e4, x, + → e2e4x+
        - Preserves castling and turn indicators
        """
        special = {self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN}
        clean_tokens = [t for t in tokens if t not in special]
        
        output = []
        modifiers = {'+', '#', 'x'}
        promotions = {'q', 'r', 'b', 'n'}
        
        for token in clean_tokens:
            # Castling and turn indicators stay as-is
            if token in ["O-O", "O-O-O", self.WHITE_TOKEN, self.BLACK_TOKEN]:
                output.append(token)
            # Promotions attach to previous move
            elif token in promotions and output and len(output[-1]) == 4:
                output[-1] += token
            # Modifiers can attach or stay separate (flexible)
            elif token in modifiers and output:
                output[-1] += token
            # Square: either start new move or complete previous
            elif len(token) == 2 and token[0] in 'abcdefgh':
                if output and len(output[-1]) == 2 and output[-1][0] in 'abcdefgh':
                    # Complete the move
                    output[-1] += token
                else:
                    # Start new move
                    output.append(token)
            else:
                output.append(token)
        
        return " ".join(output)
    
    def add_turn_indicators(self, text: str, add_white_indicator: bool = True) -> str:
        """
        Add turn indicators to help the model understand whose turn it is.
        
        Args:
            text: Game string (space-separated moves)
            add_white_indicator: If True, add [WHITE] at start (white moves first)
        
        Returns:
            Game string with turn indicators
        """
        moves = text.strip().split()
        result = []
        
        # White starts (by convention)
        is_white = add_white_indicator
        
        for move in moves:
            turn_token = self.WHITE_TOKEN if is_white else self.BLACK_TOKEN
            result.append(turn_token)
            result.append(move)
            is_white = not is_white
        
        return " ".join(result)
    
    def save_vocabulary(
        self,
        save_directory: str,
        filename_prefix: Optional[str] = None,
    ) -> tuple:
        """Save the vocabulary to a JSON file."""
        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,)
    
    @classmethod
    def build_vocab_from_iterator(cls, iterator, min_frequency=1):
        """Returns tokenizer with fixed vocabulary (doesn't depend on data)."""
        return cls()
    
    @classmethod
    def build_vocab_from_dataset(cls, **kwargs):
        """Returns tokenizer with fixed vocabulary (doesn't depend on data)."""
        return cls()
    
    @property
    def vocab_size(self) -> int:
        """Return the size of the vocabulary (79 tokens)."""
        return len(self._vocab)
    
    def get_vocab(self) -> Dict[str, int]:
        """Return the vocabulary as a dictionary."""
        return dict(self._vocab)