| """Decision Index engine for prism-system-one: one predict() call per request, refusing (never truncating).""" |
| import math |
| import os |
|
|
| import torch |
| from decision_index.engines import Engine, Unsupported |
| from modeling_prism import PrismDecisionModel |
|
|
|
|
| def _finite(response): |
| for a in response["answers"].values(): |
| values = [a["noul"]] if a["type"] == "noul" else list(a["probabilities"].values()) |
| if not all(math.isfinite(v) for v in values): |
| return False |
| return True |
|
|
|
|
| class PrismEngine(Engine): |
| name = "prism-system-one" |
| latency = ("One forward pass per request on one GPU (fp16 on T4, bf16 on Ampere and newer), " |
| "including tokenization.") |
|
|
| def __init__(self, path, device=None, **options): |
| super().__init__(**options) |
| self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| cap = torch.cuda.get_device_capability()[0] if self.device == "cuda" else 0 |
| self.amp = self.device == "cuda" |
| self.dtype = torch.bfloat16 if cap >= 8 else torch.float16 |
| if self.amp and cap < 8: |
| torch.backends.cudnn.enabled = False |
| self.model = PrismDecisionModel.from_pretrained(path, device=self.device).eval() |
| self.provenance = {"repo": "paulprt/prism-system-one", |
| "revision": os.path.basename(os.path.normpath(path)), |
| "device": self.device, |
| "dtype": str(self.dtype) if self.amp else "float32", |
| "max_length": self.model.decision_model.max_length} |
|
|
| def _predict(self, request, amp): |
| with torch.autocast("cuda", dtype=self.dtype, enabled=amp): |
| return self.model.predict([request], strict=True)[0] |
|
|
| def __call__(self, state, questions): |
| request = {"state": state if state is not None else "", "questions": questions} |
| try: |
| response = self._predict(request, self.amp) |
| if self.amp and not _finite(response): |
| response = self._predict(request, False) |
| except ValueError as e: |
| raise Unsupported(str(e)) |
| return response, None |
