"""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 /forced.jsonl (rollout format, stop_type eos) and /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 # ; the model itself closes thinking right after a sentence end: "...}." + + "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}")