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