| """Exactly paired test of the forced code-check module on recorded runs.
|
|
|
| For every recorded answer whose thinking closed, the same thinking is replayed (prompt now declares the python
|
| tool), the model is made to write a check program right after </think>, the program runs, and the model continues
|
| with the output in a new thinking turn and gives its final answer. Truncated records are copied unchanged, so the
|
| difference to the source run is due to the check alone.
|
|
|
| usage: posthoc_check.py --src results/m180_base/mini.jsonl --out results/m180_chk/mini.jsonl [--max-tokens 6144]
|
| """
|
| import argparse
|
| import json
|
| import os
|
| import sys
|
| import threading
|
| import time
|
| import urllib.request
|
| import zlib
|
| from concurrent.futures import ThreadPoolExecutor
|
|
|
| from transformers import AutoTokenizer
|
|
|
| sys.argv, _argv = sys.argv[:1] + ["--inp", "x", "--out", "x"], sys.argv
|
| _src = open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "tir.py"), encoding="utf-8").read()
|
| exec(_src[:_src.index("def solve(")])
|
| sys.argv = _argv
|
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
| from grade import last_boxed
|
|
|
| ap2 = argparse.ArgumentParser()
|
| ap2.add_argument("--src", required=True)
|
| ap2.add_argument("--out", required=True)
|
| ap2.add_argument("--items", nargs="+", default=["data/coc/math180.jsonl"])
|
| ap2.add_argument("--max-tokens", type=int, default=6144)
|
| ap2.add_argument("--system", default="verify")
|
| ap2.add_argument("--check-line", default="# Independent check of the answer: compute it directly / by brute force\n")
|
| ap2.add_argument("--workers", type=int, default=32)
|
| ap2.add_argument("--mode", default="check", choices=["check", "switch"],
|
| help="check: code check after the thinking closes; switch: if the thinking runs longer than "
|
| "--switch-at tokens, stop it there and make the model compute the answer with a program")
|
| ap2.add_argument("--switch-at", type=int, default=3072)
|
| ap2.add_argument("--more-calls", type=int, default=0, help="further tool calls the model may make after the forced one")
|
| ap2.add_argument("--switch-text", default="\n\nThis is taking long. Let me compute the answer directly with a program.\n")
|
| a2 = ap2.parse_args()
|
| ENDS = {i for i in range(len(tok)) if (t := tok.convert_ids_to_tokens(i)) is not None and (t.endswith(".") or t.endswith("Ċ"))}
|
| items = {}
|
| for p in a2.items:
|
| for l in open(p, encoding="utf-8"):
|
| r = json.loads(l)
|
| items[r["id"]] = r
|
| SAMP = {"temperature": 0.6, "top_p": 0.95, "top_k": 50, "min_p": 0.0, "samplers": ["top_k", "top_p", "temperature"]}
|
|
|
|
|
| def check(rec):
|
| try:
|
| return _check(rec)
|
| except Exception as e:
|
| return {**rec, "check_error": repr(e)[:200]}
|
|
|
|
|
| def _check(rec):
|
| c = rec["completion"]
|
| it = items[rec["id"]]
|
| think = c[len("<think>"):].split("</think>", 1)[0]
|
| if a2.mode == "check" and (rec["truncated"] or "</think>" not in c):
|
| return rec
|
| th = tok(think, add_special_tokens=False)["input_ids"]
|
| if a2.mode == "switch":
|
| if len(th) <= a2.switch_at:
|
| return rec
|
| cut = next((j for j in range(a2.switch_at, a2.switch_at // 2, -1) if th[j - 1] in ENDS), a2.switch_at)
|
| think = tok.decode(th[:cut]) + a2.switch_text
|
| th = tok(think, add_special_tokens=False)["input_ids"]
|
| th = th + [THINK_END]
|
| msgs = ([{"role": "system", "content": SYSTEMS[a2.system]}] if SYSTEMS[a2.system] else []) + \
|
| [{"role": "user", "content": it["prompt"]}]
|
| ids = tok.apply_chat_template(msgs, tools=[TOOL], tokenize=True, add_generation_prompt=True, return_dict=False)
|
| used = len(th)
|
| seed = zlib.crc32(rec["id"].encode()) % 1000003 * 7 + rec["sample"] + 7
|
| pre = tok(FORCE_PREFIX + a2.check_line, add_special_tokens=False)["input_ids"]
|
| r2 = post({"prompt": ids + th + pre, "n_predict": max(64, min(1536, a2.max_tokens - used)), "seed": seed,
|
| "cache_prompt": False, "stop": ["</parameter>"], **SAMP})
|
| used += r2["tokens_predicted"]
|
| code = a2.check_line + r2["content"]
|
| out = run_code(code)
|
| ids = ids + th + pre + tok(r2["content"].rstrip("\n") + "\n</parameter>\n</function>\n</tool_call><|im_end|>\n"
|
| "<|im_start|>tool\n<tool_response>\n" + out + "\n</tool_response><|im_end|>\n"
|
| "<|im_start|>assistant\n<think>", add_special_tokens=False)["input_ids"]
|
| log = ""
|
| for extra in range(a2.more_calls + 1):
|
| r3 = post({"prompt": ids, "n_predict": max(256, a2.max_tokens - used), "seed": seed + 1 + extra,
|
| "cache_prompt": False, "return_tokens": True, **SAMP})
|
| used += r3["tokens_predicted"]
|
| after = r3["content"]
|
| m = CALL_RE.search(after.split("</think>", 1)[1]) if "</think>" in after else None
|
| if not m or extra == a2.more_calls or used >= a2.max_tokens:
|
| break
|
| out2 = run_code(m.group(1))
|
| turn = after.split("</tool_call>", 1)[0] + "</tool_call>"
|
| log += turn.replace("</think>", "\n", 1) + "\n<tool_response>\n" + out2 + "\n</tool_response>\n"
|
| ids = ids + [t for t in r3["tokens"] if t != EOS] + tok(
|
| "<|im_end|>\n<|im_start|>tool\n<tool_response>\n" + out2 + "\n</tool_response><|im_end|>\n"
|
| "<|im_start|>assistant\n<think>", add_special_tokens=False)["input_ids"]
|
| after = log + after
|
| if "</think>" not in after:
|
| if a2.mode == "switch":
|
| return {**rec, "completion": "<think>" + think + after, "n_tokens": used, "truncated": True,
|
| "check": {"code": code, "out": out, "kept": "none"}}
|
| return {**rec, "check": {"code": code, "out": out, "kept": "original"}}
|
| t2, final = after.split("</think>", 1)
|
| if "\\boxed" in it["prompt"] and last_boxed(final) is None:
|
|
|
| last = [t for t in r3["tokens"] if t != EOS]
|
| k = last.index(THINK_END) + 1 if THINK_END in last else len(last)
|
| pr = post({"prompt": ids + last[:k] + tok("The final answer is $\\boxed{", add_special_tokens=False)["input_ids"],
|
| "n_predict": 48, "temperature": 0.0, "cache_prompt": False, "stop": ["}$", "}\n"]})
|
| final = final.rstrip() + "\n\nThe final answer is $\\boxed{" + pr["content"] + "}$"
|
| return {**rec, "completion": "<think>" + think + "\n" + FORCE_PREFIX + code + "\n</parameter>\n</function>\n"
|
| "</tool_call>\n<tool_response>\n" + out + "\n</tool_response>\n" + t2 + "</think>" + final,
|
| "n_tokens": used, "truncated": False, "check": {"code": code, "out": out, "kept": "new"}}
|
|
|
|
|
| recs = [json.loads(l) for l in open(a2.src, encoding="utf-8")]
|
| os.makedirs(os.path.dirname(a2.out), exist_ok=True)
|
| t0 = time.time()
|
| with ThreadPoolExecutor(a2.workers) as ex:
|
| res = list(ex.map(check, recs))
|
| with open(a2.out, "w", encoding="utf-8") as f:
|
| for r in res:
|
| f.write(json.dumps(r, ensure_ascii=False) + "\n")
|
| print(f"{len(res)} records, {sum('check' in r for r in res)} checked, {(time.time() - t0) / 60:.1f} min")
|
|
|
