File size: 2,658 Bytes
67ebcef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | """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 # noqa: E402
from public_tests import run_public # noqa: E402
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))
|