handler.py

from __future__ import annotations

from typing import Any, Dict

try:
    from .runtime import JNUTSBRuntime
except ImportError:  # pragma: no cover
    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)