File size: 4,772 Bytes
63ae4a0 | 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 | """Decision Index engine with a learned router between the two v6 levels: every question gets the direct slot pass first; the router head (tools/router_fit_final.py,
CKPT/router-direct-brief.npz: P(direct wrong) x P(brief right | direct wrong) from the direct pass pooled slot hidden state + confidence) decides per question whether the brief
rationale path (BriefSlotEngine) is run instead. Per-option (parallel) and non-choice questions stay direct. Escalated questions of a request are batched through the brief path together.
python -m decision_index run --engine decision_engine_v6r:RoutedSlotEngine --option model=Qwen/Qwen3.5-9B --option run=RUN --option ckpt=final [--option tau=0.02] ...
Options: those of BriefSlotEngine (context is fixed to brief for the escalated path), tau (override of the stored threshold; `inf` = never escalate, `-inf` = always).
"""
import os
import sys
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import numpy as np # noqa: E402
import torch # noqa: E402
import answer_slot as AS # noqa: E402
from decision_engine_dense import DenseSlotEngine # noqa: E402
from decision_engine_v6 import BriefSlotEngine # noqa: E402
def _lg(p):
p = np.clip(p, 1e-4, 1 - 1e-4)
return np.log(p / (1 - p))
class RoutedSlotEngine(BriefSlotEngine):
name = "dense-slot-routed"
latency = ("In-process request wall time (Hugging Face path): direct slot pass for every question, router head, and the brief path (batched rationale generation + adapted slot pass) "
"only for the questions the router escalates. Not the board latency path; see tools/vllm_requests_0f.py for decode timing.")
def __init__(self, tau=None, **options):
options["context"] = "brief"
options.setdefault("temperature", "none")
super().__init__(**options)
run, ck = options["run"], options.get("ckpt", "best")
blob = np.load(os.path.join(ck if os.path.isdir(ck) else os.path.join(run, ck), "router-direct-brief.npz"))
self.knots = blob["knots"]; self.tau = float(blob["tau"]) if tau in (None, "") else float(tau)
self.W = {p: {k: blob[f"{p}_{k}"] for k in ("mh", "sh", "mc", "sc", "w")} for p in ("w", "f")}
self.provenance.update({"router": f"direct->brief head, tau {self.tau}"})
def _score(self, pooled, conf, k):
x = _lg(np.array([conf], dtype=np.float64))
C = np.concatenate([x[:, None], np.maximum(0, x[:, None] - self.knots[None]), np.log([[float(k)]]), np.log([[float(k)]]) * x[:, None]], 1)
h = pooled.double().cpu().numpy()[None]
f = lambda M: float(np.clip((np.concatenate([(h - M["mh"]) / M["sh"], (C - M["mc"]) / M["sc"], np.ones((1, 1))], 1) @ M["w"])[0], 0, 1))
return f(self.W["w"]) * f(self.W["f"])
def __call__(self, state, questions):
answers, raw, esc = {}, {}, {}
tokens = 0
try:
for qid, q in questions.items():
keys, r = self._render(state, q)
logits, pooled = self._logits(r, None)
p = torch.softmax(logits, -1).cpu().numpy().astype(np.float64)
p = p / p.sum()
score = None
if q["type"] == "choice" and AS.slot_mode(r["qtype"], len(keys)) != "parallel":
score = self._score(pooled, float(p.max()), len(keys))
if score >= self.tau:
esc[qid] = q
raw[qid] = {"routed": "brief", "router_score": round(score, 5), "direct_conf": round(float(p.max()), 4)}
continue
probs = dict(zip(keys, (float(x) for x in p)))
answers[qid] = {"type": "choice", "choice": max(keys, key=lambda x: probs[x]), "probabilities": probs} if q["type"] == "choice" else {"type": "noul", "noul": probs["true"]}
raw[qid] = {"routed": "direct", "router_score": None if score is None else round(score, 5), "prompt_tokens": len(r["ids"]), "direct_conf": round(float(p.max()), 4)}
tokens += len(r["ids"])
except torch.OutOfMemoryError:
torch.cuda.empty_cache()
try:
from decision_index.engines.base import Unsupported
except ImportError:
from engine_base import Unsupported
raise Unsupported("out of memory on a long prompt")
if esc:
resp_b, raw_b = BriefSlotEngine.__call__(self, state, esc)
answers.update(resp_b["answers"]); tokens += resp_b["usage"]["input_tokens"]
for qid, v in raw_b.items():
raw[qid].update(v)
return {"model": self.model_id, "answers": {qid: answers[qid] for qid in questions}, "usage": {"input_tokens": tokens}}, raw
|