Update app.py

app.py

@@ -166,11 +166,9 @@ def run_update_dataset():
     )
 
 
-def get_data(rl_env, task_id, path) -> pd.DataFrame:
+def get_data(rl_env, task_id, path, selected_filters: List[str] = None) -> pd.DataFrame:
     """
-    Get data from rl_env, filter by the given task_id, and drop the Task-ID column.
-    Also drops any columns that have no data (all values are NaN) or all values are 0.0.
-    :return: filtered data as a pandas DataFrame without the Task-ID column
+    Get data from rl_env, filter by the given task_id and selected filters, and drop the Task-ID column.
     """
     csv_path = path + "/" + rl_env + ".csv"
     data = pd.read_csv(csv_path)
@@ -178,16 +176,15 @@ def get_data(rl_env, task_id, path) -> pd.DataFrame:
     # Filter the data to only include rows where the "Task-ID" column matches the given task_id
     filtered_data = data[data["Task-ID"] == task_id]
 
-    # Drop the "Task-ID" column
-    filtered_data = filtered_data.drop(columns=["Task-ID"])
+    # Apply selected filters for difficulty or pattern if provided
+    if selected_filters:
+        filtered_data = filtered_data[filtered_data['Pattern'].isin(selected_filters) | filtered_data['Difficulty'].isin(selected_filters)]
 
-    # Drop the "Task" column
-    filtered_data = filtered_data.drop(columns=["Task"])
+    # Drop the "Task-ID" and "Task" columns
+    filtered_data = filtered_data.drop(columns=["Task-ID", "Task"])
 
-    # Drop columns that have no data (all values are NaN)
+    # Drop columns that have no data (all values are NaN) or where all values are 0.0
     filtered_data = filtered_data.dropna(axis=1, how='all')
-
-    # Drop columns where all values are 0.0
     filtered_data = filtered_data.loc[:, (filtered_data != 0.0).any(axis=0)]
 
     # Convert User and Model columns to clickable links
@@ -252,18 +249,24 @@ with block:
         gr.HTML(f"<p style='text-align: center;'>Get started 🚀 on our <a href='https://github.com/hivex-research/hivex'>GitHub repository</a>!</p>")
 
     path_ = download_leaderboard_dataset()
-    # gr.Markdown(INTRODUCTION_TEXT, elem_classes="markdown-text")
     # ENVIRONMENT TABS
-    with gr.Tabs() as tabs: # elem_classes="tab-buttons"
+    with gr.Tabs() as tabs:  # elem_classes="tab-buttons"
         for env_index in range(0, len(hivex_envs)):
             hivex_env = hivex_envs[env_index]
             with gr.Tab(f"{hivex_env['title']}") as env_tabs:
-                # TASK TABS
                 for task_id in range(0, hivex_env["task_count"]):
                     task_title = convert_to_title_case(get_task(hivex_env["hivex_env"], task_id, path_))
                     with gr.TabItem(f"Task {task_id}: {task_title}"):
                         with gr.Row():
-                            data = get_data(hivex_env["hivex_env"], task_id, path_)
+                            # CheckboxGroup for Difficulty/Pattern
+                            selected_filters = gr.CheckboxGroup(
+                                choices=list(pattern_map.values()), 
+                                label="Select Difficulty/Pattern",
+                                value=list(pattern_map.values())  # Default to all selected
+                            )
+    
+                        with gr.Row():
+                            data = get_data(hivex_env["hivex_env"], task_id, path_, selected_filters=selected_filters.value)
                             row_count = len(data)  # Number of rows in the data
                             
                             gr_dataframe = gr.components.Dataframe(
@@ -272,6 +275,13 @@ with block:
                                 datatype=["markdown", "markdown"],
                                 row_count=(row_count, 'fixed')  # Set to the exact number of rows in the data
                             )
+    
+                            # Update data based on checkbox selection
+                            selected_filters.change(
+                                lambda filters: get_data(hivex_env["hivex_env"], task_id, path_, filters),
+                                inputs=[selected_filters],
+                                outputs=[gr_dataframe]
+                            )
 
 
 scheduler = BackgroundScheduler()