| from __future__ import annotations |
|
|
| from typing import Any, Dict |
|
|
| try: |
| from .runtime import JNUTSBRuntime |
| except ImportError: |
| from runtime import JNUTSBRuntime |
|
|
|
|
| class EndpointHandler: |
| """Hugging Face Inference Endpoint custom handler.""" |
|
|
| def __init__(self, model_dir: str, **kwargs: Any) -> None: |
| self.runtime = JNUTSBRuntime.from_config_dir(model_dir) |
|
|
| def __call__(self, data: Dict[str, Any]) -> Any: |
| inputs = data.get("inputs", data) |
| parameters = data.get("parameters", {}) |
| return self.runtime.predict(inputs=inputs, **parameters) |
|
|
