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

chess-gpt-eval-contrastive/mamba_module.py

@@ -87,6 +87,17 @@ class MambaPlayer:
         self.ctx = ctx
         self.device = device
 
+        self.activations = {}
+        self.hooks = []
+        
+        for i, layer in enumerate(self.model.backbone.layers):
+            self.activations[i] = {"won": [], "lost": []}
+            
+            def hook(module, input, output, layer_idx=i):
+                self.activations[layer_idx]["current"] = output.detach().cpu().numpy()
+            
+            self.hooks.append(layer.register_forward_hook(hook))
+
     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()
         #game_state = ";" + game_state
@@ -149,3 +160,28 @@ 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):
+        with open(path, "wb") as f:
+            pickle.dump(self.activations, f)
+    
+    def load_activations(self, path):
+        if os.path.exists(path):
+            with open(path, "rb") as f:
+                self.activations = pickle.load(f)
+    
+    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)
+