Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,194 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import chess
+
+MOVES_PER_SQUARE = 73
+POLICY_SIZE = 64 * MOVES_PER_SQUARE
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, channels: int) -> None:
+        super().__init__()
+        self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        out = F.relu(self.conv1(x))
+        out = self.conv2(out)
+        out = out + residual
+        return F.relu(out)
+
+
+class TinyPCN(nn.Module):
+    def __init__(self, board_channels: int = 18, policy_size: int = POLICY_SIZE) -> None:
+        """Tiny policy-value net: shared trunk plus separate heads."""
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(board_channels, 32, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv2d(32, 32, kernel_size=3, padding=1)
+        self.res_block = ResidualBlock(32)
+
+        self.policy_conv = nn.Conv2d(32, 32, kernel_size=1)
+        self.policy_fc = nn.Linear(32 * 8 * 8, policy_size)
+
+        self.value_conv = nn.Conv2d(32, 1, kernel_size=1)
+        self.value_fc1 = nn.Linear(8 * 8, 64)
+        self.value_fc2 = nn.Linear(64, 1)
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = self.res_block(x)
+
+        p = F.relu(self.policy_conv(x))
+        p = p.view(p.size(0), -1)
+        policy_logits = self.policy_fc(p)
+
+        v = F.relu(self.value_conv(x))
+        v = v.view(v.size(0), -1)
+        v = F.relu(self.value_fc1(v))
+        value = torch.tanh(self.value_fc2(v))
+
+        return policy_logits, value
+
+
+def board_to_18_planes(board: chess.Board) -> torch.FloatTensor:
+    """Return 18 AlphaZero-style planes for the given board."""
+    planes = np.zeros((18, 8, 8), dtype=np.float32)
+
+    for square, piece in board.piece_map().items():
+        row = 7 - (square // 8)
+        col = square % 8
+        color_offset = 0 if piece.color == chess.WHITE else 6
+        plane_idx = (piece.piece_type - 1) + color_offset
+        planes[plane_idx, row, col] = 1.0
+
+    planes[12, :, :] = 1.0 if board.has_kingside_castling_rights(chess.WHITE) else 0.0
+    planes[13, :, :] = 1.0 if board.has_queenside_castling_rights(chess.WHITE) else 0.0
+    planes[14, :, :] = 1.0 if board.has_kingside_castling_rights(chess.BLACK) else 0.0
+    planes[15, :, :] = 1.0 if board.has_queenside_castling_rights(chess.BLACK) else 0.0
+
+    planes[16, :, :] = 1.0 if board.turn == chess.WHITE else 0.0
+
+    if board.ep_square is not None:
+        ep_row = 7 - (board.ep_square // 8)
+        ep_col = board.ep_square % 8
+        planes[17, ep_row, ep_col] = 1.0
+
+    return torch.from_numpy(planes)
+
+
+def board_to_20_planes(board: chess.Board) -> torch.FloatTensor:
+    """Return 20 planes (18 standard plus repetition and move count)."""
+    planes18 = board_to_18_planes(board).numpy()
+    extra = np.zeros((2, 8, 8), dtype=np.float32)
+
+    try:
+        repetition = board.is_repetition()
+    except Exception:
+        repetition = False
+    extra[0, :, :] = 1.0 if repetition else 0.0
+
+    move_norm = min(board.fullmove_number / 100.0, 1.0)
+    extra[1, :, :] = float(move_norm)
+
+    planes20 = np.concatenate([planes18, extra], axis=0)
+    return torch.from_numpy(planes20)
+
+
+def encode_board(board: chess.Board, variant: str = "18") -> torch.FloatTensor:
+    if variant == "18":
+        return board_to_18_planes(board)
+    if variant == "20":
+        return board_to_20_planes(board)
+    raise ValueError("variant must be '18' or '20'")
+
+
+_RAY_OFFSETS = ((1, 0), (-1, 0), (0, 1), (0, -1), (1, 1), (1, -1), (-1, 1), (-1, -1))
+_KNIGHT_OFFSETS = ((2, 1), (1, 2), (-1, 2), (-2, 1), (-2, -1), (-1, -2), (1, -2), (2, -1))
+_PROMOTION_OFFSETS = ((1, 0), (1, 1), (1, -1), (2, 0))
+_PROMOTION_PIECES = ("q", "r", "b", "n")
+
+_move_to_index: dict[tuple[int, int, str | None], int] = {}
+_index_to_move: dict[int, tuple[int, int, str | None]] = {}
+
+
+def _init_move_tables() -> None:
+    idx = 0
+    for sq in range(64):
+        row0, col0 = divmod(sq, 8)
+
+        for dx, dy in _RAY_OFFSETS:
+            for step in range(1, 8):
+                row = row0 + dx * step
+                col = col0 + dy * step
+                if 0 <= row < 8 and 0 <= col < 8:
+                    target = row * 8 + col
+                    _move_to_index[(sq, target, None)] = idx
+                    _index_to_move[idx] = (sq, target, None)
+                idx += 1
+
+        for dx, dy in _KNIGHT_OFFSETS:
+            row = row0 + dx
+            col = col0 + dy
+            if 0 <= row < 8 and 0 <= col < 8:
+                target = row * 8 + col
+                _move_to_index[(sq, target, None)] = idx
+                _index_to_move[idx] = (sq, target, None)
+            idx += 1
+
+        for dx, dy in _PROMOTION_OFFSETS:
+            row = row0 + dx
+            col = col0 + dy
+            if 0 <= row < 8 and 0 <= col < 8:
+                target = row * 8 + col
+                for promo in _PROMOTION_PIECES:
+                    _move_to_index[(sq, target, promo)] = idx
+                    _index_to_move[idx] = (sq, target, promo)
+                    idx += 1
+            else:
+                idx += len(_PROMOTION_PIECES)
+
+        while idx % MOVES_PER_SQUARE != 0:
+            _index_to_move[idx] = None
+            idx += 1
+
+
+_init_move_tables()
+
+
+def _promotion_symbol(piece_type: int | None) -> str | None:
+    if piece_type is None:
+        return None
+    return chess.Piece(piece_type, chess.WHITE).symbol().lower()
+
+
+def encode_move(move: chess.Move, board: chess.Board) -> int:
+    from_sq = move.from_square
+    to_sq = move.to_square
+    promo_symbol = _promotion_symbol(move.promotion)
+
+    if board.color_at(move.from_square) == chess.BLACK:
+        from_sq = chess.square_mirror(from_sq)
+        to_sq = chess.square_mirror(to_sq)
+
+    key = (from_sq, to_sq, promo_symbol)
+    return _move_to_index.get(key, -1)
+
+
+def decode_move(index: int, board: chess.Board | None = None) -> chess.Move | None:
+    triple = _index_to_move.get(index)
+    if triple is None:
+        return None
+
+    from_sq, to_sq, promo = triple
+
+    if board is not None and board.turn == chess.BLACK:
+        from_sq = chess.square_mirror(from_sq)
+        to_sq = chess.square_mirror(to_sq)
+
+    promotion = chess.Piece.from_symbol(promo.upper()).piece_type if promo else None
+    return chess.Move(from_sq, to_sq, promotion=promotion)