| """The inference engine behind the server: decisions, generation, and decide-then-generate. |
| |
| One model in memory. Decisions prefill the state once with the base (adapter off), then run each |
| question suffix with the adapter on against that cache. Generation runs with the adapter off. |
| `auto` decides first and, when the router flags the question, decodes from the same state cache, |
| so escalation never prefills the state a second time. |
| |
| Requests use TypeSafe-style question types: noul (yes/no), choice (named options, at most 16) |
| and score (ordered levels, at most 6). |
| """ |
|
|
| import json |
| import threading |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| from dmm import lora |
| from dmm.artifacts import read_router, router_file, weights_file |
| from dmm.cache import StateCache |
| from dmm.head import DecisionHead, option_logits |
| from dmm.router import Router, confidence_features |
| from dmm.template import Template, chat_text, decision_text |
|
|
|
|
| def to_item(state: str, name: str, q: dict) -> dict: |
| """A TypeSafe-style question as one of our decision items.""" |
| kind = q["type"] |
| text = str(q.get("instructions") or "").strip() |
| if kind == "noul": |
| return { |
| "id": name, |
| "kind": "yesno", |
| "state": state, |
| "question": text, |
| "options": ["no", "yes"], |
| "answer": None, |
| } |
| if kind == "choice": |
| names = list(q["criteria"]) |
| if not 1 <= len(names) <= 16: |
| raise ValueError(f"{name}: choice takes 1 to 16 options, got {len(names)}") |
| opts = [n if not q["criteria"][n] else f"{n}: {q['criteria'][n]}" for n in names] |
| return { |
| "id": name, |
| "kind": "choice", |
| "state": state, |
| "question": text, |
| "options": opts, |
| "answer": None, |
| "_names": names, |
| } |
| if kind == "score": |
| levels = list(q["criteria"]) |
| if not 2 <= len(levels) <= 6: |
| raise ValueError(f"{name}: score takes 2 to 6 levels, got {len(levels)}") |
| legend = "; ".join(f"{i} = {lv}" for i, lv in enumerate(levels)) |
| return { |
| "id": name, |
| "kind": "score", |
| "state": state, |
| "question": f"{text} ({legend})", |
| "options": [str(i) for i in range(len(levels))], |
| "answer": None, |
| } |
| raise ValueError(f"{name}: unknown question type {kind!r}") |
|
|
|
|
| class Engine: |
| def __init__( |
| self, |
| model, |
| tokenizer, |
| head: DecisionHead, |
| temperatures: dict, |
| router: Router | None = None, |
| router_threshold: float = 1.0, |
| ): |
| self.model, self.tok, self.head = model, tokenizer, head |
| self.tpl = Template(tokenizer) |
| self.temperatures, self.router, self.threshold = temperatures, router, router_threshold |
| self.pad = tokenizer.pad_token_id or tokenizer.eos_token_id |
| self.lock = threading.Lock() |
| self.model.eval() |
|
|
| @classmethod |
| def from_artifacts(cls, model, tokenizer, cfg: dict, ckpt_dir: str, router_dir: str | None): |
| from train.train_decision import load_checkpoint, prepare |
|
|
| head, _ = prepare(model, tokenizer, cfg) |
| load_checkpoint(model, head, weights_file(ckpt_dir)) |
| temps = json.loads(Path(f"{ckpt_dir}/temperatures.json").read_text()) |
| router, threshold = None, 1.0 |
| if router_dir: |
| router = Router(model.config.hidden_size) |
| router.load_state_dict(read_router(router_file(router_dir))) |
| threshold = json.loads(Path(f"{router_dir}/operating_point.json").read_text())[ |
| "router_threshold" |
| ] |
| return cls(model, tokenizer, head.float(), temps, router, threshold) |
|
|
| |
|
|
| @torch.no_grad() |
| def _decide_on(self, cache: StateCache, item: dict) -> dict: |
| h = cache.run_suffix(self.tpl.encode(decision_text(item))) |
| z = option_logits(self.head(h[None].to(self.head.proj.weight.device)), [item])[0] |
| z = z.float() / self.temperatures.get(item["kind"], 1.0) |
| p = torch.softmax(z, -1).cpu().numpy().astype(np.float64) |
| score = None |
| if self.router is not None: |
| hf = h.float().cpu().numpy()[None] |
| score = float(self.router.score(hf, confidence_features([p]))[0]) |
| return { |
| "probs": p, |
| "router_score": score, |
| "escalate": score is not None and score >= self.threshold, |
| } |
|
|
| def decide(self, state: str, questions: dict[str, dict]) -> dict: |
| items = {name: to_item(state, name, q) for name, q in questions.items()} |
| with self.lock: |
| cache = StateCache(self.model, self.tpl.state_ids(next(iter(items.values())))) |
| raw = {name: self._decide_on(cache, it) for name, it in items.items()} |
| return {name: self._answer(items[name], r) for name, r in raw.items()} |
|
|
| @staticmethod |
| def _answer(item: dict, r: dict) -> dict: |
| p = r["probs"] |
| out = { |
| "type": {"yesno": "noul", "choice": "choice", "score": "score"}[item["kind"]], |
| "router_score": r["router_score"], |
| "needs_generation": r["escalate"], |
| } |
| if item["kind"] == "yesno": |
| out |= { |
| "noul": float(p[1]), |
| "probabilities": {"false": float(p[0]), "true": float(p[1])}, |
| } |
| elif item["kind"] == "choice": |
| names = item["_names"] |
| out |= { |
| "choice": names[int(np.argmax(p))], |
| "probabilities": {n: float(x) for n, x in zip(names, p, strict=True)}, |
| } |
| else: |
| out |= { |
| "score": float(np.dot(p, np.arange(len(p)))), |
| "probabilities": {str(i): float(x) for i, x in enumerate(p)}, |
| } |
| return out |
|
|
| |
|
|
| @torch.no_grad() |
| def _generate( |
| self, ids: list[int], max_new_tokens: int, cache: StateCache | None = None, streamer=None |
| ): |
| x = torch.tensor([ids], device=self.model.device) |
| kw = {"past_key_values": cache.kv} if cache is not None else {} |
| if streamer is not None: |
| kw["streamer"] = streamer |
| with lora.disabled(self.model): |
| out = self.model.generate( |
| x, |
| attention_mask=torch.ones_like(x), |
| do_sample=False, |
| max_new_tokens=max_new_tokens, |
| pad_token_id=self.pad, |
| **kw, |
| ) |
| return out[0, len(ids) :].tolist() |
|
|
| def _stream(self, ids: list[int], max_new_tokens: int, cache: StateCache | None = None): |
| """Generate as _generate does, yielding {"event": "text"} pieces as they decode. |
| |
| Returns the generated token ids (the generator's return value). Generation runs in a |
| thread feeding a streamer; the thread is always joined, even if the caller stops early, |
| so the model and the state cache are never touched by two requests at once. |
| """ |
| from transformers import TextIteratorStreamer |
|
|
| streamer = TextIteratorStreamer(self.tok, skip_prompt=True, skip_special_tokens=True) |
| result: dict = {} |
|
|
| def run(): |
| try: |
| result["tokens"] = self._generate(ids, max_new_tokens, cache, streamer) |
| except BaseException as e: |
| result["error"] = e |
| streamer.end() |
|
|
| t = threading.Thread(target=run, daemon=True) |
| t.start() |
| try: |
| for piece in streamer: |
| if piece: |
| yield {"event": "text", "text": piece} |
| finally: |
| t.join() |
| if "error" in result: |
| raise result["error"] |
| return result["tokens"] |
|
|
| def chat_stream(self, messages: list[dict], max_new_tokens: int = 512): |
| """chat(), streamed: {"event": "text"} pieces, then {"event": "done", "tokens": [...]}.""" |
| ids = self.tok.apply_chat_template( |
| messages, add_generation_prompt=True, tokenize=True, return_dict=True |
| )["input_ids"] |
| with self.lock: |
| tokens = yield from self._stream(ids, max_new_tokens) |
| yield {"event": "done", "tokens": tokens} |
|
|
| def auto_stream(self, state: str, name: str, question: dict, max_new_tokens: int = 512): |
| """auto(), streamed. Yields {"event": "decision", ...} as soon as the quick decision and |
| the router's call are known, then {"event": "text"} reasoning pieces when escalated, then |
| {"event": "done", ...} carrying exactly what auto() returns.""" |
| item = to_item(state, name, question) |
| with self.lock: |
| cache = StateCache(self.model, self.tpl.state_ids(item)) |
| r = self._decide_on(cache, item) |
| out = { |
| "decision": self._answer(item, r), |
| "escalated": r["escalate"], |
| "state_prefills": 1, |
| } |
| yield {"event": "decision", **out} |
| if r["escalate"]: |
| ids = cache.ids + self.tpl.encode(chat_text(item)) |
| tokens = yield from self._stream(ids, max_new_tokens, cache) |
| out["generation_tokens"] = tokens |
| out["generation"] = self.tok.decode(tokens, skip_special_tokens=True) |
| cache.kv.crop(len(cache)) |
| yield {"event": "done", **out} |
|
|
| def chat(self, messages: list[dict], max_new_tokens: int = 512) -> list[int]: |
| ids = self.tok.apply_chat_template( |
| messages, add_generation_prompt=True, tokenize=True, return_dict=True |
| )["input_ids"] |
| with self.lock: |
| return self._generate(ids, max_new_tokens) |
|
|
| def auto(self, state: str, name: str, question: dict, max_new_tokens: int = 512) -> dict: |
| """Decide; when the router flags it, decode from the same state cache.""" |
| item = to_item(state, name, question) |
| with self.lock: |
| cache = StateCache(self.model, self.tpl.state_ids(item)) |
| r = self._decide_on(cache, item) |
| out = { |
| "decision": self._answer(item, r), |
| "escalated": r["escalate"], |
| "state_prefills": 1, |
| } |
| if r["escalate"]: |
| ids = cache.ids + self.tpl.encode(chat_text(item)) |
| out["generation_tokens"] = self._generate(ids, max_new_tokens, cache) |
| out["generation"] = self.tok.decode( |
| out["generation_tokens"], skip_special_tokens=True |
| ) |
| cache.kv.crop(len(cache)) |
| return out |
|
|
