| """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 |
| import torch |
|
|
| import answer_slot as AS |
| from decision_engine_dense import DenseSlotEngine |
| from decision_engine_v6 import BriefSlotEngine |
|
|
|
|
| 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 |
|
|
