feat: refactor the data loading function

app.py

@@ -77,41 +77,54 @@ def restart_space():
 from dataclasses import dataclass
 import pandas as pd
 from typing import Optional
+
+
 @dataclass
 class LeaderboardDataStore:
     raw_data: Optional[list]
-    original_df_qa: Optional[pd.DataFrame]
+    raw_qa_df: Optional[pd.DataFrame]
     original_df_long_doc: Optional[pd.DataFrame]
     leaderboard_df_qa: Optional[pd.DataFrame]
     leaderboard_df_long_doc: Optional[pd.DataFrame]
     reranking_models: Optional[list]
-
-data = {}
-data["AIR-Bench_24.04"] = LeaderboardDataStore(None, None, None, None, None, None)
-data["AIR-Bench_24.04"].raw_data = get_raw_eval_results(f"{EVAL_RESULTS_PATH}/AIR-Bench_24.04")
-data["AIR-Bench_24.04"].original_df_qa = get_leaderboard_df(
-    data["AIR-Bench_24.04"].raw_data, task='qa', metric=DEFAULT_METRIC_QA)
-data["AIR-Bench_24.04"].original_df_long_doc = get_leaderboard_df(
-    data["AIR-Bench_24.04"].raw_data, task='long-doc', metric=DEFAULT_METRIC_LONG_DOC)
-print(f'raw data: {len(data["AIR-Bench_24.04"].raw_data)}')
-print(f'QA data loaded: {data["AIR-Bench_24.04"].original_df_qa.shape}')
-print(f'Long-Doc data loaded: {len(data["AIR-Bench_24.04"].original_df_long_doc)}')
-
-data["AIR-Bench_24.04"].leaderboard_df_qa = data["AIR-Bench_24.04"].original_df_qa.copy()
-# leaderboard_df_qa = leaderboard_df_qa[has_no_nan_values(df, _benchmark_cols)]
-shown_columns_qa, types_qa = get_default_cols(
-    'qa', data["AIR-Bench_24.04"].leaderboard_df_qa.columns, add_fix_cols=True)
-data["AIR-Bench_24.04"].leaderboard_df_qa = data["AIR-Bench_24.04"].leaderboard_df_qa[~data["AIR-Bench_24.04"].leaderboard_df_qa[COL_NAME_IS_ANONYMOUS]][shown_columns_qa]
-data["AIR-Bench_24.04"].leaderboard_df_qa.drop([COL_NAME_REVISION, COL_NAME_TIMESTAMP], axis=1, inplace=True)
-
-data["AIR-Bench_24.04"].leaderboard_df_long_doc = data["AIR-Bench_24.04"].original_df_long_doc.copy()
-shown_columns_long_doc, types_long_doc = get_default_cols(
-    'long-doc', data["AIR-Bench_24.04"].leaderboard_df_long_doc.columns, add_fix_cols=True)
-data["AIR-Bench_24.04"].leaderboard_df_long_doc = data["AIR-Bench_24.04"].leaderboard_df_long_doc[~data["AIR-Bench_24.04"].leaderboard_df_long_doc[COL_NAME_IS_ANONYMOUS]][shown_columns_long_doc]
-data["AIR-Bench_24.04"].leaderboard_df_long_doc.drop([COL_NAME_REVISION, COL_NAME_TIMESTAMP], axis=1, inplace=True)
-
-data["AIR-Bench_24.04"].reranking_models = sorted(list(frozenset([eval_result.reranking_model for eval_result in data["AIR-Bench_24.04"].raw_data])))
-
+    types_qa: Optional[list]
+    types_long_doc: Optional[list]
+
+
+def load_eval_results(file_path: str):
+    output = {}
+    versions = ("AIR-Bench_24.04",)
+    for version in versions:
+        output[version] = LeaderboardDataStore(None, None, None, None, None, None, None, None)
+        output[version].raw_data = get_raw_eval_results(f"{file_path}/{version}")
+        output[version].raw_qa_df = get_leaderboard_df(
+            output[version].raw_data, task='qa', metric=DEFAULT_METRIC_QA)
+        output[version].original_df_long_doc = get_leaderboard_df(
+            output[version].raw_data, task='long-doc', metric=DEFAULT_METRIC_LONG_DOC)
+        print(f'raw data: {len(output[version].raw_data)}')
+        print(f'QA data loaded: {output[version].raw_qa_df.shape}')
+        print(f'Long-Doc data loaded: {len(output[version].original_df_long_doc)}')
+
+        output[version].leaderboard_df_qa = output[version].raw_qa_df.copy()
+        # leaderboard_df_qa = leaderboard_df_qa[has_no_nan_values(df, _benchmark_cols)]
+        shown_columns_qa, types_qa = get_default_cols(
+            'qa', output[version].leaderboard_df_qa.columns, add_fix_cols=True)
+        output[version].types_qa = types_qa
+        output[version].leaderboard_df_qa = output[version].leaderboard_df_qa[~output[version].leaderboard_df_qa[COL_NAME_IS_ANONYMOUS]][shown_columns_qa]
+        output[version].leaderboard_df_qa.drop([COL_NAME_REVISION, COL_NAME_TIMESTAMP], axis=1, inplace=True)
+
+        output[version].leaderboard_df_long_doc = output[version].original_df_long_doc.copy()
+        shown_columns_long_doc, types_long_doc = get_default_cols(
+            'long-doc', output[version].leaderboard_df_long_doc.columns, add_fix_cols=True)
+        output[version].types_long_doc = types_long_doc
+        output[version].leaderboard_df_long_doc = output[version].leaderboard_df_long_doc[~output[version].leaderboard_df_long_doc[COL_NAME_IS_ANONYMOUS]][shown_columns_long_doc]
+        output[version].leaderboard_df_long_doc.drop([COL_NAME_REVISION, COL_NAME_TIMESTAMP], axis=1, inplace=True)
+
+        output[version].reranking_models = sorted(list(frozenset([eval_result.reranking_model for eval_result in output[version].raw_data])))
+    return output
+
+
+data = load_eval_results(EVAL_RESULTS_PATH)
 
 def update_metric_qa(
         metric: str,
@@ -173,9 +186,9 @@ with demo:
                             # select reranking models
                             with gr.Column():
                                 selected_rerankings = get_reranking_dropdown(data["AIR-Bench_24.04"].reranking_models)
-                        leaderboard_table = get_leaderboard_table(data["AIR-Bench_24.04"].leaderboard_df_qa, types_qa)
+                        leaderboard_table = get_leaderboard_table(data["AIR-Bench_24.04"].leaderboard_df_qa, data["AIR-Bench_24.04"].types_qa)
                         # Dummy leaderboard for handling the case when the user uses backspace key
-                        hidden_leaderboard_table_for_search = get_leaderboard_table(data["AIR-Bench_24.04"].original_df_qa, types_qa, visible=False)
+                        hidden_leaderboard_table_for_search = get_leaderboard_table(data["AIR-Bench_24.04"].raw_qa_df, data["AIR-Bench_24.04"].types_qa, visible=False)
 
                         set_listeners(
                             "qa",
@@ -212,11 +225,11 @@ with demo:
                                 selected_noreranker = get_noreranking_dropdown()
                         lb_df_retriever = data["AIR-Bench_24.04"].leaderboard_df_qa[data["AIR-Bench_24.04"].leaderboard_df_qa[COL_NAME_RERANKING_MODEL] == "NoReranker"]
                         lb_df_retriever = reset_rank(lb_df_retriever)
-                        lb_table_retriever = get_leaderboard_table(lb_df_retriever, types_qa)
+                        lb_table_retriever = get_leaderboard_table(lb_df_retriever, data["AIR-Bench_24.04"].types_qa)
                         # Dummy leaderboard for handling the case when the user uses backspace key
-                        hidden_lb_df_retriever = data["AIR-Bench_24.04"].original_df_qa[data["AIR-Bench_24.04"].original_df_qa[COL_NAME_RERANKING_MODEL] == "NoReranker"]
+                        hidden_lb_df_retriever = data["AIR-Bench_24.04"].raw_qa_df[data["AIR-Bench_24.04"].raw_qa_df[COL_NAME_RERANKING_MODEL] == "NoReranker"]
                         hidden_lb_df_retriever = reset_rank(hidden_lb_df_retriever)
-                        hidden_lb_table_retriever = get_leaderboard_table(hidden_lb_df_retriever, types_qa, visible=False)
+                        hidden_lb_table_retriever = get_leaderboard_table(hidden_lb_df_retriever, data["AIR-Bench_24.04"].types_qa, visible=False)
 
                         set_listeners(
                             "qa",
@@ -254,11 +267,11 @@ with demo:
                                 selected_rerankings_reranker = get_reranking_dropdown(reranking_models_reranker)
                             with gr.Column(scale=1):
                                 search_bar_reranker = gr.Textbox(show_label=False, visible=False)
-                        lb_table_reranker = get_leaderboard_table(lb_df_reranker, types_qa)
-                        hidden_lb_df_reranker = data["AIR-Bench_24.04"].original_df_qa[data["AIR-Bench_24.04"].original_df_qa[COL_NAME_RETRIEVAL_MODEL] == BM25_LINK]
+                        lb_table_reranker = get_leaderboard_table(lb_df_reranker, data["AIR-Bench_24.04"].types_qa)
+                        hidden_lb_df_reranker = data["AIR-Bench_24.04"].raw_qa_df[data["AIR-Bench_24.04"].raw_qa_df[COL_NAME_RETRIEVAL_MODEL] == BM25_LINK]
                         hidden_lb_df_reranker = reset_rank(hidden_lb_df_reranker)
                         hidden_lb_table_reranker = get_leaderboard_table(
-                            hidden_lb_df_reranker, types_qa, visible=False
+                            hidden_lb_df_reranker, data["AIR-Bench_24.04"].types_qa, visible=False
                         )
 
                         set_listeners(
@@ -316,12 +329,12 @@ with demo:
                                 selected_rerankings = get_reranking_dropdown(data["AIR-Bench_24.04"].reranking_models)
 
                         lb_table = get_leaderboard_table(
-                            data["AIR-Bench_24.04"].leaderboard_df_long_doc, types_long_doc
+                            data["AIR-Bench_24.04"].leaderboard_df_long_doc, data["AIR-Bench_24.04"].types_long_doc
                         )
 
                         # Dummy leaderboard for handling the case when the user uses backspace key
                         hidden_lb_table_for_search = get_leaderboard_table(
-                            data["AIR-Bench_24.04"].original_df_long_doc, types_long_doc, visible=False
+                            data["AIR-Bench_24.04"].original_df_long_doc, data["AIR-Bench_24.04"].types_long_doc, visible=False
                         )
 
                         set_listeners(
@@ -366,9 +379,9 @@ with demo:
                         ]
                         hidden_lb_db_retriever_long_doc = reset_rank(hidden_lb_db_retriever_long_doc)
                         lb_table_retriever_long_doc = get_leaderboard_table(
-                            lb_df_retriever_long_doc, types_long_doc)
+                            lb_df_retriever_long_doc, data["AIR-Bench_24.04"].types_long_doc)
                         hidden_lb_table_retriever_long_doc = get_leaderboard_table(
-                            hidden_lb_db_retriever_long_doc, types_long_doc, visible=False
+                            hidden_lb_db_retriever_long_doc, data["AIR-Bench_24.04"].types_long_doc, visible=False
                         )
 
                         set_listeners(
@@ -408,11 +421,11 @@ with demo:
                                 selected_rerankings_reranker_ldoc = get_reranking_dropdown(reranking_models_reranker_ldoc)
                             with gr.Column(scale=1):
                                 search_bar_reranker_ldoc = gr.Textbox(show_label=False, visible=False)
-                        lb_table_reranker_ldoc = get_leaderboard_table(lb_df_reranker_ldoc, types_long_doc)
+                        lb_table_reranker_ldoc = get_leaderboard_table(lb_df_reranker_ldoc, data["AIR-Bench_24.04"].types_long_doc)
                         hidden_lb_df_reranker_ldoc = data["AIR-Bench_24.04"].original_df_long_doc[data["AIR-Bench_24.04"].original_df_long_doc[COL_NAME_RETRIEVAL_MODEL] == BM25_LINK]
                         hidden_lb_df_reranker_ldoc = reset_rank(hidden_lb_df_reranker_ldoc)
                         hidden_lb_table_reranker_ldoc = get_leaderboard_table(
-                            hidden_lb_df_reranker_ldoc, types_long_doc, visible=False
+                            hidden_lb_df_reranker_ldoc, data["AIR-Bench_24.04"].types_long_doc, visible=False
                         )
 
                         set_listeners(