Update utils/utils.py

utils/utils.py

@@ -19,6 +19,7 @@ def get_questions_from_file(file_path, level):
 def get_questions_from_dataset(file_path, level):
     # Extract dataset type from file path (e.g., "gaia" or "hle")
     basename = os.path.splitext(os.path.basename(file_path))[0]
+    dataset_type = basename.replace("_validation", "")
     
     # Get space ID from environment, defaulting to "bstraehle/gaia"
     space_id = os.environ.get("SPACE_ID", "bstraehle/gaia")
@@ -26,10 +27,16 @@ def get_questions_from_dataset(file_path, level):
     username = space_id.split("/")[0]
     dataset_repo = f"{username}/validation"
     
-    # Load dataset from Hugging Face - use the basename as the split name
-    dataset = load_dataset(dataset_repo, split=basename)
+    # Load dataset from Hugging Face
+    dataset = load_dataset(dataset_repo, split="validation")
     df = dataset.to_pandas()
     
+    # Filter by dataset type using the task_id prefix
+    if dataset_type == "gaia":
+        df = df[df["task_id"].str.startswith("gaia-")]
+    elif dataset_type == "hle":
+        df = df[df["task_id"].str.startswith("hle-")]
+    
     # Filter by level if level > 0 (for GAIA benchmark)
     if level > 0:
         df = df[df["Level"] == level]