"""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