Update app.py

app.py

@@ -47,7 +47,17 @@ custom_css = """
 """
 
 # Pattern: 0 Default, 1 Grid, 2 Chain, 3 Circle, 4 Square, 5 Cross, 6 Two_Rows, 7 Field, 8 Random
-pattern_map = {0: "0: Default", 1: "1: Grid", 2: "2: Chain", 3: "3: Circle", 4: "4: Square", 5: "5: Cross", 6: "6: Two Rows", 7: "7: Field", 8: "8: Random" }
+pattern_map = {
+    0: "0: Default",
+    1: "1: Grid",
+    2: "2: Chain",
+    3: "3: Circle",
+    4: "4: Square",
+    5: "5: Cross",
+    6: "6: Two Rows",
+    7: "7: Field",
+    8: "8: Random",
+}
 
 hivex_envs = [
     {
@@ -86,7 +96,7 @@ def restart():
 def download_leaderboard_dataset():
     path = snapshot_download(repo_id=DATASET_REPO_ID, repo_type="dataset")
     return path
-    
+
 
 def get_total_models():
     total_models = 0
@@ -94,7 +104,7 @@ def get_total_models():
         model_ids = get_model_ids(hivex_env["hivex_env"])
         total_models += len(model_ids)
     return total_models
-    
+
 
 def get_model_ids(hivex_env):
     api = HfApi()
@@ -130,7 +140,11 @@ def update_leaderboard_dataset_parallel(hivex_env, path):
         row["Task"] = results["task"]["name"]
         if "pattern-id" in results["task"] or "difficulty-id" in results["task"]:
             key = "Pattern" if "pattern-id" in results["task"] else "Difficulty"
-            row[key] = pattern_map[results["task"]["pattern-id"]] if "pattern-id" in results["task"] else results["task"]["difficulty-id"]
+            row[key] = (
+                pattern_map[results["task"]["pattern-id"]]
+                if "pattern-id" in results["task"]
+                else results["task"]["difficulty-id"]
+            )
 
         results_metrics = results["metrics"]
 
@@ -146,7 +160,7 @@ def update_leaderboard_dataset_parallel(hivex_env, path):
 
     # ranked_dataframe = rank_dataframe(pd.DataFrame.from_records(data))
     ranked_dataframe = pd.DataFrame.from_records(data)
-    
+
     new_history = ranked_dataframe
     file_path = path + "/" + hivex_env + ".csv"
     new_history.to_csv(file_path, index=False)
@@ -187,7 +201,7 @@ def get_data(rl_env, task_id, path) -> pd.DataFrame:
     filtered_data = filtered_data.drop(columns=["Task"])
 
     # Drop columns that have no data (all values are NaN)
-    filtered_data = filtered_data.dropna(axis=1, how='all')
+    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)]
@@ -201,6 +215,7 @@ def get_data(rl_env, task_id, path) -> pd.DataFrame:
 
     return filtered_data
 
+
 def get_task(rl_env, task_id, path) -> str:
     """
     Get the task name from the leaderboard dataset based on the rl_env and task_id.
@@ -223,13 +238,24 @@ def get_task(rl_env, task_id, path) -> str:
 def convert_to_title_case(text: str) -> str:
     # Replace underscores with spaces
     text = text.replace("_", " ")
-    
+
     # Convert each word to title case (capitalize the first letter)
     title_case_text = text.title()
-    
+
     return title_case_text
 
 
+def get_difficulty_pattern_ids_and_key(rl_env, path):
+    csv_path = path + "/" + rl_env + ".csv"
+    data = pd.read_csv(csv_path)
+
+    key = "Pattern" if "Pattern" in data.columns else "Difficulty"
+    # Get the unique values in the "Difficulty" column
+    difficulty_pattern_ids = data[key].unique()
+
+    key, difficulty_pattern_ids
+
+
 run_update_dataset()
 
 block = gr.Blocks(css=custom_css)  # Attach the custom CSS here
@@ -247,34 +273,47 @@ with block:
                 <h1 style="font-weight: bold;">HIVEX Leaderboard</h1>
             </div>
             """
-        ) 
+        )
     with gr.Row(elem_id="header-row"):
-        gr.HTML(f"<p style='text-align: center;'>Total models: {get_total_models()}</p>")
+        gr.HTML(
+            f"<p style='text-align: center;'>Total models: {get_total_models()}</p>"
+        )
     with gr.Row(elem_id="header-row"):
-        gr.HTML(f"<p style='text-align: center;'>Get started 🚀 on our <a href='https://github.com/hivex-research/hivex'>GitHub repository</a>!</p>")
+        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:
                 # ADD CHECK BOX GROUP TO SELECT DIFFICULTY / PATTERN IDs
-                
+                dp_key, difficulty_pattern_ids = get_difficulty_pattern_ids_and_key(
+                    hivex_env["hivex_env"], path_
+                )
+                gr.CheckboxGroup(difficulty_pattern_ids, label=dp_key)
+
                 # 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_))
+                    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_)
                             row_count = len(data)  # Number of rows in the data
-                            
+
                             gr_dataframe = gr.components.Dataframe(
                                 value=data,
                                 headers=["User", "Model"],
                                 datatype=["markdown", "markdown"],
-                                row_count=(row_count, 'fixed')  # Set to the exact number of rows in the data
+                                row_count=(
+                                    row_count,
+                                    "fixed",
+                                ),  # Set to the exact number of rows in the data
                             )