deepspec/eval_datasets/convert_eval_datasets_to_jsonl.py

"""Convert supported Hugging Face eval datasets to local turns JSONL files.

Usage:
    python eval_datasets/convert_eval_datasets_to_jsonl.py openai/gsm8k \
        --output-path /tmp/gsm8k.jsonl

By default, the converter writes to eval_datasets/ and raises an error when the
target JSONL file already exists.

Supported datasets:

| Dataset | Hugging Face dataset name |
| --- | --- |
| gsm8k | openai/gsm8k |
| math500 | HuggingFaceH4/MATH-500 |
| aime24 | HuggingFaceH4/aime_2024 |
| aime25 | MathArena/aime_2025 |
| alpaca | tatsu-lab/alpaca |
| mt-bench | HuggingFaceH4/mt_bench_prompts |
| humaneval | openai/openai_humaneval |
| mbpp | google-research-datasets/mbpp |
| lbpp | CohereLabs/lbpp |
| swe-bench | princeton-nlp/SWE-bench_Lite |
| livecodebench | livecodebench/code_generation_lite |

PerfectBlend comes from mlabonne/open-perfectblend, but the checked-in
perfectblend.jsonl was built from a regenerated qwen3-4b test cache. The raw HF
download, splitting, and conversion utility lives in
scripts/data/download_and_split.py.
"""

from __future__ import annotations

import argparse
import json
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path


DEFAULT_OUTPUT_ROOT = Path(__file__).resolve().parent
REASONING_SUFFIX = (
    "\nPlease reason step by step, and put your final answer within \\boxed{}."
)


@dataclass(frozen=True)
class DatasetSpec:
    output_name: str
    hf_name: str
    config_name: str | None
    split: str
    format_turns: Callable[[dict], list[str]]
    jsonl_files: tuple[str, ...] = ()
    parquet_files: tuple[str, ...] = ()


def format_gsm8k(row: dict) -> list[str]:
    return [f"{row['question']}{REASONING_SUFFIX}"]


def format_math_problem(row: dict) -> list[str]:
    return [f"{row['problem']}{REASONING_SUFFIX}"]


def format_alpaca(row: dict) -> list[str]:
    instruction = row["instruction"]
    input_text = row["input"]
    if input_text:
        return [f"{instruction}\n\nInput:\n{input_text}"]
    return [instruction]


def format_mt_bench(row: dict) -> list[str]:
    turns = row["prompt"]
    if not isinstance(turns, list):
        raise TypeError("mt-bench field `prompt` must be a list of turns.")
    return turns


def format_humaneval(row: dict) -> list[str]:
    return [
        (
            "Write a solution to the following problem and make sure that it "
            f"passes the tests:\n```python\n{row['prompt']}\n```"
        )
    ]


def format_mbpp(row: dict) -> list[str]:
    return [row["prompt"]]


def format_lbpp(row: dict) -> list[str]:
    return [row["instruction"]]


def format_swe_bench(row: dict) -> list[str]:
    return [
        (
            f"Problem Statement:\n{row['problem_statement']}\nPlease fix the issue "
            "described above."
        )
    ]


def format_livecodebench(row: dict) -> list[str]:
    if "messages" in row:
        turns = [
            message["content"]
            for message in row["messages"]
            if message["role"] == "user"
        ]
        if not turns:
            raise ValueError("livecodebench row has no user messages.")
        return turns

    question = row["question_content"]
    starter_code = row["starter_code"]
    code_template = starter_code if starter_code else "# YOUR CODE HERE"
    format_label = (
        "Use the following code structure:"
        if starter_code
        else "Write your code in the following format:"
    )
    return [
        (
            "You are an expert Python programmer. You will be given a question "
            "(problem specification) and will generate a correct Python program "
            "that matches the specification and passes all tests. You will NOT "
            "return anything except for the program\n\n"
            f"### Question:\n{question}\n\n"
            f"### Format: {format_label}\n"
            f"```python\n{code_template}\n```\n\n"
            "### Answer: (use the provided format with backticks)"
        )
    ]


DATASET_SPECS = (
    DatasetSpec("gsm8k", "openai/gsm8k", "main", "test", format_gsm8k),
    DatasetSpec("math500", "HuggingFaceH4/MATH-500", None, "test", format_math_problem),
    DatasetSpec("aime24", "HuggingFaceH4/aime_2024", None, "train", format_math_problem),
    DatasetSpec("aime25", "MathArena/aime_2025", None, "train", format_math_problem),
    DatasetSpec("alpaca", "tatsu-lab/alpaca", None, "train", format_alpaca),
    DatasetSpec("mt-bench", "HuggingFaceH4/mt_bench_prompts", None, "train", format_mt_bench),
    DatasetSpec("humaneval", "openai/openai_humaneval", None, "test", format_humaneval),
    DatasetSpec("mbpp", "google-research-datasets/mbpp", "sanitized", "test", format_mbpp),
    DatasetSpec(
        "lbpp",
        "CohereLabs/lbpp",
        None,
        "test",
        format_lbpp,
        parquet_files=("python/test.parquet",),
    ),
    DatasetSpec("swe-bench", "princeton-nlp/SWE-bench_Lite", None, "test", format_swe_bench),
    DatasetSpec(
        "livecodebench",
        "livecodebench/code_generation_lite",
        None,
        "test",
        format_livecodebench,
        jsonl_files=(
            "test.jsonl",
            "test2.jsonl",
            "test3.jsonl",
            "test4.jsonl",
            "test5.jsonl",
            "test6.jsonl",
        ),
    ),
)
DATASET_SPECS_BY_HF_NAME = {spec.hf_name: spec for spec in DATASET_SPECS}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Convert a supported Hugging Face eval dataset to turns JSONL."
    )
    parser.add_argument(
        "dataset_name",
        nargs="?",
        help="Hugging Face dataset name, for example `openai/gsm8k`.",
    )
    parser.add_argument(
        "--config-name",
        help="Override the Hugging Face dataset config/subset name.",
    )
    parser.add_argument(
        "--split",
        help="Override the Hugging Face split.",
    )
    parser.add_argument(
        "--output-root",
        type=Path,
        default=DEFAULT_OUTPUT_ROOT,
        help="Directory for the default output file.",
    )
    parser.add_argument(
        "--output-path",
        type=Path,
        help="Exact output JSONL path. Overrides --output-root.",
    )
    parser.add_argument(
        "--max-rows",
        type=int,
        help="Convert at most this many rows; useful for smoke tests.",
    )
    parser.add_argument(
        "--list-supported",
        action="store_true",
        help="Print supported Hugging Face dataset names and exit.",
    )
    return parser.parse_args()


def print_supported_datasets() -> None:
    for spec in DATASET_SPECS:
        config = spec.config_name if spec.config_name is not None else "-"
        print(
            f"{spec.hf_name}\toutput={spec.output_name}.jsonl\t"
            f"config={config}\tsplit={spec.split}"
        )


def require_dataset_spec(dataset_name: str) -> DatasetSpec:
    if dataset_name in DATASET_SPECS_BY_HF_NAME:
        return DATASET_SPECS_BY_HF_NAME[dataset_name]

    supported = ", ".join(spec.hf_name for spec in DATASET_SPECS)
    raise ValueError(
        f"Unsupported Hugging Face dataset name: {dataset_name}. "
        f"Supported names: {supported}"
    )


def load_hf_dataset(spec: DatasetSpec, args: argparse.Namespace):
    from datasets import load_dataset

    config_name = args.config_name if args.config_name is not None else spec.config_name
    split = args.split if args.split is not None else spec.split

    load_kwargs = {"split": split}

    if config_name is None:
        return load_dataset(spec.hf_name, **load_kwargs)
    return load_dataset(spec.hf_name, config_name, **load_kwargs)


def iter_hf_jsonl_rows(spec: DatasetSpec, args: argparse.Namespace):
    from huggingface_hub import hf_hub_download

    if args.config_name is not None:
        raise ValueError(f"{spec.hf_name} does not use a Hugging Face config name.")
    split = args.split if args.split is not None else spec.split
    if split != spec.split:
        raise ValueError(f"{spec.hf_name} only supports split={spec.split!r}.")

    for filename in spec.jsonl_files:
        path = Path(
            hf_hub_download(spec.hf_name, filename=filename, repo_type="dataset")
        )
        with path.open("r", encoding="utf-8") as handle:
            for line in handle:
                line = line.strip()
                if line:
                    yield json.loads(line)


def iter_hf_parquet_rows(spec: DatasetSpec, args: argparse.Namespace):
    import pyarrow.parquet as pq
    from huggingface_hub import hf_hub_download

    if args.config_name is not None:
        raise ValueError(f"{spec.hf_name} does not use a Hugging Face config name.")
    split = args.split if args.split is not None else spec.split
    if split != spec.split:
        raise ValueError(f"{spec.hf_name} only supports split={spec.split!r}.")

    for filename in spec.parquet_files:
        path = hf_hub_download(spec.hf_name, filename=filename, repo_type="dataset")
        yield from pq.read_table(path).to_pylist()


def iter_dataset_rows(spec: DatasetSpec, args: argparse.Namespace):
    if spec.jsonl_files:
        yield from iter_hf_jsonl_rows(spec, args)
        return
    if spec.parquet_files:
        yield from iter_hf_parquet_rows(spec, args)
        return

    yield from load_hf_dataset(spec, args)


def resolve_output_path(spec: DatasetSpec, args: argparse.Namespace) -> Path:
    if args.output_path is not None:
        return args.output_path
    return args.output_root / f"{spec.output_name}.jsonl"


def validate_turns(turns: list[str], row_number: int) -> None:
    if not turns:
        raise ValueError(f"row {row_number} produced no turns.")
    for turn in turns:
        if not isinstance(turn, str) or not turn:
            raise ValueError(f"row {row_number} produced an invalid turn: {turn!r}")


def convert_dataset(spec: DatasetSpec, args: argparse.Namespace) -> tuple[Path, int]:
    output_path = resolve_output_path(spec, args)
    if output_path.exists():
        raise FileExistsError(f"Output JSONL already exists: {output_path}")

    output_path.parent.mkdir(parents=True, exist_ok=True)

    count = 0
    with output_path.open("w", encoding="utf-8") as handle:
        for row_number, row in enumerate(iter_dataset_rows(spec, args), start=1):
            if args.max_rows is not None and count >= args.max_rows:
                break

            turns = spec.format_turns(row)
            validate_turns(turns, row_number)
            handle.write(json.dumps({"turns": turns}, ensure_ascii=False) + "\n")
            count += 1

    return output_path, count


def main() -> None:
    args = parse_args()
    if args.list_supported:
        print_supported_datasets()
        return
    if args.dataset_name is None:
        raise ValueError("dataset_name is required unless --list-supported is used.")

    spec = require_dataset_spec(args.dataset_name)
    try:
        output_path, count = convert_dataset(spec, args)
    except FileExistsError as exc:
        raise SystemExit(str(exc)) from None
    print(f"wrote {spec.hf_name}: {count} rows -> {output_path}")


if __name__ == "__main__":
    main()