86 lines
3.3 KiB
Python
86 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Sesija C Phase 3b-A — Gaia2 HF dataset → JSONL dumper.
|
|
|
|
Runs from inside the ARE-installed venv (uv-managed under
|
|
external/meta-agents-research-environments/.venv); `datasets` package is
|
|
already a transitive dep of meta-agents-research-environments.
|
|
|
|
Usage:
|
|
python benchmarks/gaia2/scripts/dump-tasks.py \\
|
|
--hf-config mini \\
|
|
--hf-split validation \\
|
|
--limit 10 \\
|
|
--output benchmarks/gaia2/data/tasks-mini-10.jsonl
|
|
|
|
Emits one JSON record per line, schema matching benchmarks/gaia2/adapter.ts
|
|
`Gaia2HfTask` interface (HF dataset card schema verified 2026-04-30).
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="Dump Gaia2 HF tasks to JSONL")
|
|
parser.add_argument("--hf-dataset", default="meta-agents-research-environments/gaia2")
|
|
parser.add_argument("--hf-config", required=True, help="config name (mini/search/execution/...)")
|
|
parser.add_argument("--hf-split", default="validation")
|
|
parser.add_argument("--limit", type=int, required=True)
|
|
parser.add_argument("--output", required=True, help="Output JSONL path")
|
|
args = parser.parse_args()
|
|
|
|
# `datasets` is installed via meta-agents-research-environments pyproject
|
|
# (transitive dep). Run this script under the ARE venv.
|
|
try:
|
|
from datasets import load_dataset # type: ignore
|
|
except ImportError as e:
|
|
print(
|
|
f"ERROR: `datasets` not installed. Run from ARE venv:\n"
|
|
f" cd external/meta-agents-research-environments && uv run python {sys.argv[0]} ...\n"
|
|
f"Underlying error: {e}",
|
|
file=sys.stderr,
|
|
)
|
|
return 2
|
|
|
|
print(f"Loading {args.hf_dataset} config={args.hf_config} split={args.hf_split} ...", file=sys.stderr)
|
|
ds = load_dataset(args.hf_dataset, args.hf_config, split=args.hf_split)
|
|
total = len(ds)
|
|
print(f"Dataset has {total} examples; limiting to {args.limit}", file=sys.stderr)
|
|
|
|
output_path = Path(args.output)
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
written = 0
|
|
with output_path.open("w", encoding="utf-8") as f:
|
|
for i, record in enumerate(ds):
|
|
if i >= args.limit:
|
|
break
|
|
# HF stores `data` as a serialized JSON STRING (not a nested
|
|
# object). HF dataset card sample showed the post-parse form;
|
|
# actual on-disk format is a string. Parse here so the JSONL
|
|
# written matches benchmarks/gaia2/adapter.ts Gaia2HfTask
|
|
# interface (data: { metadata, apps, events }).
|
|
data_field = record.get("data")
|
|
if isinstance(data_field, str):
|
|
try:
|
|
record["data"] = json.loads(data_field)
|
|
except json.JSONDecodeError as e:
|
|
print(
|
|
f"WARN: record {record.get('id')} has unparseable `data` "
|
|
f"string ({e}); writing raw string",
|
|
file=sys.stderr,
|
|
)
|
|
f.write(json.dumps(record, ensure_ascii=False))
|
|
f.write("\n")
|
|
written += 1
|
|
|
|
print(f"Wrote {written} task records to {output_path}", file=sys.stderr)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|