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

chess-gpt-eval-contrastive/mamba_module.py

@@ -222,7 +222,7 @@ class MambaPlayer:
             self.activations_count[layer_idx] = {bucket: {"won": 0, "lost": 0, "current": 0}
                                          for bucket in self.move_buckets}
     
-    def apply_contrastive_activations(self, path):
+    def apply_contrastive_activations(self, path, weight=1.0):
         if os.path.exists(path):
             with open(path, "rb") as f:
                 activations_sum, activations_count = pickle.load(f)
@@ -232,7 +232,7 @@ class MambaPlayer:
                 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
-                return output + torch.from_numpy(contrastive_activations[:, :seq_len, :]).to(output.device)
+                return output + torch.from_numpy(contrastive_activations[:, :seq_len, :]).to(output.device) * weight
             
             for layer_idx in activations_sum:
                 for bucket in self.move_buckets: