Update chess-gpt-eval/main.py

chess-gpt-eval/main.py

@@ -173,6 +173,31 @@ def get_move_from_gpt_response(response: Optional[str]) -> Optional[str]:
     return first_move
 
 
+def calculate_stats(csv_file_path):
+    data = []
+    with open(csv_file_path, "r") as csv_file:
+        reader = csv.DictReader(csv_file)
+        data = list(reader)
+
+    if not data:
+        return None
+
+    stats = {
+        "wins": sum(float(row["player_one_score"]) for row in data if row["player_one_score"] == 1),
+        "draws": len(data) - sum(float(row["player_two_score"]) for row in data if row["player_one_score"] == 1) - sum(float(row["player_one_score"]) for row in data if row["player_one_score"] == 1),
+        "illegal_attempts_ratio": sum(float(row["p1_illegal_attempts"]) for row in data) / (sum(float(row["p1_illegal_attempts"]) for row in data) + sum(float(row["player_one_legal_moves"]) for row in data)),
+        "illegal_moves_ratio": sum(float(row["player_one_illegal_moves"]) for row in data) / sum(float(row["player_one_illegal_moves"]) + float(row["player_one_legal_moves"]) for row in data),
+        "avg_attempts_per_illegal": sum(float(row["p1_avg_attempts_per_illegal"]) for row in data) / len(data),
+        "avg_first_illegal_move": sum(float(row["p1_first_illegal_move_num"]) for row in data if float(row["p1_first_illegal_move_num"]) > 0) / len([row for row in data if float(row["p1_first_illegal_move_num"]) > 0]),
+        "avg_illegal_move_num": sum(float(row["p1_avg_illegal_move_num"]) for row in data if float(row["p1_avg_illegal_move_num"]) > 0) / len([row for row in data if float(row["p1_avg_illegal_move_num"]) > 0]),
+        "lost_to_illegal_ratio": len([row for row in data if row["player_one_failed_to_find_legal_move"] == "True"]) / len([row for row in data if float(row["number_of_moves"]) > 0]),
+        "avg_game_length": sum(float(row["number_of_moves"]) for row in data) / len(data),
+        "max_game_length": max(float(row["number_of_moves"]) for row in data),
+    }
+
+    return stats
+
+
 def record_results(
     board: chess.Board,
     player_one: Player,
@@ -216,8 +241,8 @@ def record_results(
         # resignation / failed move situation I didn't think of
         # -1e10 at least ensures it doesn't fail silently
         if "-" in result:
-            player_one_score = result.split("-")[0]
-            player_two_score = result.split("-")[1]
+            player_one_score = float(result.split("-")[0])
+            player_two_score = float(result.split("-")[1])
         elif result == "*": # Loss due to hitting max moves
             player_one_score = 0
             player_two_score = 1
@@ -553,6 +578,7 @@ def play_games(
         print(f"Result: {board.result()}")
         print(board)
         print()
+        game_transcript = game_state.strip()
         if game_transcript not in unique_games:
             unique_games.add(game_transcript)
             record_results(
@@ -575,6 +601,7 @@ def play_games(
                 opening_moves,
                 illegal_move_numbers
             )
+            games_saved += 1
         else:
             print("Duplicate game; not saved.")
     if isinstance(player_one, StockfishPlayer):
@@ -582,7 +609,17 @@ def play_games(
     if isinstance(player_two, StockfishPlayer):
         player_two.close()
 
-        # print(game_state)
+    stats = calculate_stats(csv_file_path)
+    if stats:
+        print("\nStatistics:")
+        for key, value in stats.items():
+            print(f"{key}: {value}")
+        
+        with open(csv_file_path, "a") as csv_file:
+            writer = csv.writer(csv_file)
+            writer.writerow([""] * len(info_dict))  # Add empty cells for existing columns
+            writer.writerow(list(stats.keys()))
+            writer.writerow(list(stats.values()))
 
 
 RUN_FOR_ANALYSIS = True
@@ -615,4 +652,16 @@ if __name__ == "__main__":
             print(f"\n\nSTARTING GAMES AGAINST LC0 LEVEL {i}\n\n")
 
             play_games(player_one, player_two, num_games, book_opening=book_opening, random_opening=random_opening, random_opening_moves=random_opening_moves)
+            
+print("\n\n\n********\nFinal Statistics:\n********\n")
+for nanogpt_player in player_ones:
+    csv_file_path = f"logs/{player_one_recording_name}_vs_{player_two_recording_name}"
+    csv_file_path = csv_file_path.replace(".", "_") # Because I'm using ckpt filenames for nanogpt models 
+    csv_file_path += ".csv"
+    
+    stats = calculate_stats(csv_file_path)
+    if stats:
+        print(f"\nStatistics for {nanogpt_player}:")
+        for key, value in stats.items():
+            print(f"{key}: {value}")
 print("\n\n\n********\nDONE!\n********\n\n\n")