"""Decision Index engine for the dense LoRA checkpoints trained by dense/train_dense.py (Qwen3.5-9B, Gemma-4-12B, ...). Usage, from the decision-index checkout (PYTHONPATH must include this directory and /mnt/datasets/pylibs): python -m decision_index pipeline --engine decision_engine_dense:DenseSlotEngine \ --option model=Qwen/Qwen3.5-9B --option run=/mnt/datasets/dense-runs/qwen35-9b-pilot --suite-dir --out runs/qwen35-9b-pilot Options (all optional except model): model run (training output dir; omit for the stock model + zero residual = zero-shot baseline) ckpt (best|final|dir) quant (none|nf4) merge (auto: merge the LoRA into the weights for bf16) temperature (auto|none|global|bucket|feat|hidden) prefix_cache (true) max_tokens (32768) Every question is answered on its own with the fixed answer-slot prompt. A row's questions share the state, so the state prefix is run once and its cache (attention KV + recurrent states) is copied per question. Prompts longer than max_tokens, or that run out of memory, are `Unsupported`; nothing is truncated. """ import copy import json 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 common import QTYPES, temp_bucket # noqa: E402 from decision_index.engines.base import Engine, Unsupported # noqa: E402 from dense import temperature as TM # noqa: E402 from dense.train_dense import DensePipeline, load_model # noqa: E402 LO, HI = 0.5, 5.0 class DenseSlotEngine(Engine): name = "dense-slot" latency = "In-process request wall time including prompt construction, one cached state prefix per row and one suffix pass per question; excludes model loading." def __init__(self, model="Qwen/Qwen3.5-9B", run=None, ckpt="best", quant="none", merge="auto", temperature="auto", prefix_cache=True, max_tokens=32768, **options): super().__init__(**options) from transformers import AutoTokenizer self.tok = AutoTokenizer.from_pretrained(model) if self.tok.pad_token_id is None: self.tok.pad_token_id = self.tok.eos_token_id self.codec = AS.SlotCodec(self.tok, state_max=None, option_max=None, template_kwargs={"enable_thinking": False}, mask_id=self.tok.pad_token_id or 0, strip_label_keys=True, autoregressive=True) lm = load_model(model, quant) lm.config.use_cache = True softcap = getattr(lm.config.get_text_config(), "final_logit_softcapping", None) ckpt_dir = None if run: from peft import PeftModel ckpt_dir = ckpt if os.path.isdir(ckpt) else os.path.join(run, ckpt) lm = PeftModel.from_pretrained(lm, os.path.join(ckpt_dir, "lora")) if merge is True or (merge == "auto" and quant == "none"): lm = lm.merge_and_unload() self.pipe = DensePipeline(lm, self.codec.cand, softcap).to("cuda").eval() if ckpt_dir: self.pipe.head.load_state_dict(torch.load(os.path.join(ckpt_dir, "head.pt"), map_location="cuda")) self.limit = int(max_tokens) self.prefix_cache = bool(prefix_cache) self.model_id = model self._load_temperature(ckpt_dir, temperature) self.provenance = { "kind": "dense-answer-slot", "repo": model, "run": run, "checkpoint": ckpt_dir, "quant": quant, "merged": bool(run and (merge is True or (merge == "auto" and quant == "none"))), "temperature": self.t_kind, "context_limit_tokens": self.limit, "prefix_cache": self.prefix_cache, "layout": AS.LAYOUT, "policy": "One chat prompt per question (state in tags, options as letters A-Z up to 26 options, one forced-' No' Yes/No slot per option beyond that, a Yes/No answer for noul); " "the answer is the next token read through the frozen LM-head rows of the answer tokens plus a trained residual; softmax over options with the fitted temperature. " "A row's questions share a cached state prefix. Prompts over the limit are refused as unsupported, never truncated; no option is filtered.", } # ------------------------------------------------------------------ temperature def _load_temperature(self, ckpt_dir, mode): self.t_kind, self.t_json, self.t_head = "none", None, None if not ckpt_dir or mode == "none": return path = os.path.join(ckpt_dir, "temperature.json") if not os.path.isfile(path): if mode != "auto": raise FileNotFoundError(f"{path} (run python -m dense.temperature fit first)") return self.t_json = json.load(open(path)) kind = self.t_json["selected"] if mode == "auto" else mode self.t_kind = kind if kind in ("feat", "hidden"): blob = torch.load(os.path.join(ckpt_dir, "temperature.pt"), map_location="cuda") if blob["kind"] != kind: raise ValueError(f"temperature.pt holds a '{blob['kind']}' head, not '{kind}'") self.t_head = TM.TempHead(hidden=blob["hidden"]).to("cuda") self.t_head.load_state_dict(blob["state"]) self.t_head.eval() def _temperature(self, logits: torch.Tensor, pooled: torch.Tensor, qtype: int, k: int) -> float: kind = self.t_kind if kind == "none": return 1.0 if kind == "global": t = self.t_json["global"] elif kind == "bucket": b = self.t_json["bucket"] t = b["temperatures"].get(temp_bucket(qtype, k), b["global"]) else: mask = torch.ones(1, k, dtype=torch.bool, device="cuda") mode = TM.MODES[AS.slot_mode(qtype, k)] feats = TM.make_features(logits[None, :k], mask, torch.tensor([qtype], device="cuda"), torch.tensor([mode], device="cuda")) t = self.t_head.log_t(feats, pooled[None]).exp().item() return float(min(HI, max(LO, t))) # ------------------------------------------------------------------ forward def _render(self, state, q): if q["type"] not in ("choice", "noul"): raise Unsupported("Unsupported question type " + str(q["type"])) keys = list(q["criteria"]) if q["type"] == "choice" else ["false", "true"] if not keys: raise Unsupported("question without options") crit = q["criteria"] if q["type"] == "choice" else (q.get("criteria") or {}) k = len(keys) text = "" if state in (None, "", {}, []) else (state if isinstance(state, str) else json.dumps(state, ensure_ascii=False, separators=(",", ":"))) item = {"qtype": QTYPES[q["type"]], "target": [1.0 / k] * k, "label": 0, "state": text, "qdef": {"type": q["type"], "instructions": q.get("instructions", ""), "criteria": crit}} r = self.codec.render_ar(item) if len(r["ids"]) > self.limit: raise Unsupported(f"prompt of {len(r['ids'])} tokens exceeds the {self.limit}-token limit") return keys, r def _tensors(self, r): dev = "cuda" t = lambda x: torch.tensor([x], device=dev) k = len(r["target"]) return t(r["slot_of_opt"]), t(r["pos_idx"]), t(r["neg_idx"]), torch.ones(1, k, dtype=torch.bool, device=dev) @torch.inference_mode() def _logits(self, r, base): """(logits (K,), pooled hidden (H,)). `base` = (state_end, cache) of the row's shared prefix, or None.""" bb = self.pipe.backbone ids = torch.tensor([r["ids"]], device="cuda") with torch.amp.autocast("cuda", dtype=torch.bfloat16): if base is None: h = bb(input_ids=ids, use_cache=False).last_hidden_state hs = h[:, r["slot_pos"]] else: end, cache = base h = bb(input_ids=ids[:, end:], past_key_values=copy.deepcopy(cache), use_cache=True).last_hidden_state hs = h[:, [p - end for p in r["slot_pos"]]] logits = self.pipe.readout(hs, *self._tensors(r))[0] return logits.float(), hs.float().mean(1)[0] @torch.inference_mode() def _base(self, r): with torch.amp.autocast("cuda", dtype=torch.bfloat16): out = self.pipe.backbone(input_ids=torch.tensor([r["ids"][:r["state_end"]]], device="cuda"), use_cache=True) return r["state_end"], out.past_key_values def __call__(self, state, questions): answers, raw = {}, {} rendered = {qid: self._render(state, q) for qid, q in questions.items()} base = None try: if self.prefix_cache and len(rendered) >= 2: first = next(iter(rendered.values()))[1] if first["state_end"] >= 128 and all(r["ids"][:first["state_end"]] == first["ids"][:first["state_end"]] for _, r in rendered.values()): base = self._base(first) for qid, q in questions.items(): keys, r = rendered[qid] logits, pooled = self._logits(r, base) t = self._temperature(logits, pooled, r["qtype"], len(keys)) p = torch.softmax(logits / t, -1).cpu().numpy().astype(np.float64) p = p / p.sum() probs = dict(zip(keys, (float(x) for x in p))) if q["type"] == "choice": answers[qid] = {"type": "choice", "choice": max(keys, key=lambda x: probs[x]), "probabilities": probs} else: answers[qid] = {"type": "noul", "noul": probs["true"]} raw[qid] = {"prompt_tokens": len(r["ids"]), "temperature": round(t, 4), "cached_prefix_tokens": base[0] if base else 0} except torch.OutOfMemoryError: torch.cuda.empty_cache() raise Unsupported("out of memory on a long prompt") finally: del base return {"model": self.model_id, "answers": answers, "usage": {"input_tokens": sum(v["prompt_tokens"] for v in raw.values())}}, raw def runtime(self): info = {"torch": torch.__version__, "device": "cuda", "cuda": torch.version.cuda, "gpu": torch.cuda.get_device_name()} try: import transformers info["transformers"] = transformers.__version__ except Exception: pass return info def synchronize(self): torch.cuda.synchronize()