202 lines
7.8 KiB
Python
202 lines
7.8 KiB
Python
"""LiveCodeBench transfer eval: contamination-safe check of the MBPP-trained
|
|
implants (no LCB training exists anywhere in the pipeline).
|
|
|
|
Uses code_generation_lite, stdin-judged problems only, filtered to contests
|
|
dated >= --since (default 2026-01-01, past the base model's training data).
|
|
Labeling pass (direct vs terse-plan, greedy, frozen model) defines the
|
|
descriptive hard bucket, exactly as for HumanEval; then arms at k grid.
|
|
|
|
Verifier: run the program on each test's stdin, compare whitespace-normalized
|
|
stdout. Caps: --cap problems, 8 test cases per problem, 15s per case.
|
|
"""
|
|
|
|
import argparse
|
|
import base64
|
|
import json
|
|
import os
|
|
import pickle
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import zlib
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
from loop_common import BandLooper, MergeAdapter
|
|
from prep_mbpp import batch_generate, extract_code
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
from jlens.core import load_model # noqa: E402
|
|
|
|
OUT = Path(os.environ.get("LOOP_OUT",
|
|
Path(__file__).resolve().parent.parent / "results-loop"))
|
|
|
|
DIRECT = ("Solve this competitive programming problem in Python. Read from "
|
|
"standard input, write to standard output. Return ONLY the complete "
|
|
"program in a ```python code block.\n\n{q}")
|
|
PLAN = ("First write a very brief plan: at most 4 short bullet lines. Then "
|
|
"solve this competitive programming problem in Python (stdin -> "
|
|
"stdout), the complete program in a ```python code block.\n\n{q}")
|
|
|
|
|
|
def decode_private(s):
|
|
try:
|
|
return json.loads(s)
|
|
except (json.JSONDecodeError, TypeError):
|
|
return json.loads(pickle.loads(zlib.decompress(base64.b64decode(s))))
|
|
|
|
|
|
def load_items(since, cap):
|
|
# script-loader is dead in current `datasets`; read the raw jsonl shards
|
|
from huggingface_hub import hf_hub_download
|
|
rows = []
|
|
for shard in ("test.jsonl", "test2.jsonl", "test3.jsonl",
|
|
"test4.jsonl", "test5.jsonl", "test6.jsonl"):
|
|
try:
|
|
p = hf_hub_download("livecodebench/code_generation_lite", shard,
|
|
repo_type="dataset")
|
|
except Exception as e: # noqa: BLE001
|
|
print(f"({shard}: {type(e).__name__})", flush=True)
|
|
continue
|
|
with open(p) as f:
|
|
rows += [json.loads(l) for l in f]
|
|
items = []
|
|
for r in rows:
|
|
if str(r["contest_date"])[:10] < since:
|
|
continue
|
|
tests = json.loads(r["public_test_cases"])
|
|
try:
|
|
tests += decode_private(r["private_test_cases"])
|
|
except Exception: # noqa: BLE001
|
|
pass
|
|
tests = [t for t in tests if t.get("testtype") == "stdin"][:8]
|
|
if not tests:
|
|
continue
|
|
items.append({"task_id": r["question_id"],
|
|
"date": str(r["contest_date"])[:10],
|
|
"q": r["question_content"], "tests": tests,
|
|
"difficulty": r.get("difficulty", "")})
|
|
items.sort(key=lambda it: it["date"], reverse=True) # newest first
|
|
items = items[:cap]
|
|
items.sort(key=lambda it: it["task_id"])
|
|
print(f"LCB: {len(items)} stdin problems, {items[0]['date'] if items else '-'}"
|
|
f" .. {max(it['date'] for it in items) if items else '-'}", flush=True)
|
|
return items
|
|
|
|
|
|
def run_case(code, inp, expected, timeout=15):
|
|
try:
|
|
r = subprocess.run([sys.executable, "-c", code], input=inp,
|
|
capture_output=True, text=True, timeout=timeout)
|
|
except (subprocess.TimeoutExpired, OSError):
|
|
return False
|
|
if r.returncode != 0:
|
|
return False
|
|
got = "\n".join(" ".join(l.split()) for l in r.stdout.strip().splitlines())
|
|
exp = "\n".join(" ".join(l.split()) for l in expected.strip().splitlines())
|
|
return got == exp
|
|
|
|
|
|
def passes(code, item):
|
|
if not code:
|
|
return False
|
|
return all(run_case(code, t["input"], t["output"]) for t in item["tests"])
|
|
|
|
|
|
def prompt(tok, item, tmpl):
|
|
return tok.apply_chat_template(
|
|
[{"role": "user", "content": tmpl.format(q=item["q"])}],
|
|
tokenize=False, add_generation_prompt=True)
|
|
|
|
|
|
@torch.no_grad()
|
|
def loop_eval(looper, adapter, tok, items, k, batch=4, max_new=700,
|
|
feedforward=False):
|
|
codes = []
|
|
for i in range(0, len(items), batch):
|
|
torch.cuda.empty_cache()
|
|
chunk = items[i : i + batch]
|
|
enc = tok([prompt(tok, it, DIRECT) for it in chunk],
|
|
return_tensors="pt", padding=True,
|
|
add_special_tokens=False).to("cuda")
|
|
gen = looper.generate_frozen_prompt(adapter, tok, enc["input_ids"], k,
|
|
max_new_tokens=max_new,
|
|
attention_mask=enc["attention_mask"],
|
|
feedforward=feedforward)
|
|
for j in range(len(chunk)):
|
|
codes.append(extract_code(
|
|
tok.decode(gen[j, enc["input_ids"].shape[1]:],
|
|
skip_special_tokens=True)))
|
|
with ThreadPoolExecutor(8) as ex:
|
|
return list(ex.map(lambda ci: passes(ci[0], ci[1]),
|
|
zip(codes, items)))
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--adapter", default=None)
|
|
ap.add_argument("--tag", default="lcb")
|
|
ap.add_argument("--ks", default="0,4")
|
|
ap.add_argument("--since", default="2024-07-01",
|
|
help="newest LCB shard ends 2025-04; gemma-4 cutoff is "
|
|
"undisclosed, so this is newest-available, not "
|
|
"provably post-cutoff")
|
|
ap.add_argument("--cap", type=int, default=150)
|
|
ap.add_argument("--feedforward", action="store_true")
|
|
args = ap.parse_args()
|
|
ks = [int(x) for x in args.ks.split(",")]
|
|
|
|
items = load_items(args.since, args.cap)
|
|
model, tok = load_model(dtype=torch.bfloat16)
|
|
tok.padding_side = "left"
|
|
looper = BandLooper(model)
|
|
adapter = MergeAdapter(
|
|
d=model.config.get_text_config().hidden_size).cuda()
|
|
if args.adapter:
|
|
adapter.load_state_dict(torch.load(args.adapter, map_location="cuda"))
|
|
adapter.eval()
|
|
|
|
lab_path = OUT / "lcb_labels.json"
|
|
if lab_path.exists():
|
|
labels = json.load(open(lab_path))
|
|
else:
|
|
plans = batch_generate(model, tok,
|
|
[prompt(tok, it, PLAN) for it in items],
|
|
max_new_tokens=1100, batch_size=4)
|
|
with ThreadPoolExecutor(8) as ex:
|
|
plan_ok = list(ex.map(
|
|
lambda gi: passes(extract_code(gi[0]), gi[1]),
|
|
zip(plans, items)))
|
|
labels = {it["task_id"]: bool(ok) for it, ok in zip(items, plan_ok)}
|
|
json.dump(labels, open(lab_path, "w"), indent=1)
|
|
print(f"plan-solvable: {sum(labels.values())}/{len(items)}", flush=True)
|
|
|
|
res = {"tag": args.tag, "since": args.since, "n": len(items), "ks": {}}
|
|
k0_ok = None
|
|
for k in ks:
|
|
t0 = time.time()
|
|
oks = loop_eval(looper, adapter, tok, items, k,
|
|
feedforward=args.feedforward)
|
|
if k == 0:
|
|
k0_ok = oks
|
|
hard = [i for i, it in enumerate(items)
|
|
if k0_ok and not k0_ok[i] and labels.get(it["task_id"])]
|
|
acc = sum(oks) / len(items)
|
|
hacc = (sum(oks[i] for i in hard) / len(hard)) if hard else None
|
|
res["ks"][k] = {"acc": acc, "hard_n": len(hard), "hard_acc": hacc,
|
|
"per_item": [{"task_id": it["task_id"],
|
|
"ok": bool(o)}
|
|
for it, o in zip(items, oks)]}
|
|
print(f"k={k}: pass@1={acc:.3f} hard({len(hard)})="
|
|
f"{hacc if hacc is None else round(hacc, 3)} "
|
|
f"({time.time()-t0:.0f}s)", flush=True)
|
|
|
|
json.dump(res, open(OUT / f"eval_lcb_{args.tag}.json", "w"), indent=1)
|
|
print("wrote", OUT / f"eval_lcb_{args.tag}.json")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|