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

chess-gpt-eval-contrastive/mamba_module.py

@@ -110,7 +110,7 @@ class MambaPlayer:
             else:
                 self.linear_probes = {}
         if update_contrastive or update_linear:
-            linear_size = self.model.config.d_model * self.max_seq_len
+            linear_size = self.model.config.d_model * 8 #self.model.config.d_model * self.max_seq_len
             for i, layer in enumerate(self.model.backbone.layers):
                 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)),
@@ -184,6 +184,7 @@ class MambaPlayer:
                     have_non_space = True
                 input_ids = torch.cat([input_ids, next_token_id.unsqueeze(0)], dim=1)
 
+        self.seq_len = input_ids[0].size(dim=0)
         model_response = self.decode(input_ids[0].tolist())
         model_response = model_response[len(game_state):].split(";")[0]
         return model_response
@@ -344,7 +345,7 @@ class MambaPlayer:
         for layer_idx in self.linear_probes:
             for bucket in self.move_buckets:
                 if self.activations_count[layer_idx][bucket]['current'] > 0:
-                    X = torch.from_numpy(self.activations_sum[layer_idx][bucket]['current']).float().flatten(1) #/ self.activations_count[layer_idx][bucket]['current']).float()
+                    X = torch.from_numpy(self.activations_sum[layer_idx][bucket]['current']).float().flatten(1)[:self.seq_len][-8:] #/ self.activations_count[layer_idx][bucket]['current']).float()
                     for probe_type in ['q_value', 'q_value_delta', 'material_balance']:
                         y = torch.tensor(self.linear_probe_targets[layer_idx][bucket][probe_type]).float().unsqueeze(1)
                         if len(y) > 0:
@@ -375,7 +376,7 @@ class MambaPlayer:
         self.move_num = board.fullmove_number
         bucket = next(b for b in self.move_buckets if self.move_num <= b)
         for layer_idx in self.linear_probes:
-            X = torch.from_numpy(self.activations_sum[layer_idx][bucket]['current']).float().flatten(1)
+            X = torch.from_numpy(self.activations_sum[layer_idx][bucket]['current']).float().flatten(1)[:self.seq_len][-8:]
             for probe_type in ['q_value', 'q_value_delta', 'material_balance']:
                 target = torch.tensor(self.linear_probe_targets[layer_idx][bucket][probe_type]).float().item()
                 probe = self.linear_probes[layer_idx][probe_type]