Add col slector

app.py

@@ -135,6 +135,9 @@ def get_leaderboard_df(merge_values: bool = True):
 
     # Trim minimath column names
     df.columns = [c.replace("_level_", "_l") for c in df.columns]
+
+    # Trim AIMO column names
+    df.columns = [c.replace("Aimo_", "") for c in df.columns]
     return df
 
 
@@ -153,6 +156,16 @@ def update_table(search_query):
     return df
 
 
+def filter_columns(cols):
+    index_cols = list(leaderboard_df.columns[:2])
+    new_cols = index_cols + cols
+    df = get_leaderboard_df()
+    df = df.copy()[new_cols]
+    # Drop rows with NaN values
+    df = df.copy().dropna(how="all", axis=0, subset=[c for c in df.columns if c in cols])
+    return df
+
+
 leaderboard_df = get_leaderboard_df()
 
 demo = gr.Blocks()
@@ -168,6 +181,10 @@ with demo:
                 label="Merge evals",
                 info="Merge evals for the same model. If there are duplicates, we display the largest one.",
             )
+        with gr.Row():
+            cols_bar = gr.CheckboxGroup(
+                choices=list(leaderboard_df.columns[2:]), show_label=False, info="Select columns to display"
+            )
         with gr.Group():
             leaderboard_df = get_leaderboard_df()
             leaderboard_table = gr.Dataframe(
@@ -179,6 +196,7 @@ with demo:
         with gr.Row():
             refresh_button = gr.Button("Refresh")
 
+    cols_bar.change(filter_columns, inputs=[cols_bar], outputs=[leaderboard_table])
     merge_values.change(refresh, inputs=[merge_values], outputs=[leaderboard_table])
     search_bar.submit(update_table, inputs=[search_bar], outputs=[leaderboard_table])
     refresh_button.click(refresh, inputs=[], outputs=[leaderboard_table])