diff --git a/requirements.txt b/requirements.txt index b305160..32f5243 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,4 +3,5 @@ scikit-learn torch pandas pandas-stubs -matplotlib \ No newline at end of file +matplotlib +safetensors \ No newline at end of file diff --git a/src/inference/simpleChessAI.py b/src/inference/simpleChessAI.py new file mode 100644 index 0000000..3662e99 --- /dev/null +++ b/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.") \ No newline at end of file