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

chess-gpt-eval-contrastive/mamba_module.py

@@ -160,18 +160,39 @@ class MambaPlayer:
     def get_config(self) -> dict:
         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"])
-    
     def save_activations(self, path):
+        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"])
+            }
+        
         with open(path, "wb") as f:
-            pickle.dump(self.activations, f)
+            pickle.dump((activations_sum, activations_count), f)
+        
+        self.activations = {}
     
-    def load_activations(self, path):
+    def apply_contrastive_activations(self, path):
         if os.path.exists(path):
             with open(path, "rb") as f:
-                self.activations = pickle.load(f)
+                activations_sum, activations_count = pickle.load(f)
+            
+            for layer_idx in activations_sum:
+                won_activations = activations_sum[layer_idx]["won"] / activations_count[layer_idx]["won"]
+                lost_activations = activations_sum[layer_idx]["lost"] / activations_count[layer_idx]["lost"]
+                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)
     
     def apply_contrastive_activations(self):
         for layer_idx, layer_activations in self.activations.items():