filter nb_shots

app.py

@@ -70,13 +70,14 @@ def update_table(
     columns: list,
     phenotypes: list,
     metrics: list,
+    nb_shots: list,
     type_query: list,
     precision_query: str,
     size_query: list,
     show_deleted: bool,
     query: str,
 ):
-    filtered_df = filter_models(hidden_df, type_query, size_query, precision_query, show_deleted)
+    filtered_df = filter_models(hidden_df, type_query, size_query, precision_query, show_deleted, nb_shots)
     filtered_df = filter_queries(query, filtered_df)
     df = select_columns(filtered_df, columns, phenotypes, metrics)
     return df
@@ -124,8 +125,7 @@ def filter_queries(query: str, filtered_df: pd.DataFrame) -> pd.DataFrame:
 
 
 def filter_models(
-    df: pd.DataFrame, type_query: list, size_query: list, precision_query: list, show_deleted: bool
-) -> pd.DataFrame:
+    df: pd.DataFrame, type_query: list, size_query: list, precision_query: list, show_deleted: bool, nb_shots: list) -> pd.DataFrame:
     # Show all models
     if show_deleted:
         filtered_df = df
@@ -135,7 +135,9 @@ def filter_models(
     type_emoji = [t[0] for t in type_query]
     filtered_df = filtered_df.loc[df[AutoEvalColumn.model_type_symbol.name].isin(type_emoji)]
     filtered_df = filtered_df.loc[df[AutoEvalColumn.precision.name].isin(precision_query + ["None"])]
-
+    if -1 not in nb_shots:
+        filtered_df = filtered_df.loc[df[AutoEvalColumn.nb_shots.name].isin(nb_shots)]
+    
     numeric_interval = pd.IntervalIndex(sorted([NUMERIC_INTERVALS[s] for s in size_query]))
     params_column = pd.to_numeric(df[AutoEvalColumn.params.name], errors="coerce")
     mask = params_column.apply(lambda x: any(numeric_interval.contains(x)))
@@ -201,8 +203,8 @@ with demo:
                     with gr.Column(min_width=320):
                         filter_nb_shots = gr.CheckboxGroup(
                             label="Number of shots",
-                            choices=["Zero-shot", "10-shot", "All"],
-                            value=["Zero-shot", "10-shot", "All"],
+                            choices=[("Zero-shot", 0), ("10-shot", 10), ("All", -1)],
+                            value=[-1],
                             interactive=True,
                             elem_id="filter-nb-shots",
                         )
@@ -272,7 +274,7 @@ with demo:
                 ],
                 leaderboard_table,
             )
-            for selector in [shown_phenotypes, shown_metrics, shown_columns, filter_columns_type, filter_columns_precision, filter_columns_size, deleted_models_visibility]:
+            for selector in [shown_phenotypes, shown_metrics, shown_columns, filter_columns_type, filter_columns_precision, filter_columns_size, deleted_models_visibility, filter_nb_shots]:
                 selector.change(
                     update_table,
                     [
@@ -280,6 +282,7 @@ with demo:
                         shown_columns,
                         shown_phenotypes,
                         shown_metrics,
+                        filter_nb_shots,
                         filter_columns_type,
                         filter_columns_precision,
                         filter_columns_size,