| """Commit-on-consensus (CoC), streaming version: an inference-time halting mechanism.
|
|
|
| The main reasoning is ONE streamed request with exactly the sampling settings and seed of gen.py, so without a
|
| commit the output is the ordinary generation. Every --chunk generated tokens (while the model is still thinking) a
|
| probe runs on a separate slot: the reasoning so far is closed at the last sentence end with the model's own closing
|
| pattern ("...." + </think>) and the model writes its final answer greedily. When --agree consecutive probes give the
|
| same extracted answer, the reasoning is committed: the stream is dropped and the output becomes
|
| reasoning[:cut] + </think> + that probe's answer. Code / instruction-following tasks are never probed.
|
|
|
| --shadow: on consensus keep generating to the end and store the committed answer next to the full one, so one run
|
| gives both outcomes on the same trajectory (the full arm is the ordinary baseline generation).
|
| --record-all: never stop probing (with --shadow), store every probe for offline simulation of other settings.
|
|
|
| usage: gen_coc2.py --inp FILE.. --out OUT.jsonl [--max-tokens 16384] [--chunk 1024] [--agree 2] [--shadow]
|
| """
|
| import argparse
|
| import json
|
| import os
|
| import queue
|
| 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
|
|
|
| ap = argparse.ArgumentParser()
|
| ap.add_argument("--inp", nargs="+", required=True)
|
| ap.add_argument("--out", required=True)
|
| ap.add_argument("--n", type=int, default=1)
|
| ap.add_argument("--max-tokens", type=int, default=16384)
|
| ap.add_argument("--chunk", type=int, default=1024)
|
| ap.add_argument("--agree", type=int, default=2)
|
| ap.add_argument("--probe-tokens", type=int, default=400)
|
| ap.add_argument("--shadow", action="store_true")
|
| ap.add_argument("--free-probes", action="store_true", help="let the model write its answer freely (slow)")
|
| ap.add_argument("--fallback", action="store_true",
|
| help="if the reasoning hits the token limit, close it and answer instead of being cut off "
|
| "(short-answer tasks: forced short answer within the last 64 tokens; other tasks: thinking is "
|
| "closed --free-reserve tokens before the limit and the answer is generated normally)")
|
| ap.add_argument("--free-reserve", type=int, default=1536)
|
| ap.add_argument("--commit-tasks", nargs="*", default=None,
|
| help="tasks where early commitment is allowed (default: all probed tasks; the released module uses "
|
| "mmlupro bbh, chosen on dev: numeric math answers get revised late and commits hurt there)")
|
| ap.add_argument("--record-all", action="store_true")
|
| ap.add_argument("--workers", type=int, default=24, help="main slots 0..W-1, probe slots W..2W-1 (server -np 2W)")
|
| ap.add_argument("--temperature", type=float, default=0.6)
|
| ap.add_argument("--top-p", type=float, default=0.95)
|
| ap.add_argument("--top-k", type=int, default=50)
|
| ap.add_argument("--seed", type=int, default=0)
|
| ap.add_argument("--port", type=int, default=8080)
|
| ap.add_argument("--tokenizer", default="models/LFM3-1.7B")
|
| args = ap.parse_args()
|
|
|
| THINK_END = 124902
|
| MATH_SUFFIX = "The final answer is $\\boxed{"
|
| URL = f"http://localhost:{args.port}/completion"
|
| tok = AutoTokenizer.from_pretrained(args.tokenizer)
|
| MATH_SUFFIX_IDS = tok(MATH_SUFFIX, add_special_tokens=False)["input_ids"]
|
|
|
| SHORT = {"math": (MATH_SUFFIX, 48, ["}$", "}\n"], "}$"), "mmlupro": ("The answer is (", 3, [")"], ")."),
|
| "bbh": ("The answer is ", 24, ["\n", ".\n"], ".")}
|
| SHORT["math500"] = SHORT["aime"] = SHORT["gsm8k"] = SHORT["math"]
|
| SHORT_IDS = {k: tok(v[0], add_special_tokens=False)["input_ids"] for k, v in SHORT.items()}
|
| 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_)
|
|
|
| items = []
|
| for path in args.inp:
|
| items += [json.loads(l) for l in open(path, encoding="utf-8")]
|
| done = set()
|
| if os.path.exists(args.out):
|
| done = {(json.loads(l)["id"], json.loads(l)["sample"]) for l in open(args.out, encoding="utf-8")}
|
| jobs = [(it, s) for it in items for s in range(args.n) if (it["id"], s) not in done]
|
| print(len(jobs), "jobs", flush=True)
|
| os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
|
| out_f = open(args.out, "a", encoding="utf-8")
|
| lock = threading.Lock()
|
| slots = queue.Queue()
|
| for w in range(args.workers):
|
| slots.put(w)
|
| stats = {"n": 0, "commit": 0, "t0": time.time(), "tok": 0}
|
|
|
|
|
| def post(body):
|
| for attempt in range(5):
|
| 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 Exception:
|
| if attempt == 4:
|
| raise
|
| time.sleep(2)
|
|
|
|
|
| def extract(task, text):
|
| a = _extract(task, text)
|
| if a is None or "..." in a or a in ("?", "") or len(a) > 80:
|
| return None
|
| return a
|
|
|
|
|
| def _extract(task, text):
|
| a = _extract(task, text)
|
| if a is None or "..." in a or a in ("?", "") or len(a) > 80:
|
| return None
|
| return a
|
|
|
|
|
| def _extract(task, text):
|
| if task == "mmlupro":
|
| m = re.findall(r"answer is:?\s*\(?\**\s*([A-J])\b", text) or re.findall(r"\\boxed\{\(?([A-J])\)?\}", text)
|
| return m[-1] if m else None
|
| if 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()
|
| return a or None
|
| b = last_boxed(text)
|
| if b is None and task == "gsm8k":
|
| nums = re.findall(r"-?\d[\d,]*\.?\d*", text)
|
| b = nums[-1] if nums else None
|
| return re.sub(r"\s+", "", b) if b else None
|
|
|
|
|
| def snap(gen, lo):
|
| for j in range(len(gen), lo, -1):
|
| if gen[j - 1] in ENDS:
|
| return j
|
| return None
|
|
|
|
|
| def run(job):
|
| it, s = job
|
| w = slots.get()
|
| try:
|
| ids = tok.apply_chat_template([{"role": "user", "content": it["prompt"]}], tokenize=True,
|
| add_generation_prompt=True, return_dict=False)
|
| task = it["task"]
|
| probing = not task.startswith("code") and task not in ("humaneval", "mbpp", "ifeval")
|
| short = task in SHORT and not args.free_probes
|
| if args.commit_tasks is not None and task not in args.commit_tasks:
|
| probing = False
|
| body = {"prompt": ids, "n_predict": args.max_tokens, "temperature": args.temperature, "top_p": args.top_p,
|
| "top_k": args.top_k, "min_p": 0.0, "repeat_penalty": 1.0, "presence_penalty": 0.0,
|
| "frequency_penalty": 0.0, "dry_multiplier": 0.0, "xtc_probability": 0.0, "typical_p": 1.0,
|
| "seed": args.seed * 1000003 + zlib.crc32(it["id"].encode()) % 1000003 * 7 + s,
|
| "cache_prompt": False, "samplers": ["top_k", "top_p", "temperature"], "return_tokens": True,
|
| "stream": True, "id_slot": w}
|
| req = urllib.request.Request(URL, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"})
|
| resp = urllib.request.urlopen(req, timeout=3600)
|
| gen, final, next_probe, prev_cut = [], None, args.chunk, 0
|
| probe_q, probe_res = queue.Queue(), []
|
|
|
| def prober():
|
| while True:
|
| job_ = probe_q.get()
|
| if job_ is None:
|
| return
|
| cut_, snap_ = job_
|
| if short:
|
| suf, n_new, stops, close = SHORT[task]
|
| pr = post({"prompt": ids + snap_ + [THINK_END] + SHORT_IDS[task], "n_predict": n_new,
|
| "temperature": 0.0, "cache_prompt": False, "id_slot": args.workers + w,
|
| "return_tokens": True, "stop": stops})
|
| 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
|
| toks = SHORT_IDS[task] + pr["tokens"]
|
| else:
|
| pr = post({"prompt": ids + snap_ + [THINK_END], "n_predict": args.probe_tokens, "temperature": 0.0,
|
| "cache_prompt": False, "id_slot": args.workers + w, "return_tokens": True})
|
| text, toks = pr["content"], pr["tokens"]
|
| probe_res.append({"cut": cut_, "ans": extract(task, text), "text": text, "toks": toks})
|
|
|
| def fallback_answer(gen_):
|
| """answer within the budget when the reasoning was cut off by the limit"""
|
| if short:
|
| cut_ = snap(gen_[: args.max_tokens - 64], 0)
|
| if cut_ is None:
|
| return None
|
| 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, "id_slot": args.workers + w,
|
| "return_tokens": True, "stop": stops})
|
| 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 {"cut": cut_, "text": text_, "n": len(SHORT_IDS[task]) + len(pr["tokens"])}
|
| cut_ = snap(gen_[: args.max_tokens - args.free_reserve], 0)
|
| if cut_ is None:
|
| return None
|
| pr = post({"prompt": ids + gen_[:cut_] + [THINK_END], "n_predict": args.free_reserve,
|
| "temperature": args.temperature, "top_p": args.top_p, "top_k": args.top_k, "min_p": 0.0,
|
| "seed": body["seed"] + 17, "cache_prompt": False, "id_slot": args.workers + w,
|
| "samplers": ["top_k", "top_p", "temperature"], "return_tokens": True})
|
| return {"cut": cut_, "text": pr["content"], "n": len(pr["tokens"])}
|
|
|
| pth = threading.Thread(target=prober, daemon=True)
|
| pth.start()
|
| commit, checked, closed = None, 0, False
|
| while True:
|
| line = resp.readline()
|
| if not line:
|
| break
|
| line = line.decode("utf-8").strip()
|
| if not line.startswith("data:"):
|
| continue
|
| ev = json.loads(line[5:])
|
| gen += ev.get("tokens", [])
|
| if ev.get("stop"):
|
| final = ev
|
| break
|
| if THINK_END in ev.get("tokens", []):
|
| closed = True
|
| if probing and not closed and (commit is None or args.record_all) and len(gen) >= next_probe:
|
| next_probe += args.chunk
|
| cut = snap(gen, prev_cut)
|
| if cut:
|
| prev_cut = cut
|
| probe_q.put((cut, list(gen[:cut])))
|
|
|
| while commit is None and checked < len(probe_res):
|
| checked += 1
|
| tail = probe_res[:checked][-args.agree:]
|
| if len(tail) == args.agree and tail[0]["ans"] is not None and all(p["ans"] == tail[0]["ans"] for p in tail):
|
| commit = tail[-1]
|
| if commit is not None and not args.shadow:
|
| resp.close()
|
| break
|
| probe_q.put(None)
|
| if commit is None or args.shadow:
|
| pth.join(timeout=900)
|
|
|
| while commit is None and checked < len(probe_res):
|
| checked += 1
|
| tail = probe_res[:checked][-args.agree:]
|
| if len(tail) == args.agree and tail[0]["ans"] is not None and all(p["ans"] == tail[0]["ans"] for p in tail):
|
|
|
| commit = tail[-1]
|
| if gen and gen[-1] == 124900:
|
| gen = gen[:-1]
|
| full_text = "<think>" + tok.decode(gen)
|
| rec = {"id": it["id"], "task": task, "sample": s}
|
| truncated = bool(final) and final.get("stop_type") == "limit"
|
| if final is None and commit is None:
|
| truncated = len(gen) >= args.max_tokens
|
| fb = fallback_answer(gen) if (args.fallback and commit is None and truncated) else None
|
| coc = None
|
| if commit is not None:
|
| coc = {"completion": "<think>" + tok.decode(gen[:commit["cut"]]) + "</think>" + commit["text"],
|
| "n_tokens": commit["cut"] + len(commit["toks"]), "cut": commit["cut"]}
|
| elif fb is not None:
|
| coc = {"completion": "<think>" + tok.decode(gen[:fb["cut"]]) + "</think>" + fb["text"],
|
| "n_tokens": fb["cut"] + fb["n"], "cut": fb["cut"]}
|
| rec["fallback"] = True
|
| if args.shadow:
|
| rec.update(completion=full_text, n_tokens=len(gen),
|
| truncated=bool(final) and final.get("stop_type") == "limit",
|
| stop_type=final.get("stop_type") if final else None)
|
| if coc:
|
| rec.update(coc_completion=coc["completion"], coc_n_tokens=coc["n_tokens"], coc_cut=coc["cut"])
|
| if args.record_all:
|
| rec["probe_log"] = [{k: p[k] for k in ("cut", "ans", "text")} | {"n": len(p["toks"])}
|
| for p in probe_res]
|
| elif coc:
|
| rec.update(completion=coc["completion"], n_tokens=coc["n_tokens"], truncated=False, stop_type="commit",
|
| coc_cut=coc["cut"])
|
| else:
|
| rec.update(completion=full_text, n_tokens=len(gen),
|
| truncated=bool(final) and final.get("stop_type") == "limit",
|
| stop_type=final.get("stop_type") if final else None, coc_cut=None)
|
| rec["probes"] = len(probe_res)
|
| with lock:
|
| out_f.write(json.dumps(rec, ensure_ascii=False) + "\n")
|
| out_f.flush()
|
| stats["n"] += 1
|
| stats["commit"] += commit is not None
|
| stats["tok"] += len(gen)
|
| if stats["n"] % 20 == 0:
|
| el = time.time() - stats["t0"]
|
| print(f"{stats['n']}/{len(jobs)} commits {stats['commit']} {stats['tok'] / el:.0f} gen tok/s "
|
| f"{el / 60:.1f} min", flush=True)
|
| finally:
|
| slots.put(w)
|
|
|
|
|
| with ThreadPoolExecutor(args.workers) as ex:
|
| list(ex.map(run, jobs))
|
| print("done", stats, flush=True)
|
|
|
