Update app.py

app.py

@@ -2,6 +2,7 @@ import io
 import traceback
 from typing import List
 
+import chess
 import chess.pgn
 import chess.svg
 import gradio as gr
@@ -10,12 +11,12 @@ 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
+from transformers import (AutoModelForCausalLM, GPT2LMHeadModel,
+                          PreTrainedTokenizerFast)
 
 checkpoint_name = "austindavis/chess-gpt2-uci-8x8x512"
 
+
 class UciTokenizer(PreTrainedTokenizerFast):
     _PAD_TOKEN: str
     _UNK_TOKEN: str
@@ -40,14 +41,15 @@ class UciTokenizer(PreTrainedTokenizerFast):
     ):
         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)
+        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()
@@ -58,8 +60,8 @@ class UciTokenizer(PreTrainedTokenizerFast):
             pair=None,
             special_tokens=[(bos_token, 1)],
         )
-        slow_tokenizer.post_processor=post_proc
-        
+        slow_tokenizer.post_processor = post_proc
+
         super().__init__(
             tokenizer_object=slow_tokenizer,
             unk_token=self._UNK_TOKEN,
@@ -84,14 +86,13 @@ class UciTokenizer(PreTrainedTokenizerFast):
 
             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]
+                tokens_str = [self.itos.get(xi, self._UNK_TOKEN)
+                              for xi in token_ids]
                 moves = self._process_str_tokens(tokens_str)
 
                 return " ".join(moves)
-            
-            
 
         self._decode = _decode
 
@@ -100,32 +101,45 @@ class UciTokenizer(PreTrainedTokenizerFast):
 
     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"""
+    """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))
+            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))
+            zip(
+                ["<pad>", "<s>", "</s>", "<unk>"] +
+                chess.SQUARE_NAMES + list("qrbn"),
+                range(72),
+            )
         )
     }
 
-    id2square:List[int] = [None]*4 + list(range(64))+[None]*4
+    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`    
+    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
 
@@ -147,7 +161,8 @@ class UciTileTokenizer(UciTokenizer):
         pre_tokenizer = pre_tokenizers.Sequence(
             [
                 pre_tokenizers.Whitespace(),
-                pre_tokenizers.Split(pattern=pattern, behavior="merged_with_previous"),
+                pre_tokenizers.Split(pattern=pattern,
+                                     behavior="merged_with_previous"),
             ]
         )
         return pre_tokenizer
@@ -175,40 +190,39 @@ class UciTileTokenizer(UciTokenizer):
 
         moves.append(next_move)
         return moves
-    
+
+
 def setup_app(model: GPT2LMHeadModel):
     """
-    Configures a Gradio App to use the GPT model for move generation. 
+    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: 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):
+        "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':
+        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. ':
+        if input[0] == "[" or input[:3] == "1. ":
             pgn = io.StringIO(input)
             game = chess.pgn.read_game(pgn)
             board.reset()
@@ -218,8 +232,10 @@ def setup_app(model: GPT2LMHeadModel):
                 board.push(move)
                 node.add_variation(move)
 
-            return chess.svg.board(board=board,lastmove=move), ""#str(node.root()).split(']')[-1].strip()
-
+            return (
+                chess.svg.board(board=board, lastmove=move),
+                "",
+            )  # str(node.root()).split(']')[-1].strip()
 
         try:
             move = chess.Move.from_uci(input)
@@ -232,22 +248,35 @@ def setup_app(model: GPT2LMHeadModel):
 
                 # 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]
-                
+                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
+                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]
@@ -260,40 +289,61 @@ def setup_app(model: GPT2LMHeadModel):
                         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()
-                
+                        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)]
+                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)
+                    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}"
-        
+                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}"
+            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")
+    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']],
+        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",
-        share=True
+        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"
@@ -302,8 +352,9 @@ def setup_app(model: GPT2LMHeadModel):
 
     return iface
 
+
 model: GPT2LMHeadModel = AutoModelForCausalLM.from_pretrained(checkpoint_name)
 model.requires_grad_(False)
 
 iface = setup_app(model)
-iface.launch()
+iface.launch(share=True)