Spaces:
Running
Running
| import csv | |
| 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 = {} | |
| with open(model_info_file, 'r', newline='', encoding='utf-8') as file: | |
| csv_reader = csv.reader(file) | |
| header = next(csv_reader) | |
| for row in csv_reader: | |
| if len(row) == 8: # Ensure that all columns are present in the row | |
| name, upload_date, description, parameter_count, creator, result_path, license, link = row | |
| model_info[name] = { | |
| "upload_date": upload_date, | |
| "description": description, | |
| "parameter_count": parameter_count, | |
| "creator": creator, | |
| "result_path": result_path, | |
| "license": license, | |
| "link": link | |
| } | |
| with open(elo_rating_pkl, "rb") as fin: | |
| elo_rating_results = pickle.load(fin) | |
| anony_elo_rating_results = elo_rating_results["anony"] | |
| full_elo_rating_results = elo_rating_results["full"] | |
| anony_leaderboard_data = anony_elo_rating_results["leaderboard_table_df"] | |
| full_leaderboard_data = full_elo_rating_results["leaderboard_table_df"] | |
| fields = ["key", "Model", "Arena Elo rating (anony)", "Arena Elo rating (full)", "license", "creator", "link"] | |
| # set creator and license to empty for now | |
| all_models = anony_leaderboard_data.index.tolist() | |
| for model in all_models: | |
| if not model in model_info: | |
| model_info[model] = {} | |
| model_info[model]["license"] = "N/A" | |
| model_info[model]["creator"] = "N/A" | |
| model_info[model]["link"] = "N/A" | |
| print(f"Model {model} not found in model_info.json") | |
| #continue # 没有 model_info 不显示 | |
| model_info[model]["Model"] = model | |
| model_info[model]["key"] = model | |
| if model in anony_leaderboard_data.index: | |
| model_info[model]["Arena Elo rating (anony)"] = anony_leaderboard_data.loc[model, "rating"] | |
| else: | |
| model_info[model]["Arena Elo rating (anony)"] = 0 | |
| if model in full_elo_rating_results["leaderboard_table_df"].index: | |
| model_info[model]["Arena Elo rating (full)"] = full_leaderboard_data.loc[model, "rating"] | |
| else: | |
| model_info[model]["Arena Elo rating (full)"] = 0 | |
| # if model in anony_leaderboard_data.index: | |
| # model_info[model]["Arena Elo rating"] = anony_leaderboard_data.loc[model, "rating"] | |
| # else: | |
| # model_info[model]["Arena Elo rating"] = 0 | |
| final_model_info = {} | |
| for model in model_info: | |
| if "Model" in model_info[model]: | |
| final_model_info[model] = model_info[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 (anony)"], ascending=False) | |
| df.to_csv(output_csv, index=False) | |
| print("Leaderboard data saved to", output_csv) | |
| print(df) | |
| if __name__ == "__main__": | |
| fire.Fire(main) |