SLM-RAG-Arena / utils /leaderboard.py
Haoguang Cai
add UI and data processing
8a142a6
raw
history blame
2.83 kB
import os
import pandas as pd
import random
from .models import model_names
def load_leaderboard_data():
"""
Loads the leaderboard data from the leaderboard CSV file.
Returns the data in a format compatible with the application.
"""
# Initialize the results structure
results = {"wins": {}, "losses": {}, "ties": {}, "votes": 0}
try:
# Define the path to the CSV file for leaderboard
csv_path = os.path.join('utils', 'arena_df_leaderboard.csv')
# Check if the file exists and load it
if os.path.exists(csv_path):
df = pd.read_csv(csv_path)
# Process the data into our structure
for _, row in df.iterrows():
model = row['model']
results["wins"][model] = row['wins']
results["losses"][model] = row['losses']
results["ties"][model] = row['ties']
# Calculate total votes
for model in results["wins"].keys():
results["votes"] += results["wins"][model] + results["losses"][model] + results["ties"][model] // 2
else:
# If file doesn't exist, pre-populate with some data
for model in model_names:
results["wins"][model] = random.randint(0, 10)
results["losses"][model] = random.randint(0, 10)
results["ties"][model] = random.randint(0, 5)
# Calculate total votes
for model in model_names:
results["votes"] += results["wins"][model] + results["losses"][model] + results["ties"][model] // 2
return results
except Exception as e:
print(f"Error loading leaderboard data: {e}")
# Return the initialized structure if file can't be loaded
return results
def save_leaderboard_data(results):
"""
Saves the current leaderboard results back to the CSV file.
Parameters:
- results: The results dictionary containing wins, losses, ties, and votes
"""
try:
# Define the path to the CSV file
csv_path = os.path.join('utils', 'arena_df_leaderboard.csv')
# Convert the results dictionary to a DataFrame
data = []
for model in results["wins"].keys():
data.append({
'model': model,
'wins': results["wins"].get(model, 0),
'losses': results["losses"].get(model, 0),
'ties': results["ties"].get(model, 0)
})
df = pd.DataFrame(data)
# Save to CSV
df.to_csv(csv_path, index=False)
print(f"Leaderboard data saved successfully to {csv_path}")
except Exception as e:
print(f"Error saving leaderboard data: {e}")