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

chess-gpt-eval-contrastive/mamba_module.py

@@ -89,16 +89,20 @@ class MambaPlayer:
         self.ctx = ctx
         self.device = device
 
+        self.move_num = 0
         self.hooks = []
         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] = {"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}
+            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):
@@ -106,8 +110,9 @@ class MambaPlayer:
                 else:
                     tensor_output = output
                 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
+                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))
 
@@ -167,6 +172,7 @@ class MambaPlayer:
             return None
 
     def get_move(self, board: str, game_state: str, temperature: float) -> str:
+        self.move_num = game_state.count('.')
         completion = self.get_mamba_response(game_state, temperature, 8, self.vocab_size)
         return self.get_move_from_response(completion)
 
@@ -175,9 +181,10 @@ class MambaPlayer:
 
     def update_activations(self, result):
         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] += 1
-
+            for bucket in self.move_buckets:
+                self.activations_sum[layer_idx][bucket][result] += self.activations_sum[layer_idx][bucket]["current"]
+                self.activations_count[layer_idx][bucket][result] += 1
+    
     def save_activations(self, path):
         if os.path.exists(path):
             with open(path, "rb") as f:
@@ -187,38 +194,47 @@ class MambaPlayer:
             activations_count = {}
         
         for layer_idx in self.activations_sum:
-            for category in ["won", "lost"]:
+            for bucket in self.move_buckets:
+                if self.activations_count[layer_idx][bucket]["current"] == 0:
+                    continue
                 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] = 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]
-                activations_count[layer_idx][category] += self.activations_count[layer_idx][category]
+                if bucket not in activations_sum[layer_idx]:
+                    activations_sum[layer_idx][bucket] = {}
+                    activations_count[layer_idx][bucket] = {}
+                for category in ["won", "lost"]:
+                    if category not in activations_sum[layer_idx][bucket]:
+                        activations_sum[layer_idx][bucket][category] = np.zeros((1, self.max_seq_len, self.model.config.d_model))
+                        activations_count[layer_idx][bucket][category] = 0
+                    
+                    activations_sum[layer_idx][bucket][category] += self.activations_sum[layer_idx][bucket][category]
+                    activations_count[layer_idx][bucket][category] += self.activations_count[layer_idx][bucket][category]
         
         with open(path, "wb") as f:
             pickle.dump((activations_sum, activations_count), f)
         
         for layer_idx in self.activations_sum:
-            self.activations_sum[layer_idx]["current"].fill(0)
-            self.activations_count[layer_idx]["current"] = 0
+            for bucket in self.move_buckets:
+                self.activations_sum[layer_idx][bucket]["current"].fill(0)
+                self.activations_count[layer_idx][bucket]["current"] = 0
     
     def apply_contrastive_activations(self, path):
         if os.path.exists(path):
             with open(path, "rb") as 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"]
+            def hook(module, input, output, layer_idx, bucket):
+                seq_len = output.shape[1]
+                won_activations = activations_sum[layer_idx][bucket]["won"] / activations_count[layer_idx][bucket]["won"]
+                lost_activations = activations_sum[layer_idx][bucket]["lost"] / activations_count[layer_idx][bucket]["lost"]
                 contrastive_activations = won_activations - lost_activations
-                
-                def hook(module, input, output):
-                    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)
+                return output + torch.from_numpy(contrastive_activations[:, :seq_len, :]).to(output.device)
+            
+            for layer_idx in activations_sum:
+                for bucket in self.move_buckets:
+                    self.hooks.append(self.model.backbone.layers[layer_idx].register_forward_hook(
+                        lambda module, input, output, layer_idx=layer_idx, bucket=bucket: hook(module, input, output, layer_idx, bucket)
+                    ))