ChessBot / usage_example.py
Maxlegrec's picture
Upload ChessBot Chess model
311cfe1 verified
raw
history blame
1.14 kB
"""
Example usage of the ChessBot Chess Model
This model can be used without installing any external packages except:
- torch
- transformers
- chess (python-chess)
- numpy
"""
import torch
import sys
sys.path.append("./") # Add the model directory to path
from modeling_chessbot import ChessBotModel, ChessBotConfig
# Load the model
config = ChessBotConfig()
model = ChessBotModel.from_pretrained("./")
# Alternative: You can also try AutoModel (may require additional setup)
# from transformers import AutoModel
# model = AutoModel.from_pretrained("./", trust_remote_code=True)
# Example usage
fen = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
# Get the best move
move = model.get_move_from_fen_no_thinking(fen, T=0.1, device=device)
print(f"Best move: {move}")
# Get move probabilities
probs = model.get_move_from_fen_no_thinking(fen, T=0.1, device=device, return_probs=True)
top_moves = sorted(probs.items(), key=lambda x: x[1], reverse=True)[:5]
print("Top 5 moves:")
for move, prob in top_moves:
print(f" {move}: {prob:.4f}")