File size: 6,498 Bytes
508f87a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1cd8369
508f87a
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
"""
Chess Position Evaluator (V7 HULK) - Inference Example
Author: Ahmed Darwish (@engdarwish)
https://huggingface.co/engdarwish/chess-position-evaluator

Minimal, runnable example: download the weights, rebuild the exact
SE-ResNet-20 dual-head architecture, encode a FEN position, and get
a value estimate + ranked legal moves.

Install:
    pip install torch python-chess huggingface_hub numpy
"""

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import chess
from huggingface_hub import hf_hub_download

REPO_ID = "engdarwish/chess-position-evaluator"
MOVE_SPACE = 64 * 64 * 5

PROMO_MAP = {None: 0, chess.QUEEN: 1, chess.ROOK: 2, chess.BISHOP: 3, chess.KNIGHT: 4}
PROMO_MAP_INV = {v: k for k, v in PROMO_MAP.items() if k is not None}


def uci_to_idx(uci_str: str) -> int:
    move = chess.Move.from_uci(uci_str)
    promo = PROMO_MAP.get(move.promotion, 0)
    return (move.from_square * 64 + move.to_square) * 5 + promo


def idx_to_uci(idx: int) -> str:
    promo_type = idx % 5
    rest = idx // 5
    to_sq = rest % 64
    from_sq = rest // 64
    promo_piece = PROMO_MAP_INV.get(promo_type)
    return chess.Move(from_sq, to_sq, promotion=promo_piece).uci()


PIECE_IDX = {
    chess.PAWN: 0, chess.KNIGHT: 1, chess.BISHOP: 2,
    chess.ROOK: 3, chess.QUEEN: 4, chess.KING: 5,
}


def fen_to_tensor(fen: str) -> np.ndarray:
    board = chess.Board(fen)
    t = np.zeros((18, 8, 8), dtype=np.float32)
    for sq in chess.SQUARES:
        piece = board.piece_at(sq)
        if piece:
            r = 7 - (sq >> 3)
            c = sq & 7
            ch = PIECE_IDX[piece.piece_type]
            t[ch if piece.color else ch + 6, r, c] = 1.0
    if board.turn == chess.WHITE:
        t[12] = 1.0
    t[13] = float(board.has_kingside_castling_rights(chess.WHITE))
    t[14] = float(board.has_queenside_castling_rights(chess.WHITE))
    t[15] = float(board.has_kingside_castling_rights(chess.BLACK))
    t[16] = float(board.has_queenside_castling_rights(chess.BLACK))
    if board.ep_square is not None:
        t[17, 7 - (board.ep_square >> 3), board.ep_square & 7] = 1.0
    return t


class SEBlock(nn.Module):
    def __init__(self, channels, reduction=16):
        super().__init__()
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channels, channels // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(channels // reduction, channels, bias=False),
            nn.Sigmoid(),
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y


class SEResBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv1 = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(channels)
        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(channels)
        self.se = SEBlock(channels)

    def forward(self, x):
        r = x
        out = F.relu(self.bn1(self.conv1(x)), inplace=True)
        out = self.bn2(self.conv2(out))
        out = self.se(out)
        return F.relu(out + r, inplace=True)


class ChessPositionEvaluator(nn.Module):
    """
    Chess Position Evaluator - V7 HULK (SE-ResNet-20, Dual Head)
    Input  : (B, 18, 8, 8)
    Value  : scalar in [-1, 1]   (tanh, from White's perspective)
    Policy : logits over 20,480 moves (includes underpromotion)
    """

    def __init__(self, in_channels=18, num_filters=256, num_res_blocks=20, policy_size=MOVE_SPACE):
        super().__init__()
        self.input_block = nn.Sequential(
            nn.Conv2d(in_channels, num_filters, 3, padding=1, bias=False),
            nn.BatchNorm2d(num_filters),
            nn.ReLU(inplace=True),
        )
        self.tower = nn.Sequential(*[SEResBlock(num_filters) for _ in range(num_res_blocks)])
        self.value_head = nn.Sequential(
            nn.Conv2d(num_filters, 32, 1, bias=False),
            nn.BatchNorm2d(32), nn.ReLU(inplace=True),
            nn.Flatten(),
            nn.Linear(32 * 8 * 8, 256), nn.ReLU(inplace=True),
            nn.Dropout(0.3),
            nn.Linear(256, 1), nn.Tanh(),
        )
        self.policy_head = nn.Sequential(
            nn.Conv2d(num_filters, 32, 1, bias=False),
            nn.BatchNorm2d(32), nn.ReLU(inplace=True),
            nn.Flatten(),
            nn.Linear(32 * 8 * 8, 1024), nn.ReLU(inplace=True),
            nn.Dropout(0.3),
            nn.Linear(1024, policy_size),
        )

    def forward(self, x):
        x = self.input_block(x)
        x = self.tower(x)
        value = self.value_head(x).squeeze(-1)
        policy = self.policy_head(x)
        return value, policy


def load_model(device: str = "cpu") -> ChessPositionEvaluator:
    weights_path = hf_hub_download(repo_id=REPO_ID, filename="model_weights.pt")
    model = ChessPositionEvaluator().to(device)
    state_dict = torch.load(weights_path, map_location=device)
    model.load_state_dict(state_dict)
    model.eval()
    return model


@torch.no_grad()
def get_best_move(model: ChessPositionEvaluator, fen: str, top_k: int = 5, device: str = "cpu"):
    """Returns (value, [(uci_move, probability), ...]) using legal-move masking."""
    board = chess.Board(fen)
    t = fen_to_tensor(fen)
    inp = torch.tensor(t, dtype=torch.float32).unsqueeze(0).to(device)
    value, policy_logits = model(inp)
    logits = policy_logits[0].float().cpu()

    mask = torch.full((MOVE_SPACE,), float("-inf"))
    legal_moves = list(board.legal_moves)
    for move in legal_moves:
        idx = uci_to_idx(move.uci())
        if 0 <= idx < MOVE_SPACE:
            mask[idx] = 0.0

    probs = torch.softmax(logits + mask, dim=0)
    top_indices = torch.topk(probs, min(top_k, len(legal_moves))).indices
    return value.item(), [(idx_to_uci(i.item()), probs[i].item()) for i in top_indices]


if __name__ == "__main__":
    model = load_model()
    print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")

    positions = [
        ("Starting position", chess.STARTING_FEN),
        ("After 1.e4", "rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR b KQkq e3 0 1"),
    ]
    for name, fen in positions:
        val, moves = get_best_move(model, fen, top_k=3)
        print(f"\n{name}")
        print(f"  Value: {val:.4f}")
        for uci, prob in moves:
            print(f"    {uci}  ({prob * 100:.1f}%)")