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

chess-gpt-eval-contrastive/mamba_module.py

@@ -89,18 +89,21 @@ class MambaPlayer:
         self.ctx = ctx
         self.device = device
 
-        self.activations = {}
         self.hooks = []
+        self.activations_sum = {}
+        self.activations_count = {}
         
         for i, layer in enumerate(self.model.backbone.layers):
-            self.activations[i] = {"won": [], "lost": []}
+            self.activations_sum[i] = {"won": 0, "lost": 0}
+            self.activations_count[i] = {"won": 0, "lost": 0}
             
             def hook(module, input, output, layer_idx=i):
                 if isinstance(output, tuple):
                     tensor_output = output[0]
                 else:
                     tensor_output = output
-                self.activations[layer_idx]["current"] = tensor_output.detach().cpu().numpy()
+                self.activations_sum[layer_idx]["current"] += tensor_output.detach().cpu().numpy()
+                self.activations_count[layer_idx]["current"] += 1
             
             self.hooks.append(layer.register_forward_hook(hook))
 
@@ -167,27 +170,37 @@ class MambaPlayer:
         return {"model": self.model_name}
 
     def update_activations(self, result):
-        for layer_idx in self.activations:
-            self.activations[layer_idx][result].append(self.activations[layer_idx]["current"])
+        for layer_idx in self.activations_sum:
+            self.activations_sum[layer_idx][result] += self.activations_sum[layer_idx]["current"]
+            self.activations_count[layer_idx][result] += self.activations_count[layer_idx]["current"]
+            self.activations_sum[layer_idx]["current"] = 0
+            self.activations_count[layer_idx]["current"] = 0
 
     def save_activations(self, path):
-        activations_sum = {}
-        activations_count = {}
+        if os.path.exists(path):
+            with open(path, "rb") as f:
+                activations_sum, activations_count = pickle.load(f)
+        else:
+            activations_sum = {}
+            activations_count = {}
         
-        for layer_idx, layer_activations in self.activations.items():
-            activations_sum[layer_idx] = {
-                "won": np.sum(layer_activations["won"], axis=0),
-                "lost": np.sum(layer_activations["lost"], axis=0)
-            }
-            activations_count[layer_idx] = {
-                "won": len(layer_activations["won"]),
-                "lost": len(layer_activations["lost"])
-            }
+        for layer_idx in self.activations_sum:
+            for category in ["won", "lost"]:
+                if layer_idx not in activations_sum:
+                    activations_sum[layer_idx] = {}
+                    activations_count[layer_idx] = {}
+                if category not in activations_sum[layer_idx]:
+                    activations_sum[layer_idx][category] = 0
+                    activations_count[layer_idx][category] = 0
+                
+                activations_sum[layer_idx][category] += self.activations_sum[layer_idx][category]
+                activations_count[layer_idx][category] += self.activations_count[layer_idx][category]
         
         with open(path, "wb") as f:
             pickle.dump((activations_sum, activations_count), f)
-        
-        self.activations = {}
+
+        self.activations_sum = {}
+        self.activations_count = {}
     
     def apply_contrastive_activations(self, path):
         if os.path.exists(path):
@@ -203,16 +216,5 @@ class MambaPlayer:
                     return output + torch.from_numpy(contrastive_activations).to(output.device)
                 
                 self.hooks[layer_idx] = self.model.backbone.layers[layer_idx].register_forward_hook(hook)
-    
-    def apply_contrastive_activations(self):
-        for layer_idx, layer_activations in self.activations.items():
-            if len(layer_activations["won"]) > 0 and len(layer_activations["lost"]) > 0:
-                won_activations = np.mean(layer_activations["won"], axis=0)
-                lost_activations = np.mean(layer_activations["lost"], axis=0)
-                contrastive_activations = won_activations - lost_activations
-                
-                def hook(module, input, output):
-                    return output + torch.from_numpy(contrastive_activations).to(output.device)
-                
-                self.hooks[layer_idx] = self.model.backbone.layers[layer_idx].register_forward_hook(hook)
+