code/decision_engine_dense.py

"""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  # 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 <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.",
        }

    # ------------------------------------------------------------------ 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()