File size: 4,625 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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | """Code read-out rescue for the deliberation module.
When the thinking of a code task hits the budget, the module closes it and lets the model answer. On hard problems the
model then writes an editorial-style "Solution Explanation" with pseudo-code and runs out of space before the program
(LCB v6: 37 of 112 budget read-outs had no Python block at all). Rescue: if the answer generated at the budget contains
no program that compiles, generate it again with the code block already opened ("```python" right after </think>).
Other records are copied unchanged, so the comparison with the source run is exactly paired. Tokens of the discarded
answer stay counted.
usage: code_rescue.py --src results/lcb_dm/lcb.jsonl --items data/lcb/lcb6.jsonl --out results/lcb_dmr/lcb.jsonl
"""
import argparse
import json
import os
import sys
import time
import urllib.request
import zlib
from concurrent.futures import ThreadPoolExecutor
from transformers import AutoTokenizer
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from grade import extract_code # noqa: E402
THINK_END = 124902
PREFILL = "```python\n"
CODE_TASKS = {"code_io", "humaneval", "mbpp", "code_fn", "lcb"}
SAMPLING = {"temperature": 0.6, "top_p": 0.95, "top_k": 50, "min_p": 0.0, "repeat_penalty": 1.0,
"samplers": ["top_k", "top_p", "temperature"]}
def runnable(text):
code = extract_code(text or "")
if not code:
return False
try:
compile(code, "<answer>", "exec")
return True
except Exception:
return False
def close_block(content):
text = PREFILL + content
return text if "```" in content else text.rstrip() + "\n```"
def post(port, body):
for attempt in range(5):
try:
req = urllib.request.Request(f"http://localhost:{port}/completion", data=json.dumps(body).encode(),
headers={"Content-Type": "application/json"})
return json.loads(urllib.request.urlopen(req, timeout=3600).read())
except Exception:
if attempt == 4:
raise
time.sleep(5)
if __name__ == "__main__":
ap = argparse.ArgumentParser()
ap.add_argument("--src", required=True)
ap.add_argument("--items", nargs="+", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--reserve", type=int, default=4096)
ap.add_argument("--workers", type=int, default=32)
ap.add_argument("--port", type=int, default=8080)
ap.add_argument("--tokenizer", default="models/LFM3-1.7B")
args = ap.parse_args()
tok = AutoTokenizer.from_pretrained(args.tokenizer)
pre_ids = tok(PREFILL, add_special_tokens=False)["input_ids"]
items = {}
for path in args.items:
for l in open(path, encoding="utf-8"):
r = json.loads(l)
items[r["id"]] = r
def process(rec):
it = items[rec["id"]]
c = rec["completion"]
if it["task"] not in CODE_TASKS or rec.get("module") != "fallback" or "</think>" not in c:
return rec
thinking, final = c[len("<think>"):].split("</think>", 1)
if runnable(final):
return rec
ids = tok.apply_chat_template([{"role": "user", "content": it["prompt"]}], tokenize=True,
add_generation_prompt=True, return_dict=False)
gen = tok(thinking, add_special_tokens=False)["input_ids"]
seed = zlib.crc32(rec["id"].encode()) % 1000003 * 7 + rec.get("sample", 0) + 31
r = post(args.port, {"prompt": ids + gen + [THINK_END] + pre_ids, "n_predict": args.reserve, "seed": seed,
"cache_prompt": False, "return_tokens": True, **SAMPLING})
text = close_block(r["content"])
return {**rec, "completion": "<think>" + thinking + "</think>" + text,
"n_tokens": rec["n_tokens"] + len(pre_ids) + len(r["tokens"]), "module": "fallback_rescue",
"rescued": runnable(text)}
recs = [json.loads(l) for l in open(args.src, encoding="utf-8")]
t0 = time.time()
with ThreadPoolExecutor(args.workers) as ex:
out = list(ex.map(process, recs))
os.makedirs(os.path.dirname(args.out), exist_ok=True)
with open(args.out, "w", encoding="utf-8") as f:
for r in out:
f.write(json.dumps(r, ensure_ascii=False) + "\n")
res = [r for r in out if r.get("module") == "fallback_rescue"]
print(f"{len(out)} records, {len(res)} read-outs without a runnable program regenerated, "
f"{sum(r['rescued'] for r in res)} now have one, {(time.time() - t0) / 60:.1f} min")
|