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

chess-gpt-eval-contrastive/mamba_module.py

@@ -7,6 +7,7 @@ from contextlib import nullcontext
 import numpy as np
 from functools import partial
 import chess
+from sklearn.linear_model import LinearRegression
 
 BASE_DIR = "mamba/"
 
@@ -100,10 +101,10 @@ class MambaPlayer:
             self.activations_count = {}
         if update_linear:
             if linear_probe_path and os.path.exists(linear_probe_path):
-                self.linear_probes = torch.load(linear_probe_data_path)
+                with open(linear_probe_path, 'rb') as f:
+                    self.linear_probes = pickle.load(f)
             else:
                 self.linear_probes = {}
-            self.linear_probe_targets = {}
         if update_contrastive or update_linear:
             for i, layer in enumerate(self.model.backbone.layers):
                 self.activations_sum[i] = {bucket: {"won": np.zeros((1, self.max_seq_len, self.model.config.d_model)),
@@ -131,7 +132,8 @@ class MambaPlayer:
                             'q_value_delta': torch.nn.Linear(self.model.config.d_model, 1),
                             'material_balance': torch.nn.Linear(self.model.config.d_model, 1)
                         }
-                    self.linear_probe_targets[i] = {bucket: {'q_value': [], 'q_value_delta': [], 'material_balance': []} for bucket in self.move_buckets}
+            if update_linear:
+                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}
 
     def get_mamba_response(self, game_state: str, temperature: float, max_new_tokens: int, top_k: int):
         game_state = game_state.split("\n\n")[-1].strip()
@@ -307,16 +309,16 @@ class MambaPlayer:
             for bucket in self.move_buckets:
                 if self.activations_count[layer_idx][bucket]['current'] > 0:
                     X = self.activations_sum[layer_idx][bucket]['current'] / self.activations_count[layer_idx][bucket]['current']
-                    X = torch.from_numpy(X).float()
                     for probe_type in ['q_value', 'q_value_delta', 'material_balance']:
-                        y = torch.tensor(self.linear_probe_targets[layer_idx][bucket][probe_type]).float().unsqueeze(1)
+                        y = self.linear_probe_targets[layer_idx][bucket][probe_type]
                         self.linear_probes[layer_idx][probe_type].fit(X, y)
-    
+
         # Reset linear_probe_targets after training
         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}
-    
+
     def save_linear_probe_data(self, path):
-        torch.save(self.linear_probes, path)
+        with open(path, 'wb') as f:
+            pickle.dump(self.linear_probes, f)
 
     def evaluate_linear_probes(self, board: chess.Board, game_state: str):
         self.move_num = game_state.count('.')