Spaces:
Running
Running
| import random | |
| import json | |
| class EloRank: | |
| def __init__(self, initial_rating=1000, k_factor=32): | |
| """ | |
| Initialize the EloRank class. | |
| :param initial_rating: Initial ELO rating for each model. | |
| :param k_factor: The K-factor that determines the sensitivity of rating changes. | |
| """ | |
| self.ratings = {} | |
| self.initial_rating = initial_rating | |
| self.k_factor = k_factor | |
| self.wins = {} | |
| def add_model(self, model_id): | |
| """ | |
| Add a new model with the initial rating. | |
| :param model_id: Unique identifier for the model. | |
| """ | |
| self.ratings[model_id] = self.initial_rating | |
| self.wins[model_id] = 0 | |
| def record_match(self, winner, loser): | |
| """ | |
| Update the ratings based on a match result. | |
| :param winner: Model ID of the winner. | |
| :param loser: Model ID of the loser. | |
| """ | |
| rating_winner = self.ratings[winner] | |
| rating_loser = self.ratings[loser] | |
| expected_winner = self.expected_score(rating_winner, rating_loser) | |
| expected_loser = self.expected_score(rating_loser, rating_winner) | |
| self.ratings[winner] += self.k_factor * (1 - expected_winner) | |
| self.ratings[loser] += self.k_factor * (0 - expected_loser) | |
| # Update win count | |
| self.wins[winner] += 1 | |
| def expected_score(self, rating_a, rating_b): | |
| """ | |
| Calculate the expected score for a model. | |
| :param rating_a: Rating of model A. | |
| :param rating_b: Rating of model B. | |
| :return: Expected score. | |
| """ | |
| return 1 / (1 + 10 ** ((rating_b - rating_a) / 400)) | |
| def get_rating(self, model_id): | |
| """ | |
| Get the current rating of a model. | |
| :param model_id: Unique identifier for the model. | |
| :return: Current rating of the model. | |
| """ | |
| return self.ratings.get(model_id, None) | |
| def get_wins(self, model_id): | |
| """ | |
| Get the number of wins of a model. | |
| :param model_id: Unique identifier for the model. | |
| :return: Number of wins of the model. | |
| """ | |
| return self.wins.get(model_id, 0) | |
| def get_top_models(self, n=2): | |
| """ | |
| Get the top N models by rating. | |
| :param n: Number of top models to retrieve. | |
| :return: List of model IDs of the top models. | |
| """ | |
| return sorted(self.ratings, key=self.ratings.get, reverse=True)[:n] | |
| def sample_next_match(self): | |
| """ | |
| Sample the next match based on the probability proportional to the current rating. | |
| This approach helps accelerate the convergence of ranking. | |
| :return: Tuple of two model IDs for the next match. | |
| """ | |
| model_ids = list(self.ratings.keys()) | |
| probabilities = [self.ratings[model_id] for model_id in model_ids] | |
| total_rating = sum(probabilities) | |
| probabilities = [rating / total_rating for rating in probabilities] | |
| # Sample two different models for the next match | |
| next_match = random.choices(model_ids, probabilities, k=2) | |
| while next_match[0] == next_match[1]: | |
| next_match = random.choices(model_ids, probabilities, k=2) | |
| return tuple(next_match) | |
| def process_match_records(self, file_path): | |
| """ | |
| Process match records from a JSON file and update ratings and win counts accordingly. | |
| :param file_path: Path to the JSON file containing match records. | |
| """ | |
| with open(file_path, 'r') as file: | |
| match_records = json.load(file) | |
| for record in match_records: | |
| winner = record['winner'] | |
| model_1 = record['model_1'] | |
| model_2 = record['model_2'] | |
| # Add models if they are not already added | |
| if model_1 not in self.ratings: | |
| self.add_model(model_1) | |
| if model_2 not in self.ratings: | |
| self.add_model(model_2) | |
| # Record the match result | |
| if winner == model_1: | |
| self.record_match(model_1, model_2) | |
| elif winner == model_2: | |
| self.record_match(model_2, model_1) | |
| # # Example Usage | |
| # e = EloRank() | |
| # e.add_model('model_A') | |
| # e.add_model('model_B') | |
| # e.add_model('model_C') | |
| # e.record_match('model_A', 'model_B') | |
| # print(e.get_rating('model_A')) # Should be greater than the initial rating | |
| # print(e.get_rating('model_B')) # Should be less than the initial rating | |
| # print(e.get_top_models(2)) # Get the top 2 models | |
| # print(e.sample_next_match()) # Sample the next match based on ratings | |
| # # Process match records from a JSON file | |
| # e.process_match_records('match_records.json') | |
| # print(e.get_wins('model_A')) # Get the number of wins for model_A | |