Refactor search_leaderboard to accept columns to show and filter results accordingly

app.py

@@ -41,31 +41,36 @@ CITATION_BUTTON_TEXT = """
 """
 
 df = None
+original_columns_order = None
 
-def search_leaderboard(model_name):
+def search_leaderboard(model_name, columns_to_show):
     if len(model_name.strip()) == 0:
-        return df
+        return df.loc[:, columns_to_show]
     
     threshold = 95  # You can adjust this value to make the search more or less strict
+    filtered_df = df.copy()
 
     def calculate_similarity(row):
         similarity = fuzz.partial_ratio(model_name.lower(), row["Model"].lower())
         return similarity if similarity >= threshold else 0
 
     # Add a new column for similarity scores
-    df["similarity"] = df.apply(calculate_similarity, axis=1)
+    filtered_df["similarity"] = filtered_df.apply(calculate_similarity, axis=1)
 
     # Filter and sort the dataframe
-    filtered_df = df[df["similarity"] > 0].sort_values('similarity', ascending=False)
+    filtered_df = filtered_df[filtered_df["similarity"] > 0].sort_values('similarity', ascending=False)
 
     # Remove the similarity column before returning
     filtered_df = filtered_df.drop('similarity', axis=1)
 
+    # Filter the columns to show
+    filtered_df = filtered_df.loc[:, columns_to_show]
+
     return filtered_df
 
 
 def main():
-    global df
+    global df, original_columns_order
     df = load_retrieval_results()
     df[["Model"]] = df[["Model"]].applymap(lambda x: f'<a href="https://huggingface.co/{x}" target="_blank">{x}</a>')
     df.drop(columns=["Revision", "Precision", "Task"], inplace=True)
@@ -103,11 +108,11 @@ def main():
                         # Submit the search box and the leaderboard
                         search_box_retrieval.input(
                             search_leaderboard,
-                            inputs=[search_box_retrieval],
+                            inputs=[search_box_retrieval, columns_to_show_input],
                             outputs=retrieval_leaderboard
                         )
                         columns_to_show_input.select(
-                            lambda columns: gr.update(value=df[[col for col in df.columns if col in columns]]),
+                            lambda columns: df.loc[:, columns],
                             inputs=columns_to_show_input,
                             outputs=retrieval_leaderboard
                         )