File size: 9,401 Bytes
c69aaec | 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 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | """Decision Index engine for kev decision models (`--engine kev.bench:KevEngine`).
Both the base model (no checkpoint: the readout is Qwen's own logits for the answer codes, a
zero-shot in kev's prompt format) and trained checkpoints run through the same fixed prompt.
Every question of a request is scored in token-budgeted batches; requests longer than
`max_tokens` are refused (`Unsupported`), never truncated, as the kit's rules require.
"""
import hashlib
import math
from pathlib import Path
from typing import Any, cast
import torch
from decision_index.engines import Engine, Unsupported
from kev.decide import CapacityError, decide
from kev.model import BASE_MODEL, DecisionModel
class KevEngine(Engine):
name = "kev"
latency = "In-process request wall time: prompt rendering, tokenization and one forward pass per question batch."
def __init__(self, checkpoint: str | None = None, base_model: str = BASE_MODEL, temperature: float | None = None,
max_tokens: int = 131072, token_budget: int = 131072, batch_size: int = 64, device: str | None = None,
label: str | None = None, **options: Any) -> None:
super().__init__(**options)
self.model = DecisionModel(checkpoint=checkpoint, base_model=base_model, device=device)
self.temperature = self.model.temperature if temperature is None else float(temperature)
if not math.isfinite(self.temperature) or self.temperature <= 0:
raise ValueError("Temperature must be positive and finite")
self.max_tokens, self.token_budget, self.batch_size = int(max_tokens), int(token_budget), int(batch_size)
self.label = label or (Path(checkpoint).resolve().parent.parent.name if checkpoint else "qwen3.8-27b-base-zeroshot")
config = Path(checkpoint) / "decision_config.json" if checkpoint else None
self.provenance = {
"base_model": self.model.base_model, "revision": self.model.revision,
"checkpoint": str(Path(checkpoint).resolve()) if checkpoint else None,
"checkpoint_config_sha256": hashlib.sha256(config.read_bytes()).hexdigest() if config else None,
"temperature": self.temperature, "max_tokens": self.max_tokens, "prompt": "kev.model.decision_messages",
}
def runtime(self) -> dict[str, object]:
if not torch.cuda.is_available():
return {"device": "cpu", "torch": torch.__version__}
properties = torch.cuda.get_device_properties(0)
return {"device": properties.name, "memory_gib": round(properties.total_memory / 2**30), "torch": torch.__version__,
"hip": torch.version.hip}
def synchronize(self) -> None:
if torch.cuda.is_available():
torch.cuda.synchronize()
def __call__(self, state: object, questions: dict[str, dict[str, Any]]) -> tuple[dict[str, object], None]:
for question in questions.values():
if question["type"] not in ("choice", "noul"):
raise Unsupported(f"question type {question['type']!r}")
try:
distributions, input_tokens = decide(self.model, state, questions, temperature=self.temperature,
max_tokens=self.max_tokens, token_budget=self.token_budget,
batch_size=self.batch_size)
except CapacityError as error:
raise Unsupported(str(error)) from error
answers = {key: self._answer(questions[key], values) for key, values in distributions.items()}
return {"model": self.label, "answers": answers, "usage": {"input_tokens": input_tokens}}, None
@staticmethod
def _answer(question: dict[str, Any], values: list[float]) -> dict[str, object]:
total = sum(values)
values = [value / total for value in values]
if question["type"] == "noul":
return {"type": "noul", "noul": values[1]}
keys = list(question["criteria"])
return {"type": "choice", "choice": keys[max(range(len(keys)), key=values.__getitem__)],
"probabilities": dict(zip(keys, values, strict=True))}
class MainJevEngine(Engine):
"""Decision Index engine for main-jev-v0-style checkpoints (`--engine kev.bench:MainJevEngine`).
Imports the checkpoint's own `decision.py` (prompt `Context/Question/Options ... Answer: (`,
label-row readout) so scoring matches the author's `predict.py`. Choice options are rendered as
"key: description" (or the key alone); noul questions become [false, true] options.
"""
name = "main-jev"
latency = "In-process request wall time: prompt encoding and one forward pass per question batch."
def __init__(self, model: str = "/shared/model-cache/main-jev-v0", batch_size: int = 16, device: str | None = None,
**options: Any) -> None:
super().__init__(**options)
import importlib.util
import json
import sys
from transformers import AutoModelForCausalLM, AutoTokenizer
directory = Path(model)
self.config = json.loads((directory / "decision_config.json").read_text())
sys.dont_write_bytecode = True # never write __pycache__ into someone else's model directory
spec = importlib.util.spec_from_file_location("main_jev_decision", directory / "decision.py")
assert spec is not None and spec.loader is not None
self.decision = importlib.util.module_from_spec(spec)
spec.loader.exec_module(self.decision)
self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
self.tokenizer = AutoTokenizer.from_pretrained(directory, local_files_only=True)
dtype = torch.bfloat16 if self.device.type == "cuda" else torch.float32
lm = AutoModelForCausalLM.from_pretrained(directory, dtype=dtype, attn_implementation="sdpa", local_files_only=True)
self.model = self.decision.DecisionModel(lm, self.config["label_ids"]).to(self.device).eval()
self.collate = self.decision.Collator(self.tokenizer.pad_token_id)
self.batch_size = int(batch_size)
self.max_options = len(self.config["label_ids"])
self.provenance = {"model": str(directory.resolve()), "name": self.config.get("name"),
"decision_config_sha256": hashlib.sha256((directory / "decision_config.json").read_bytes()).hexdigest(),
"decision_py_sha256": hashlib.sha256((directory / "decision.py").read_bytes()).hexdigest(),
"max_length": self.config["max_length"], "prompt": self.config.get("prompt_format")}
def runtime(self) -> dict[str, object]:
return KevEngine.runtime(self)
def synchronize(self) -> None:
if torch.cuda.is_available():
torch.cuda.synchronize()
@staticmethod
def _text(value: object) -> str:
from decision_index.engines.base import text
return str(text(value)) if value is not None else ""
def __call__(self, state: object, questions: dict[str, dict[str, Any]]) -> tuple[dict[str, object], None]:
encoded = []
for key, question in questions.items():
if question["type"] == "choice":
options = [str(name) if description is None else f"{name}: {self._text(description)}"
for name, description in question["criteria"].items()]
elif question["type"] == "noul":
criteria = question.get("criteria") or {}
options = [self._text(criteria.get("false")) or "No", self._text(criteria.get("true")) or "Yes"]
else:
raise Unsupported(f"question type {question['type']!r}")
if len(options) > self.max_options:
raise Unsupported(f"declared limit of {self.max_options} options")
row = {"context": self._text(state), "question": self._text(question.get("instructions")), "options": options,
"target": [1 / len(options)] * len(options), "id": key, "source": "decision-index", "kind": question["type"],
"key": key}
item = self.decision.encode(row, self.tokenizer, self.config["label_ids"], self.config["max_length"])
if item is None:
raise Unsupported(f"request exceeds the model's {self.config['max_length']}-token input limit")
encoded.append(item)
answers: dict[str, object] = {}
input_tokens = 0
for start in range(0, len(encoded), self.batch_size):
chunk = encoded[start:start + self.batch_size]
batch = {name: tensor.to(self.device) for name, tensor in self.collate(chunk).items()}
batch.pop("targets")
input_tokens += int(batch["attention_mask"].sum())
with torch.inference_mode(), torch.autocast(self.device.type, dtype=torch.bfloat16, enabled=self.device.type == "cuda"):
probabilities = self.model(**batch)["logits"].softmax(-1).float().cpu().tolist()
for item, values in zip(chunk, probabilities, strict=True):
answers[item["id"]] = KevEngine._answer(questions[item["id"]], values[:item["nopts"]])
return {"model": str(self.config.get("name")), "answers": {key: answers[key] for key in questions},
"usage": {"input_tokens": input_tokens}}, None
|