Upload jobs/occ_humaneval_harness.py

jobs/occ_humaneval_harness.py

@@ -0,0 +1,212 @@
+#!/usr/bin/env python3
+"""
+OCC HumanEval Benchmark using BigCode Evaluation Harness + Qwen3-Coder-30B-A3B-Instruct.
+
+Strategy:
+  Pass 1: Generate with 128 tokens, run tests
+  Pass 2: Only on failures, regenerate with 1024 tokens
+  Merge: pass@1 = (pass1_passed + pass2_added_passed) / total
+
+This uses the standard evaluation harness (not custom extraction).
+The harness handles completion-format prompts, stop tokens, and code extraction correctly.
+"""
+
+import subprocess
+import json
+import os
+import sys
+import time
+from pathlib import Path
+
+# === CONFIG ===
+MODEL = "Qwen/Qwen3-Coder-30B-A3B-Instruct"
+MODEL_ARGS = f"pretrained={MODEL},trust_remote_code=True,dtype=bfloat16"
+TOTAL_PROBLEMS = 164
+SHORT_TOKENS = 128
+LONG_TOKENS = 1024
+HARNESS_REPO = "https://github.com/bigcode-project/bigcode-evaluation-harness.git"
+HARNESS_DIR = "/app/bigcode-evaluation-harness"
+RESULTS_DIR = "/app/occ_humaneval_results"
+
+
+def log(msg):
+    print(f"[OCC] {msg}", flush=True)
+
+
+def run_cmd(cmd, cwd=None, check=True):
+    log(f"  $ {' '.join(cmd[:6])}{'...' if len(cmd) > 6 else ''}")
+    result = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True)
+    if result.returncode != 0 and check:
+        log(f"  FAILED (rc={result.returncode})")
+        log(f"  STDERR: {result.stderr[-1500:]}")
+    return result.returncode == 0, result.stdout, result.stderr
+
+
+def setup():
+    """Clone harness, install deps."""
+    if not os.path.exists(HARNESS_DIR):
+        log("Cloning bigcode-evaluation-harness...")
+        ok, _, _ = run_cmd(["git", "clone", HARNESS_REPO, HARNESS_DIR])
+        if not ok:
+            sys.exit(1)
+
+    log("Installing harness + deps...")
+    # Install harness
+    run_cmd(["pip", "install", "-e", HARNESS_DIR], check=False)
+    # Ensure accelerate
+    run_cmd(["pip", "install", "accelerate"], check=False)
+    log("Setup complete.")
+
+
+def run_eval(max_tokens, label):
+    """Run harness with given token limit. Returns (metrics_dict, generations_dict)."""
+    os.makedirs(RESULTS_DIR, exist_ok=True)
+    gen_path = f"{RESULTS_DIR}/{label}_generations.json"
+    met_path = f"{RESULTS_DIR}/{label}_metrics.json"
+
+    cmd = [
+        "accelerate", "launch", "main.py",
+        "--model", "hf",
+        "--model_args", MODEL_ARGS,
+        "--tasks", "humaneval",
+        "--max_length_generation", str(max_tokens),
+        "--batch_size", "1",
+        "--n_samples", "1",
+        "--allow_code_execution",
+        "--save_generations",
+        "--generations_path", gen_path,
+        "--metric_output_path", met_path,
+        "--temperature", "0.0",  # greedy
+    ]
+
+    log(f"Running {label} (max_tokens={max_tokens})...")
+    t0 = time.time()
+    ok, stdout, stderr = run_cmd(cmd, cwd=HARNESS_DIR, check=False)
+    elapsed = time.time() - t0
+    log(f"  Elapsed: {elapsed:.0f}s")
+
+    if not ok:
+        log(f"  FAILED. Dumping last 2000 chars of output:")
+        log(f"  STDOUT: {stdout[-2000:]}")
+        log(f"  STDERR: {stderr[-2000:]}")
+        return None, None
+
+    metrics = None
+    if os.path.exists(met_path):
+        with open(met_path) as f:
+            metrics = json.load(f)
+
+    generations = None
+    if os.path.exists(gen_path):
+        with open(gen_path) as f:
+            generations = json.load(f)
+
+    # Parse pass@1
+    if metrics:
+        for key in metrics:
+            if "pass@1" in metrics[key]:
+                log(f"  pass@1: {metrics[key]['pass@1']:.4f}")
+
+    return metrics, generations
+
+
+def count_passed(generations):
+    """Count how many problems passed tests."""
+    passed = []
+    failed = []
+    for task_id, val in generations.items():
+        # val is typically [code_string, test_result_bool]
+        if isinstance(val, list) and len(val) >= 2:
+            if val[1]:
+                passed.append(task_id)
+            else:
+                failed.append(task_id)
+        else:
+            failed.append(task_id)
+    return passed, failed
+
+
+def main():
+    log("=" * 60)
+    log(f"OCC HumanEval — {MODEL}")
+    log(f"Tiered: {SHORT_TOKENS} tokens → test → {LONG_TOKENS} tokens on failures")
+    log("=" * 60)
+
+    setup()
+
+    # --- Pass 1: Short generation ---
+    log("\n--- PASS 1: SHORT ---")
+    m1, g1 = run_eval(SHORT_TOKENS, "pass1_short")
+
+    if g1 is None:
+        log("CRITICAL: Pass 1 failed. Aborting.")
+        sys.exit(1)
+
+    passed_1, failed_1 = count_passed(g1)
+    n_pass1 = len(passed_1)
+    n_fail1 = len(failed_1)
+    log(f"Pass 1: {n_pass1} passed, {n_fail1} failed ({n_pass1/len(g1)*100:.1f}%)")
+    pass1_tokens = len(g1) * SHORT_TOKENS
+
+    # --- Pass 2: Long generation on failures ---
+    pass2_tokens = 0
+    new_passes = 0
+
+    if n_fail1 > 0:
+        log(f"\n--- PASS 2: LONG ({LONG_TOKENS} tokens) on {n_fail1} failures ---")
+        # Run full eval with long tokens — simpler and correct.
+        # We only count problems that were in the failure set.
+        m2, g2 = run_eval(LONG_TOKENS, "pass2_long")
+
+        if g2:
+            passed_2, failed_2 = count_passed(g2)
+            new_pass_set = set(passed_2) & set(failed_1)
+            still_fail_set = set(failed_2) & set(failed_1)
+            new_passes = len(new_pass_set)
+            still_failed = len(still_fail_set)
+            log(f"Pass 2: {new_passes} newly passed, {still_failed} still failed")
+            pass2_tokens = n_fail1 * LONG_TOKENS
+        else:
+            log("Pass 2 failed — keeping Pass 1 results only.")
+    else:
+        log("\n--- PASS 2: SKIPPED (all passed) ---")
+
+    # --- Merge results ---
+    final_passed = n_pass1 + new_passes
+    final_pass_at_1 = final_passed / TOTAL_PROBLEMS
+    total_tokens = pass1_tokens + pass2_tokens
+    baseline_tokens = TOTAL_PROBLEMS * LONG_TOKENS
+    token_savings_pct = (1.0 - total_tokens / baseline_tokens) * 100 if baseline_tokens > 0 else 0
+
+    log("\n" + "=" * 60)
+    log("FINAL RESULTS")
+    log("=" * 60)
+    log(f"  OCC pass@1:    {final_pass_at_1:.4f} ({final_passed}/{TOTAL_PROBLEMS})")
+    log(f"  Pass 1 only:   {n_pass1/TOTAL_PROBLEMS:.4f} ({n_pass1}/{TOTAL_PROBLEMS})")
+    log(f"  Tokens used:   {total_tokens} (OCC) vs {baseline_tokens} (baseline)")
+    log(f"  Savings:       {token_savings_pct:.1f}%")
+    log(f"  New from P2:   {new_passes}")
+    log(f"  Still failed:  {TOTAL_PROBLEMS - final_passed}")
+
+    # Save final report
+    results = {
+        "model": MODEL,
+        "config": {"short_tokens": SHORT_TOKENS, "long_tokens": LONG_TOKENS, "total_problems": TOTAL_PROBLEMS},
+        "pass1": {"pass_at_1": n_pass1 / TOTAL_PROBLEMS, "passed": n_pass1, "failed": n_fail1, "tokens": pass1_tokens},
+        "pass2": {"newly_passed": new_passes, "tokens": pass2_tokens},
+        "occ_combined": {"pass_at_1": final_pass_at_1, "total_passed": final_passed, "total_tokens": total_tokens},
+        "baseline_long_only": {"tokens": baseline_tokens, "savings_pct": token_savings_pct},
+    }
+
+    report_path = f"{RESULTS_DIR}/occ_results.json"
+    with open(report_path, "w") as f:
+        json.dump(results, f, indent=2)
+
+    log(f"\nResults saved to {report_path}")
+
+    # Also write a summary to stdout for easy parsing
+    print(f"\nOCC_SUMMARY: pass@1={final_pass_at_1:.4f} tokens={total_tokens} baseline={baseline_tokens} savings={token_savings_pct:.1f}%")
+
+
+if __name__ == "__main__":
+    main()