Create modeling_chessnet.py

modeling_chessnet.py

@@ -0,0 +1,270 @@
+"""
+Blackwell Knight - ChessNet Modeling File
+Custom chess neural network for the TobiasLogic Chess Arena.
+
+Architecture: ResNet backbone with policy + value heads.
+Board input: 19 planes of 8x8.
+Move output: Flat index over 20,480 possible moves.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import chess
+
+# =====================================================================
+# BOARD ENCODING
+# =====================================================================
+
+def encode_board(board):
+    """
+    Encode a chess.Board into a (19, 8, 8) float32 tensor.
+    
+    Planes:
+      0-5:   White pieces (Pawn, Knight, Bishop, Rook, Queen, King)
+      6-11:  Black pieces (Pawn, Knight, Bishop, Rook, Queen, King)
+      12:    Side to move (1.0 = white, 0.0 = black)
+      13-16: Castling rights (WK, WQ, BK, BQ)
+      17:    En passant target square
+      18:    Halfmove clock (normalized by /100)
+    """
+    planes = np.zeros((19, 64), dtype=np.float32)
+
+    for square, piece in board.piece_map().items():
+        idx = piece.piece_type - 1 + (0 if piece.color == chess.WHITE else 6)
+        planes[idx, square] = 1.0
+
+    planes[12].fill(1.0 if board.turn == chess.WHITE else 0.0)
+    planes[13].fill(1.0 if board.has_kingside_castling_rights(chess.WHITE) else 0.0)
+    planes[14].fill(1.0 if board.has_queenside_castling_rights(chess.WHITE) else 0.0)
+    planes[15].fill(1.0 if board.has_kingside_castling_rights(chess.BLACK) else 0.0)
+    planes[16].fill(1.0 if board.has_queenside_castling_rights(chess.BLACK) else 0.0)
+
+    if board.ep_square is not None:
+        planes[17, board.ep_square] = 1.0
+
+    planes[18].fill(min(board.halfmove_clock / 100.0, 1.0))
+
+    return planes.reshape(19, 8, 8)
+
+
+# =====================================================================
+# MOVE ENCODING / DECODING
+# =====================================================================
+
+POLICY_SIZE = 20480  # 64 * 64 * 5
+
+def get_move_index(move):
+    """Convert a chess.Move to a flat index (0-20479)."""
+    promo_idx = 0
+    if move.promotion:
+        promo_idx = {
+            chess.KNIGHT: 1,
+            chess.BISHOP: 2,
+            chess.ROOK: 3,
+            chess.QUEEN: 4,
+        }.get(move.promotion, 0)
+    return (move.from_square * 64 + move.to_square) * 5 + promo_idx
+
+
+def index_to_move(idx, board):
+    """
+    Convert a flat index back to a chess.Move.
+    If the decoded move is illegal, returns None.
+    """
+    from_sq = (idx // 5) // 64
+    to_sq = (idx // 5) % 64
+    promo_idx = idx % 5
+
+    promo = None
+    if promo_idx == 1:
+        promo = chess.KNIGHT
+    elif promo_idx == 2:
+        promo = chess.BISHOP
+    elif promo_idx == 3:
+        promo = chess.ROOK
+    elif promo_idx == 4:
+        promo = chess.QUEEN
+
+    move = chess.Move(from_sq, to_sq, promotion=promo)
+
+    # Safety: check if the move is actually legal
+    if move in board.legal_moves:
+        return move
+    return None
+
+
+# =====================================================================
+# MODEL ARCHITECTURE
+# =====================================================================
+
+class ResidualBlock(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)
+
+    def forward(self, x):
+        identity = x
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += identity
+        return F.relu(out)
+
+
+class ChessNet(nn.Module):
+    """
+    Blackwell Knight ChessNet.
+    
+    ~18.5M parameters.
+    Input:  (batch, 19, 8, 8) board planes.
+    Output: policy_logits (batch, 20480), value (batch, 1) in [-1, 1].
+    """
+
+    def __init__(self, config=None):
+        super().__init__()
+
+        if config is None:
+            config = {}
+
+        in_planes = config.get("input_planes", 19)
+        num_blocks = config.get("num_blocks", 10)
+        num_channels = config.get("num_channels", 128)
+        policy_channels = config.get("policy_channels", 32)
+        value_channels = config.get("value_channels", 8)
+        policy_size = config.get("policy_size", POLICY_SIZE)
+
+        self.stem = nn.Sequential(
+            nn.Conv2d(in_planes, num_channels, 3, padding=1, bias=False),
+            nn.BatchNorm2d(num_channels),
+            nn.ReLU(),
+        )
+
+        self.blocks = nn.Sequential(
+            *[ResidualBlock(num_channels) for _ in range(num_blocks)]
+        )
+
+        self.policy_head = nn.Sequential(
+            nn.Conv2d(num_channels, policy_channels, 1, bias=False),
+            nn.BatchNorm2d(policy_channels),
+            nn.ReLU(),
+            nn.Flatten(),
+            nn.Linear(policy_channels * 8 * 8, policy_size),
+        )
+
+        self.value_head = nn.Sequential(
+            nn.Conv2d(num_channels, value_channels, 1, bias=False),
+            nn.BatchNorm2d(value_channels),
+            nn.ReLU(),
+            nn.Flatten(),
+            nn.Linear(value_channels * 8 * 8, 64),
+            nn.ReLU(),
+            nn.Linear(64, 1),
+            nn.Tanh(),
+        )
+
+    def forward(self, x):
+        x = self.stem(x)
+        x = self.blocks(x)
+        policy_logits = self.policy_head(x)
+        value = self.value_head(x)
+        return policy_logits, value
+
+
+# =====================================================================
+# INFERENCE / PLAY CHESS
+# =====================================================================
+
+class BlackwellKnightPlayer:
+    """
+    Arena-compatible player wrapper.
+    
+    Usage:
+        player = BlackwellKnightPlayer("akshit/blackwell-knight-v1")
+        move = player.choose_move(board)
+    """
+
+    def __init__(self, model_repo, device=None, temperature=0.0):
+        from huggingface_hub import hf_hub_download
+        from safetensors.torch import load_file
+        import json
+
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.temperature = temperature
+
+        # Load config
+        config_path = hf_hub_download(model_repo, "config.json")
+        with open(config_path, "r") as f:
+            self.config = json.load(f)
+
+        # Build model
+        arch = self.config.get("model_architecture", {})
+        self.model = ChessNet({
+            "input_planes": self.config.get("board_encoding", {}).get("input_planes", 19),
+            "num_blocks": arch.get("num_blocks", 10),
+            "num_channels": arch.get("num_channels", 128),
+            "policy_channels": arch.get("policy_channels", 32),
+            "value_channels": arch.get("value_channels", 8),
+            "policy_size": self.config.get("move_encoding", {}).get("policy_size", POLICY_SIZE),
+        })
+
+        # Load weights
+        weights_path = hf_hub_download(model_repo, "model.safetensors")
+        state_dict = load_file(weights_path)
+        self.model.load_state_dict(state_dict)
+        self.model.to(self.device)
+        self.model.eval()
+
+    def choose_move(self, board):
+        """
+        Given a chess.Board, return the best legal move.
+        """
+        x = encode_board(board)
+        x_tensor = torch.tensor(x, dtype=torch.float32).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            policy_logits, value = self.model(x_tensor)
+
+        # Get legal moves
+        legal_moves = list(board.legal_moves)
+        if not legal_moves:
+            return None
+
+        legal_indices = [get_move_index(m) for m in legal_moves]
+
+        # Mask illegal moves
+        mask = torch.full((POLICY_SIZE,), float("-inf"), device=self.device)
+        mask[legal_indices] = 0.0
+
+        logits = policy_logits[0] + mask
+
+        if self.temperature > 0:
+            probs = torch.softmax(logits / self.temperature, dim=-1)
+            idx = torch.multinomial(probs, 1).item()
+        else:
+            idx = torch.argmax(logits).item()
+
+        move = index_to_move(idx, board)
+
+        # Fallback: if something went wrong, pick a random legal move
+        if move is None or move not in legal_moves:
+            import random
+            move = random.choice(legal_moves)
+
+        return move
+
+    def evaluate(self, board):
+        """
+        Return the value estimate for the current position.
+        Positive = good for side to move.
+        """
+        x = encode_board(board)
+        x_tensor = torch.tensor(x, dtype=torch.float32).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            _, value = self.model(x_tensor)
+
+        return value.item()