gaia benchmark class and entrypoint script

Author: ollmer

.gitignore

@@ -167,4 +167,6 @@ _sandbox.py
 results/
 
 # gradio
-.gradio/
+.gradio/
+
+outputs/

scripts/run_gaia.py

@@ -0,0 +1,13 @@
+from agentlab.agents.tapeagent import TapeAgentArgs
+from agentlab.benchmarks.gaia import GaiaBenchmark
+from agentlab.experiments.study import make_study
+
+exp_dir = "./outputs/gaia/debug1"
+agent_args = TapeAgentArgs("gaia_agent")
+study = make_study(
+    benchmark=GaiaBenchmark(split="validation", exp_dir=exp_dir),
+    agent_args=[agent_args],
+    comment="Gaia eval",
+)
+
+study.run(n_jobs=1)

src/agentlab/benchmarks/gaia.py

@@ -1,6 +1,7 @@
 import os
-from typing import Literal
+from typing import Any, Literal
 
+import bgym
 import datasets
 from tapeagents.environment import ContainerExecutor
 from tapeagents.tools.browser import Browser
@@ -12,29 +13,48 @@
 from agentlab.benchmarks.multitool_gym import MultiToolGym
 
 
+class GaiaBenchmark(bgym.Benchmark):
+    name = "gaia"
+    split: Literal["test", "validation"]
+    exp_dir: str
+
+    high_level_action_set_args = None
+    is_multi_tab = False
+    supports_parallel_seeds = False
+    backends = ["gaia"]
+    env_args_list = None
+    task_metadata = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.env_args_list = []
+        dataset = datasets.load_dataset("gaia-benchmark/GAIA", "2023_all")[self.split]
+        for task in dataset:
+            task_dir = os.path.join(self.name, task["task_id"])
+            env_args = GaiaGymArgs(task=task, exp_dir=task_dir)
+            self.env_args_list.append(env_args)
+
+
 class GaiaGym(MultiToolGym):
     task: dict
     exp_dir: str
 
 
 class GaiaGymArgs(AbstractEnvArgs):
-    task_id: str
+    task: dict[str, Any]
     split: Literal["test", "validation"]
     exp_dir: str
     viewport_chars: int = 64000
 
     def make_env(self) -> GaiaGym:
         self.init_code_sandbox()
-        dataset = datasets.load_dataset("gaia-benchmark/GAIA", "2023_all")
-        tasks_by_id = {task["task_id"]: task for task in dataset[self.split]}
-        task = tasks_by_id[self.task_id]
         tools = [
             WebSearch(),
             VideoReader(self.exp_dir),
             Browser(self.exp_dir, viewport_chars=self.viewport_chars),
             CodeExecutor(self.exp_dir),
         ]
-        env = GaiaGym(tools=tools, task=task, exp_dir=self.exp_dir)
+        env = GaiaGym(tools=tools, task=self.task, exp_dir=self.exp_dir)
         return env
 
     def init_code_sandbox(self) -> None: