| """Decision Index interface for JADE.""" |
| from .engine import JadeEngine, UnsupportedInput |
|
|
|
|
| class DecisionIndexEngine(JadeEngine): |
| name = "jade" |
| latency = "In-process vLLM wall time including rendering; model initialization excluded." |
|
|
| def __call__(self, state, questions): |
| from decision_index.engines.base import Unsupported |
|
|
| try: |
| return super().__call__(state, questions) |
| except UnsupportedInput as exc: |
| raise Unsupported(str(exc)) from exc |
|
|
| def close(self): |
| del self.model |
|
|
| def warmup(self): |
| questions = { |
| "warmup": { |
| "type": "choice", |
| "instructions": "Which color is named?", |
| "criteria": {"red": "red", "blue": "blue"}, |
| } |
| } |
| for _ in range(2): |
| self("The color is red.", questions) |
|
|
| def runtime(self): |
| import torch |
| import vllm |
|
|
| return { |
| "torch": torch.__version__, |
| "vllm": vllm.__version__, |
| "cuda": torch.version.cuda, |
| "gpu": torch.cuda.get_device_name(), |
| } |
|
|
| def synchronize(self): |
| import torch |
|
|
| torch.cuda.synchronize() |
|
|
