| """Candidate selection for code: K independent answers (each already through the deliberation fallback) are run on |
| the statement's public tests; the kept answer is the first candidate that passes all of them, else the candidate |
| passing the most (ties: the earliest sample). Only public tests are used, never the hidden grading tests. |
| The output keeps one record per (problem, sample 0) so it is graded like a single answer; n_tokens = sum over the |
| candidates actually generated (all K are counted, the cost of selection). |
| |
| usage: code_select.py --src results/lcb_dm/lcb.jsonl results/lcb_dm2/lcb.jsonl ... --items data/lcb/lcb6.jsonl |
| --out results/lcb_sel/lcb.jsonl [--k 4] |
| """ |
| import argparse |
| import collections |
| import json |
| import os |
| import sys |
| from concurrent.futures import ThreadPoolExecutor |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from grade import extract_code |
| from public_tests import run_public |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("--src", nargs="+", required=True) |
| ap.add_argument("--items", nargs="+", required=True) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--k", type=int, default=4) |
| a = ap.parse_args() |
| items = {} |
| for p in a.items: |
| for l in open(p, encoding="utf-8"): |
| r = json.loads(l) |
| items[r["id"]] = r |
| cands = collections.defaultdict(list) |
| for p in a.src: |
| for l in open(p, encoding="utf-8"): |
| r = json.loads(l) |
| cands[r["id"]].append(r) |
|
|
|
|
| def kind(it): |
| return "lcb" if it["task"] == "lcb" else "humaneval" if it["task"] == "humaneval" else \ |
| "mbpp" if it["task"] in ("mbpp", "code_fn") else "code_io" |
|
|
|
|
| def pick(pid): |
| rs = sorted(cands[pid], key=lambda r: r["sample"])[: a.k] |
| it = items[pid] |
| best, best_key = rs[0], (-1, 0) |
| for i, r in enumerate(rs): |
| fin = r["completion"].split("</think>", 1)[1] if "</think>" in r["completion"] else "" |
| code = extract_code(fin) |
| ok, tot = (run_public(kind(it), code, it["prompt"], it)[:2] if code else (0, 1)) |
| if (ok, -i) > best_key: |
| best, best_key = r, (ok, -i) |
| if tot and ok == tot: |
| break |
| return {**best, "sample": 0, "n_tokens": sum(r["n_tokens"] for r in rs), "picked": best["sample"], |
| "public_passed": best_key[0]} |
|
|
|
|
| with ThreadPoolExecutor(12) as ex: |
| out = list(ex.map(pick, list(cands))) |
| os.makedirs(os.path.dirname(a.out), exist_ok=True) |
| with open(a.out, "w", encoding="utf-8") as f: |
| for r in out: |
| f.write(json.dumps(r, ensure_ascii=False) + "\n") |
| print(len(out), "problems; picked sample distribution:", collections.Counter(r["picked"] for r in out)) |
|
|
