File size: 4,605 Bytes
cd34530 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | #!/usr/bin/env python3
"""The merge: bonsai thinks, decider-4b judges, MetaCog wires them into one model.
thinker: bonsai (bjev llama.cpp, OpenAI chat) .... http://localhost:8010
judge: decider-4b (/v1/systemone shim) ......... http://localhost:8008
Usage:
.venv/bin/python merge.py --problem "What is 17*23?"
.venv/bin/python merge.py --problems problems.jsonl --out results.jsonl
"""
from __future__ import annotations
import argparse
import json
import sys
import time
import re
def clean_answer(text: str) -> str:
"""Strip <think>...</think> reasoning blocks the thinker may leak."""
return re.sub(r"<think>.*?</think>\s*", "", text, flags=re.S).strip()
from metacog import Config, MetaCog, OpenAICompatThinker, SystemOneJudge
import os
THINKER_URL = os.environ.get("THINKER_URL", "http://localhost:8010")
THINKER_MODEL = os.environ.get("THINKER_MODEL", "bonsai")
JUDGE_URL = os.environ.get("JUDGE_URL", "http://localhost:8008")
def build_thinker_judge():
thinker = OpenAICompatThinker(THINKER_URL, model=THINKER_MODEL)
# non-TypeSafe base_url -> no auth header; the local shim is open
judge = SystemOneJudge(base_url=JUDGE_URL)
return thinker, judge
def build(n_paths: int = 8, no_cascade: bool = False) -> MetaCog:
thinker, judge = build_thinker_judge()
config = Config(
mode="best_of_n",
strategy="noul",
n_paths=n_paths,
greedy_anchor=True,
cascade_confidence=None if no_cascade else 0.95,
max_tokens=2048,
temperature=0.8,
)
return MetaCog(thinker, judge, config)
def main() -> int:
ap = argparse.ArgumentParser(description="bonsai + decider-4b MetaCog merge")
src = ap.add_mutually_exclusive_group(required=True)
src.add_argument("--problem")
src.add_argument("--problems")
ap.add_argument("--n-paths", type=int, default=8)
ap.add_argument("--no-cascade", action="store_true")
ap.add_argument("--stop-confidence", type=float, default=0.80,
help="race loop early-stop bar (default 0.80)")
ap.add_argument("--no-race", action="store_true",
help="use the old sequential MetaCog best_of_n instead of the race loop")
ap.add_argument("--out", default="runs/merge_results.jsonl")
args = ap.parse_args()
if args.problem:
problems = [args.problem]
else:
problems = [json.loads(l)["problem"] for l in open(args.problems)
if l.strip()]
if not problems:
print("no problems", file=sys.stderr)
return 1
mc = build(args.n_paths, args.no_cascade)
use_race = not args.no_race
thinker = judge = None
if use_race:
from race import race
thinker, judge = build_thinker_judge()
t0 = time.time()
with open(args.out, "w") as out:
for i, problem in enumerate(problems):
t1 = time.time()
if use_race:
r = race(problem, thinker, judge, n_paths=args.n_paths,
stop_confidence=args.stop_confidence)
answer, finished = clean_answer(r.answer), r.finished
conf, toks, jcalls, dt = (r.confidence, r.thinker_tokens,
r.judge_calls, r.seconds)
extra = {"early_stopped": r.early_stopped,
"winner_source": r.winner_source,
"paths_judged": r.paths_judged}
else:
result = mc.run(problem)
dt = time.time() - t1
rnd = result.trace.rounds[-1] if result.trace.rounds else None
answer, finished = clean_answer(result.answer), result.finished
conf = round(rnd.verdict.confidence, 4) if rnd else 0.0
toks, jcalls = result.trace.thinker_tokens, result.trace.judge_calls
extra = {}
row = {
"i": i,
"problem": problem,
"answer": answer,
"finished": finished,
"judge_confidence": round(conf, 4),
"thinker_tokens": toks,
"judge_calls": jcalls,
"seconds": round(dt, 2),
**extra,
}
out.write(json.dumps(row) + "\n")
out.flush()
print(f"[{i+1}/{len(problems)}] conf={conf:.2f} "
f"{dt:.1f}s: {answer[:120]}", flush=True)
print(f"done: {len(problems)} problems in {time.time()-t0:.0f}s -> {args.out}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|