added support of inputs parameters that are recognised by the API.
Models are now loaded in separate endpoints for the inputs to be easier to recognise
This commit is contained in:
parent
900c58ffcb
commit
7bd84c8570
17 changed files with 163 additions and 128 deletions
|
@ -1,7 +1,9 @@
|
|||
import abc
|
||||
import gc
|
||||
import typing
|
||||
from pathlib import Path
|
||||
|
||||
from source import api
|
||||
from source.manager import ModelManager
|
||||
|
||||
|
||||
|
@ -10,13 +12,13 @@ class BaseModel(abc.ABC):
|
|||
Represent a model.
|
||||
"""
|
||||
|
||||
def __init__(self, manager: ModelManager, configuration: dict, path: Path):
|
||||
def __init__(self, manager: ModelManager, configuration: dict[str, typing.Any], path: Path):
|
||||
# the environment directory of the model
|
||||
self.path = path
|
||||
# the model manager
|
||||
self.manager = manager
|
||||
# the mimetype of the model responses
|
||||
self.response_mimetype: str = configuration.get("response_mimetype", "application/json")
|
||||
self.output_type: str = configuration.get("output_type", "application/json")
|
||||
|
||||
self._loaded = False
|
||||
|
||||
|
@ -101,13 +103,11 @@ class BaseModel(abc.ABC):
|
|||
"""
|
||||
Unload the model
|
||||
Do not call manually, use `unload` instead.
|
||||
:return:
|
||||
"""
|
||||
|
||||
def infer(self, payload: dict) -> str | bytes:
|
||||
def infer(self, *args, **kwargs) -> typing.Iterator[bytes]:
|
||||
"""
|
||||
Infer our payload through the model within the model manager
|
||||
:param payload: the payload to give to the model
|
||||
:return: the response of the model
|
||||
"""
|
||||
|
||||
|
@ -115,12 +115,11 @@ class BaseModel(abc.ABC):
|
|||
self.load()
|
||||
|
||||
# model specific inference part
|
||||
return self._infer(payload)
|
||||
return self._infer(*args, **kwargs)
|
||||
|
||||
@abc.abstractmethod
|
||||
def _infer(self, payload: dict) -> str | bytes:
|
||||
def _infer(self, *args, **kwargs) -> typing.Iterator[bytes]:
|
||||
"""
|
||||
Infer our payload through the model
|
||||
:param payload: the payload to give to the model
|
||||
:return: the response of the model
|
||||
"""
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue