Update chess-gpt-eval-contrastive/mamba_module.py

chess-gpt-eval-contrastive/mamba_module.py

@@ -358,12 +358,13 @@ class MambaPlayer:
         }, step=self.wandb_step)
         torch.save(self.linear_probes, path)
 
-    def evaluate_linear_probes(self, board: chess.Board, game_state: str):
-        self.move_num = game_state.count('.')
+    def evaluate_linear_probes(self, board: chess.Board):
+        self.move_num = board.fullmove_number
         bucket = next(b for b in self.move_buckets if self.move_num <= b)
         for layer_idx in self.linear_probes:
-            X = torch.cat(self.activations_sum[layer_idx][bucket]['current'], dim=0)
+            X = torch.from_numpy(self.activations_sum[layer_idx][bucket]['current']).float()
             for probe_type in ['q_value', 'q_value_delta', 'material_balance']:
+                target = torch.tensor(self.linear_probe_targets[layer_idx][bucket][probe_type]).float().item() #.unsqueeze(1)
                 probe = self.linear_probes[layer_idx][probe_type]
-                prediction = probe(X)
-                print(f"Layer {layer_idx}, {probe_type}: {prediction.item()}")
+                prediction = probe(X).item()
+                print(f"Layer {layer_idx}, {probe_type}: {prediction} vs {target}")