| """Apply the deliberation module to already generated reasoning (exactly what it would do online).
|
|
|
| The module never changes the reasoning before it acts, and its probes are greedy (deterministic given the prefix),
|
| so running it over recorded trajectories reproduces its online behaviour on those same trajectories. This gives a
|
| fully paired comparison with the recorded baseline.
|
|
|
| Module (settings chosen on dev data only):
|
| * every --every tokens of thinking: close at the last sentence end and read a short forced answer
|
| * commit when --agree successive answers are identical
|
| * if the reasoning hit the token limit: answer within the budget instead of being cut off
|
| (short-answer tasks: forced short answer at the limit; code / instruction tasks: thinking closed
|
| --free-reserve tokens before the limit, answer generated with the normal sampling settings)
|
| Output: results/<out_tag>/*.jsonl in gen.py format (+ "module": commit/fallback/none) -> grade.py, compare.py
|
|
|
| usage: posthoc_coc.py --src results/orig/main.jsonl results/orig/aime.jsonl --out-tag orig_dm
|
| """
|
| import argparse
|
| import json
|
| import os
|
| import re
|
| import sys
|
| import threading
|
| import time
|
| import urllib.request
|
| import zlib
|
| from concurrent.futures import ThreadPoolExecutor
|
|
|
| from transformers import AutoTokenizer
|
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
| from grade import last_boxed, load_suite
|
|
|
| ap = argparse.ArgumentParser()
|
| ap.add_argument("--src", nargs="+", required=True)
|
| ap.add_argument("--out-tag", required=True)
|
| ap.add_argument("--max-tokens", type=int, default=16384)
|
| ap.add_argument("--every", type=int, default=2048)
|
| ap.add_argument("--agree", type=int, default=3)
|
| ap.add_argument("--free-reserve", type=int, default=1536)
|
| ap.add_argument("--commit-tasks", nargs="*", default=["mmlupro", "bbh"],
|
| help="tasks where early commitment is allowed (dev: it hurts numeric math at a 16k budget)")
|
| ap.add_argument("--only-tasks", nargs="*", default=None, help="process only these tasks (others copied unchanged)")
|
| ap.add_argument("--workers", type=int, default=32)
|
| ap.add_argument("--port", type=int, default=8080)
|
| ap.add_argument("--tokenizer", default="models/LFM3-1.7B")
|
| ap.add_argument("--items", default=None, help="glob of item files (default: the eval suite)")
|
| args = ap.parse_args()
|
|
|
| THINK_END = 124902
|
| URL = f"http://localhost:{args.port}/completion"
|
| tok = AutoTokenizer.from_pretrained(args.tokenizer)
|
| 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_)
|
| MATH = ("The final answer is $\\boxed{", 48, ["}$", "}\n"], "}$")
|
| SHORT = {"math500": MATH, "aime": MATH, "gsm8k": MATH, "math": MATH,
|
| "mmlupro": ("The answer is (", 3, [")"], ")."), "bbh": ("The answer is ", 24, ["\n", ".\n"], ".")}
|
| SHORT_IDS = {k: tok(v[0], add_special_tokens=False)["input_ids"] for k, v in SHORT.items()}
|
| items = load_suite(args.items) if args.items else load_suite()
|
| lock = threading.Lock()
|
| stats = {"n": 0, "commit": 0, "fallback": 0, "t0": time.time()}
|
|
|
|
|
| def post(body):
|
| for attempt in range(6):
|
| try:
|
| req = urllib.request.Request(URL, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"})
|
| return json.loads(urllib.request.urlopen(req, timeout=3600).read())
|
| except urllib.error.HTTPError as e:
|
| if e.code == 400:
|
| return None
|
| time.sleep(3)
|
| except Exception:
|
| time.sleep(3)
|
| return None
|
|
|
|
|
| def extract(task, text):
|
| if task == "mmlupro":
|
| m = re.findall(r"answer is:?\s*\(?\**\s*([A-J])\b", text)
|
| a = m[-1] if m else None
|
| elif task == "bbh":
|
| m = re.findall(r"answer is:?\s*(.+)", text, re.I)
|
| a = m[-1] if m else (last_boxed(text) or "")
|
| a = re.sub(r"\\boxed\{(.*)\}", r"\1", a).strip().strip("*$ .").lower() or None
|
| else:
|
| b = last_boxed(text)
|
| a = re.sub(r"\s+", "", b) if b else None
|
| if a is None or "..." in a or a in ("?", "") or len(a) > 80:
|
| return None
|
| return a
|
|
|
|
|
| def snap(gen, k, lo=0):
|
| for j in range(min(k, len(gen)), lo, -1):
|
| if gen[j - 1] in ENDS:
|
| return j
|
| return None
|
|
|
|
|
| def short_probe(ids, gen, cut, task):
|
| suf, n_new, stops, close = SHORT[task]
|
| pr = post({"prompt": ids + gen[:cut] + [THINK_END] + SHORT_IDS[task], "n_predict": n_new, "temperature": 0.0,
|
| "cache_prompt": False, "return_tokens": True, "stop": stops})
|
| if pr is None:
|
| return None, None, 0
|
| c = pr["content"]
|
| if close == "}$":
|
| depth = 1
|
| for j, ch in enumerate(c):
|
| depth += (ch == "{") - (ch == "}")
|
| if depth == 0:
|
| c = c[:j]
|
| break
|
| c = c.strip().rstrip(".").strip() if close == "." else c.strip()
|
| text = suf + c + close
|
| return text, extract(task, text), len(SHORT_IDS[task]) + len(pr["tokens"])
|
|
|
|
|
| def run(rec):
|
| it = items[rec["id"]]
|
| task = it["task"]
|
| ids = tok.apply_chat_template([{"role": "user", "content": it["prompt"]}], tokenize=True,
|
| add_generation_prompt=True, return_dict=False)
|
| body = rec["completion"][len("<think>"):]
|
| thinking = body.split("</think>", 1)[0]
|
| gen = tok(thinking, add_special_tokens=False)["input_ids"]
|
| out = dict(rec)
|
| out["module"] = "none"
|
| if args.only_tasks is not None and task not in args.only_tasks:
|
| return out
|
| if task in SHORT and task in args.commit_tasks:
|
| prev, answers = 0, []
|
| for k in range(args.every, len(gen), args.every):
|
| cut = snap(gen, k, prev)
|
| if cut is None:
|
| continue
|
| prev = cut
|
| text, ans, n = short_probe(ids, gen, cut, task)
|
| answers.append(ans)
|
| tail = answers[-args.agree:]
|
| if len(tail) == args.agree and tail[0] is not None and all(a == tail[0] for a in tail):
|
| out.update(completion="<think>" + tok.decode(gen[:cut]) + "</think>" + text, n_tokens=cut + n,
|
| truncated=False, module="commit", module_cut=cut)
|
| break
|
| if out["module"] == "none" and rec.get("truncated"):
|
| if task in SHORT:
|
| cut = snap(gen, args.max_tokens - 64)
|
| if cut:
|
| text, ans, n = short_probe(ids, gen, cut, task)
|
| if text is not None:
|
| out.update(completion="<think>" + tok.decode(gen[:cut]) + "</think>" + text, n_tokens=cut + n,
|
| truncated=False, module="fallback", module_cut=cut)
|
| else:
|
| cut = snap(gen, args.max_tokens - args.free_reserve)
|
| if cut:
|
| seed = zlib.crc32(rec["id"].encode()) % 1000003 * 7 + rec.get("sample", 0) + 17
|
| pr = post({"prompt": ids + gen[:cut] + [THINK_END], "n_predict": args.free_reserve,
|
| "temperature": 0.6, "top_p": 0.95, "top_k": 50, "min_p": 0.0, "seed": seed,
|
| "cache_prompt": False, "samplers": ["top_k", "top_p", "temperature"], "return_tokens": True})
|
| if pr is not None:
|
| out.update(completion="<think>" + tok.decode(gen[:cut]) + "</think>" + pr["content"],
|
| n_tokens=cut + len(pr["tokens"]), truncated=pr.get("stop_type") == "limit",
|
| module="fallback", module_cut=cut)
|
| with lock:
|
| stats["n"] += 1
|
| stats["commit"] += out["module"] == "commit"
|
| stats["fallback"] += out["module"] == "fallback"
|
| if stats["n"] % 250 == 0:
|
| print(f"{stats['n']} done commits {stats['commit']} fallbacks {stats['fallback']} "
|
| f"{(time.time() - stats['t0']) / 60:.1f} min", flush=True)
|
| return out
|
|
|
|
|
| for src in args.src:
|
| recs = [json.loads(l) for l in open(src, encoding="utf-8")]
|
| dst = os.path.join("results", args.out_tag, os.path.basename(src))
|
| os.makedirs(os.path.dirname(dst), exist_ok=True)
|
| done = set()
|
| if os.path.exists(dst):
|
| done = {(json.loads(l)["id"], json.loads(l)["sample"]) for l in open(dst, encoding="utf-8")}
|
| todo = [r for r in recs if (r["id"], r["sample"]) not in done]
|
| print(src, len(todo), "records to process", flush=True)
|
| with open(dst, "a", encoding="utf-8") as f, ThreadPoolExecutor(args.workers) as ex:
|
| for out in ex.map(run, todo):
|
| f.write(json.dumps(out, ensure_ascii=False) + "\n")
|
| f.flush()
|
| print("done", stats, flush=True)
|
|
|
