Update modeling_chessbot.py

modeling_chessbot.py

@@ -652,6 +652,34 @@ class ChessBotModel(ChessBotPreTrainedModel):
             selected_move = policy_index[torch.multinomial(probs, 1).item()]
         
         return selected_move
+        
+ def get_position_value(self, fen, device="cuda"):
+        """
+        Get the value evaluation for a given FEN position.
+        Returns the value vector [black_win_prob, draw_prob, white_win_prob]
+        """
+        x = torch.from_numpy(fen_to_tensor(fen)).to(device).to(torch.float32)
+        x = x.view(1, 1, 8, 8, 19)
+        
+        # Forward pass through the model to get value
+        with torch.no_grad():
+            # We need to run through the model layers to get to value_head
+            b, seq_len, _, _, emb = x.size()
+            x_processed = x.view(b * seq_len, 64, emb)
+            x_processed = self.linear1(x_processed)
+            x_processed = F.gelu(x_processed)
+            x_processed = self.layernorm1(x_processed)
+            x_processed = self.ma_gating(x_processed)
+            
+            pos_enc = self.positional(x_processed)
+            for i in range(self.num_layers):
+                x_processed = self.layers[i](x_processed, pos_enc)
+            
+            value_logits = self.value_head_q(x_processed)
+            value_logits = value_logits.view(b, seq_len, 3)
+            value_logits = torch.softmax(value_logits, dim=-1)
+             
+        return value_logits.squeeze()  # Remove batch and sequence dimensions
 
     def get_batch_position_values(self, fens, device="cuda"):
         """