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

chess-gpt-eval-contrastive/mamba_module.py

@@ -12,6 +12,7 @@ import torch.nn as nn
 import torch.optim as optim
 import wandb
 import math
+import json
 
 BASE_DIR = "mamba/"
 
@@ -376,13 +377,33 @@ class MambaPlayer:
     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)
+        
+        # Create a dictionary to store the statistics for the current move
+        probe_stats = {probe_type: {layer_idx: {self.move_num: None} for layer_idx in self.linear_probes} for probe_type in ['q_value', 'q_value_delta', 'material_balance']}
+        
         for layer_idx in self.linear_probes:
             X = torch.from_numpy(self.activations_sum[layer_idx][bucket]['current']).float().flatten(1)
             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()
                 probe = self.linear_probes[layer_idx][probe_type]
-                #probe.eval()
                 prediction = probe(X).item()
-                if probe_type == 'material_balance':
-                    print(f"Layer {layer_idx}, {probe_type}: {int(prediction)} vs {int(target)}")
+                #print(f"Layer {layer_idx}, {probe_type}: {prediction} vs {target}")
+                
+                # Calculate the percentage accuracy based on the probe type
+                if probe_type == 'q_value':
+                    accuracy = 1 - abs(prediction - target) / 2  # Q-value range: -1 to 1
+                elif probe_type == 'q_value_delta':
+                    accuracy = 1 - abs(prediction - target) / 4  # Q-value delta range: -2 to 2
+                else:  # material_balance
+                    max_range = 35  # Adjust this value based on the expected range of material balance
+                    accuracy = 1 - min(abs(prediction - target) / max_range, 1)
+                
+                # Store the accuracy in the probe_stats dictionary for the current move
+                probe_stats[probe_type][layer_idx][self.move_num] = accuracy
+        
         self.linear_probe_targets = {i: {bucket: {'q_value': [], 'q_value_delta': [], 'material_balance': []} for bucket in self.move_buckets} for i in self.linear_probes}
+        
+        # Append the probe_stats to the file
+        with open('probe_stats.json', 'a') as f:
+            json.dump(probe_stats, f)
+            f.write('\n')  # Add a newline separator between moves