Refactor retrieval and reranking result loading to support customizable sorting and column dropping

app.py

@@ -110,12 +110,12 @@ def main():
     global retrieval_df, reranking_df
 
     # Prepare retrieval dataframe
-    retrieval_df = load_retrieval_results(prepare_for_display=True)
+    retrieval_df = load_retrieval_results(True, "Web Search Dataset (Overall Score)", ["Revision", "Precision", "Task"])
     retrieval_columns_to_show = ["Model", "Web Search Dataset (Overall Score)", "Model Size (in Millions)", "Embedding Dimension", "Max Tokens", "Num Likes"]
     retrieval_cols = retrieval_df.columns.tolist()  # cache columns
 
     # Prepare reranking dataframe
-    reranking_df = load_reranking_results(prepare_for_display=True)
+    reranking_df = load_reranking_results(True, sort_col="Overall Score")
     reranking_columns_to_show = reranking_df.columns.tolist()
     reranking_cols = reranking_df.columns.tolist()  # cache columns
 

utils.py

@@ -30,21 +30,22 @@ def load_json_results(file_path: Path, prepare_for_display=False, sort_col=None,
             df.sort_values(sort_col, ascending=False, inplace=True)
     return df
 
-def load_retrieval_results(prepare_for_display=False):
+def load_retrieval_results(prepare_for_display=False, sort_col=None, drop_cols=None):
     dataframe_path = results_dir / "retrieval_results.json"
     return load_json_results(
         dataframe_path,
         prepare_for_display=prepare_for_display,
-        sort_col="Web Search Dataset (Overall Score)",
-        drop_cols=["Revision", "Precision", "Task"]
+        sort_col=sort_col,
+        drop_cols=drop_cols
     )
 
-def load_reranking_results(prepare_for_display=False):
+def load_reranking_results(prepare_for_display=False, sort_col=None, drop_cols=None):
     dataframe_path = results_dir / "reranking_results.json"
     return load_json_results(
         dataframe_path,
         prepare_for_display=prepare_for_display,
-        sort_col="Overall Score"
+        sort_col=sort_col,
+        drop_cols=drop_cols
     )
 
 def get_model_info(model_id, verbose=False):