Upload chess-gpt-eval-contrastive/mamba_module.py with huggingface_hub

chess-gpt-eval-contrastive/mamba_module.py

@@ -321,10 +321,10 @@ class MambaPlayer:
 
     def train_linear_probes(self):
         def get_lr(it):
-            warmup_iters = 25 * 43
+            warmup_iters = 150 * 43
             lr_decay_iters = 5000 * 43
-            learning_rate = 0.025
-            min_lr = 0.0001
+            learning_rate = 0.000015
+            min_lr = 0.000001
             # 1) linear warmup for warmup_iters steps
             if it < warmup_iters:
                 return learning_rate * it / warmup_iters
@@ -365,7 +365,7 @@ class MambaPlayer:
         self.linear_probe_targets = {i: {bucket: {'q_value': [], 'q_value_delta': [], 'material_balance': []} for bucket in self.move_buckets} for i in self.linear_probes}
 
     def save_linear_probe_data(self, path):
-        self.linear_save_ct += 1
+        self.linear_save_ct += 25
         wandb.log({
             "etc/games": self.linear_save_ct
         }, step=self.wandb_step)
@@ -382,4 +382,4 @@ class MambaPlayer:
                 #probe.eval()
                 prediction = probe(X).item()
                 print(f"Layer {layer_idx}, {probe_type}: {prediction} vs {target}")
-        self.linear_probe_targets = {i: {bucket: {'q_value': [], 'q_value_delta': [], 'material_balance': []} for bucket in self.move_buckets} for i in self.linear_probes}
+        self.linear_probe_targets = {i: {bucket: {'q_value': [], 'q_value_delta': [], 'material_balance': []} for bucket in self.move_buckets} for i in self.linear_probes}