File size: 3,313 Bytes
d10ad42 d664391 d10ad42 | 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 | """In-process engine for the Decision Index kit (`--engine maincode_jev_serve.engine:MaincodeJevEngine`).
Optional: needs the `decision-index` package installed alongside. Same scoring as the server.
Example: python -m decision_index run --engine maincode_jev_serve.engine:MaincodeJevEngine \
--option checkpoint=./matilda-jev-v1 --rows sample.jsonl.gz --out runs/x
"""
import hashlib
import math
from pathlib import Path
from typing import Any
import torch
from decision_index.engines import Engine, Unsupported
from maincode_jev_serve.config import gpu_limits
from maincode_jev_serve.decide import CapacityError, decide
from maincode_jev_serve.model import DecisionModel
class MaincodeJevEngine(Engine): # type: ignore[misc]
name = "maincode-jev"
latency = "In-process request wall time: prompt rendering, tokenization and one forward pass per question batch."
def __init__(self, checkpoint: str, max_tokens: int | None = None, token_budget: int | None = None, batch_size: int = 64,
device: str | None = None, **options: Any) -> None:
super().__init__(**options)
self.model = DecisionModel(checkpoint=checkpoint, device=device)
auto_context, auto_budget = gpu_limits()
self.max_tokens, self.token_budget, self.batch_size = int(max_tokens or auto_context), int(token_budget or auto_budget), batch_size
config = Path(checkpoint) / "decision_config.json"
self.provenance = {"checkpoint": str(Path(checkpoint).resolve()), "revision": self.model.revision,
"checkpoint_config_sha256": hashlib.sha256(config.read_bytes()).hexdigest(),
"temperature": self.model.temperature, "max_tokens": self.max_tokens}
self.label = Path(checkpoint).resolve().name
if not math.isfinite(self.model.temperature) or self.model.temperature <= 0:
raise ValueError("Checkpoint temperature must be positive and finite")
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.model.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: dict[str, object] = {}
for key, values in distributions.items():
if questions[key]["type"] == "noul":
answers[key] = {"type": "noul", "noul": values[1]}
else:
keys = list(questions[key]["criteria"])
answers[key] = {"type": "choice", "choice": keys[max(range(len(keys)), key=values.__getitem__)],
"probabilities": dict(zip(keys, values, strict=True))}
return {"model": self.label, "answers": answers, "usage": {"input_tokens": input_tokens}}, None
|