File size: 7,488 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 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | """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 # reuse tir.py's helpers without running it
_src = open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "tir.py"), encoding="utf-8").read()
exec(_src[:_src.index("def solve(")]) # noqa: S102 (constants, TOOL, SYSTEMS, run_code, post, CALL_RE)
sys.argv = _argv
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from grade import last_boxed # noqa: E402
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): # the model may run more code (e.g. after an error) before answering
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: # the post-check turn ran out: keep the original answer (switch: no answer)
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:
# after a tool turn the model often answers in prose: read the answer out in the requested form
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")
|