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

chess-gpt-eval-contrastive/mamba_module.py

@@ -90,11 +90,14 @@ class MambaPlayer:
         self.device = device
 
         self.hooks = []
+        self.max_seq_len = 1536
+        
         self.activations_sum = {}
         self.activations_count = {}
-        
         for i, layer in enumerate(self.model.backbone.layers):
-            self.activations_sum[i] = {"won": 0, "lost": 0, "current": 0}
+            self.activations_sum[i] = {"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))}
             self.activations_count[i] = {"won": 0, "lost": 0, "current": 0}
             
             def hook(module, input, output, layer_idx=i):
@@ -102,7 +105,8 @@ class MambaPlayer:
                     tensor_output = output[0]
                 else:
                     tensor_output = output
-                self.activations_sum[layer_idx]["current"] += tensor_output.detach().cpu().numpy()
+                seq_len = tensor_output.shape[1]
+                self.activations_sum[layer_idx]["current"][:, :seq_len, :] += tensor_output.detach().cpu().numpy()
                 self.activations_count[layer_idx]["current"] += 1
             
             self.hooks.append(layer.register_forward_hook(hook))
@@ -188,7 +192,7 @@ class MambaPlayer:
                     activations_sum[layer_idx] = {}
                     activations_count[layer_idx] = {}
                 if category not in activations_sum[layer_idx]:
-                    activations_sum[layer_idx][category] = 0
+                    activations_sum[layer_idx][category] = np.zeros((1, self.max_seq_len, self.model.config.d_model))
                     activations_count[layer_idx][category] = 0
                 
                 activations_sum[layer_idx][category] += self.activations_sum[layer_idx][category]
@@ -196,12 +200,10 @@ class MambaPlayer:
         
         with open(path, "wb") as f:
             pickle.dump((activations_sum, activations_count), f)
-
-        self.activations_sum = {}
-        self.activations_count = {}
-        for i, layer in enumerate(self.model.backbone.layers):
-            self.activations_sum[i] = {"won": 0, "lost": 0, "current": 0}
-            self.activations_count[i] = {"won": 0, "lost": 0, "current": 0}
+        
+        for layer_idx in self.activations_sum:
+            self.activations_sum[layer_idx]["current"].fill(0)
+            self.activations_count[layer_idx]["current"] = 0
     
     def apply_contrastive_activations(self, path):
         if os.path.exists(path):
@@ -214,7 +216,8 @@ class MambaPlayer:
                 contrastive_activations = won_activations - lost_activations
                 
                 def hook(module, input, output):
-                    return output + torch.from_numpy(contrastive_activations).to(output.device)
+                    seq_len = output.shape[1]
+                    return output + torch.from_numpy(contrastive_activations[:, :seq_len, :]).to(output.device)
                 
                 self.hooks[layer_idx] = self.model.backbone.layers[layer_idx].register_forward_hook(hook)