File size: 5,645 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 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | #!/usr/bin/env python3
"""Heuristic 1.1 race loop: parallel thought paths, judged as they finish.
Replaces the sequential best_of_n + judge-everything-at-the-end loop:
1. The greedy anchor (temperature 0) is generated and judged FIRST. If the
judge scores it >= stop_confidence and it finished, it wins immediately --
no sampled paths are ever generated.
2. Otherwise n_paths-1 sampled paths are generated IN PARALLEL
(ThreadPoolExecutor over the thinker's HTTP endpoint).
3. Each path is judged the moment it completes. The first FINISHED path to
score >= stop_confidence wins on the spot; the race is cut off at that
instant -- queued paths are cancelled and race() returns immediately.
(Honest cutoff semantics: threads already running a blocking HTTP
generation can't be killed, so they finish silently in the background
and their results are discarded. Nothing waits for them.)
4. If nothing clears the bar, the best judged finished path wins. Finished
paths always outrank truncated ones, whatever the score.
Only the thinker's .generate() and the judge's .score() from the metacog
package are used -- the MetaCog controller is not touched.
Caveat: decider-4b's confidence scale is not calibrated (values above 1.0
have been observed), so stop_confidence is a knob to measure, not a true
probability. The 0.95 default in the MetaCog repo was measured on real Jev
(226/229 correct at >= 0.95); 0.80 here is the operator's directive and is
untested on decider-4b.
"""
from __future__ import annotations
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
@dataclass
class RaceResult:
answer: str
finished: bool
confidence: float
thinker_tokens: int = 0
judge_calls: int = 0
paths_submitted: int = 0
paths_judged: int = 0
early_stopped: bool = False
winner_source: str = "sample" # "greedy" | "sample"
seconds: float = 0.0
def _judge_one(judge, problem: str, text: str) -> float:
"""Judge's P(correct) for a single path; 0.0 if the judge call fails."""
try:
return float(judge.score(problem, [text]).confidence)
except Exception:
return 0.0
def race(
problem: str,
thinker,
judge,
n_paths: int = 8,
stop_confidence: float = 0.80,
max_tokens: int = 2048,
temperature: float = 0.8,
max_workers: int = 8,
) -> RaceResult:
t0 = time.time()
lock = threading.Lock()
thinker_tokens = 0
judge_calls = 0
def gen(**kw):
nonlocal thinker_tokens
gens = thinker.generate(problem, "", max_tokens=max_tokens, **kw)
with lock:
for g in gens:
thinker_tokens += g.tokens or 0
return gens
def judge_of(text: str) -> float:
nonlocal judge_calls
with lock:
judge_calls += 1
return _judge_one(judge, problem, text)
# --- 1. greedy anchor first -------------------------------------------
greedy = gen(n=1, temperature=0.0)[0]
g_conf = judge_of(greedy.text)
pool = [(greedy.text, greedy.finished, g_conf, "greedy")]
if greedy.finished and g_conf >= stop_confidence:
return RaceResult(
answer=greedy.text, finished=True, confidence=g_conf,
thinker_tokens=thinker_tokens, judge_calls=judge_calls,
paths_submitted=1, paths_judged=1,
early_stopped=True, winner_source="greedy",
seconds=time.time() - t0,
)
# --- 2. sampled paths in parallel, judged as they finish ---------------
n_sampled = max(n_paths - 1, 0)
early = False
winner = None
# NOTE: not a `with` block on purpose. Exiting a ThreadPoolExecutor
# context manager WAITS for every worker, which would defeat the early
# stop. Instead we shut down without waiting: queued futures are
# cancelled, race() returns at once, and any thread already inside a
# blocking thinker.generate() finishes on its own in the background
# (its result is discarded). The thinker's generate() is MetaCog's
# library code and is deliberately not modified, so true GPU-side
# abort of an in-flight generation isn't available -- this is a
# stop-waiting cutoff, which is what cuts user-facing latency.
ex = ThreadPoolExecutor(max_workers=max(1, min(max_workers, n_sampled or 1)))
try:
futs = [ex.submit(gen, n=1, temperature=temperature)
for _ in range(n_sampled)]
for fut in as_completed(futs):
try:
g = fut.result()[0]
except Exception:
continue # a dead path just never gets judged
conf = judge_of(g.text)
pool.append((g.text, g.finished, conf, "sample"))
if g.finished and conf >= stop_confidence:
winner = (g.text, True, conf, "sample")
early = True
break
finally:
ex.shutdown(wait=False, cancel_futures=True)
# --- 3. no early stop: best judged finished path wins -------------------
if winner is None:
finished = [p for p in pool if p[1]]
text, fin, conf, src = max(finished or pool, key=lambda p: p[2])
winner = (text, fin, conf, src)
text, fin, conf, src = winner
with lock:
toks, calls = thinker_tokens, judge_calls
return RaceResult(
answer=text, finished=fin, confidence=conf,
thinker_tokens=toks, judge_calls=calls,
paths_submitted=1 + n_sampled, paths_judged=len(pool),
early_stopped=early, winner_source=src,
seconds=time.time() - t0,
)
|