Update app.py

app.py

@@ -1,128 +1,307 @@
-!pip install transformers chess torch 
+import io
+import traceback
+from typing import List
 
-# %%
-from transformers import AutoModelForCausalLM, GPT2LMHeadModel
-import torch
-import chess
 import chess.pgn
 import chess.svg
 import gradio as gr
-from modeling.utils import uci_to_board
-from modeling.uci_tokenizers import UciTileTokenizer
 import numpy as np
-import io
-import traceback
-# %%
-checkpoint_name = "austindavis/chess-gpt2-uci-8x8x512"
-tokenizer = UciTileTokenizer()
-model: GPT2LMHeadModel = AutoModelForCausalLM.from_pretrained(checkpoint_name)
-model.requires_grad_(False)
+import tokenizers
+import torch
+from tokenizers import models, pre_tokenizers, processors
+from torch import Tensor as TT
+from transformers import AutoModelForCausalLM, GPT2LMHeadModel, PreTrainedTokenizerFast
+
+import chess
 
-# %%
-# Initialize the chess board
-board = chess.Board()
-game:chess.pgn.GameNode = chess.pgn.Game()
+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"""
 
-game.headers["Event"] = "Example"
+    def __init__(
+        self,
+        stoi,
+        itos,
+        pad_token,
+        unk_token,
+        bos_token,
+        eos_token,
+        name_or_path,
+    ):
+        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
 
-generate_kwargs = {
-                "max_new_tokens": 3,
-                "num_return_sequences": 10,
-                "temperature": 0.5, 
-                "output_scores": True,
-                "output_logits": True,
-                "return_dict_in_generate": True
-                }
+        # Define the model
+        tok_model = models.WordLevel(vocab=self.stoi, unk_token=self._UNK_TOKEN)
 
-def make_move(input:str, node=game, board = board):
-    # check for reset
-    if input.lower() == 'reset':
-        board.reset()
-        node.root().variations.clear()
-        return chess.svg.board(board=board), "New game!"
-    
-    # check for pgn
-    if input[0] == '[' or input[:3] == '1. ':
-        pgn = io.StringIO(input)
-        game = chess.pgn.read_game(pgn)
-        board.reset()
-        node.root().variations.clear()
-
-        for move in game.mainline_moves():
-            board.push(move)
-            node.add_variation(move)
-
-        return chess.svg.board(board=board,lastmove=move), ""#str(node.root()).split(']')[-1].strip()
-
-
-    try:
-        move = chess.Move.from_uci(input)
-        if move in board.legal_moves:
-            board.push(move)
-
-            while node.next() is not None:
-                node = node.next()
-            node = node.add_variation(move)
-
-            # get computer's move
-
-            prefix = ' '.join([x.uci() for x in board.move_stack])
-            encoding = tokenizer(text=prefix,
-                return_tensors='pt', 
-                )['input_ids']
-
-            output = model.generate(encoding, **generate_kwargs) # [b,p,v]
-            new_tokens = tokenizer.batch_decode(output.sequences[:,-3:])
-            unique_moves, unique_indices = np.unique([x[:4] if ' ' in x else x for x in new_tokens], return_index=True)
-            unique_indices = torch.Tensor(list(unique_indices)).to(dtype=torch.int)
-            logits = torch.stack(output.logits) # [token, batch, vocab]
-            logits = logits[:,unique_indices]  # [token, batch, vocab]
+        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,
+        )
+
+        # Override the decode behavior to ensure spaces are correctly handled
+        def _decode(
+            token_ids: int | List[int],
+            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()
             
-            # select moves based on mean logit value for tokens 1 and 2
-            logit_priority_order = logits.max(dim=-1).values.T[:,:2].mean(-1).topk(len(unique_indices)).indices
-            priority_ordered_moves = unique_moves[logit_priority_order]
+            if isinstance(token_ids, list):
+                tokens_str = [self.itos.get(xi, self._UNK_TOKEN) for xi in token_ids]
+                moves = self._process_str_tokens(tokens_str)
+
+                return " ".join(moves)
             
-            # if there's only 1 option, we have to pack it back into a list
-            if isinstance(priority_ordered_moves, str):
-                priority_ordered_moves = [priority_ordered_moves]
-
-            # test if any moves are valid
-            for uci in priority_ordered_moves:
-                move = chess.Move.from_uci(uci)
-                if move in board.legal_moves:
-                    board.push(move)
-                    while node.next() is not None:
-                        node = node.next()
-                    node = node.add_variation(move)
-                    return chess.svg.board(board=board,lastmove=move), ""#str(node.root()).split(']')[-1].strip()
             
-            # no moves are valid
-            bad_from_tiles = [chess.parse_square(x) for x in [x[:2] for x in unique_moves]]
-            bad_to_tiles = [chess.parse_square(x) for x in [x[2:] for x in unique_moves]]
-            arrows = [chess.svg.Arrow(tail, head, color="red") for (tail, head) in zip(bad_from_tiles, bad_to_tiles)]
-            return chess.svg.board(board=board,arrows=arrows), '|'.join(unique_moves)
-        else:
-            return chess.svg.board(board=board,lastmove=move), f"Illegal move:  {input}"
+
+        self._decode = _decode
+
+    def _init_pretokenizer(self) -> pre_tokenizers.PreTokenizer:
+        raise NotImplementedError
+
+    def _process_str_tokens(self, tokens_str: list[str]) -> 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"""
+    stoi = {
+        tok: idx
+        for tok, idx in list(
+            zip(["<pad>", "<s>", "</s>", "<unk>"] + chess.SQUARE_NAMES + list("qrbn"), range(72))
+        )
+    }
+    
+    itos = {
+        idx: tok
+        for tok, idx in list(
+            zip(["<pad>", "<s>", "</s>", "<unk>"] + chess.SQUARE_NAMES + list("qrbn"), range(72))
+        )
+    }
+
+    id2square:List[int] = [None]*4 + list(range(64))+[None]*4
+    """
+    List mapping token IDs to squares on the chess board. Order is file then row, i.e.: 
+    `A1, B1, C1, ..., F8, G8, H8`    
+    """
+    
+    def get_id2square_list(self) -> List[int]:
+        return self.id2square
+
+    def __init__(self):
+
+        super().__init__(
+            self.stoi,
+            self.itos,
+            pad_token="<pad>",
+            unk_token="<unk>",
+            bos_token="<s>",
+            eos_token="</s>",
+            name_or_path="austindavis/uci_tile_tokenizer",
+        )
+
+    def _init_pretokenizer(self):
+        # Pre-tokenizer to split input into UCI moves
+        pattern = tokenizers.Regex(r"\d")
+        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):
+        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:
+                moves.append(next_move + token)
+                continue
+
+            # handle regular tokens
+            if len(next_move) == 4:
+                moves.append(next_move)
+                next_move = token
+            else:
+                next_move += token
+
+        moves.append(next_move)
+        return moves
     
-    except ValueError:
-        return chess.svg.board(board=board), f"Invalid UCI format:  {uci} {list(unique_moves)}"
-    except Exception:
-        return chess.svg.board(board=board), traceback.format_exc()
-
-# Define the Gradio interface
-iface = gr.Interface(
-    fn=make_move,
-    inputs="text",
-    outputs=["html", "text"],
-    examples=[['e2e4'], ['d2d4'], ['Reset']],
-    title="Play Versus ChessGPT",
-    description="Enter moves in UCI notation (e.g., e2e4 for pawn from e2 to e4). Enter 'reset' to restart the game.",
-    allow_flagging='never',
-)
-
-# Launch the Gradio app
-iface.launch()
-
-# %%
+def setup_app(model: GPT2LMHeadModel):
+    """
+    Configures a Gradio App to use the GPT model for move generation. 
+    The model must be compatible with a UciTileTokenizer.
+    """
+    tokenizer = UciTileTokenizer()
+
+    # Initialize the chess board
+    board = chess.Board()
+    game:chess.pgn.GameNode = chess.pgn.Game()
+
+
+
+    game.headers["Event"] = "Example"
+
+    generate_kwargs = {
+                    "max_new_tokens": 3,
+                    "num_return_sequences": 10,
+                    "temperature": 0.5, 
+                    "output_scores": True,
+                    "output_logits": True,
+                    "return_dict_in_generate": True
+                    }
+
+    def make_move(input:str, node=game, board = board):
+        # check for reset
+        if input.lower() == 'reset':
+            board.reset()
+            node.root().variations.clear()
+            return chess.svg.board(board=board), "New game!"
+        
+        # check for pgn
+        if input[0] == '[' or input[:3] == '1. ':
+            pgn = io.StringIO(input)
+            game = chess.pgn.read_game(pgn)
+            board.reset()
+            node.root().variations.clear()
+
+            for move in game.mainline_moves():
+                board.push(move)
+                node.add_variation(move)
+
+            return chess.svg.board(board=board,lastmove=move), ""#str(node.root()).split(']')[-1].strip()
+
+
+        try:
+            move = chess.Move.from_uci(input)
+            if move in board.legal_moves:
+                board.push(move)
+
+                while node.next() is not None:
+                    node = node.next()
+                node = node.add_variation(move)
+
+                # get computer's move
+
+                prefix = ' '.join([x.uci() for x in board.move_stack])
+                encoding = tokenizer(text=prefix,
+                    return_tensors='pt', 
+                    )['input_ids']
+
+                output = model.generate(encoding, **generate_kwargs) # [b,p,v]
+                new_tokens = tokenizer.batch_decode(output.sequences[:,-3:])
+                unique_moves, unique_indices = np.unique([x[:4] if ' ' in x else x for x in new_tokens], return_index=True)
+                unique_indices = torch.Tensor(list(unique_indices)).to(dtype=torch.int)
+                logits = torch.stack(output.logits) # [token, batch, vocab]
+                logits = logits[:,unique_indices]  # [token, batch, vocab]
+                
+                # select moves based on mean logit value for tokens 1 and 2
+                logit_priority_order = logits.max(dim=-1).values.T[:,:2].mean(-1).topk(len(unique_indices)).indices
+                priority_ordered_moves = unique_moves[logit_priority_order]
+                
+                # if there's only 1 option, we have to pack it back into a list
+                if isinstance(priority_ordered_moves, str):
+                    priority_ordered_moves = [priority_ordered_moves]
+
+                # test if any moves are valid
+                for uci in priority_ordered_moves:
+                    move = chess.Move.from_uci(uci)
+                    if move in board.legal_moves:
+                        board.push(move)
+                        while node.next() is not None:
+                            node = node.next()
+                        node = node.add_variation(move)
+                        return chess.svg.board(board=board,lastmove=move), "".join(str(node.root()).split("]")[-1]).strip()
+                
+                # no moves are valid
+                bad_from_tiles = [chess.parse_square(x) for x in [x[:2] for x in unique_moves]]
+                bad_to_tiles = [chess.parse_square(x) for x in [x[2:] for x in unique_moves]]
+                arrows = [chess.svg.Arrow(tail, head, color="red") for (tail, head) in zip(bad_from_tiles, bad_to_tiles)]
+                checks = None
+                if board.is_check():
+                    checks = board.pieces(chess.PIECE_TYPES[-1],board.turn).pop()
+
+                return chess.svg.board(board=board,arrows=arrows, check=checks), '|'.join(unique_moves)
+            else:
+                return chess.svg.board(board=board,lastmove=move), f"Illegal move:  {input}"
+        
+        except chess.InvalidMoveError:
+            return chess.svg.board(board=board), f"Invalid UCI format:  {input}"
+        except Exception:
+            return chess.svg.board(board=board), traceback.format_exc()
+
+    input_box = gr.Textbox(None,placeholder="Enter your move in UCI format")
+
+    # Define the Gradio interface
+    iface = gr.Interface(
+        fn=make_move,
+        inputs=input_box,
+        outputs=["html", "text"],
+        examples=[['e2e4'], ['d2d4'], ['Reset']],
+        title="Play Versus ChessGPT",
+        description="Enter moves in UCI notation (e.g., e2e4 for pawn from e2 to e4). Enter 'reset' to restart the game.",
+        allow_flagging='never',
+        submit_btn = "Move",
+        stop_btn = "Stop",
+        clear_btn = "Clear w/o reset",
+    )
+
+    iface.output_components[0].label = "Board"
+    iface.output_components[0].show_label = True
+    iface.output_components[1].label = "Move Sequence"
+
+    return iface
+
+checkpoint_name = "austindavis/gpt2-lichess-uci-202306"
+model: GPT2LMHeadModel = AutoModelForCausalLM.from_pretrained(checkpoint_name)
+model.requires_grad_(False)
+
+iface = setup_app(model)
+iface.launch()