| """Commit-on-consensus (CoC) generation: an inference-time halting mechanism.
|
|
|
| The reasoning is generated in chunks (--chunk tokens, KV cache of the slot is reused). After every chunk, while the
|
| model is still thinking, a probe closes the thinking at the last sentence end (the model's own closing pattern:
|
| "...." + </think>) and lets the model write its final answer greedily on a separate slot. When two consecutive probes
|
| give the same extracted answer, the reasoning is committed there: the output becomes reasoning[:cut] + </think> +
|
| that probe's answer. Code tasks are never probed (no cheap answer equivalence).
|
|
|
| --no-probe runs the identical chunked generation without probes: the control for measuring CoC's own contribution
|
| (both runs are token-for-token identical until CoC commits).
|
|
|
| Output records have the same format as gen.py, plus "coc_cut" (None if not committed).
|
| usage: gen_coc.py --inp data/quick/mini.jsonl --out results/coc/mini.jsonl --max-tokens 6144 [--no-probe]
|
| """
|
| 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("--max-tokens", type=int, default=16384)
|
| ap.add_argument("--chunk", type=int, default=1024)
|
| ap.add_argument("--agree", type=int, default=2, help="consecutive identical probe answers needed to commit")
|
| ap.add_argument("--probe-tokens", type=int, default=400)
|
| ap.add_argument("--no-probe", action="store_true")
|
| ap.add_argument("--record-all", action="store_true",
|
| help="shadow mode that never stops probing: stores every probe (cut, answer, text) for offline "
|
| "simulation of other chunk/agreement settings")
|
| ap.add_argument("--shadow", action="store_true",
|
| help="on consensus, record the committed answer but keep generating to the end: one run yields both "
|
| "outcomes on the same trajectory (exact pairing) and also training data")
|
| 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, EOS = 124902, 124900
|
| 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_)
|
|
|
| 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"] for l in open(args.out, encoding="utf-8")}
|
| items = [it for it in items if it["id"] not in done]
|
| print(len(items), "items", 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(f"http://localhost:{args.port}/completion", 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):
|
| """answer identity used for the consensus test"""
|
| if task in ("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(it):
|
| 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"]
|
| probe = not args.no_probe and not task.startswith("code") and task not in ("humaneval", "mbpp", "ifeval")
|
| base_seed = args.seed * 1000003 + zlib.crc32(it["id"].encode()) % 1000003 * 7
|
| gen, stop, last_ans, prev_cut, commit, n_chunk, n_probe = [], None, None, 0, None, 0, 0
|
| plog = []
|
| while True:
|
| n_pred = min(args.chunk, args.max_tokens - len(gen))
|
| if n_pred <= 0:
|
| stop = "limit"
|
| break
|
| res = post({"prompt": ids + gen, "n_predict": n_pred, "temperature": args.temperature,
|
| "top_p": args.top_p, "top_k": args.top_k, "min_p": 0.0, "repeat_penalty": 1.0,
|
| "seed": base_seed + 7919 * n_chunk, "cache_prompt": True, "id_slot": w,
|
| "samplers": ["top_k", "top_p", "temperature"], "return_tokens": True})
|
| n_chunk += 1
|
| gen += res["tokens"]
|
| if res.get("stop_type") == "eos" or (res["tokens"] and res["tokens"][-1] == EOS):
|
| stop = "eos"
|
| break
|
| if not probe or THINK_END in gen:
|
| continue
|
| cut = snap(gen, prev_cut)
|
| if cut is None:
|
| continue
|
| prev_cut = cut
|
| pr = post({"prompt": ids + gen[:cut] + [THINK_END], "n_predict": args.probe_tokens, "temperature": 0.0,
|
| "cache_prompt": False, "id_slot": args.workers + w, "return_tokens": True})
|
| n_probe += 1
|
| ans = extract(task, pr["content"])
|
| if args.record_all:
|
| plog.append({"cut": cut, "ans": ans, "text": pr["content"], "n": len(pr["tokens"])})
|
| continue
|
| if ans is not None and ans == last_ans:
|
| commit = (cut, pr)
|
| if not args.shadow:
|
| break
|
| probe = False
|
| continue
|
| last_ans = ans
|
| shadow = {}
|
| if commit and args.shadow:
|
| cut, pr = commit
|
| shadow = {"coc_completion": "<think>" + tok.decode(gen[:cut]) + "</think>" + pr["content"],
|
| "coc_n_tokens": cut + len(pr["tokens"]), "coc_tokens": gen[:cut] + [THINK_END] + pr["tokens"]}
|
| commit = None
|
| if commit:
|
| cut, pr = commit
|
| text = tok.decode(gen[:cut]) + "</think>" + pr["content"]
|
| n_tok, trunc, stop_type = cut + len(pr["tokens"]), False, "commit"
|
| else:
|
| if gen and gen[-1] == EOS:
|
| gen = gen[:-1]
|
| text = tok.decode(gen)
|
| n_tok, trunc, stop_type = len(gen), stop == "limit", stop
|
| rec = {"id": it["id"], "task": task, "sample": 0, "completion": "<think>" + text, "n_tokens": n_tok,
|
| "truncated": trunc, "stop_type": stop_type, "coc_cut": commit[0] if commit else None,
|
| "probes": n_probe, "gen_tokens": len(gen), **shadow}
|
| if args.record_all:
|
| rec["probe_log"] = plog
|
| if shadow:
|
| rec["coc_cut"] = shadow["coc_n_tokens"]
|
| with lock:
|
| out_f.write(json.dumps(rec, ensure_ascii=False) + "\n")
|
| out_f.flush()
|
| stats["n"] += 1
|
| stats["commit"] += bool(commit) or bool(shadow)
|
| stats["tok"] += len(gen)
|
| if stats["n"] % 20 == 0:
|
| el = time.time() - stats["t0"]
|
| print(f"{stats['n']}/{len(items)} 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, items))
|
| print("done", stats, flush=True)
|
|
|
