Update app.py

app.py

@@ -249,13 +249,21 @@ 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()
+    if "Pattern" in data.columns:
+        key = "Pattern"
+        difficulty_pattern_ids = data[key].unique()
+    elif "Difficulty" in data.columns:
+        key = "Difficulty"
+        difficulty_pattern_ids = data[key].unique()
+    else:
+        # Handle the case where neither 'Pattern' nor 'Difficulty' columns exist
+        key = None
+        difficulty_pattern_ids = []
 
     return key, difficulty_pattern_ids
 
 
+
 run_update_dataset()
 
 block = gr.Blocks(css=custom_css)  # Attach the custom CSS here
@@ -291,10 +299,10 @@ with block:
             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([str(dp_id) for dp_id in difficulty_pattern_ids], label=dp_key)
+                if dp_key is not None:
+                    gr.CheckboxGroup([str(dp_id) for dp_id in difficulty_pattern_ids], label=dp_key)
+                else:
+                    gr.HTML("<p>No difficulty or pattern data available for this environment.</p>")
 
                 # TASK TABS
                 for task_id in range(0, hivex_env["task_count"]):