Views
No views yet
| Variant | d_model | Layers | Heads | Parameters | Checkpoint |
|---|---|---|---|---|---|
| Small | 256 | 8 | 8 | ~19M | small/checkpoint.pt |
| Mid | 512 | 16 | 8 | ~100M | mid/checkpoint.pt |
| Large | 1024 | 40 | 32 | ~500M | large/checkpoint.pt |
1import torch
2from artoria import ChessTokenizer, GrandmasterChessModel, ChessModelConfig
3import json
4
5# Load config
6with open("small/config.json") as f:
7 config = ChessModelConfig(**json.load(f))
8
9tokenizer = ChessTokenizer()
10config.num_classes = tokenizer.num_actions
11
12model = GrandmasterChessModel(config)
13checkpoint = torch.load("small/checkpoint.pt", map_location="cpu")
14model.load_state_dict(checkpoint["model_state_dict"])
15model.eval()
16
17# Predict move from FEN
18fen = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
19tokens = tokenizer.tokenize(fen).unsqueeze(0)
20
21with torch.no_grad():
22 logits, value = model(tokens)
23
24# Get best legal move
25import chess
26board = chess.Board(fen)
27best_move, best_prob = None, -1
28probs = torch.softmax(logits[0], dim=0)
29for m in board.legal_moves:
30 idx = tokenizer.action_to_class(m.uci())
31 if idx != -1 and probs[idx].item() > best_prob:
32 best_move, best_prob = m.uci(), probs[idx].item()
33
34print(f"Best move: {best_move} (confidence: {best_prob:.4f}, eval: {value.item():.4f})")