| |
| """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) |
| |
| 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()) |
|
|
