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

chess-gpt-eval-contrastive/main.py

@@ -597,7 +597,7 @@ def play_games(
             #print(f"|{game_state}|")
             #print(f"{current_move_num}", end=" ")
 
-            if update_linear:
+            if update_linear or eval_linear:
                 prev_q_value = evaluate_position(board.fen(), player_two.backend)
             (
                 game_state,
@@ -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 or update_linear:
+            if update_activations or update_linear or eval_linear:
                 player_one.update_activations("current")
             if (
                 board.is_game_over()
@@ -619,12 +619,15 @@ def play_games(
             ):
                 break
 
-            if update_linear:
+            if update_linear or eval_linear:
                 curr_q_value = evaluate_position(board.fen(), player_two.backend)
                 q_value_delta = curr_q_value - prev_q_value
                 material_bal = material_balance(board)
                 player_one.update_linear_probe_targets(curr_q_value, q_value_delta, material_bal)
-                player_one.train_linear_probes()
+                if update_linear:
+                    player_one.train_linear_probes()
+                if eval_linear:
+                    player_one.evaluate_linear_probes(board)
                 player_one.update_activations("reset")
             
             (
@@ -744,8 +747,8 @@ save_activations_every = 25
 contrastive_weight = 0.8
 
 linear_path="linear.pkl"
-update_linear = True
-eval_linear = False
+update_linear = False
+eval_linear = True
 if __name__ == "__main__":
     for nanogpt_player in player_ones:
         i = 0
@@ -753,7 +756,7 @@ if __name__ == "__main__":
 #        for rm in range(5, 36, 5):
         for i in [0]: # [3] #range(11):
 #        for wgt in [0.005, 0.01, 0.025, 0.05]:
-            num_games = 5000
+            num_games = 500
             # player_one = GPTPlayer(model="gpt-3.5-turbo-instruct")
             # player_one = LocalLlamaPlayer(model_name="meta-llama/Llama-2-7b-hf")
             # player_one = LocalLoraLlamaPlayer("meta-llama/Llama-2-7b-hf", "/workspace/axolotl/lora2-out")