Update elo.py

elo.py

@@ -14,25 +14,21 @@ def calculate_elo(old_rating, opponent_rating, score, k_factor=32):
     expected_score = 1 / (1 + 10 ** ((opponent_rating - old_rating) / 400))
     new_rating = old_rating + k_factor * (score - expected_score)
     return new_rating
-
+    
 def update_elo_ratings(ratings_dataset, winner, loser, k_factor=32):
-    """
-    Update ELO ratings for two players in a Hugging Face dataset.
-
-    :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 as a Hugging Face dataset.
-    """
-    # Convert the Hugging Face dataset to a pandas DataFrame for easier manipulation
+    # Convert the Hugging Face dataset to a pandas DataFrame
     ratings_df = pd.DataFrame(ratings_dataset)
 
+    # Check and add new players if they don't exist in the dataset
+    for player in [winner, loser]:
+        if player not in ratings_df['bot_name'].values:
+            new_player = {'bot_name': player, 'elo_rating': 1200}
+            ratings_df = ratings_df.append(new_player, ignore_index=True)
+
     # Extract old ratings
     winner_old_rating = ratings_df.loc[ratings_df['bot_name'] == winner, 'elo_rating'].iloc[0]
     loser_old_rating = ratings_df.loc[ratings_df['bot_name'] == loser, 'elo_rating'].iloc[0]
 
-
     # Calculate new ratings
     winner_new_rating = calculate_elo(winner_old_rating, loser_old_rating, 1, k_factor)
     loser_new_rating = calculate_elo(loser_old_rating, winner_old_rating, 0, k_factor)
@@ -44,4 +40,4 @@ def update_elo_ratings(ratings_dataset, winner, loser, k_factor=32):
     # Convert the DataFrame back to a Hugging Face dataset
     updated_ratings_dataset = Dataset.from_pandas(ratings_df)
 
-    return updated_ratings_dataset
+    return updated_ratings_dataset