Upload jobs/occ_humaneval_v2.py

jobs/occ_humaneval_v2.py

@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+"""
+OCC HumanEval — Direct completion-format evaluation with Qwen3-Coder-30B-A3B-Instruct.
+Uses completion format (not chat), stop-token trimming, prompt+completion concatenation.
+Two-pass: 128 tokens → test → 1024 tokens on failures. Saves real token counts.
+"""
+import json, time, sys, torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from datasets import load_dataset
+
+MODEL = "Qwen/Qwen3-Coder-30B-A3B-Instruct"
+SHORT, LONG = 128, 1024
+STOP = ["\nclass", "\ndef", "\n#", "\nif __name__", "\nprint("]
+
+def log(msg): print(f"[OCC] {msg}", flush=True)
+
+log(f"Loading {MODEL}...")
+tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
+tok.pad_token = tok.eos_token
+model = AutoModelForCausalLM.from_pretrained(MODEL, trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto")
+log(f"Loaded on {model.device}")
+
+ds = load_dataset("openai/openai_humaneval", split="test")
+problems = list(ds)
+log(f"Loaded {len(problems)} problems")
+
+def generate(prompt, mx):
+    inp = tok(prompt, return_tensors="pt", truncation=True, max_length=2048).to(model.device)
+    ilen = inp.input_ids.shape[1]
+    with torch.no_grad():
+        out = model.generate(**inp, max_new_tokens=mx, do_sample=False, pad_token_id=tok.eos_token_id)
+    gen = tok.decode(out[0][ilen:], skip_special_tokens=True)
+    for sw in STOP:
+        idx = gen.find(sw)
+        if idx != -1: gen = gen[:idx]
+    return gen
+
+def clean_body(gen):
+    lines = gen.split("\n")
+    while lines and not lines[0].strip(): lines = lines[1:]
+    while lines and not lines[-1].strip(): lines = lines[:-1]
+    return "\n".join(lines)
+
+def test_problem(prompt, gen, test, entry):
+    body = clean_body(gen)
+    if not body.strip(): return False, "empty"
+    code = prompt.rstrip() + "\n" + body + "\n\n" + test
+    try: compile(code, "<t>", "exec")
+    except SyntaxError as e: return False, f"SyntaxErr:{e}"
+    try:
+        ns = {}
+        exec(code, ns)
+        if "check" in ns: ns["check"](ns.get(entry))
+        return True, "OK"
+    except AssertionError: return False, "AssertErr"
+    except Exception as e: return False, f"{type(e).__name__}:{str(e)[:60]}"
+
+def run_eval(probs, mx, label):
+    log(f"\n--- {label} ---")
+    ok, fail, tok_total = [], [], 0
+    for i, p in enumerate(probs):
+        t0 = time.time()
+        gen = generate(p["prompt"], mx)
+        elapsed = time.time() - t0
+        ntok = len(tok.encode(gen))
+        tok_total += ntok
+        passed, reason = test_problem(p["prompt"], gen, p["test"], p["entry_point"])
+        if passed:
+            ok.append(p["task_id"])
+            log(f"  [{i+1:3d}] {p['task_id']}: PASS ({ntok}tok {elapsed:.1f}s)")
+        else:
+            fail.append(p["task_id"])
+            log(f"  [{i+1:3d}] {p['task_id']}: FAIL ({ntok}tok) — {reason[:80]}")
+    acc = len(ok) / len(probs)
+    log(f"{label}: {len(ok)}/{len(probs)} ({acc:.4f}), {tok_total} tokens")
+    return ok, fail, tok_total, acc
+
+# Pass 1
+p1_ok, p1_fail, p1_tok, p1_acc = run_eval(problems, SHORT, "Pass1-SHORT")
+
+# Pass 2
+p2_new = 0; p2_tok = 0
+if p1_fail:
+    fail_set = set(p1_fail)
+    fp = [p for p in problems if p["task_id"] in fail_set]
+    p2_ok, _, p2_tok, _ = run_eval(fp, LONG, "Pass2-LONG")
+    p2_new = len(p2_ok)
+else:
+    log("\nPass2: SKIPPED (100% on pass 1)")
+
+final = len(p1_ok) + p2_new
+final_acc = final / len(problems)
+tt = p1_tok + p2_tok
+base = len(problems) * LONG
+sav = (1 - tt/base)*100
+
+log(f"\n{'='*60}")
+log(f"Qwen3-Coder-30B-A3B-Instruct HumanEval OCC")
+log(f"P1({SHORT}tok): {len(p1_ok)} passed, P2(+{LONG}tok): +{p2_new}")
+log(f"OCC pass@1={final_acc:.4f} ({final}/{len(problems)}), {tt} tok")
+log(f"Baseline={base} tok, Savings={sav:.1f}%")
+log(f"{'='*60}")
+
+res = {"model":MODEL,"pass1_ok":len(p1_ok),"pass2_new":p2_new,
+       "occ_pass1":final_acc,"occ_tokens":tt,"baseline_tokens":base,"savings_pct":sav}
+json.dump(res, open("/app/occ_humaneval_results.json","w"), indent=2)
+log("Saved to /app/occ_humaneval_results.json")