File size: 9,368 Bytes
63ae4a0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | """Decision Index engine for the v6 checkpoints (dense/train_dense.py --contexts): every question gets a short stock-model rationale, then the slot is read with the adapter ON.
Per request: (1) the rationale prompts of ALL questions of the request are built (dense/contexts.py: system line + question prompt + "Rationale:\\n") and generated together in
length-sorted batches with the adapter DISABLED (greedy, at most `brief_cap` tokens, stop <|im_end|>), (2) the generated text gets the same clean-up as the training rationales
(trailing bare-answer line dropped), (3) every question is rendered in the `brief` layout (rationale then "\\n\\nAnswer: ", per-option items put it before the slot block) and the
slot is read with the adapter ON: option logits = frozen LM-head rows of the answer tokens + trained residual, softmax over options (optional fitted temperature).
This is the accuracy path (Hugging Face generate); the latency path is tools/vllm_requests_0f.py (vLLM) with the same layout.
python -m decision_index run --engine decision_engine_v6:BriefSlotEngine --option model=Qwen/Qwen3.5-9B --option run=RUN --option ckpt=best|final --suite-dir ... --rows ... --out ...
Options: model, run, ckpt (best|final|dir), context (brief|direct: direct = the v5 behaviour, no rationale), brief_cap (64), gen_tokens_budget (padded tokens per generate batch, 24000),
temperature (auto|none|...: same files as decision_engine_dense), max_tokens (32768).
"""
import os
import sys
import time
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 # noqa: E402
from decision_engine_dense import DenseSlotEngine, _flag # noqa: E402
try:
from decision_index.engines.base import Unsupported # noqa: E402
except ImportError: # Decision Index kit not installed
from engine_base import Unsupported # noqa: E402
from dense.contexts import Contexts # noqa: E402
class BriefSlotEngine(DenseSlotEngine):
name = "dense-slot-brief"
latency = ("In-process request wall time including prompt construction, the batched stock-model rationale generation of all questions of the request (Hugging Face generate, "
"adapter disabled) and one adapted slot pass per question; excludes model loading.")
def __init__(self, model="Qwen/Qwen3.5-9B", run=None, ckpt="best", context="brief", parallel_context="brief", brief_cap=64, gen_tokens_budget=24000, temperature="auto", max_tokens=32768, **options):
if not run:
raise ValueError("BriefSlotEngine needs a trained run (option run=...)")
self.temp_name = f"temperature-{context}" # fitted by tools/fit_temperature_v6.py
super().__init__(model=model, run=run, ckpt=ckpt, quant="none", merge=False, temperature=temperature, prefix_cache=False, max_tokens=max_tokens, **options)
self.context = str(context)
self.parallel_context = str(parallel_context) # context of per-option ('parallel' slot) questions: brief | direct (v6b: trained without a rationale for these families)
if self.context not in ("brief", "direct"):
raise ValueError("context must be brief or direct")
self.brief_cap, self.gen_budget = int(brief_cap), int(gen_tokens_budget)
self.ctx = Contexts(self.tok, self.codec, brief_cap=self.brief_cap)
self.peft = self.pipe.peft
self.eos = self.tok.convert_tokens_to_ids("<|im_end|>")
self.provenance.update({"context": self.context, "brief_cap": self.brief_cap, "generation": "stock model (adapter disabled), greedy, Hugging Face generate",
"policy": self.provenance["policy"] + " Before the slot is read, a stock-model rationale of at most %d tokens is generated for every question." % self.brief_cap})
# ------------------------------------------------------------------ generation (adapter off)
@torch.inference_mode()
def _generate(self, prompts):
"""Greedy rationales for a list of token-id prompts; length-sorted batches under a padded-token budget."""
out = [None] * len(prompts)
order = sorted(range(len(prompts)), key=lambda i: len(prompts[i]))
s = 0
pad = self.tok.pad_token_id
while s < len(order):
e = s + 1
while e < len(order) and (e - s + 1) * len(prompts[order[e]]) <= self.gen_budget:
e += 1
idx = order[s:e]
L = max(len(prompts[i]) for i in idx)
ids = torch.full((len(idx), L), pad, dtype=torch.long)
att = torch.zeros((len(idx), L), dtype=torch.long)
for j, i in enumerate(idx):
p = prompts[i]
ids[j, L - len(p):] = torch.tensor(p)
att[j, L - len(p):] = 1
base = self.peft.get_base_model()
with self.peft.disable_adapter(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
gen = base.generate(input_ids=ids.cuda(), attention_mask=att.cuda(), max_new_tokens=self.brief_cap, do_sample=False, eos_token_id=self.eos, pad_token_id=pad, use_cache=True)
for j, i in enumerate(idx):
g = gen[j, L:].tolist()
if self.eos in g:
g = g[:g.index(self.eos)]
while g and g[-1] == pad:
g.pop()
out[i] = self.ctx.clean_brief(g)
s = e
return out
def _item(self, state, q):
"""The codec item of a question (as DenseSlotEngine._render, but returns the item)."""
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 __import__("json").dumps(state, ensure_ascii=False, separators=(",", ":")))
return keys, {"qtype": QTYPES[q["type"]], "target": [1.0 / k] * k, "label": 0, "state": text, "source": "index",
"qdef": {"type": q["type"], "instructions": q.get("instructions", ""), "criteria": crit}}
def __call__(self, state, questions):
if self.context == "direct":
return super().__call__(state, questions)
answers, raw = {}, {}
items = {qid: self._item(state, q) for qid, q in questions.items()}
if self.parallel_context == "direct": # per-option questions skip the rationale: they are answered by the plain slot path
par = {qid for qid, (_, it) in items.items() if AS.slot_mode(it["qtype"], len(it["target"])) == "parallel"}
if par:
resp_p, raw_p = DenseSlotEngine.__call__(self, state, {qid: questions[qid] for qid in par})
rest = {qid: q for qid, q in questions.items() if qid not in par}
resp_r, raw_r = self.__call__(state, rest) if rest else ({"model": self.model_id, "answers": {}, "usage": {"input_tokens": 0}}, {})
resp_p["answers"].update(resp_r["answers"]); resp_p["usage"]["input_tokens"] += resp_r["usage"]["input_tokens"]
return resp_p, {**raw_p, **raw_r}
qids = list(items)
try:
gen_prompts = [self.ctx.gen_prompt(items[qid][1]) for qid in qids]
for qid, p in zip(qids, gen_prompts):
if len(p) + self.brief_cap > self.limit:
raise Unsupported(f"prompt of {len(p)} tokens exceeds the {self.limit}-token limit")
t0 = time.perf_counter()
rats = self._generate(gen_prompts)
t_gen = time.perf_counter() - t0
for qid, rat in zip(qids, rats):
keys, item = items[qid]
ex = {**item, "rationale": {"brief": {"tokens": rat, "n": len(rat)}}}
r = self.ctx.render(ex, "brief")
if len(r["ids"]) > self.limit:
raise Unsupported(f"prompt of {len(r['ids'])} tokens exceeds the {self.limit}-token limit")
logits, pooled = self._logits(r, None)
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)))
q = questions[qid]
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"]), "rationale_tokens": len(rat), "rationale": self.tok.decode(rat), "temperature": round(t, 4), "generation_s": round(t_gen, 3)}
except torch.OutOfMemoryError:
torch.cuda.empty_cache()
raise Unsupported("out of memory on a long prompt")
return {"model": self.model_id, "answers": answers, "usage": {"input_tokens": sum(v["prompt_tokens"] for v in raw.values())}}, raw
|