Simple chess AI

requirements.txt

@@ -4,3 +4,4 @@ torch
 pandas
 pandas-stubs
 matplotlib
+safetensors

src/inference/simpleChessAI.py

@@ -0,0 +1,85 @@
+import torch
+import chess
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+class LightweightFlyInference:
+    def __init__(self, model_name="lazy-guy12/chess-llama"):
+        self.model_name = model_name
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        print(f"Auto-downloading and loading '{self.model_name}' from Hugging Face...")
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(self.model_name).to(self.device)
+        self.model.eval()
+        
+        # Pre-cache vocabulary token IDs to valid python-chess Move objects
+        self._build_vocab_move_cache()
+        print("Model successfully loaded and ready for inference.")
+
+    def _build_vocab_move_cache(self):
+        """Maps chess-llama's custom vocabulary tokens to valid python-chess Move objects."""
+        self.vocab_map = {}
+        for idx in range(len(self.tokenizer)):
+            token_str = self.tokenizer.decode([idx]).strip()
+            token_str = token_str.split()[0] if token_str.split() else ""
+            try:
+                if len(token_str) >= 4:
+                    move = chess.Move.from_uci(token_str)
+                    self.vocab_map[idx] = move
+            except ValueError:
+                continue
+
+    def forward(self, board: chess.Board) -> chess.Move:
+        """
+        Evaluates board state using model logits and selects the highest-probability 
+        legal move from the model's vocabulary distribution.
+        """
+        move_uci_list = [move.uci() for move in board.move_stack]
+        
+        # Handle FEN initialization check if move stack is empty
+        if not move_uci_list and board.fen().startswith("rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR"):
+            move_uci_list = ["e2e4"]
+
+        prompt = " ".join(move_uci_list)
+        if prompt:
+            prompt += " "
+
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+
+        with torch.no_grad():
+            # Get next-token logits directly from the model forward pass
+            outputs = self.model(**inputs)
+            next_token_logits = outputs.logits[0, -1, :] # Shape: (vocab_size,)
+
+        # Sort vocabulary indices by model confidence score (descending order)
+        sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
+
+        # Iterate through the model's predictions from most confident to least confident,
+        # and pick the first one that is completely legal for the current board state.
+        for idx in sorted_indices.tolist():
+            if idx in self.vocab_map:
+                move = self.vocab_map[idx]
+                if move in board.legal_moves:
+                    return move
+
+        # Ultimate fallback to the first legal move if no vocabulary match is found
+        print("Warning: No valid move found in model vocabulary. Falling back to first legal move.")
+        legal_moves = list(board.legal_moves)
+        if legal_moves:
+            return legal_moves[0]
+        
+        return None
+
+if __name__ == "__main__":
+    inference = LightweightFlyInference()
+
+    board = chess.Board("rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR b KQkq - 0 1")
+    best_move = inference.forward(board)
+    
+    if best_move:
+        print(f"Recommended Move (UCI): {best_move.uci()}")
+        print(f"Recommended Move (SAN): {board.san(best_move)}")
+        board.push(best_move)
+        print("\nUpdated Chess Board:\n", board)
+    else:
+        print("No legal moves available.")