3DGen-Arena / arena_elo /elo_rating /generate_leaderboard.py
ZhangYuhan's picture
update leaderboard
d75a844
raw
history blame
3.92 kB
import fire
import json
import pandas as pd
import pickle
def main(
model_info_file: str,
elo_rating_pkl: str,
output_csv: str
):
model_info = json.load(open(model_info_file))
with open(elo_rating_pkl, "rb") as fin:
elo_rating_results = pickle.load(fin)
# Model, Dim Elo rating (anony), Arena Elo rating (anony), Link, Orgnization
model_ratings = model_info
fields = ["key", "Model"]
for dim, dim_results in elo_rating_results.items():
anony_elo_rating_results = dim_results["anony"]
full_elo_rating_results = dim_results["full"]
anony_leaderboard_data = anony_elo_rating_results["leaderboard_table_df"]
full_leaderboard_data = full_elo_rating_results["leaderboard_table_df"]
fields += [f"{dim} Elo rating"]
all_models = anony_leaderboard_data.index.tolist()
for model in all_models:
if not model in model_ratings:
# set Organization and license to empty
model_ratings[model] = {}
model_ratings[model]["Organization"] = "N/A"
model_ratings[model]["Link"] = "N/A"
model_ratings[model]["Model"] = model
model_ratings[model]["key"] = model
if model in anony_leaderboard_data.index:
model_ratings[model][f"{dim} Elo rating"] = anony_leaderboard_data.loc[model, "rating"]
else:
model_ratings[model][f"{dim} Elo rating"] = 0
if "Arena Elo rating" not in model_ratings[model].keys():
model_ratings[model]["Arena Elo rating"] = 0
model_ratings[model]["Arena Elo rating"] += model_ratings[model][f"{dim} Elo rating"]
## Anony
# if model in anony_leaderboard_data.index:
# model_ratings[model][f"{dim} Elo rating (anony)"] = anony_leaderboard_data.loc[model, "rating"]
# else:
# model_ratings[model][f"{dim} Elo rating (anony)"] = 0
# if "Arena Elo rating (anony)" not in model_ratings[model].keys():
# model_ratings[model]["Arena Elo rating (anony)"] = 0
# model_ratings[model]["Arena Elo rating (anony)"] += model_ratings[model][f"{dim} Elo rating (anony)"]
## Anony + Named
# if model in full_elo_rating_results["leaderboard_table_df"].index:
# model_ratings[model][f"{dim} Elo rating (full)"] = full_leaderboard_data.loc[model, "rating"]
# else:
# model_ratings[model][f"{dim} Elo rating (full)"] = 0
# if "Arena Elo rating (full)" not in model_ratings[model].keys():
# model_ratings[model]["Arena Elo rating (full)"] = 0
# model_ratings[model]["Arena Elo rating (full)"] += model_ratings[model][f"{dim} Elo rating (full)"]
fields += ["Arena Elo rating", "Link", "Organization"]
# fields += ["Arena Elo rating (anony)", "Arena Elo rating (full)", "Link", "Organization"]
final_model_info = {}
print(model_ratings)
for model in model_ratings:
if "Model" in model_ratings[model]:
# model_ratings[model]["Arena Elo rating (anony)"] /= 5
# model_ratings[model]["Arena Elo rating (full)"] /= 5
model_ratings[model]["Arena Elo rating"] /= 5
final_model_info[model] = model_ratings[model]
model_info = final_model_info
exclude_keys = ['starting_from']
for key in exclude_keys:
for model in model_info:
if key in model_info[model]:
del model_info[model][key]
df = pd.DataFrame(model_info).T
df = df[fields]
# sort by anony rating
df = df.sort_values(by=["Arena Elo rating"], ascending=False)
df.to_csv(output_csv, index=False)
print("Leaderboard data saved to", output_csv)
print(df)
if __name__ == "__main__":
fire.Fire(main)