#!/usr/bin/env python3 """ Download or export benchmark datasets into benchmarks/data. Uses the HuggingFace `datasets` library for reliable full downloads. Falls back to the REST API or builtins if `datasets` is not installed. """ from __future__ import annotations import argparse import gzip import json import re import urllib.parse import urllib.request from dataclasses import asdict, dataclass from pathlib import Path from typing import Any, Callable from benchmarks.suites.aider import _BUILTIN_PROBLEMS as _AIDER_BUILTINS from benchmarks.suites.aime import _BUILTIN_PROBLEMS as _AIME_BUILTINS from benchmarks.suites.bfcl import _BUILTIN_PROBLEMS as _BFCL_BUILTINS from benchmarks.suites.bigbench import _BUILTIN_PROBLEMS as _BIGBENCH_BUILTINS from benchmarks.suites.codeforces import _BUILTIN_PROBLEMS as _CODEFORCES_BUILTINS from benchmarks.suites.gpqa import _BUILTIN_PROBLEMS as _GPQA_BUILTINS from benchmarks.suites.gsm8k import _BUILTIN_PROBLEMS as _GSM8K_BUILTINS from benchmarks.suites.hle import _BUILTIN_PROBLEMS as _HLE_BUILTINS from benchmarks.suites.humaneval import _BUILTIN_PROBLEMS as _HUMANEVAL_BUILTINS from benchmarks.suites.ifeval import _BUILTIN_PROBLEMS as _IFEVAL_BUILTINS from benchmarks.suites.livecodebench import _BUILTIN_PROBLEMS as _LIVECODEBENCH_BUILTINS from benchmarks.suites.math_bench import _BUILTIN_PROBLEMS as _MATH_BUILTINS from benchmarks.suites.mbpp import _BUILTIN_PROBLEMS as _MBPP_BUILTINS from benchmarks.suites.mmmlu import _BUILTIN_PROBLEMS as _MMMLU_BUILTINS from benchmarks.suites.mmlu_pro import _BUILTIN_PROBLEMS as _MMLU_PRO_BUILTINS from benchmarks.suites.swe_bench import _BUILTIN_PROBLEMS as _SWE_BUILTINS from benchmarks.suites.tau2 import _BUILTIN_PROBLEMS as _TAU2_BUILTINS HUMANEVAL_GZ_URL = "https://raw.githubusercontent.com/openai/human-eval/master/data/HumanEval.jsonl.gz" DEFAULT_DATA_DIR = Path(__file__).resolve().parent / "data" # Legacy REST API support (fallback only) HF_DATASET_VIEWER_BASE = "https://datasets-server.huggingface.co" JsonFetcher = Callable[[str, dict[str, object], dict[str, str] | None, float], object] @dataclass class DownloadResult: suite: str rows: int path: str source: str note: str = "" def fetch_bytes(url: str, timeout: float, headers: dict[str, str] | None = None) -> bytes: request = urllib.request.Request(url, headers=headers or {}) with urllib.request.urlopen(request, timeout=timeout) as response: return response.read() def fetch_json( endpoint: str, params: dict[str, object], headers: dict[str, str] | None, timeout: float, ) -> object: query = urllib.parse.urlencode(params, doseq=True) url = f"{HF_DATASET_VIEWER_BASE}/{endpoint}" if query: url = f"{url}?{query}" raw = fetch_bytes(url, timeout=timeout, headers=headers) return json.loads(raw.decode("utf-8")) def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> int: path.parent.mkdir(parents=True, exist_ok=True) with open(path, "w", encoding="utf-8") as handle: for row in rows: handle.write(json.dumps(row, ensure_ascii=False) + "\n") return len(rows) # --------------------------------------------------------------------------- # HuggingFace `datasets` library helpers # --------------------------------------------------------------------------- def _load_hf_dataset( dataset_name: str, config: str | None = None, split: str = "test", ) -> list[dict[str, Any]]: """Load a dataset using the HuggingFace `datasets` library.""" from datasets import load_dataset # type: ignore[import-untyped] kwargs: dict[str, Any] = {} if config: kwargs["name"] = config ds = load_dataset(dataset_name, split=split, **kwargs) return [dict(row) for row in ds] # type: ignore[union-attr] def _try_load_hf( dataset_name: str, config: str | None = None, split_preference: tuple[str, ...] = ("test", "validation", "train"), ) -> list[dict[str, Any]]: """Try loading with preferred splits, falling back through the list.""" for split in split_preference: try: rows = _load_hf_dataset(dataset_name, config=config, split=split) if rows: print(f" Loaded {len(rows)} rows from {dataset_name} [{split}]") return rows except (ValueError, KeyError): continue raise ValueError(f"No valid split found for {dataset_name}") def _fetch_hf_rows( dataset: str, *, config_preference: tuple[str, ...] = (), split_preference: tuple[str, ...] = ("test", "validation", "train"), json_fetcher: JsonFetcher = fetch_json, timeout: float = 60.0, headers: dict[str, str] | None = None, ) -> list[dict[str, Any]]: """Legacy REST-API fetcher (kept for backward compatibility with tests).""" splits_payload = json_fetcher("splits", {"dataset": dataset}, headers, timeout) splits = list((splits_payload or {}).get("splits", [])) # type: ignore[assignment] if not splits: return [] chosen: dict[str, Any] | None = None for config_name in config_preference: for split_name in split_preference: chosen = next( ( item for item in splits if item.get("config") == config_name and item.get("split") == split_name ), None, ) if chosen is not None: break if chosen is not None: break if chosen is None: for split_name in split_preference: chosen = next((item for item in splits if item.get("split") == split_name), None) if chosen is not None: break if chosen is None: chosen = splits[0] rows: list[dict[str, Any]] = [] offset = 0 length = 100 while True: payload = json_fetcher( "rows", { "dataset": chosen["dataset"], "config": chosen["config"], "split": chosen["split"], "offset": offset, "length": length, }, headers, timeout, ) batch = [item.get("row", {}) for item in (payload or {}).get("rows", [])] # type: ignore[union-attr] rows.extend(batch) total = int((payload or {}).get("num_rows_total", len(rows))) # type: ignore[union-attr] offset += len(batch) if not batch or offset >= total: break return rows # --------------------------------------------------------------------------- # Answer extraction helpers # --------------------------------------------------------------------------- def _extract_gsm8k_answer(text: str) -> str: if "####" in text: text = text.split("####", 1)[1] numbers = re.findall(r"-?\d[\d,]*\.?\d*", text.replace("$", "")) if numbers: return numbers[-1].replace(",", "") return text.strip().replace(",", "") def _extract_math_answer(solution: str) -> str: boxed_fraction = re.search(r"\\boxed\{\\frac\{([^}]+)\}\{([^}]+)\}\}", solution, flags=re.DOTALL) if boxed_fraction: return f"{boxed_fraction.group(1).strip()}/{boxed_fraction.group(2).strip()}" boxed = re.search(r"\\boxed\{([^{}]+)\}", solution, flags=re.DOTALL) value = boxed.group(1) if boxed else solution value = value.strip() value = value.replace("\\frac{", "").replace("}{", "/").replace("}", "") value = value.replace("$", "").replace(",", "").strip() fraction = re.search(r"-?\d+\s*/\s*-?\d+", value) if fraction: return fraction.group(0).replace(" ", "") numbers = re.findall(r"-?\d+(?:/\d+)?(?:\.\d+)?", value) return numbers[-1] if numbers else value # --------------------------------------------------------------------------- # Individual dataset downloaders (using `datasets` library) # --------------------------------------------------------------------------- def _download_humaneval(output_path: Path, *, timeout: float) -> DownloadResult: raw = fetch_bytes(HUMANEVAL_GZ_URL, timeout=timeout) if raw[:2] == b"\x1f\x8b": raw = gzip.decompress(raw) lines = [json.loads(line) for line in raw.decode("utf-8").splitlines() if line.strip()] rows = [ { "task_id": item["task_id"], "prompt": item["prompt"], "canonical_solution": item.get("canonical_solution", ""), "test": item["test"], "entry_point": item["entry_point"], } for item in lines ] count = _write_jsonl(output_path, rows) return DownloadResult("humaneval", count, str(output_path), "official") def _download_gsm8k( output_path: Path, *, timeout: float, json_fetcher: JsonFetcher | None = None, ) -> DownloadResult: if json_fetcher is not None: # Legacy path for tests rows = _fetch_hf_rows( "openai/gsm8k", config_preference=("main",), split_preference=("test",), json_fetcher=json_fetcher, timeout=timeout, ) else: rows = _try_load_hf("openai/gsm8k", config="main", split_preference=("test",)) normalized = [ { "id": f"gsm8k-{index + 1:04d}", "question": row["question"], "answer": _extract_gsm8k_answer(str(row["answer"])), } for index, row in enumerate(rows) ] count = _write_jsonl(output_path, normalized) return DownloadResult("gsm8k", count, str(output_path), "official") def _download_mbpp(output_path: Path, *, timeout: float) -> DownloadResult: try: rows = _try_load_hf("google-research-datasets/mbpp", config="sanitized", split_preference=("test", "validation")) except Exception: rows = _try_load_hf("google-research-datasets/mbpp", config="full", split_preference=("test", "validation")) normalized = [ { "task_id": row.get("task_id", index + 1), "text": row.get("text") or row.get("prompt") or "", "code": row.get("code", ""), "test_list": row.get("test_list") or row.get("test_setup_code", []), } for index, row in enumerate(rows) ] count = _write_jsonl(output_path, normalized) return DownloadResult("mbpp", count, str(output_path), "official") def _download_math(output_path: Path, *, timeout: float) -> DownloadResult: rows = _try_load_hf("DigitalLearningGmbH/MATH-lighteval", split_preference=("test", "train")) normalized = [ { "id": row.get("problem_id", f"math-{index + 1:04d}"), "problem": row.get("problem", ""), "answer": _extract_math_answer(str(row.get("solution", row.get("answer", "")))), "subject": row.get("type", row.get("subject", "unknown")), "level": row.get("level", 0), } for index, row in enumerate(rows) ] count = _write_jsonl(output_path, normalized) return DownloadResult("math", count, str(output_path), "official") def _download_mmlu_pro(output_path: Path, *, timeout: float) -> DownloadResult: rows = _try_load_hf("TIGER-Lab/MMLU-Pro", split_preference=("test", "validation")) letters = "ABCDEFGHIJ" normalized = [] for index, row in enumerate(rows): answer_raw = row.get("answer", "") if isinstance(answer_raw, int) and answer_raw < len(letters): answer = letters[answer_raw] else: answer = str(answer_raw) normalized.append({ "id": f"mmlu-pro-{index + 1:04d}", "subject": row.get("category", row.get("subject", "unknown")), "question": row.get("question", ""), "choices": row.get("options", row.get("choices", [])), "answer": answer, }) count = _write_jsonl(output_path, normalized) return DownloadResult("mmlu-pro", count, str(output_path), "official") def _download_gpqa(output_path: Path, *, timeout: float) -> DownloadResult: rows = _try_load_hf("Idavidrein/gpqa", config="gpqa_diamond", split_preference=("train",)) normalized = [] for index, row in enumerate(rows): choices = [ row.get("Correct Answer", ""), row.get("Incorrect Answer 1", ""), row.get("Incorrect Answer 2", ""), row.get("Incorrect Answer 3", ""), ] normalized.append({ "id": f"gpqa-{index + 1:04d}", "subject": row.get("Subdomain", row.get("domain", "science")), "question": row.get("Question", ""), "choices": choices, "answer": "A", }) count = _write_jsonl(output_path, normalized) return DownloadResult("gpqa-diamond", count, str(output_path), "official") def _download_bigbench_hard(output_path: Path, *, timeout: float) -> DownloadResult: from datasets import load_dataset # type: ignore[import-untyped] configs = [ "boolean_expressions", "causal_judgement", "date_understanding", "disambiguation_qa", "dyck_languages", "formal_fallacies", "geometric_shapes", "hyperbaton", "logical_deduction_three_objects", "logical_deduction_five_objects", "logical_deduction_seven_objects", "movie_recommendation", "multistep_arithmetic_two", "navigate", "object_counting", "penguins_in_a_table", "reasoning_about_colored_objects", "ruin_names", "salient_translation_error_detection", "snarks", "sports_understanding", "temporal_sequences", "tracking_shuffled_objects_three_objects", "tracking_shuffled_objects_five_objects", "tracking_shuffled_objects_seven_objects", "web_of_lies", "word_sorting", ] all_rows: list[dict[str, Any]] = [] for config in configs: try: ds = load_dataset("lukaemon/bbh", config, split="test") for row in ds: row_dict = dict(row) # type: ignore[arg-type] row_dict["task"] = config all_rows.append(row_dict) except Exception: continue print(f" Loaded {len(all_rows)} rows from lukaemon/bbh [{len(configs)} tasks]") normalized = [] for index, row in enumerate(all_rows): target = row.get("target", row.get("answer", "")) normalized.append({ "id": f"bbh-{index + 1:05d}", "task": row.get("task", "unknown"), "question": row.get("input", row.get("question", "")), "choices": [], "answer": str(target).strip(), }) count = _write_jsonl(output_path, normalized) return DownloadResult("bigbench-hard", count, str(output_path), "official") def _download_mmmlu(output_path: Path, *, timeout: float) -> DownloadResult: rows = _try_load_hf("openai/MMMLU", split_preference=("test", "validation")) letters = "ABCD" normalized = [] for index, row in enumerate(rows): answer_raw = row.get("answer", "") if isinstance(answer_raw, int) and answer_raw < len(letters): answer = letters[answer_raw] else: answer = str(answer_raw) normalized.append({ "id": f"mmmlu-{index + 1:04d}", "language": row.get("language", "en"), "subject": row.get("subject", "unknown"), "question": row.get("question", ""), "choices": row.get("choices", row.get("options", [])), "answer": answer, }) count = _write_jsonl(output_path, normalized) return DownloadResult("mmmlu", count, str(output_path), "official") def _download_hle(output_path: Path, *, timeout: float) -> DownloadResult: rows = _try_load_hf("cais/hle", split_preference=("test", "validation", "train")) normalized = [] for index, row in enumerate(rows): entry: dict[str, Any] = { "id": f"hle-{index + 1:04d}", "subject": row.get("category", row.get("subject", "general")), "question": row.get("question", ""), "answer": str(row.get("answer", "")), } if row.get("choices") or row.get("options"): entry["answer_type"] = "multiple_choice" entry["choices"] = row.get("choices", row.get("options", [])) else: entry["answer_type"] = "exact" normalized.append(entry) count = _write_jsonl(output_path, normalized) return DownloadResult("hle", count, str(output_path), "official") def _export_builtin(output_path: Path, suite: str, rows: list[dict[str, Any]], *, source: str = "builtin", note: str = "") -> DownloadResult: count = _write_jsonl(output_path, rows) return DownloadResult(suite, count, str(output_path), source, note) def _builtin_rows(suite: str) -> list[dict[str, Any]]: mapping: dict[str, list[dict[str, Any]]] = { "humaneval": list(_HUMANEVAL_BUILTINS), "mbpp": list(_MBPP_BUILTINS), "gsm8k": list(_GSM8K_BUILTINS), "math": list(_MATH_BUILTINS), "swe-bench": list(_SWE_BUILTINS), "aider": list(_AIDER_BUILTINS), "livecodebench": list(_LIVECODEBENCH_BUILTINS), "aime": list(_AIME_BUILTINS), "ifeval": list(_IFEVAL_BUILTINS), "bfcl": list(_BFCL_BUILTINS), "mmlu-pro": list(_MMLU_PRO_BUILTINS), "gpqa-diamond": list(_GPQA_BUILTINS), "bigbench-hard": list(_BIGBENCH_BUILTINS), "mmmlu": list(_MMMLU_BUILTINS), "hle": list(_HLE_BUILTINS), "tau2": list(_TAU2_BUILTINS), "codeforces": list(_CODEFORCES_BUILTINS), } return list(mapping[suite]) def prepare_suite( suite: str, *, data_dir: Path, force: bool, builtin_only: bool, official_only: bool, timeout: float, ) -> DownloadResult: output_map = { "mmlu-pro": "mmlu_pro.jsonl", "gpqa-diamond": "gpqa.jsonl", "bigbench-hard": "bigbench_hard.jsonl", "mmmlu": "mmmlu.jsonl", "hle": "hle.jsonl", "tau2": "tau2.jsonl", "codeforces": "codeforces.jsonl", } output_path = data_dir / output_map.get(suite, f"{suite}.jsonl") if output_path.exists() and not force: lines = [line for line in output_path.read_text(encoding="utf-8").splitlines() if line.strip()] return DownloadResult(suite, len(lines), str(output_path), "existing") data_dir.mkdir(parents=True, exist_ok=True) builtin_only_suites = {"swe-bench", "aider", "livecodebench", "aime", "ifeval", "bfcl", "tau2", "codeforces"} official_downloaders = { "humaneval": _download_humaneval, "gsm8k": _download_gsm8k, "mbpp": _download_mbpp, "math": _download_math, "mmlu-pro": _download_mmlu_pro, "gpqa-diamond": _download_gpqa, "bigbench-hard": _download_bigbench_hard, "mmmlu": _download_mmmlu, "hle": _download_hle, } if suite in builtin_only_suites or builtin_only: return _export_builtin(output_path, suite, _builtin_rows(suite)) downloader = official_downloaders.get(suite) if downloader is None: return _export_builtin(output_path, suite, _builtin_rows(suite)) try: return downloader(output_path, timeout=timeout) except Exception as exc: if official_only: raise note = f"official download failed: {exc}" print(f" WARNING: {note}") print(f" Falling back to {len(_builtin_rows(suite))} built-in problems.") print(f" To get full data, install `datasets`: pip install datasets") return _export_builtin( output_path, suite, _builtin_rows(suite), source="builtin-fallback", note=note, ) def _write_manifest(data_dir: Path, results: list[DownloadResult]) -> Path: manifest_path = data_dir / "manifest.json" payload = { "generated_at": __import__("time").strftime("%Y-%m-%dT%H:%M:%S"), "results": [asdict(item) for item in results], } manifest_path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") return manifest_path def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description="Download benchmark datasets for claw-code-agent.") parser.add_argument("--suite", action="append", default=[], help="Suite to prepare. Can be repeated.") parser.add_argument("--all", action="store_true", help="Prepare all known suites.") parser.add_argument("--list", action="store_true", help="List known suites.") parser.add_argument("--data-dir", type=Path, default=DEFAULT_DATA_DIR, help="Output data directory.") parser.add_argument("--force", action="store_true", help="Overwrite existing files.") parser.add_argument("--builtin-only", action="store_true", help="Skip official downloads and export builtins only.") parser.add_argument("--official-only", action="store_true", help="Do not fall back to builtins.") parser.add_argument("--timeout", type=float, default=60.0, help="Network timeout in seconds.") return parser def main() -> None: parser = build_parser() args = parser.parse_args() known = [ "humaneval", "mbpp", "gsm8k", "math", "swe-bench", "aider", "livecodebench", "aime", "ifeval", "bfcl", "mmlu-pro", "gpqa-diamond", "bigbench-hard", "mmmlu", "hle", "tau2", "codeforces", ] if args.list: for name in known: print(name) return suites = list(args.suite) if args.all: suites = known if not suites: parser.error("specify --suite or --all") results = [ prepare_suite( suite, data_dir=args.data_dir, force=args.force, builtin_only=args.builtin_only, official_only=args.official_only, timeout=args.timeout, ) for suite in suites ] manifest = _write_manifest(args.data_dir, results) print(f"Wrote {len(results)} suite files to {args.data_dir}") print(f"Manifest: {manifest}") if __name__ == "__main__": main()