| """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 |
|
|
| 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") |
|
|
