Update utils/utils.py

utils/utils.py

@@ -39,22 +39,20 @@ def get_dataset(dataset_type, level):
     dataset_repo = f"{username}/validation"
     dataset = load_dataset(dataset_repo, split="validation")
     df = dataset.to_pandas()
-    
+
     if dataset_type == DATASET_TYPE_GAIA:
         print(f"filtering for dataset_type={dataset_type}")
-        df = df[df["task_id"].str.startswith("gaia-")]
+        df = df[df["Level"].isin([1, 2, 3])]
     elif dataset_type == DATASET_TYPE_HLE:
         print(f"filtering for dataset_type={dataset_type}")
-        df = df[df["task_id"].str.startswith("hle-")]
+        df = df[df["Level"] == 0]
     
-    if level > 0:
-        df = df[df["Level"] == level]
-
-    print(df)
+    df = df[df["Level"] == level]
     
     result=[]
     
     for _, row in df.iterrows():
+        print(row)
         result.append([row["Question"], row["Final answer"], row["file_name"]])
     
     return result