agents/uci_tokenizers.py

from typing import List

import chess
import tiktoken
import tokenizers
from tokenizers import models, pre_tokenizers, processors
from torch import Tensor as TT
from transformers import PreTrainedTokenizerFast
from transformers.tokenization_utils_fast import BatchEncoding


def getTiktokenizer() -> tiktoken.Encoding:
    """
    Defines a tiktoken-based BPE encoder for UCI chess moves. This
    tokenizer effectively tokenizes UCI moves by the square names.
    One notable variation is that promotions must be in upper-case.

    Vocabulary:
    Special Tokens (4): "\<|pad|\>", "\<|startoftext|\>", "\<|endoftext|\>", "\<|unknown|\>"
    Square Tokens (64): a1 through h8
    Promote Tokens (4): Q, B, R, N
    UNUSED (8120): Need 8192-4-64-4=8120 unused tokens of the form <|unused####|>
    """
    special_tokens = ["<|pad|>", "<|startoftext|>", "<|endoftext|>", "<|unknown|>"]
    unused_tokens = [f"<|unused{i:04d}" for i in range(8120)]
    chess_vocab = special_tokens + chess.SQUARE_NAMES + list("QBRN") + unused_tokens
    mergeable_ranks = {k.encode():v for (v,k) in enumerate(chess_vocab)}
    chess_pat_str = r'[a-h][1-8]|[QBRN]'

    enc = tiktoken.Encoding(
        name="chess_enc",
        pat_str=chess_pat_str, # or \d|\s
        mergeable_ranks=mergeable_ranks,
        special_tokens={k:v for (v,k) in enumerate(special_tokens)},
    )

    return enc


class UciTokenizer(PreTrainedTokenizerFast):
    _PAD_TOKEN: str
    _UNK_TOKEN: str
    _EOS_TOKEN: str
    _BOS_TOKEN: str
    

    stoi: dict[str, int]
    """Integer to String mapping"""

    itos: dict[int, str]
    """String to Integer Mapping. This is the vocab"""

    def __init__(
        self,
        stoi,
        itos,
        pad_token,
        unk_token,
        bos_token,
        eos_token,
        name_or_path,
        **kwargs
    ):
        self.stoi = stoi
        self.itos = itos
        
        self._PAD_TOKEN = pad_token
        self._UNK_TOKEN = unk_token
        self._EOS_TOKEN = eos_token
        self._BOS_TOKEN = bos_token

        # Define the model
        tok_model = models.WordLevel(vocab=self.stoi, unk_token=self._UNK_TOKEN)

        slow_tokenizer = tokenizers.Tokenizer(tok_model)
        slow_tokenizer.pre_tokenizer = self._init_pretokenizer()

        # post processing adds special tokens unless explicitly ignored
        post_proc = processors.TemplateProcessing(
            single=f"{bos_token} $0",
            pair=None,
            special_tokens=[(bos_token, 1)],
        )
        slow_tokenizer.post_processor=post_proc
        
        super().__init__(
            tokenizer_object=slow_tokenizer,
            unk_token=self._UNK_TOKEN,
            bos_token=self._BOS_TOKEN,
            eos_token=self._EOS_TOKEN,
            pad_token=self._PAD_TOKEN,
            name_or_path=name_or_path,
            **kwargs
        )

        # Override the decode behavior to ensure spaces are correctly handled
        def _decode(
            token_ids: int | List[int] | dict | TT,
            skip_special_tokens=False,
            clean_up_tokenization_spaces=False,
        ) -> int | List[int]:

            if isinstance(token_ids, int):
                return self.itos.get(token_ids, self._UNK_TOKEN)

            if isinstance(token_ids, dict):
                token_ids = token_ids["input_ids"]

            if isinstance(token_ids, TT):
                token_ids = token_ids.tolist()
            
            if isinstance(token_ids, list):
                tokens_str = [self.itos.get(xi, self._UNK_TOKEN) for xi in token_ids]
                processed_tokens = self._process_str_tokens(tokens_str)

                return " ".join(processed_tokens)
            
            raise ValueError(f"Unknown input type to decode() for argument 'token_ids'. Received: {type(token_ids)} ")
            

        self._decode = _decode

    def _init_pretokenizer(self) -> pre_tokenizers.PreTokenizer:
        raise NotImplementedError

    def _process_str_tokens(self, tokens_str: list[str], return_player_ids: bool) -> list[str]:
        raise NotImplementedError
    
    def get_id2square_list() -> list[int]:
        raise NotImplementedError


class UciTileTokenizer(UciTokenizer):
    """ Uci tokenizer converting start/end tiles and promotion types each into individual tokens"""

    SPECIAL_TOKENS = ["<|pad|>", "<|startoftext|>", "<|endoftext|>", "<|unknown|>"]

    stoi = {
        tok: idx
        for tok, idx in list(
            zip(SPECIAL_TOKENS + chess.SQUARE_NAMES + list("QRBN"), range(72))
        )
    }

    itos = {
        idx: tok
        for tok, idx in list(
            zip(SPECIAL_TOKENS + chess.SQUARE_NAMES + list("QRBN"), range(72))
        )
    }

    id2square:List[int] = list(range(4,68))
    """
    List mapping token IDs to squares on the chess board. Order is file then rank, i.e.: 
    `A1, B1, C1, ..., F8, G8, H8`    
    """

    def get_id2square_list(self) -> List[int]:
        return self.id2square

    def __init__(self, **kwargs):
        super().__init__(
            self.stoi,
            self.itos,
            pad_token="<|pad|>",
            unk_token="<|unknown|>",
            bos_token="<|startoftext|>",
            eos_token="<|endoftext|>",
            name_or_path="austindavis/uci_tile_tokenizer",
            clean_up_tokenization_spaces=False,
            **kwargs
        )

    def _init_pretokenizer(self):
        # Pre-tokenizer to split input into UCI moves
        pattern = tokenizers.Regex(r"\d|[QBRN]")
        pre_tokenizer = pre_tokenizers.Sequence(
            [
                pre_tokenizers.Whitespace(),
                pre_tokenizers.Split(pattern=pattern, behavior="merged_with_previous"),
            ]
        )
        return pre_tokenizer

    def _process_str_tokens(self, token_str: list[str]):
        moves = []
        next_move = ""
        for token in token_str:

            # skip special tokens
            if token in self.all_special_tokens:
                continue

            # handle promotions
            if len(token) == 1:
                next_move += token
                continue

            # handle regular tokens if there's room
            if len(next_move) < 4:
                next_move += token
                continue

            moves.append(next_move)
            next_move = token

        moves.append(next_move)
        return moves

    @staticmethod
    def compute_players(encoding: BatchEncoding, according_to='output'):
        """
        Determines which player (white=True, black=False) is associated with each token in the sequence. 
        This method works based on chess move sequences tokenized using the UciTileTokenizer.

        # Parameters:
        ----------
        **`encoding`** : BatchEncoding
            Tokenized input of a chess game, where each token represents a move or special token.
            
        **`according_to`** : str (optional, default='output')
            Specifies the perspective for associating players:
            - 'output': Returns the player whose next move is predicted by the sequence (the output move).
            - Otherwise: Returns the player associated with the input tokens (i.e., which player made each move).

        # Returns:
        -------
        List[bool]
            A list of boolean values indicating the player for each token:
            - True for white (player 1), 
            - False for black (player 2).
            
            The list length corresponds to the number of tokens in the sequence, including special tokens if any.

        # Example Usage: 
        ```
        >>> tok = UciTileTokenizer()
        >>> encoding = tok('e2e4 d7d5 e4d5 e7e6 d5e6 d8g5 e6e7 g5f6 e7f8Q')
        >>> print(encoding['input_ids'])
        [1, 16, 32, 55, 39, 32, 39, 56, 48, 39, 48, 63, 42, 48, 56, 42, 49, 56, 65, 68]
        >>> tok.compute_players(encoding)
        [True, True, False, False, True, True, False, False, True, True, False, False, True, True, False, False, True, True, True, False]
        >>> tok.compute_players(encoding, according_to='input')        
        [True, True, True, False, False, True, True, False, False, True, True, False, False, True, True, False, False, True, True, True]
        ```

        # Notes:
        -------
        This method does not rely on board position calculations. Therefore, when
        using `according_to='output'`, it cannot reliably predict which player is
        responsible for selecting the final token of the sequence. For instance,
        if a pawn is moved to the back rank (e.g., 'e7e8'), then white must select
        the promotion class on the next token; however, this algorithm will predict 
        that black is responsible for selecting the next token instead of white. 
        """
        
        return [UciTileTokenizer._compute_players_single(encoding[i].ids) for i in range(len(encoding['input_ids']))]
            
        
    
    @staticmethod
    def _compute_players_single(input_ids: list[int], according_to: str='output'):
        players = [] if according_to == "output" else [True]
        current_player = False
        num_tokens_in_ply = 0
        has_specials = False
        
        for i, token_id in enumerate(input_ids):
            if token_id == 1:
                has_specials = True
                continue
            
            if num_tokens_in_ply == 0:
                # check if promotion OR unknown token ID
                if token_id > 67 or token_id == 3: 
                    players.append(current_player)
                    num_tokens_in_ply = 0
                else:
                    num_tokens_in_ply += 1
                    current_player = not current_player
                    players.append(current_player)
            elif num_tokens_in_ply == 1:
                num_tokens_in_ply = 0
                players.append(current_player)
            else:
                raise ValueError("Illegal move sequence")

        if according_to == "output":
            # anticipate what output should be based on the final input token
            # see notes for more detail
            if num_tokens_in_ply == 0:
                if token_id > 67:
                    players.append(not current_player)
                else:
                    players.append(current_player)
            else:
                players.append(current_player)

        return players if has_specials else players[1:]

if __name__ == "__main__":
    tok = UciTileTokenizer()
    encoding = tok('e2e4Q b7b8N e2e7 a1',add_special_tokens=True)
    print(f"{encoding['input_ids']=}\n{tok.compute_players(encoding,  according_to='output')=}")
    print(f"{encoding['input_ids']=}\n{tok.compute_players(encoding, according_to='input')=}")

    encoding = tok('e2e4Q b7b8N e2e7 a1',add_special_tokens=False)
    print(f"{encoding['input_ids']=}\n{tok.compute_players(encoding,  according_to='output')=}")
    print(f"{encoding['input_ids']=}\n{tok.compute_players(encoding, according_to='input')=}")

    encoding = tok('e2e4 d7d5 e4d5 e7e6 d5e6 d8g5 e6e7 g5f6 e7f8Q')
    print(encoding['input_ids'])
    print(tok.compute_players(encoding))
    print(tok.compute_players(encoding, according_to='input'))