| """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 <suite-0.3> --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 |
| import torch |
|
|
| import answer_slot as AS |
| from common import QTYPES, temp_bucket |
| from decision_index.engines.base import Engine, Unsupported |
| from dense import temperature as TM |
| from dense.train_dense import DensePipeline, load_model |
|
|
| 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 <text> 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.", |
| } |
|
|
| |
| 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))) |
|
|
| |
| 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() |
|
|
