fix: fix filtering

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 
-from app.utils import add_rank_and_format, filter_models, get_refresh_function
+from app.utils import add_rank_and_format, filter_models, get_refresh_function, deprecated_get_refresh_function
 from data.deprecated_model_handler import DeprecatedModelHandler
 from data.model_handler import ModelHandler
 
@@ -109,8 +109,8 @@ def main():
 
                 def update_data_1(metric, search_term, selected_columns):
                     model_handler.get_vidore_data(metric)
-                    data = deprecated_model_handler.render_df(metric, benchmark_version=1)
-                    data = add_rank_and_format(data, benchmark_version=1)
+                    data = model_handler.render_df(metric, benchmark_version=1)
+                    data = add_rank_and_format(data, benchmark_version=1, selected_columns=selected_columns)
                     data = filter_models(data, search_term)
                     if selected_columns:
                         data = data[["Rank", "Model", "Model Size (Million Parameters)", "Average"] + selected_columns]
@@ -193,8 +193,8 @@ def main():
 
                 def update_data_2(metric, search_term, selected_columns):
                     model_handler.get_vidore_data(metric)
-                    data = deprecated_model_handler.render_df(metric, benchmark_version=2)
-                    data = add_rank_and_format(data, benchmark_version=2)
+                    data = model_handler.render_df(metric, benchmark_version=2)
+                    data = add_rank_and_format(data, benchmark_version=2, selected_columns=selected_columns)
                     data = filter_models(data, search_term)
                     # data = remove_duplicates(data)  # Add this line
                     if selected_columns:
@@ -305,56 +305,56 @@ def main():
                 Refer to the [ColPali paper](https://arxiv.org/abs/2407.01449) for details on metrics, tasks and models.
                 """
                 )
-                datasets_columns_1 = list(deprecated_data_benchmark_1.columns[3:])
+                deprecated_datasets_columns_1 = list(deprecated_data_benchmark_1.columns[3:])
 
                 with gr.Row():
-                    metric_dropdown_1 = gr.Dropdown(choices=METRICS, value=initial_metric, label="Select Metric")
-                    research_textbox_1 = gr.Textbox(
+                    deprecated_metric_dropdown_1 = gr.Dropdown(choices=METRICS, value=initial_metric, label="Select Metric")
+                    deprecated_research_textbox_1 = gr.Textbox(
                         placeholder="🔍 Search Models... [press enter]",
                         label="Filter Models by Name",
                     )
-                    column_checkboxes_1 = gr.CheckboxGroup(
-                        choices=datasets_columns_1, value=datasets_columns_1, label="Select Columns to Display"
+                    deprecated_column_checkboxes_1 = gr.CheckboxGroup(
+                        choices=deprecated_datasets_columns_1, value=deprecated_datasets_columns_1, label="Select Columns to Display"
                     )
 
                 with gr.Row():
-                    datatype_1 = ["number", "markdown"] + ["number"] * (deprecated_num_datasets_1 + 1)
-                    dataframe_1 = gr.Dataframe(deprecated_data_benchmark_1, datatype=datatype_1, type="pandas")
+                    deprecated_datatype_1 = ["number", "markdown"] + ["number"] * (deprecated_num_datasets_1 + 1)
+                    deprecated_dataframe_1 = gr.Dataframe(deprecated_data_benchmark_1, datatype=deprecated_datatype_1, type="pandas")
 
-                def update_data_1(metric, search_term, selected_columns):
+                def deprecated_update_data_1(metric, search_term, selected_columns):
                     deprecated_model_handler.get_vidore_data(metric)
                     data = deprecated_model_handler.render_df(metric, benchmark_version=1)
-                    data = add_rank_and_format(data, benchmark_version=1)
+                    data = add_rank_and_format(data, benchmark_version=1, selected_columns=selected_columns)
                     data = filter_models(data, search_term)
                     # data = remove_duplicates(data)  # Add this line
                     if selected_columns:
-                        data = data[["Rank", "Model", "Model Size (Million Parameters)", "Average"] + selected_columns]
+                        data = data[["Rank", "Model", "Average"] + selected_columns]
                     return data
 
                 with gr.Row():
-                    refresh_button_1 = gr.Button("Refresh")
-                    refresh_button_1.click(
-                        get_refresh_function(deprecated_model_handler, benchmark_version=1),
-                        inputs=[metric_dropdown_1],
-                        outputs=dataframe_1,
+                    deprecated_refresh_button_1 = gr.Button("Refresh")
+                    deprecated_refresh_button_1.click(
+                        deprecated_get_refresh_function(deprecated_model_handler, benchmark_version=1),
+                        inputs=[deprecated_metric_dropdown_1],
+                        outputs=deprecated_dataframe_1,
                         concurrency_limit=20,
                     )
 
                 # Automatically refresh the dataframe when the dropdown value changes
-                metric_dropdown_1.change(
-                    get_refresh_function(deprecated_model_handler, benchmark_version=1),
-                    inputs=[metric_dropdown_1],
-                    outputs=dataframe_1,
+                deprecated_metric_dropdown_1.change(
+                    deprecated_get_refresh_function(deprecated_model_handler, benchmark_version=1),
+                    inputs=[deprecated_metric_dropdown_1],
+                    outputs=deprecated_dataframe_1,
                 )
-                research_textbox_1.submit(
-                    lambda metric, search_term, selected_columns: update_data_1(metric, search_term, selected_columns),
-                    inputs=[metric_dropdown_1, research_textbox_1, column_checkboxes_1],
-                    outputs=dataframe_1,
+                deprecated_research_textbox_1.submit(
+                    lambda metric, search_term, selected_columns: deprecated_update_data_1(metric, search_term, selected_columns),
+                    inputs=[deprecated_metric_dropdown_1, deprecated_research_textbox_1, deprecated_column_checkboxes_1],
+                    outputs=deprecated_dataframe_1,
                 )
-                column_checkboxes_1.change(
-                    lambda metric, search_term, selected_columns: update_data_1(metric, search_term, selected_columns),
-                    inputs=[metric_dropdown_1, research_textbox_1, column_checkboxes_1],
-                    outputs=dataframe_1,
+                deprecated_column_checkboxes_1.change(
+                    lambda metric, search_term, selected_columns: deprecated_update_data_1(metric, search_term, selected_columns),
+                    inputs=[deprecated_metric_dropdown_1, deprecated_research_textbox_1, deprecated_column_checkboxes_1],
+                    outputs=deprecated_dataframe_1,
                 )
 
                 gr.Markdown(
@@ -398,38 +398,38 @@ def main():
                 Refer to the [ColPali paper](https://arxiv.org/abs/2407.01449) for details on metrics and models.
                 """
                 )
-                datasets_columns_2 = list(deprecated_data_benchmark_2.columns[3:])
+                deprecated_datasets_columns_2 = list(deprecated_data_benchmark_2.columns[3:])
 
                 with gr.Row():
-                    metric_dropdown_2 = gr.Dropdown(choices=METRICS, value=initial_metric, label="Select Metric")
-                    research_textbox_2 = gr.Textbox(
+                    deprecated_metric_dropdown_2 = gr.Dropdown(choices=METRICS, value=initial_metric, label="Select Metric")
+                    deprecated_research_textbox_2 = gr.Textbox(
                         placeholder="🔍 Search Models... [press enter]",
                         label="Filter Models by Name",
                     )
-                    column_checkboxes_2 = gr.CheckboxGroup(
-                        choices=datasets_columns_2, value=datasets_columns_2, label="Select Columns to Display"
+                    deprecated_column_checkboxes_2 = gr.CheckboxGroup(
+                        choices=deprecated_datasets_columns_2, value=deprecated_datasets_columns_2, label="Select Columns to Display"
                     )
 
                 with gr.Row():
-                    datatype_2 = ["number", "markdown"] + ["number"] * (deprecated_num_datasets_2 + 1)
-                    dataframe_2 = gr.Dataframe(deprecated_data_benchmark_2, datatype=datatype_2, type="pandas")
+                    deprecated_datatype_2 = ["number", "markdown"] + ["number"] * (deprecated_num_datasets_2 + 1)
+                    deprecated_dataframe_2 = gr.Dataframe(deprecated_data_benchmark_2, datatype=deprecated_datatype_2, type="pandas")
 
-                def update_data_2(metric, search_term, selected_columns):
+                def deprecated_update_data_2(metric, search_term, selected_columns):
                     deprecated_model_handler.get_vidore_data(metric)
                     data = deprecated_model_handler.render_df(metric, benchmark_version=2)
-                    data = add_rank_and_format(data, benchmark_version=2)
+                    data = add_rank_and_format(data, benchmark_version=2, selected_columns=selected_columns)
                     data = filter_models(data, search_term)
                     # data = remove_duplicates(data)  # Add this line
                     if selected_columns:
-                        data = data[["Rank", "Model", "Model Size (Million Parameters)", "Average"] + selected_columns]
+                        data = data[["Rank", "Model", "Average"] + selected_columns]
                     return data
 
                 with gr.Row():
-                    refresh_button_2 = gr.Button("Refresh")
-                    refresh_button_2.click(
-                        get_refresh_function(deprecated_model_handler, benchmark_version=2),
-                        inputs=[metric_dropdown_2],
-                        outputs=dataframe_2,
+                    deprecated_refresh_button_2 = gr.Button("Refresh")
+                    deprecated_refresh_button_2.click(
+                        deprecated_get_refresh_function(deprecated_model_handler, benchmark_version=2),
+                        inputs=[deprecated_metric_dropdown_2],
+                        outputs=deprecated_dataframe_2,
                         concurrency_limit=20,
                     )
 
@@ -442,20 +442,20 @@ def main():
                     )
 
                 # Automatically refresh the dataframe when the dropdown value changes
-                metric_dropdown_2.change(
-                    get_refresh_function(deprecated_model_handler, benchmark_version=2),
-                    inputs=[metric_dropdown_2],
-                    outputs=dataframe_2,
+                deprecated_metric_dropdown_2.change(
+                    deprecated_get_refresh_function(deprecated_model_handler, benchmark_version=2),
+                    inputs=[deprecated_metric_dropdown_2],
+                    outputs=deprecated_dataframe_2,
                 )
-                research_textbox_2.submit(
-                    lambda metric, search_term, selected_columns: update_data_2(metric, search_term, selected_columns),
-                    inputs=[metric_dropdown_2, research_textbox_2, column_checkboxes_2],
-                    outputs=dataframe_2,
+                deprecated_research_textbox_2.submit(
+                    lambda metric, search_term, selected_columns: deprecated_update_data_2(metric, search_term, selected_columns),
+                    inputs=[deprecated_metric_dropdown_2, deprecated_research_textbox_2, deprecated_column_checkboxes_2],
+                    outputs=deprecated_dataframe_2,
                 )
-                column_checkboxes_2.change(
-                    lambda metric, search_term, selected_columns: update_data_2(metric, search_term, selected_columns),
-                    inputs=[metric_dropdown_2, research_textbox_2, column_checkboxes_2],
-                    outputs=dataframe_2,
+                deprecated_column_checkboxes_2.change(
+                    lambda metric, search_term, selected_columns: deprecated_update_data_2(metric, search_term, selected_columns),
+                    inputs=[deprecated_metric_dropdown_2, deprecated_research_textbox_2, deprecated_column_checkboxes_2],
+                    outputs=deprecated_dataframe_2,
                 )
 
                 gr.Markdown(

app/utils.py

@@ -17,20 +17,23 @@ def make_clickable_model(model_name, link=None):
     return f'<a target="_blank" style="text-decoration: underline" href="{link}">{desanitized_model_name}</a>'
 
 
-def add_rank(df, benchmark_version=1):
+def add_rank(df, benchmark_version=1, selected_columns=None):
         df.fillna(0.0, inplace=True)
-        cols_to_rank = [
-            col
-            for col in df.columns
-            if col
-            not in [
-                "Model",
-                "Model Size (Million Parameters)",
-                "Memory Usage (GB, fp32)",
-                "Embedding Dimensions",
-                "Max Tokens",
+        if selected_columns is None:
+            cols_to_rank = [
+                col
+                for col in df.columns
+                if col
+                not in [
+                    "Model",
+                    "Model Size (Million Parameters)",
+                    "Memory Usage (GB, fp32)",
+                    "Embedding Dimensions",
+                    "Max Tokens",
+                ]
             ]
-        ]
+        else:
+            cols_to_rank = selected_columns
 
         if len(cols_to_rank) == 1:
             df.sort_values(cols_to_rank[0], ascending=False, inplace=True)
@@ -45,10 +48,10 @@ def add_rank(df, benchmark_version=1):
         return df
 
 
-def add_rank_and_format(df, benchmark_version=1):
+def add_rank_and_format(df, benchmark_version=1, selected_columns=None):
     df = df.reset_index()
     df = df.rename(columns={"index": "Model"})
-    df = add_rank(df, benchmark_version)
+    df = add_rank(df, benchmark_version, selected_columns)
     df["Model"] = df["Model"].apply(make_clickable_model)
     # df = remove_duplicates(df)
     return df
@@ -71,6 +74,15 @@ def get_refresh_function(model_handler, benchmark_version):
 
     return _refresh
 
+def deprecated_get_refresh_function(model_handler, benchmark_version):
+    def _refresh(metric):
+        model_handler.get_vidore_data(metric)
+        data_task_category = model_handler.render_df(metric, benchmark_version)
+        df = add_rank_and_format(data_task_category, benchmark_version)
+        return df
+
+    return _refresh
+
 
 def filter_models(data, search_term):
     if search_term: