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

chess-gpt-eval-contrastive/mamba_module.py

@@ -11,7 +11,7 @@ import chess
 BASE_DIR = "mamba/"
 
 class MambaPlayer:
-    def __init__(self, model_name: str, move_num_in_gamestate: bool=False):
+    def __init__(self, model_name: str, move_num_in_gamestate: bool=False, update_contrastive: bool=False, update_linear: bool=False, linear_probe_path: str=None):
         self.model_name = model_name
         self.move_num_in_gamestate = move_num_in_gamestate
         # -----------------------------------------------------------------------------
@@ -95,27 +95,43 @@ class MambaPlayer:
         self.max_seq_len = 1536
         self.move_buckets = [10, 20, 30, 40, float('inf')]
         
-        self.activations_sum = {}
-        self.activations_count = {}
-        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)),
-                                                "lost": np.zeros((1, self.max_seq_len, self.model.config.d_model)),
-                                                "current": np.zeros((1, self.max_seq_len, self.model.config.d_model))}
-                                       for bucket in self.move_buckets}
-            self.activations_count[i] = {bucket: {"won": 0, "lost": 0, "current": 0}
-                                         for bucket in self.move_buckets}
-            
-            def hook(module, input, output, layer_idx=i):
-                if isinstance(output, tuple):
-                    tensor_output = output[0]
-                else:
-                    tensor_output = output
-                seq_len = tensor_output.shape[1]
-                bucket = next(b for b in self.move_buckets if self.move_num <= b)
-                self.activations_sum[layer_idx][bucket]["current"][:, :seq_len, :] += tensor_output.detach().cpu().numpy()
-                self.activations_count[layer_idx][bucket]["current"] += 1
-            
-            self.hooks.append(layer.register_forward_hook(hook))
+        if update_contrastive or update_linear:
+            self.activations_sum = {}
+            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)
+            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)),
+                                                    "lost": np.zeros((1, self.max_seq_len, self.model.config.d_model)),
+                                                    "current": np.zeros((1, self.max_seq_len, self.model.config.d_model))}
+                                           for bucket in self.move_buckets}
+                self.activations_count[i] = {bucket: {"won": 0, "lost": 0, "current": 0}
+                                             for bucket in self.move_buckets}
+                
+                def hook(module, input, output, layer_idx=i):
+                    if isinstance(output, tuple):
+                        tensor_output = output[0]
+                    else:
+                        tensor_output = output
+                    seq_len = tensor_output.shape[1]
+                    bucket = next(b for b in self.move_buckets if self.move_num <= b)
+                    self.activations_sum[layer_idx][bucket]["current"][:, :seq_len, :] += tensor_output.detach().cpu().numpy()
+                    self.activations_count[layer_idx][bucket]["current"] += 1
+                
+                self.hooks.append(layer.register_forward_hook(hook))
+                if update_linear:
+                    if not linear_probe_path or not os.path.exists(linear_probe_path):
+                        self.linear_probes[i] = {
+                            'q_value': torch.nn.Linear(self.model.config.d_model, 1),
+                            '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}
 
     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()
@@ -270,3 +286,36 @@ class MambaPlayer:
                 self.hooks.append(self.model.backbone.layers[layer_idx].register_forward_hook(
                     lambda module, input, output, layer_idx=layer_idx: hook(module, input, output, layer_idx)
                 ))
+
+    def update_linear_probe_targets(self, curr_q_value, q_value_delta, material_bal):
+        bucket = next(b for b in self.move_buckets if self.move_num <= b)
+        for layer_idx in self.linear_probe_targets:
+            self.linear_probe_targets[layer_idx][bucket]['q_value'].append(curr_q_value)
+            self.linear_probe_targets[layer_idx][bucket]['q_value_delta'].append(q_value_delta)
+            self.linear_probe_targets[layer_idx][bucket]['material_balance'].append(material_bal)
+
+    def train_linear_probes(self):
+        for layer_idx in self.linear_probes:
+            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)
+                        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)
+
+    def evaluate_linear_probes(self, board: chess.Board, game_state: str):
+        self.move_num = game_state.count('.')
+        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)
+            for probe_type in ['q_value', 'q_value_delta', 'material_balance']:
+                probe = self.linear_probes[layer_idx][probe_type]
+                prediction = probe(X)
+                print(f"Layer {layer_idx}, {probe_type}: {prediction.item()}")