Update elo.py

elo.py

@@ -1,3 +1,6 @@
+import pandas as pd
+from datasets import Dataset
+
 def calculate_elo(old_rating, opponent_rating, score, k_factor=32):
     """
     Calculate the new ELO rating for a player.
@@ -12,26 +15,32 @@ def calculate_elo(old_rating, opponent_rating, score, k_factor=32):
     new_rating = old_rating + k_factor * (score - expected_score)
     return new_rating
 
-def update_elo_ratings(ratings, winner, loser, k_factor=32):
+def update_elo_ratings(ratings_dataset, winner, loser, k_factor=32):
     """
-    Update ELO ratings for two players.
+    Update ELO ratings for two players in a Hugging Face dataset.
 
-    :param ratings: A dictionary of current ELO ratings.
-    :param winner: The model name of the winning player.
-    :param loser: The model name of the losing player.
+    :param ratings_dataset: A Hugging Face dataset of current ELO ratings.
+    :param winner: The name of the winning player.
+    :param loser: The name of the losing player.
     :param k_factor: The K-factor used in ELO rating (default is 32).
-    :return: Updated ELO ratings.
+    :return: Updated ELO ratings as a Hugging Face dataset.
     """
-    winner_old_rating = ratings[winner]
-    loser_old_rating = ratings[loser]
+    # Convert the Hugging Face dataset to a pandas DataFrame for easier manipulation
+    ratings_df = pd.DataFrame(ratings_dataset)
+
+    # Extract old ratings
+    winner_old_rating = ratings_df.loc[ratings_df == winner, 'rating'].iloc[0]
+    loser_old_rating = ratings_df.loc[ratings_df == loser, 'rating'].iloc[0]
 
-    # Winner's new rating
+    # Calculate new ratings
     winner_new_rating = calculate_elo(winner_old_rating, loser_old_rating, 1, k_factor)
-    # Loser's new rating
     loser_new_rating = calculate_elo(loser_old_rating, winner_old_rating, 0, k_factor)
 
-    # Update the ratings
-    ratings[winner] = winner_new_rating
-    ratings[loser] = loser_new_rating
+    # Update the DataFrame
+    ratings_df.loc[ratings_df == winner, 'rating'] = winner_new_rating
+    ratings_df.loc[ratings_df == loser, 'rating'] = loser_new_rating
+
+    # Convert the DataFrame back to a Hugging Face dataset
+    updated_ratings_dataset = Dataset.from_pandas(ratings_df)
 
-    return ratings
+    return updated_ratings_dataset