run.py

import argparse
import json
import os
import subprocess
import sys
from pathlib import Path

import yaml
from tabulate import tabulate

RESULTS_DIR = Path("evalplus_results")
VALID_DATASETS = {"humaneval", "mbpp"}


def load_config(path: Path) -> dict:
    with open(path) as f:
        config = yaml.safe_load(f)

    if not config.get("models"):
        raise ValueError("No models specified in config")

    datasets = config.get("datasets", ["humaneval"])
    if isinstance(datasets, str):
        datasets = [datasets]
    for ds in datasets:
        if ds not in VALID_DATASETS:
            raise ValueError(f"Invalid dataset '{ds}', must be one of {VALID_DATASETS}")

    return config


def run_evalplus(
    model: str,
    dataset: str,
    base_url: str,
    max_tokens: int | None = None,
    timeout: int | None = None,
) -> int:
    env = os.environ.copy()
    env["OPENAI_API_KEY"] = env.get("API_KEY", "none")
    if timeout:
        env["OPENAI_TIMEOUT"] = str(timeout)

    cmd = [
        sys.executable,
        "-m",
        "evalplus.evaluate",
        "--model",
        model,
        "--dataset",
        dataset,
        "--backend",
        "openai",
        "--base-url",
        base_url,
        "--greedy",
        "--root",
        str(RESULTS_DIR),
    ]

    if max_tokens:
        cmd.extend(["--max-tokens", str(max_tokens)])

    print(f"  $ {' '.join(cmd)}")
    result = subprocess.run(cmd, env=env)
    return result.returncode


def find_eval_results(dataset: str) -> list[Path]:
    dataset_dir = RESULTS_DIR / dataset
    if not dataset_dir.exists():
        return []
    return sorted(dataset_dir.glob("*.eval_results.json"))


def parse_model_name(filename: str) -> str:
    name = filename.replace(".eval_results", "")
    name = name.replace("_openai_temp_0.0", "")
    name = name.replace("--", "/")
    return name


def report(dataset: str):
    results_files = find_eval_results(dataset)
    if not results_files:
        print("\nNo results found. Run a benchmark first.")
        return

    rows = []
    for path in results_files:
        with open(path) as f:
            data = json.load(f)

        pass_at_k = data.get("pass_at_k", {})
        if not pass_at_k:
            continue

        base_score = pass_at_k.get("base", {}).get("pass@1", 0)
        plus_score = pass_at_k.get("plus", {}).get("pass@1", 0)
        total = len(data.get("eval", {}))
        model_name = parse_model_name(path.stem)

        rows.append(
            [
                model_name,
                total,
                f"{base_score * 100:.1f}%",
                f"{plus_score * 100:.1f}%",
            ]
        )

    rows.sort(key=lambda r: r[3], reverse=True)

    headers = ["Model", "Tasks", f"{dataset} pass@1", f"{dataset}+ pass@1"]
    print(f"\n{'=' * 60}")
    print(f"RESULTS — {dataset}")
    print(f"{'=' * 60}")
    print(tabulate(rows, headers=headers, tablefmt="grid"))


def main():
    parser = argparse.ArgumentParser(description="Benchmark LLMs with EvalPlus")
    parser.add_argument("--config", type=Path, default=Path("config.yaml"))
    parser.add_argument("--models", nargs="+", help="Override models from config")
    parser.add_argument("--dataset", nargs="+", help="Override datasets from config")
    parser.add_argument("--report-only", action="store_true")
    args = parser.parse_args()

    config = load_config(args.config)
    datasets = args.dataset or config.get("datasets", ["humaneval"])
    if isinstance(datasets, str):
        datasets = [datasets]
    base_url = config["server"]["base_url"]
    bench_cfg = config.get("benchmark", {})
    max_tokens = bench_cfg.get("max_tokens")
    timeout = bench_cfg.get("timeout")
    models = args.models or config["models"]

    if not args.report_only:
        for dataset in datasets:
            for model in models:
                print(f"\n{'=' * 60}")
                print(f"Benchmarking: {model} on {dataset}")
                print(f"{'=' * 60}")
                rc = run_evalplus(model, dataset, base_url, max_tokens, timeout)
                if rc != 0:
                    print(f"  WARNING: evalplus exited with code {rc}")

    for dataset in datasets:
        report(dataset)


if __name__ == "__main__":
    main()