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

chess-gpt-eval-contrastive/main.py

@@ -610,7 +610,7 @@ def play_games(
             if illegal_moves_one != 0:
                 player_one_legal_moves -= 1
                 illegal_move_numbers.append(board.fullmove_number)
-            if update_activations:
+            if update_activations or update_linear:
                 player_one.update_activations("current")
             if (
                 board.is_game_over()
@@ -677,6 +677,9 @@ def play_games(
             )
             games_saved += 1
 
+            if update_linear:
+                player_one.train_linear_probes()
+                
             if update_activations:
                 if player_one_resignation or player_one_failed_to_find_legal_move:
                     player_one.update_activations("lost")
@@ -690,6 +693,8 @@ def play_games(
         
                 if games_saved % save_activations_every == 0:
                     player_one.save_activations(activations_path)
+            elif update_linear:
+                player_one.update_activations("reset")
 
             if update_linear and games_saved % save_activations_every == 0:
                 player_one.save_linear_probe_data(linear_path)
@@ -756,7 +761,7 @@ if __name__ == "__main__":
             player_one_recording_name = nanogpt_player
 #            player_one = NanoGptPlayer(model_name=player_one_recording_name, move_num_in_gamestate=move_num_in_gamestate)
             #player_one_recording_name = f"xformer_rdm_{rm}"
-            player_one = MambaPlayer(model_name=player_one_recording_name, move_num_in_gamestate=move_num_in_gamestate)
+            player_one = MambaPlayer(model_name=player_one_recording_name, move_num_in_gamestate=move_num_in_gamestate, update_contrastive=update_activations, update_linear=update_linear, linear_probe_path=linear_path)
             player_one_recording_name = f'contrastive_weights_rdm/mamba_rdm_wgt_{wgt}' #f'contrastive_rdm/mamba_rdm_{rm}'
             if apply_activations:
                 player_one.apply_contrastive_activations(path=activations_path, weight=wgt)