Create app.py

app.py

@@ -0,0 +1,126 @@
+# %%
+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/gpt2-lichess-uci-201601"
+tokenizer = UciTileTokenizer()
+model: GPT2LMHeadModel = AutoModelForCausalLM.from_pretrained(checkpoint_name)
+model.requires_grad_(False)
+
+# %%
+# 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), ""#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}"
+    
+    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()
+
+# %%