File size: 6,813 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 | """Commitment bootstrap: turn truncated (budget-exhausted) reasoning into learning signal.
For every truncated math rollout, the reasoning is cut every --step tokens, thinking is closed by force and the
policy answers greedily. Each cut yields a candidate "stop here" continuation that is graded against the reference:
* the earliest cut whose forced answer is correct -> forced sample with reward 1 ("you could have stopped here")
* the latest cut before it whose answer is wrong -> forced sample with reward 0 ("stopping here is premature")
If no cut is correct, the last wrong cut is added with reward 0. Forced samples join the GRPO group of their prompt,
so advantages are computed against the policy's own full-length attempts.
Writes <out_dir>/forced.jsonl (rollout format, stop_type eos) and <out_dir>/forced_graded.jsonl.
usage: bc_probe.py --rollouts R --graded G --items P --out-dir D [--step 1024]"""
import argparse, json, sys, threading, urllib.request
from concurrent.futures import ThreadPoolExecutor
from transformers import AutoTokenizer
sys.path.insert(0, __file__.rsplit("\\", 1)[0].rsplit("/", 1)[0])
from grade import grade_math # noqa: E402
ap = argparse.ArgumentParser()
for a in ["--rollouts", "--graded", "--items", "--out-dir"]:
ap.add_argument(a, required=True)
ap.add_argument("--step", type=int, default=1024)
ap.add_argument("--stable", type=int, default=1,
help="a cut counts as 'answer known' only if it and the next stable-1 cuts are all correct")
ap.add_argument("--tokenizer", default="models/LFM3-1.7B")
ap.add_argument("--from-probes", default=None, help="probes.jsonl of an earlier run with the same cuts")
ap.add_argument("--unpaired-neg", action="store_true",
help="also add a negative forced stop for rollouts whose answer is never right (unbalances the signal)")
args = ap.parse_args()
tok = AutoTokenizer.from_pretrained(args.tokenizer)
items = {json.loads(l)["id"]: json.loads(l) for l in open(args.items, encoding="utf-8")}
recs = [json.loads(l) for l in open(args.rollouts, encoding="utf-8")]
recs = [r for r in recs if r["truncated"] and items[r["id"]]["task"] == "math" and "tokens" in r]
THINK_END = 124902 # </think>; the model itself closes thinking right after a sentence end: "...}." + </think> + "The"
SUFFIX = "The final answer is $\\boxed{"
suffix_ids = [THINK_END] + tok(SUFFIX, add_special_tokens=False)["input_ids"]
ENDS = set()
for i_ in range(len(tok)):
t_ = tok.convert_ids_to_tokens(i_)
if t_ is not None and (t_.endswith(".") or t_.endswith("\u010a")):
ENDS.add(i_)
def snap(tokens, k, lo):
"""last sentence/paragraph end at or before k (cut right after it), strictly after lo; None if there is none"""
for j in range(k, lo, -1):
if tokens[j - 1] in ENDS:
return j
return None
prompts = {i: tok.apply_chat_template([{"role": "user", "content": items[i]["prompt"]}], tokenize=True,
add_generation_prompt=True, return_dict=False) for i in {r["id"] for r in recs}}
jobs = []
for r in recs:
prev = 0
for k in range(args.step, len(r["tokens"]), args.step):
j = snap(r["tokens"], k, prev)
if j:
jobs.append((r, j))
prev = j
print(len(recs), "truncated math rollouts,", len(jobs), "probes", flush=True)
res = {}
lock = threading.Lock()
def run(job):
r, k = job
body = {"prompt": prompts[r["id"]] + r["tokens"][:k] + suffix_ids, "n_predict": 48, "temperature": 0.0,
"cache_prompt": False, "stop": ["}$", "}\n", "$"]}
try:
out = json.loads(urllib.request.urlopen(urllib.request.Request(
"http://localhost:8080/completion", data=json.dumps(body).encode(),
headers={"Content-Type": "application/json"}), timeout=900).read())["content"]
except urllib.error.HTTPError as e: # e.g. prompt + cut exceeds the slot context: skip this cut
print("probe skipped:", e.code, flush=True)
return
depth, cut = 1, len(out)
for i, ch in enumerate(out):
depth += (ch == "{") - (ch == "}")
if depth == 0:
cut = i
break
pred = out[:cut].strip()
ok = bool(pred) and grade_math("\\boxed{" + pred + "}", items[r["id"]]["answer"])
with lock:
res[(r["id"], r["sample"], k)] = (pred, ok)
if args.from_probes: # reuse saved probe results (no server needed)
for l in open(args.from_probes, encoding="utf-8"):
q = json.loads(l)
res[(q["id"], q["sample"], q["k"])] = (q["pred"], q["ok"])
else:
with ThreadPoolExecutor(32) as ex:
list(ex.map(run, jobs))
with open(f"{args.out_dir}/probes.jsonl", "w", encoding="utf-8") as f:
f.writelines(json.dumps({"id": i, "sample": s_, "k": k, "pred": v[0], "ok": v[1]}) + "\n" for (i, s_, k), v in res.items())
forced, graded, n_pos, n_neg, n_any = [], [], 0, 0, 0
for r in recs:
cuts = sorted(k for (i, s, k) in res if i == r["id"] and s == r["sample"])
oks = [res[(r["id"], r["sample"], k)] for k in cuts]
first_ok = next((j for j in range(len(oks)) if all(ok for _, ok in oks[j: j + args.stable])), None)
picks = []
if first_ok is not None:
n_any += 1
picks.append((cuts[first_ok], 1.0))
wrong_before = [j for j in range(first_ok) if not oks[j][1]]
if wrong_before:
picks.append((cuts[wrong_before[-1]], 0.0))
elif cuts and args.unpaired_neg:
picks.append((cuts[-1], 0.0)) # off by default: 'no answer anywhere' is not a premature-stop signal
for k, rew in picks:
pred = res[(r["id"], r["sample"], k)][0]
tail = [THINK_END] + tok(SUFFIX + pred + "}$", add_special_tokens=False)["input_ids"]
sid = 1000 + r["sample"] * 100 + k // 256
forced.append({"id": r["id"], "task": "math", "sample": sid, "tokens": r["tokens"][:k] + tail,
"stop_type": "eos", "truncated": False, "forced": True, "cut": k, "loss_from": k})
graded.append({"id": r["id"], "task": "math", "sample": sid, "correct": rew, "forced": True})
n_pos += rew == 1.0
n_neg += rew == 0.0
with open(f"{args.out_dir}/forced.jsonl", "w", encoding="utf-8") as f:
f.writelines(json.dumps(x) + "\n" for x in forced)
with open(f"{args.out_dir}/forced_graded.jsonl", "w", encoding="utf-8") as f:
f.writelines(json.dumps(x) + "\n" for x in graded)
print(f"truncated rollouts whose forced answer is right at some cut: {n_any}/{len(recs)}; "
f"forced samples: {n_pos} positive, {n_neg} negative", flush=True)
probe_acc = sum(ok for _, ok in res.values()) / max(1, len(res))
print(f"probe accuracy over all cuts: {probe_acc:.3f}")
|