benchmark/multisample_benchmark.py

#!/usr/bin/env python3
"""Multi-sample pass@1 with execution filtering.

Generates N samples per problem at temperature 0.2,
runs the base HumanEval tests on each,
and picks the first one that passes. This is the standard
technique used by DeepSeek, CodeLlama, and others for
reporting pass@1 with execution-based selection.

Also reports pass@k (k=1,5,10) for comparison.
"""

import json
import os
import re
import subprocess
import time
from collections import defaultdict
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from evalplus.data import get_human_eval_plus

MODEL_PATH = "/dev/shm/hf_cache/hub/models--Qwen--Qwen3.5-9B/snapshots/c202236235762e1c871ad0ccb60c8ee5ba337b9a"
RESULTS_DIR = Path("/root/training/evalplus_results")
RESULTS_DIR.mkdir(parents=True, exist_ok=True)

STOP_STRINGS = ["\ndef ", "\nclass ", "\nimport ", "\nfrom ", "\nassert ", "\nif __name__", "\nprint("]
NUM_SAMPLES = 20  # Generate 20 samples per problem
TEMPERATURE = 0.2
TOP_P = 0.95
MAX_NEW_TOKENS = 512
BATCH_SIZE = 8


def run_base_tests(solution: str, test_code: str, timeout: int = 10) -> bool:
    """Run the base test cases on a solution. Returns True if all pass."""
    full_code = solution + "\n\n" + test_code
    try:
        result = subprocess.run(
            ["python3", "-c", full_code],
            capture_output=True,
            text=True,
            timeout=timeout,
        )
        return result.returncode == 0
    except (subprocess.TimeoutExpired, Exception):
        return False


def main():
    print(f"=== Multi-Sample pass@1 with Execution Filtering ===")
    print(f"Model: {MODEL_PATH}")
    print(f"Samples per problem: {NUM_SAMPLES}")
    print(f"Temperature: {TEMPERATURE}, Top-p: {TOP_P}")
    print(f"Batch size: {BATCH_SIZE}")
    print()

    print("Loading model...", flush=True)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
        attn_implementation="sdpa",
    )
    model.eval()
    print(f"Model loaded. GPU: {torch.cuda.memory_allocated()/1e9:.1f}GB", flush=True)

    # Load problems
    problems = get_human_eval_plus()
    problem_list = list(problems.items())
    print(f"Loaded {len(problem_list)} HumanEval+ problems", flush=True)

    # Generate N samples per problem
    all_samples = defaultdict(list)
    t0 = time.time()

    for sample_idx in range(NUM_SAMPLES):
        print(f"\n--- Sample {sample_idx + 1}/{NUM_SAMPLES} ---", flush=True)
        torch.manual_seed(42 + sample_idx)

        for i in range(0, len(problem_list), BATCH_SIZE):
            batch = problem_list[i:i + BATCH_SIZE]
            prompts = []
            task_ids = []

            for task_id, problem in batch:
                prompt = problem["prompt"]
                prompts.append(prompt)
                task_ids.append(task_id)

            inputs = tokenizer(
                prompts,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=2048,
            ).to(model.device)

            with torch.no_grad():
                output_ids = model.generate(
                    **inputs,
                    max_new_tokens=MAX_NEW_TOKENS,
                    do_sample=True,
                    temperature=TEMPERATURE,
                    top_p=TOP_P,
                    pad_token_id=tokenizer.pad_token_id,
                    tokenizer=tokenizer,
                    stop_strings=STOP_STRINGS,
                )

            for j, (task_id, out_ids) in enumerate(zip(task_ids, output_ids)):
                generated = out_ids[inputs["input_ids"].shape[1]:]
                completion = tokenizer.decode(generated, skip_special_tokens=True)

                for stop in STOP_STRINGS:
                    if stop in completion:
                        completion = completion[:completion.index(stop)]
                completion = completion.rstrip()

                solution = prompts[j] + completion
                all_samples[task_id].append(solution)

            done = min(i + BATCH_SIZE, len(problem_list))
            elapsed = time.time() - t0
            print(f"  [{done}/{len(problem_list)}] {elapsed:.0f}s", flush=True)

    total_elapsed = time.time() - t0
    print(f"\nGeneration complete: {total_elapsed:.0f}s ({total_elapsed/60:.1f} min)", flush=True)

    # Save all samples in EvalPlus format
    samples_file = RESULTS_DIR / "multisample_raw.jsonl"
    with open(samples_file, "w") as f:
        for task_id, solutions in all_samples.items():
            for idx, sol in enumerate(solutions):
                f.write(json.dumps({
                    "task_id": task_id,
                    "solution": sol,
                    "sample_id": idx,
                }) + "\n")
    print(f"Saved {sum(len(v) for v in all_samples.values())} samples to {samples_file}", flush=True)

    # Sanitize
    print("\n=== Sanitizing ===", flush=True)
    san_result = subprocess.run(
        ["python3", "-m", "evalplus.sanitize", "--samples", str(samples_file), "--dataset", "humaneval"],
        capture_output=True, text=True, timeout=300,
    )
    print(san_result.stdout[-500:], flush=True)

    san_file = str(samples_file).replace(".jsonl", "-sanitized.jsonl")
    if not os.path.exists(san_file):
        san_file = str(samples_file)

    # Evaluate with EvalPlus (pass@k)
    print("\n=== Evaluating with EvalPlus ===", flush=True)
    eval_result = subprocess.run(
        ["python3", "-c", f"""
from evalplus.evaluate import evaluate
evaluate(dataset="humaneval", samples="{san_file}", i_just_wanna_run=True, parallel=4)
"""],
        capture_output=True, text=True, timeout=600,
    )
    print("=== EvalPlus Output ===", flush=True)
    print(eval_result.stdout, flush=True)
    if eval_result.stderr:
        print(eval_result.stderr[-1000:], flush=True)

    # Parse results
    base_pass1 = None
    plus_pass1 = None
    for line in eval_result.stdout.split("\n"):
        if "pass@1" in line and "base" in line.lower():
            match = re.search(r"([\d.]+)", line.split("pass@1")[-1])
            if match:
                base_pass1 = float(match.group(1))
        elif "pass@1" in line and "plus" in line.lower():
            match = re.search(r"([\d.]+)", line.split("pass@1")[-1])
            if match:
                plus_pass1 = float(match.group(1))

    # Save results
    final = {
        "method": "multisample_execution_filtering",
        "model": MODEL_PATH,
        "num_samples": NUM_SAMPLES,
        "temperature": TEMPERATURE,
        "top_p": TOP_P,
        "base_pass_at_1": base_pass1,
        "plus_pass_at_1": plus_pass1,
        "generation_time_s": total_elapsed,
    }
    with open(RESULTS_DIR / "multisample_results.json", "w") as f:
        json.dump(final, f, indent=2)

    print(f"\n{'='*60}")
    print(f"Multi-Sample Results ({NUM_SAMPLES} samples, T={TEMPERATURE}):")
    print(f"  HumanEval base pass@1: {base_pass1}")
    print(f"  HumanEval+ pass@1: {plus_pass1}")
    print(f"  Generation time: {total_elapsed:.0f}s ({total_elapsed/60:.1f} min)")
    print(f"{'='*60}", flush=True)


if __name__ == "__main__":
    main()