kev/bench.py

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