run_bcb_experiment.py

"""
Run BigCodeBench samples from a JSONL file against the live evaluator
in batches, then report which tasks pass, fail, or timeout.

Usage:
    python run_bcb_experiment.py \
        --filename /tmp/bigcodebench-instruct-codegen_results-sanitized_calibrated.jsonl \
        --batch-size 50

The results-sanitized_calibrated.jsonl file already has the fields the
/evaluate/ endpoint expects (task_id, res_id, test, solution, entry_point,
code_prompt).  The requests.jsonl can also be used — this script detects its
nested structure and extracts the needed fields from doc + doc.choices.
"""

import argparse
import json
import sys
import time

import httpx

DEFAULT_URL = "http://localhost:7860/evaluate/"
TIMEOUT = 600  # per-batch timeout in seconds


def load_jsonl(path):
    with open(path) as f:
        return [json.loads(line) for line in f if line.strip()]


def normalize_sample(raw):
    """
    Accept either format:
      - results-sanitized_calibrated.jsonl  (flat: task_id, res_id, test, solution, ...)
      - requests.jsonl  (nested: doc.task_id, doc.test, doc.choices[0], ...)
    Return a dict suitable for the /evaluate/ endpoint.
    """
    if "solution" in raw and "test" in raw:
        # Already in the right shape
        return {
            "task_id": raw["task_id"],
            "res_id": raw.get("res_id", 0),
            "test": raw["test"],
            "solution": raw["solution"],
            "entry_point": raw["entry_point"],
            "code_prompt": raw.get("code_prompt", ""),
        }

    # requests.jsonl format — extract from doc + choices
    doc = raw.get("doc", {})
    choices = doc.get("choices", [])
    solution = choices[0] if choices else doc.get("canonical_solution", "")
    # Prepend code_prompt to the choice body to form a complete solution
    code_prompt = doc.get("code_prompt", "")
    full_solution = code_prompt + solution

    return {
        "task_id": doc["task_id"],
        "res_id": raw.get("idx", 0),
        "test": doc["test"],
        "solution": full_solution,
        "entry_point": doc["entry_point"],
        "code_prompt": code_prompt,
    }


def send_batch(client, url, samples, batch_idx, calibrate, retries=3, backoff=30):
    for attempt in range(1, retries + 1):
        label = f"  Batch {batch_idx}" + (f" (attempt {attempt})" if attempt > 1 else "")
        print(f"{label}: sending {len(samples)} samples ...", end=" ", flush=True)
        try:
            t0 = time.time()
            resp = client.post(
                url,
                json=samples,
                params={"calibrate": calibrate, "min_time_limit": 1},
            )
            elapsed = time.time() - t0
        except (httpx.ReadTimeout, httpx.ConnectTimeout, httpx.ConnectError,
                httpx.ReadError, httpx.RemoteProtocolError) as exc:
            print(f"connection error: {exc}")
            if attempt < retries:
                wait = backoff * attempt
                print(f"    Waiting {wait}s before retry ...")
                time.sleep(wait)
            continue

        if resp.status_code == 200:
            print(f"OK ({elapsed:.1f}s)")
            return resp.json()

        print(f"HTTP {resp.status_code} ({elapsed:.1f}s)")
        if attempt < retries:
            wait = backoff * attempt
            print(f"    Waiting {wait}s before retry ...")
            time.sleep(wait)

    print(f"  Batch {batch_idx}: all {retries} attempts failed")
    return None


def main():
    parser = argparse.ArgumentParser(description="Run BCB experiments against live evaluator")
    parser.add_argument("--filename", required=True, help="Path to JSONL file")
    parser.add_argument("--batch-size", type=int, default=50, help="Samples per request (default 50)")
    parser.add_argument("--limit", type=int, default=0, help="Max samples to evaluate (0 = all)")
    parser.add_argument("--calibrate", action="store_true", default=False,
                        help="Enable calibrate mode (prepend code_prompt+'pass')")
    parser.add_argument("--url", type=str, default=DEFAULT_URL, help=f"Evaluate endpoint URL (default {DEFAULT_URL})")
    parser.add_argument("--output", type=str, default=None, help="Save detailed results to JSONL")
    parser.add_argument("--retries", type=int, default=3, help="Retries per batch on failure (default 3)")
    parser.add_argument("--backoff", type=int, default=30, help="Base backoff seconds between retries (default 30)")
    parser.add_argument("--delay", type=int, default=2, help="Delay seconds between batches (default 2)")
    args = parser.parse_args()

    raw_data = load_jsonl(args.filename)
    samples = [normalize_sample(r) for r in raw_data]
    if args.limit > 0:
        samples = samples[: args.limit]

    evaluate_url = args.url

    print(f"Loaded {len(samples)} samples from {args.filename}")
    print(f"Endpoint: {evaluate_url}")
    print(f"Batch size: {args.batch_size}, calibrate: {args.calibrate}\n")

    all_results = []  # flat list of per-sample result dicts
    batches = [samples[i : i + args.batch_size] for i in range(0, len(samples), args.batch_size)]

    skipped_samples = []

    with httpx.Client(timeout=TIMEOUT) as client:
        for idx, batch in enumerate(batches, 1):
            result = send_batch(client, evaluate_url, batch, idx, args.calibrate,
                                retries=args.retries, backoff=args.backoff)
            if result is None:
                skipped_samples.extend(batch)
                continue
            for task_results in result["eval"].values():
                all_results.extend(task_results)
            if idx < len(batches) and args.delay > 0:
                time.sleep(args.delay)

    # --- Summary ---
    passed = [r for r in all_results if r["status"] == "pass"]
    failed = [r for r in all_results if r["status"] == "fail"]
    timed_out = [r for r in all_results if r["status"] == "timeout"]

    print("\n" + "=" * 65)
    print(f"{'RESULTS SUMMARY':^65}")
    print("=" * 65)
    print(f"  Total evaluated : {len(all_results)}")
    print(f"  Passed          : {len(passed)}")
    print(f"  Failed          : {len(failed)}")
    print(f"  Timed out       : {len(timed_out)}")
    if all_results:
        print(f"  pass@1          : {len(passed) / len(all_results):.4f}")
    if skipped_samples:
        print(f"  Skipped (server errors): {len(skipped_samples)}")
    print("=" * 65)

    # Show failed task IDs
    if failed:
        print(f"\nFailed tasks ({len(failed)}):")
        for r in sorted(failed, key=lambda x: x["task_id"]):
            detail_keys = list(r.get("details", {}).keys())
            short = ", ".join(detail_keys[:3])
            if len(detail_keys) > 3:
                short += f" ... (+{len(detail_keys) - 3} more)"
            print(f"  {r['task_id']:<30} failing: {short}")

    if timed_out:
        print(f"\nTimed-out tasks ({len(timed_out)}):")
        for r in sorted(timed_out, key=lambda x: x["task_id"]):
            print(f"  {r['task_id']}")

    # Save detailed output
    if args.output:
        with open(args.output, "w") as f:
            for r in all_results:
                f.write(json.dumps(r) + "\n")
        print(f"\nDetailed results saved to {args.output}")


if __name__ == "__main__":
    main()