kev/decide.py

"""One decision request -> probability distributions, shared by the benchmark engine and the server.

Every question of a request is scored against the same state in token-budgeted, length-sorted
batches (one forward pass per batch). Nothing is truncated: a request that does not fit raises
CapacityError, whose message carries the standard capacity markers.
"""

from collections.abc import Mapping, Sequence
from typing import Any, cast

import torch

from kev.model import MAX_OPTIONS, DecisionModel, options
from kev.train import length_estimate, microbatches
from kev.types import Example, ImageInput


class CapacityError(ValueError):
    pass


def decide(model: DecisionModel, state: object, questions: Mapping[str, Mapping[str, Any]], *, temperature: float,
           max_tokens: int, token_budget: int, batch_size: int, images: Sequence[ImageInput] = ()) -> tuple[dict[str, list[float]], int]:
    """Normalized probabilities per question key (in option order) and the input tokens used."""
    rows: list[Example] = []
    for key, question in questions.items():
        count = len(options(cast(Any, question))[0])
        if count > MAX_OPTIONS:
            raise CapacityError(f"at most {MAX_OPTIONS} options per choice question are supported ({count} given)")
        row = {"state": state, "question": dict(question), "id": key, "suite": "", "family": key, "label": "", "target": "",
               "source": {}}
        if images:
            row["images"] = list(images)
        rows.append(cast(Example, row))
    distributions: dict[str, list[float]] = {}
    input_tokens = 0
    for batch in microbatches(sorted(rows, key=length_estimate), batch_size, token_budget):
        try:
            prepared = model.prepare(batch, max_length=max_tokens)
        except ValueError as error:
            if "token limit" in str(error):
                raise CapacityError(f"request exceeds the maximum context length of {max_tokens} tokens") from error
            raise
        input_tokens += prepared.input_tokens
        with torch.inference_mode():
            probabilities = (model(prepared) / temperature).softmax(-1).float().cpu().tolist()
        for item, values, count in zip(batch, probabilities, prepared.counts, strict=True):
            total = sum(values[:count])
            distributions[item["id"]] = [value / total for value in values[:count]]
    return {key: distributions[key] for key in questions}, input_tokens