Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,756 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .attn_map import apm_map, apm_out
+import math
+from .encoding_simple import encode_fen_to_tensor, encode_moves_to_tensor
+from .vocab import policy_index
+from typing import Union, List, Optional
+import bulletchess
+import numpy as np
+
+class Gating(nn.Module):
+    def __init__(self, features_shape, additive=True, init_value=None):
+        super(Gating, self).__init__()
+        self.additive = additive
+        if init_value is None:
+            init_value = 0 if self.additive else 1
+        
+        self.gate = nn.Parameter(torch.full(features_shape, float(init_value)))
+        if not self.additive:
+            self.gate.register_hook(lambda grad: torch.clamp(grad, min=0))
+
+    def forward(self, x):
+        if self.additive:
+            return x + self.gate
+        else:
+            return x * self.gate
+
+def ma_gating(x, in_features):
+    x = Gating(in_features, additive=False)(x)
+    x = Gating(in_features, additive=True)(x)
+    return x
+
+class RMSNorm(nn.Module):
+    def __init__(self, in_features, scale=True):
+        super(RMSNorm, self).__init__()
+        self.scale = scale
+        if self.scale:
+            self.gamma = nn.Parameter(torch.ones(in_features))
+
+    def forward(self, x):
+        rms = torch.sqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-5)
+        x_normalized = x / rms
+        if self.scale:
+            return x_normalized * self.gamma
+        return x_normalized
+
+class ApplyAttentionPolicyMap(nn.Module):
+    def __init__(self):
+        super(ApplyAttentionPolicyMap, self).__init__()
+        # Register as buffers so they move with the model when .to(device) is called
+        # Use same names as before for backward compatibility with saved models
+        self.register_buffer('fc1', torch.from_numpy(apm_map).float())
+        self.register_buffer('idx', torch.from_numpy(apm_out).long())
+
+    def forward(self, logits, pp_logits):
+        logits = torch.cat([logits.reshape(-1, 64 * 64),
+                            pp_logits.reshape(-1, 8 * 24)],
+                           dim=1)
+        
+        batch_size = logits.size(0)
+        idx = self.idx.unsqueeze(0).expand(batch_size, -1)
+        
+        return torch.gather(logits, 1, idx)
+
+class Mish(nn.Module):
+    def __init__(self):
+        super(Mish, self).__init__()
+
+    def forward(self, x):
+        return x * torch.tanh(F.softplus(x))
+
+class CustomMHA(nn.Module):
+    def __init__(self, emb_size, d_model, num_heads, dropout=0.0, use_bias_qkv=True, use_bias_out=True,
+                 use_smolgen=True, smol_hidden_channels=32, smol_hidden_sz=256, smol_gen_sz=256, smol_activation='swish'):
+        super(CustomMHA, self).__init__()
+        assert d_model % num_heads == 0
+        self.emb_size = emb_size
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.head_dim = d_model // num_heads
+        self.wq = nn.Linear(emb_size, d_model, bias=use_bias_qkv)
+        self.wk = nn.Linear(emb_size, d_model, bias=use_bias_qkv)
+        self.wv = nn.Linear(emb_size, d_model, bias=use_bias_qkv)
+        self.out_proj = nn.Linear(d_model, emb_size, bias=use_bias_out)
+        self.attn_dropout = nn.Dropout(dropout)
+        # Optional Smolgen components
+        self.smol_compress = None
+        self.smol_hidden1 = None
+        self.smol_hidden1_ln = None
+        self.smol_gen_from = None
+        self.smol_gen_from_ln = None
+        self.smol_weight_gen = None
+        if use_smolgen:
+            self.smol_compress = nn.Linear(emb_size, smol_hidden_channels, bias=False)
+            self.smol_hidden1 = nn.Linear(64 * smol_hidden_channels, smol_hidden_sz, bias=True)
+            self.smol_hidden1_ln = nn.LayerNorm(smol_hidden_sz, eps=1e-3)
+            self.smol_gen_from = nn.Linear(smol_hidden_sz, num_heads * smol_gen_sz, bias=True)
+            self.smol_gen_from_ln = nn.LayerNorm(num_heads * smol_gen_sz, eps=1e-3)
+            self.smol_weight_gen = nn.Linear(smol_gen_sz, 64 * 64, bias=False)
+        self.smol_activation = smol_activation
+
+    def _shape(self, x):
+        b, l, _ = x.shape
+        return x.view(b, l, self.num_heads, self.head_dim).transpose(1, 2)
+
+    def forward(self, x, return_attn=False):
+        # x: (B, L, emb_size)
+        q = self.wq(x)
+        k = self.wk(x)
+        v = self.wv(x)
+        q = self._shape(q)  # (B, H, L, D)
+        k = self._shape(k)
+        v = self._shape(v)
+        scale = torch.sqrt(torch.tensor(self.head_dim, dtype=x.dtype, device=x.device))
+        attn_logits = torch.matmul(q, k.transpose(-2, -1)) / scale
+        # Add Smolgen weights if present
+        smol_w = None
+        if self.smol_compress is not None:
+            b, l, _ = x.shape
+            compressed = self.smol_compress(x)  # (B, L, hidden_channels)
+            compressed = compressed.reshape(b, l * compressed.shape[-1])  # (B, 64*hidden_channels)
+            hidden_pre = self.smol_hidden1(compressed)
+            hidden = F.silu(hidden_pre) if self.smol_activation == 'swish' else F.silu(hidden_pre)
+            hidden_ln = self.smol_hidden1_ln(hidden)
+            gen_from_pre = self.smol_gen_from(hidden_ln)
+            gen_from_act = F.silu(gen_from_pre) if self.smol_activation == 'swish' else F.silu(gen_from_pre)
+            gen_from = self.smol_gen_from_ln(gen_from_act)
+            gen_from = gen_from.view(b, self.num_heads, -1)  # (B, H, gen_sz)
+            smol_w = self.smol_weight_gen(gen_from)  # (B, H, 64*64)
+            smol_w = smol_w.view(b, self.num_heads, l, l)
+            attn_logits = attn_logits + smol_w
+        # Numerically stable softmax matching TF (float32, subtract max)
+        attn_logits = attn_logits - attn_logits.max(dim=-1, keepdim=True)[0]
+        attn_weights = torch.exp(attn_logits)
+        attn_weights = attn_weights / attn_weights.sum(dim=-1, keepdim=True)
+        attn_weights = self.attn_dropout(attn_weights)
+        attn_output = torch.matmul(attn_weights, v)  # (B, H, L, D)
+        attn_output = attn_output.transpose(1, 2).contiguous().view(x.size(0), x.size(1), self.d_model)
+        out = self.out_proj(attn_output)
+        if return_attn:
+            return out, attn_weights, smol_w, attn_logits
+        return out
+
+class FFN(nn.Module):
+    def __init__(self, emb_size, dff, activation=Mish(), omit_other_biases=False):
+        super(FFN, self).__init__()
+        self.dense1 = nn.Linear(emb_size, dff, bias=not omit_other_biases)
+        self.activation = activation
+        self.dense2 = nn.Linear(dff, emb_size, bias=not omit_other_biases)
+
+    def forward(self, x):
+        x = self.dense1(x)
+        x = self.activation(x)
+        x = self.dense2(x)
+        return x
+
+class EncoderLayer(nn.Module):
+    def __init__(self, emb_size, d_model, num_heads, dff, dropout_rate, encoder_layers, skip_first_ln=False, encoder_rms_norm=False, omit_qkv_biases=False, omit_other_biases=False,
+                 use_smolgen=True, smol_hidden_channels=32, smol_hidden_sz=256, smol_gen_sz=256, smol_activation='swish'):
+        super(EncoderLayer, self).__init__()
+        self.mha = CustomMHA(emb_size, d_model, num_heads, dropout=dropout_rate, use_bias_qkv=not omit_qkv_biases, use_bias_out=not omit_other_biases,
+                             use_smolgen=use_smolgen, smol_hidden_channels=smol_hidden_channels, smol_hidden_sz=smol_hidden_sz, smol_gen_sz=smol_gen_sz, smol_activation=smol_activation)
+        self.ffn = FFN(emb_size, dff, omit_other_biases=omit_other_biases)
+        
+        self.norm1 = RMSNorm(emb_size) if encoder_rms_norm else nn.LayerNorm(emb_size, eps=0.001)
+        self.norm2 = RMSNorm(emb_size) if encoder_rms_norm else nn.LayerNorm(emb_size, eps=0.001)
+        
+        self.dropout1 = nn.Dropout(dropout_rate)
+        self.dropout2 = nn.Dropout(dropout_rate)
+        
+        self.alpha = (2. * encoder_layers)**-0.25
+        self.skip_first_ln = skip_first_ln
+
+    def forward(self, x):
+        attn_output = self.mha(x)
+        attn_output = self.dropout1(attn_output)
+        
+        out1 = x + attn_output * self.alpha
+        if not self.skip_first_ln:
+            out1 = self.norm1(out1)
+        ffn_output = self.ffn(out1)
+        ffn_output = self.dropout2(ffn_output)
+        
+        out2 = self.norm2(out1 + ffn_output * self.alpha)
+        return out2
+
+class PolicyHead(nn.Module):
+    def __init__(self, pol_embedding_size, policy_d_model, opponent=False):
+        super(PolicyHead, self).__init__()
+        self.opponent = opponent
+        self.wq = nn.Linear(pol_embedding_size, policy_d_model)
+        self.wk = nn.Linear(pol_embedding_size, policy_d_model)
+        self.ppo = nn.Linear(policy_d_model, 4, bias=False)
+        self.apply_map = ApplyAttentionPolicyMap()
+
+    def forward(self, x):
+        if self.opponent:
+            x = torch.flip(x, [1])
+
+        queries = self.wq(x)
+        keys = self.wk(x)
+
+        matmul_qk = torch.matmul(queries, keys.transpose(-2, -1))
+        
+        dk = torch.sqrt(torch.tensor(keys.shape[-1], dtype=keys.dtype, device=keys.device))
+        promotion_keys = keys[:, -8:, :]
+        promotion_offsets = self.ppo(promotion_keys).transpose(-2,-1) * dk
+        promotion_offsets = promotion_offsets[:, :3, :] + promotion_offsets[:, 3:4, :]
+
+        n_promo_logits = matmul_qk[:, -16:-8, -8:]
+        q_promo_logits = (n_promo_logits + promotion_offsets[:, 0:1, :]).unsqueeze(3)
+        r_promo_logits = (n_promo_logits + promotion_offsets[:, 1:2, :]).unsqueeze(3)
+        b_promo_logits = (n_promo_logits + promotion_offsets[:, 2:3, :]).unsqueeze(3)
+        promotion_logits = torch.cat([q_promo_logits, r_promo_logits, b_promo_logits], axis=3).reshape(-1, 8, 24)
+
+        promotion_logits = promotion_logits / dk
+        policy_attn_logits = matmul_qk / dk
+
+        return self.apply_map(policy_attn_logits, promotion_logits)
+
+class ValueHead(nn.Module):
+    def __init__(self, embedding_size, val_embedding_size, default_activation=Mish()):
+        super(ValueHead, self).__init__()
+        self.embedding = nn.Linear(embedding_size, val_embedding_size)
+        self.activation = default_activation
+        self.flatten = nn.Flatten()
+        self.dense1 = nn.Linear(val_embedding_size * 64, 128)
+        self.dense2 = nn.Linear(128, 3)
+
+    def forward(self, x):
+        x = self.embedding(x)
+        x = self.activation(x)
+        x = self.flatten(x)
+        x = self.dense1(x)
+        x = self.activation(x)
+        x = self.dense2(x)
+        return x
+
+class BT4(nn.Module):
+    def __init__(self, embedding_size=1024, embedding_dense_sz=512, encoder_layers=15, encoder_d_model=1024, encoder_heads=32, encoder_dff=1536, dropout_rate=0.0, pol_embedding_size=1024, policy_d_model=1024, val_embedding_size=128, default_activation=Mish(),
+                 use_smolgen=True, smol_hidden_channels=32, smol_hidden_sz=256, smol_gen_sz=256, smol_activation='swish'):
+        super(BT4, self).__init__()
+        self.embedding_dense_sz = embedding_dense_sz
+        # DeepNorm alpha used in embedding residual; default uses provided encoder_layers
+        self.deepnorm_alpha = (2. * encoder_layers) ** -0.25
+        
+        self.embedding_preprocess = nn.Linear(64*12, 64*self.embedding_dense_sz)
+        self.embedding = nn.Linear(112 + self.embedding_dense_sz, embedding_size)
+        nn.init.xavier_uniform_(self.embedding.weight) # Explicitly set initializer
+        nn.init.zeros_(self.embedding.bias)
+
+        self.embedding_ln = nn.LayerNorm(embedding_size, eps=0.001)
+        
+        self.gating_mult = Gating((64, embedding_size), additive=False)
+        self.gating_add = Gating((64, embedding_size), additive=True)
+
+        self.embedding_ffn = FFN(embedding_size, encoder_dff)
+        self.embedding_ffn_ln = nn.LayerNorm(embedding_size, eps=0.001)
+        
+        self.encoder_layers_list = nn.ModuleList([
+            EncoderLayer(embedding_size, encoder_d_model, encoder_heads, encoder_dff, dropout_rate, encoder_layers,
+                         use_smolgen=use_smolgen, smol_hidden_channels=smol_hidden_channels, smol_hidden_sz=smol_hidden_sz, smol_gen_sz=smol_gen_sz, smol_activation=smol_activation)
+            for _ in range(encoder_layers)
+        ])
+        
+        self.policy_embedding = nn.Linear(embedding_size, pol_embedding_size)
+        self.policy_head = PolicyHead(pol_embedding_size, policy_d_model)
+        self.value_head_winner = ValueHead(embedding_size, val_embedding_size)
+        self.value_head_q = ValueHead(embedding_size, val_embedding_size)
+        self.activation = default_activation
+        
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            # Keras' glorot_normal is equivalent to PyTorch's xavier_normal_
+            nn.init.xavier_normal_(module.weight)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+
+    def forward(self, x):
+        # x shape: (batch, 112, 8, 8)
+        flow = x.permute(0, 2, 3, 1).reshape(-1, 64, 112)
+        
+        pos_info = flow[..., :12]
+        pos_info_flat = pos_info.reshape(-1, 64 * 12)
+        
+        pos_info_processed = self.embedding_preprocess(pos_info_flat)
+        pos_info = pos_info_processed.reshape(-1, 64, self.embedding_dense_sz)
+        
+        flow = torch.cat([flow, pos_info], dim=-1)
+        
+        flow = self.embedding(flow)
+        
+        flow = self.activation(flow)
+        
+        flow = self.embedding_ln(flow)
+
+        flow = self.gating_mult(flow)
+        flow = self.gating_add(flow)
+        
+        ffn_dense1_pre = self.embedding_ffn.dense1(flow)
+        ffn_dense1 = self.embedding_ffn.activation(ffn_dense1_pre)
+        ffn_output = self.embedding_ffn.dense2(ffn_dense1)
+        
+        residual = flow + ffn_output * self.deepnorm_alpha
+        flow = self.embedding_ffn_ln(residual)
+        
+        for i, layer in enumerate(self.encoder_layers_list):
+            flow = layer(flow)
+        
+        policy_tokens = self.policy_embedding(flow)
+        policy_tokens = self.activation(policy_tokens)
+        
+        policy_logits = self.policy_head(policy_tokens)
+        
+        value_winner = self.value_head_winner(flow)
+        value_q = self.value_head_q(flow)
+
+        return policy_logits, value_winner, value_q
+    
+    def get_move_from_fen_no_thinking(self, fen_or_moves: Union[str, List[str]], T: float, device: str = None, **kwargs) -> str:
+        """
+        Predict a move from a FEN position or move history without thinking/search.
+        
+        Args:
+            fen_or_moves: Either a FEN string representing the chess position, or a list of UCI moves
+            T: Temperature for sampling (0.0 = deterministic/argmax, >0.0 = stochastic)
+            device: Device to run the model on (if None, uses model's device)
+        
+        Returns:
+            UCI move string (e.g., 'e2e4')
+        """
+        # Detect device from model if not provided
+        if device is None:
+            device = next(self.parameters()).device
+        else:
+            device = torch.device(device)
+        
+        # Determine if input is FEN string or list of moves
+        if isinstance(fen_or_moves, str):
+            # FEN string input
+            fen = fen_or_moves
+            is_black_to_move = fen.split()[1] == 'b'
+            input_tensor_112, legal_moves_mask = encode_fen_to_tensor(fen)
+            castling_rights = fen.split()[2] if len(fen.split()) > 2 else ""
+        elif isinstance(fen_or_moves, list):
+            # List of UCI moves input
+            move_history = fen_or_moves
+            input_tensor_112, legal_moves_mask = encode_moves_to_tensor(move_history)
+            # Create board to check if black is to move and for castling rights
+            board = bulletchess.Board()
+            for mv in move_history:
+                move = bulletchess.Move.from_uci(mv)
+                board.apply(move)
+            is_black_to_move = (board.turn == bulletchess.BLACK)
+            fen_parts = board.fen().split()
+            castling_rights = fen_parts[2] if len(fen_parts) > 2 else ""
+        else:
+            raise ValueError("Input must be a FEN string or a list of UCI moves")
+        
+        input_tensor_112 = input_tensor_112.to(device, non_blocking=True)
+        
+        self.eval()
+        with torch.inference_mode():
+            policy_logits,_,_ = self.forward(input_tensor_112)
+        
+        # Apply legal moves mask without in-place ops (inference tensor)
+        logits0 = policy_logits[0] + torch.from_numpy(legal_moves_mask).to(policy_logits.device)
+        
+        # Check if return_probs is requested
+        return_probs = kwargs.get('return_probs', False)
+        
+        if return_probs:
+            # Return probabilities dictionary
+            scaled_logits = logits0 / T if T > 0 else logits0
+            probs = F.softmax(scaled_logits, dim=0)
+            probs_dict = {}
+            for idx, move in enumerate(policy_index):
+                prob_val = probs[idx].item()
+                if prob_val > 1e-6:  # Only include moves with non-negligible probability
+                    probs_dict[move] = prob_val
+            return probs_dict
+        
+        if T == 0.0:
+            # Deterministic: return best move
+            best_move_idx = torch.argmax(logits0).item()
+            uci_move = policy_index[best_move_idx]
+        else:
+            # Stochastic sampling with temperature
+            # Apply temperature scaling
+            scaled_logits = logits0 / T
+            # Apply softmax to get probabilities
+            probs = F.softmax(scaled_logits, dim=0)
+            # Sample from the distribution
+            move_idx = torch.multinomial(probs, 1).item()
+            uci_move = policy_index[move_idx]
+        
+        # If black is to move, the board was mirrored during encoding, so we need to mirror the move back
+        # Mirror ranks: 1↔8, 2↔7, 3↔6, 4↔5 (keep file letters the same)
+        if is_black_to_move:
+            def mirror_rank(rank_char):
+                rank = int(rank_char)
+                return str(9 - rank)
+            
+            # UCI format: e2e4, e7e8q, etc.
+            if len(uci_move) >= 4:
+                from_file = uci_move[0]
+                from_rank = uci_move[1]
+                to_file = uci_move[2]
+                to_rank = uci_move[3]
+                promo = uci_move[4:] if len(uci_move) > 4 else ""
+                
+                uci_move = from_file + mirror_rank(from_rank) + to_file + mirror_rank(to_rank) + promo
+        
+        # Convert castling moves from king-to-rook-square format to standard castling format
+        # Only if castling rights are available (check FEN castling rights)
+        # Check and convert white castling moves
+        if uci_move == "e1h1" and "K" in castling_rights:
+            uci_move = "e1g1"
+        elif uci_move == "e1a1" and "Q" in castling_rights:
+            uci_move = "e1c1"
+        # Check and convert black castling moves
+        elif uci_move == "e8h8" and "k" in castling_rights:
+            uci_move = "e8g8"
+        elif uci_move == "e8a8" and "q" in castling_rights:
+            uci_move = "e8c8"
+        
+        return uci_move
+    
+    def get_best_move_value(self, fen_or_moves: Union[str, List[str]], T: float = 0.0, device: str = None) -> tuple:
+        """
+        Get the best move and its value using value analysis.
+        
+        Args:
+            fen_or_moves: Either a FEN string representing the chess position, or a list of UCI moves
+            T: Temperature for sampling (0.0 = deterministic/argmax, >0.0 = stochastic)
+            device: Device to run the model on (if None, uses model's device)
+        
+        Returns:
+            Tuple of (best_move, value) where value is the position evaluation
+        """
+        # Detect device from model if not provided
+        if device is None:
+            device = next(self.parameters()).device
+        else:
+            device = torch.device(device)
+        
+        # Determine if input is FEN string or list of moves
+        if isinstance(fen_or_moves, str):
+            fen = fen_or_moves
+            is_black_to_move = fen.split()[1] == 'b'
+            input_tensor_112, legal_moves_mask = encode_fen_to_tensor(fen)
+            castling_rights = fen.split()[2] if len(fen.split()) > 2 else ""
+        elif isinstance(fen_or_moves, list):
+            move_history = fen_or_moves
+            input_tensor_112, legal_moves_mask = encode_moves_to_tensor(move_history)
+            board = bulletchess.Board()
+            for mv in move_history:
+                move = bulletchess.Move.from_uci(mv)
+                board.apply(move)
+            is_black_to_move = (board.turn == bulletchess.BLACK)
+            fen_parts = board.fen().split()
+            castling_rights = fen_parts[2] if len(fen_parts) > 2 else ""
+        else:
+            raise ValueError("Input must be a FEN string or a list of UCI moves")
+        
+        input_tensor_112 = input_tensor_112.to(device, non_blocking=True)
+        
+        self.eval()
+        with torch.inference_mode():
+            policy_logits, _, value_q = self.forward(input_tensor_112)
+        
+        # Apply legal moves mask
+        logits0 = policy_logits[0] + torch.from_numpy(legal_moves_mask).to(policy_logits.device)
+        
+        # Get best move
+        if T == 0.0:
+            best_move_idx = torch.argmax(logits0).item()
+        else:
+            scaled_logits = logits0 / T
+            probs = F.softmax(scaled_logits, dim=0)
+            move_idx = torch.multinomial(probs, 1).item()
+            best_move_idx = move_idx
+        
+        uci_move = policy_index[best_move_idx]
+        
+        # Mirror move if black is to move
+        if is_black_to_move:
+            def mirror_rank(rank_char):
+                rank = int(rank_char)
+                return str(9 - rank)
+            
+            if len(uci_move) >= 4:
+                from_file = uci_move[0]
+                from_rank = uci_move[1]
+                to_file = uci_move[2]
+                to_rank = uci_move[3]
+                promo = uci_move[4:] if len(uci_move) > 4 else ""
+                uci_move = from_file + mirror_rank(from_rank) + to_file + mirror_rank(to_rank) + promo
+        
+        # Convert castling moves
+        if uci_move == "e1h1" and "K" in castling_rights:
+            uci_move = "e1g1"
+        elif uci_move == "e1a1" and "Q" in castling_rights:
+            uci_move = "e1c1"
+        elif uci_move == "e8h8" and "k" in castling_rights:
+            uci_move = "e8g8"
+        elif uci_move == "e8a8" and "q" in castling_rights:
+            uci_move = "e8c8"
+        
+        # Get value (softmax over value_q)
+        value_probs = F.softmax(value_q[0], dim=0)
+        value = value_probs.cpu().numpy()
+        
+        return uci_move, value
+    
+    def get_position_value(self, fen_or_moves: Union[str, List[str]], device: str = None) -> np.ndarray:
+        """
+        Get position evaluation using value_q.
+        
+        Args:
+            fen_or_moves: Either a FEN string representing the chess position, or a list of UCI moves
+            device: Device to run the model on (if None, uses model's device)
+        
+        Returns:
+            Array of [black_win, draw, white_win] probabilities
+        """
+        # Detect device from model if not provided
+        if device is None:
+            device = next(self.parameters()).device
+        else:
+            device = torch.device(device)
+        
+        # Determine if input is FEN string or list of moves
+        if isinstance(fen_or_moves, str):
+            input_tensor_112, _ = encode_fen_to_tensor(fen_or_moves)
+        elif isinstance(fen_or_moves, list):
+            input_tensor_112, _ = encode_moves_to_tensor(fen_or_moves)
+        else:
+            raise ValueError("Input must be a FEN string or a list of UCI moves")
+        
+        input_tensor_112 = input_tensor_112.to(device, non_blocking=True)
+        
+        self.eval()
+        with torch.inference_mode():
+            _, _, value_q = self.forward(input_tensor_112)
+        
+        # Apply softmax to get probabilities [black_win, draw, white_win]
+        value_probs = F.softmax(value_q[0], dim=0)
+        return value_probs.cpu().numpy()
+    
+    def batch_get_moves_from_fens(self, fens: List[str], T: float, device: str = None, use_fp16: bool = False) -> List[str]:
+        """
+        Get moves for multiple FEN positions using batched inference.
+        
+        Args:
+            fens: List of FEN strings representing chess positions
+            T: Temperature for sampling (0.0 = deterministic/argmax, >0.0 = stochastic)
+            device: Device to run the model on (if None, uses model's device)
+        
+        Returns:
+            List of UCI move strings
+        """
+        if not fens:
+            return []
+        
+        # Detect device from model if not provided
+        if device is None:
+            device = next(self.parameters()).device
+        else:
+            device = torch.device(device)
+        
+        batch_size = len(fens)
+        
+        # Batch encode all FENs
+        input_tensors = []
+        legal_moves_masks = []
+        is_black_to_move_list = []
+        castling_rights_list = []
+        
+        for fen in fens:
+            input_tensor, legal_mask = encode_fen_to_tensor(fen)
+            input_tensors.append(input_tensor.squeeze(0))  # Remove batch dim
+            legal_moves_masks.append(legal_mask)
+            is_black_to_move_list.append(fen.split()[1] == 'b')
+            castling_rights_list.append(fen.split()[2] if len(fen.split()) > 2 else "")
+        
+        # Stack into batch tensor: (batch_size, 112, 8, 8)
+        batch_tensor = torch.stack(input_tensors).to(device, non_blocking=True)
+        if use_fp16 and device.type == 'cuda':
+            batch_tensor = batch_tensor.half()
+        
+        # Run batched inference
+        self.eval()
+        with torch.inference_mode():
+            if use_fp16 and device.type == 'cuda':
+                with torch.autocast(device_type='cuda', dtype=torch.float16):
+                    policy_logits,_,_ = self.forward(batch_tensor)
+            else:
+                policy_logits,_,_ = self.forward(batch_tensor)
+        
+        # Process each position in the batch
+        moves = []
+        for i in range(batch_size):
+            # Apply legal moves mask
+            logits = policy_logits[i] + torch.from_numpy(legal_moves_masks[i]).to(policy_logits.device, dtype=policy_logits.dtype)
+            
+            # Sample move
+            if T == 0.0:
+                best_move_idx = torch.argmax(logits).item()
+                uci_move = policy_index[best_move_idx]
+            else:
+                scaled_logits = logits / T
+                probs = F.softmax(scaled_logits, dim=0)
+                move_idx = torch.multinomial(probs, 1).item()
+                uci_move = policy_index[move_idx]
+            
+            # Mirror move if black is to move
+            if is_black_to_move_list[i]:
+                def mirror_rank(rank_char):
+                    rank = int(rank_char)
+                    return str(9 - rank)
+                
+                if len(uci_move) >= 4:
+                    from_file = uci_move[0]
+                    from_rank = uci_move[1]
+                    to_file = uci_move[2]
+                    to_rank = uci_move[3]
+                    promo = uci_move[4:] if len(uci_move) > 4 else ""
+                    
+                    uci_move = from_file + mirror_rank(from_rank) + to_file + mirror_rank(to_rank) + promo
+            
+            # Convert castling moves
+            castling_rights = castling_rights_list[i]
+            if uci_move == "e1h1" and "K" in castling_rights:
+                uci_move = "e1g1"
+            elif uci_move == "e1a1" and "Q" in castling_rights:
+                uci_move = "e1c1"
+            elif uci_move == "e8h8" and "k" in castling_rights:
+                uci_move = "e8g8"
+            elif uci_move == "e8a8" and "q" in castling_rights:
+                uci_move = "e8c8"
+            
+            moves.append(uci_move)
+        
+        return moves
+    
+    def batch_get_moves_from_move_lists(self, move_lists: List[List[str]], T: float, device: str = None, use_fp16: bool = False, fens: Optional[List[str]] = None):
+        """
+        Get moves for multiple move histories using batched inference.
+        
+        Args:
+            move_lists: List of move sequences, where each sequence is a list of UCI moves
+            T: Temperature for sampling (0.0 = deterministic/argmax, >0.0 = stochastic)
+            device: Device to run the model on (if None, uses model's device)
+            fens: Optional list of FEN strings that represent the board state prior to 
+                  applying the corresponding move list. When provided, each move history
+                  is applied starting from the supplied FEN instead of the standard initial position.
+        
+        Returns:
+            List of UCI move strings
+        """
+        if not move_lists:
+            return []
+        
+        # Detect device from model if not provided
+        if device is None:
+            device = next(self.parameters()).device
+        else:
+            device = torch.device(device)
+        
+        batch_size = len(move_lists)
+        
+        if fens is not None and len(fens) != len(move_lists):
+            raise ValueError("Length of fens must match length of move_lists when provided.")
+        
+        # Batch encode all move histories
+        input_tensors = []
+        legal_moves_masks = []
+        is_black_to_move_list = []
+        castling_rights_list = []
+        
+        for idx, move_history in enumerate(move_lists):
+            starting_fen = fens[idx] if fens is not None else None
+            input_tensor, legal_mask = encode_moves_to_tensor(move_history, starting_fen=starting_fen)
+            input_tensors.append(input_tensor.squeeze(0))  # Remove batch dim
+            legal_moves_masks.append(legal_mask)
+            
+            board = bulletchess.Board.from_fen(starting_fen) if starting_fen is not None else bulletchess.Board()
+            for mv in move_history:
+                move = bulletchess.Move.from_uci(mv)
+                board.apply(move)
+            is_black_to_move_list.append(board.turn == bulletchess.BLACK)
+            fen_parts = board.fen().split()
+            castling_rights_list.append(fen_parts[2] if len(fen_parts) > 2 else "")
+        
+        # Stack into batch tensor: (batch_size, 112, 8, 8)
+        batch_tensor = torch.stack(input_tensors).to(device, non_blocking=True)
+        if use_fp16 and device.type == 'cuda':
+            batch_tensor = batch_tensor.half()
+        
+        # Run batched inference
+        self.eval()
+        with torch.inference_mode():
+            if use_fp16 and device.type == 'cuda':
+                with torch.autocast(device_type='cuda', dtype=torch.float16):
+                    policy_logits,_,_ = self.forward(batch_tensor)
+            else:
+                policy_logits,_,_ = self.forward(batch_tensor)
+        
+        # Process each position in the batch
+        moves = []
+        for i in range(batch_size):
+            # Apply legal moves mask
+            logits = policy_logits[i] + torch.from_numpy(legal_moves_masks[i]).to(policy_logits.device, dtype=policy_logits.dtype)
+            
+            # Sample move
+            if T == 0.0:
+                best_move_idx = torch.argmax(logits).item()
+                uci_move = policy_index[best_move_idx]
+            else:
+                scaled_logits = logits / T
+                probs = F.softmax(scaled_logits, dim=0)
+                move_idx = torch.multinomial(probs, 1).item()
+                uci_move = policy_index[move_idx]
+            
+            # Mirror move if black is to move
+            if is_black_to_move_list[i]:
+                def mirror_rank(rank_char):
+                    rank = int(rank_char)
+                    return str(9 - rank)
+                
+                if len(uci_move) >= 4:
+                    from_file = uci_move[0]
+                    from_rank = uci_move[1]
+                    to_file = uci_move[2]
+                    to_rank = uci_move[3]
+                    promo = uci_move[4:] if len(uci_move) > 4 else ""
+                    
+                    uci_move = from_file + mirror_rank(from_rank) + to_file + mirror_rank(to_rank) + promo
+            
+            # Convert castling moves
+            castling_rights = castling_rights_list[i]
+            if uci_move == "e1h1" and "K" in castling_rights:
+                uci_move = "e1g1"
+            elif uci_move == "e1a1" and "Q" in castling_rights:
+                uci_move = "e1c1"
+            elif uci_move == "e8h8" and "k" in castling_rights:
+                uci_move = "e8g8"
+            elif uci_move == "e8a8" and "q" in castling_rights:
+                uci_move = "e8c8"
+            
+            moves.append(uci_move)
+        return moves