124 lines
3.1 KiB
Python
124 lines
3.1 KiB
Python
import abc
|
|
import gc
|
|
import typing
|
|
from pathlib import Path
|
|
|
|
from source.manager import ModelManager
|
|
|
|
|
|
class BaseModel(abc.ABC):
|
|
"""
|
|
Represent a model.
|
|
"""
|
|
|
|
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.output_type: str = configuration.get("output_type", "application/json")
|
|
|
|
self._loaded = False
|
|
|
|
def __repr__(self):
|
|
return f"<{self.__class__.__name__}: {self.name}>"
|
|
|
|
@property
|
|
def name(self):
|
|
"""
|
|
Get the name of the model
|
|
:return: the name of the model
|
|
"""
|
|
|
|
return self.path.name
|
|
|
|
def get_information(self):
|
|
"""
|
|
Get information about the model
|
|
:return: information about the model
|
|
"""
|
|
|
|
return {
|
|
"name": self.name,
|
|
"output_type": self.output_type,
|
|
}
|
|
|
|
def load(self) -> None:
|
|
"""
|
|
Load the model within the model manager
|
|
"""
|
|
|
|
# if we are already loaded, stop
|
|
if self._loaded:
|
|
return
|
|
|
|
# check if we are the current loaded model
|
|
if self.manager.current_loaded_model is not self:
|
|
# unload the previous model
|
|
if self.manager.current_loaded_model is not None:
|
|
self.manager.current_loaded_model.unload()
|
|
|
|
# model specific loading
|
|
self._load()
|
|
|
|
# declare ourselves as the currently loaded model
|
|
self.manager.current_loaded_model = self
|
|
|
|
# mark the model as loaded
|
|
self._loaded = True
|
|
|
|
@abc.abstractmethod
|
|
def _load(self):
|
|
"""
|
|
Load the model
|
|
Do not call manually, use `load` instead.
|
|
"""
|
|
|
|
def unload(self) -> None:
|
|
"""
|
|
Unload the model within the model manager
|
|
"""
|
|
|
|
# if we are not already loaded, stop
|
|
if not self._loaded:
|
|
return
|
|
|
|
# if we were the currently loaded model of the manager, demote ourselves
|
|
if self.manager.current_loaded_model is self:
|
|
self.manager.current_loaded_model = None
|
|
|
|
# model specific unloading part
|
|
self._unload()
|
|
|
|
# force the garbage collector to clean the memory
|
|
gc.collect()
|
|
|
|
# mark the model as unloaded
|
|
self._loaded = False
|
|
|
|
@abc.abstractmethod
|
|
def _unload(self):
|
|
"""
|
|
Unload the model
|
|
Do not call manually, use `unload` instead.
|
|
"""
|
|
|
|
def infer(self, **kwargs) -> typing.Iterator[bytes]:
|
|
"""
|
|
Infer our payload through the model within the model manager
|
|
:return: the response of the model
|
|
"""
|
|
|
|
# make sure we are loaded before an inference
|
|
self.load()
|
|
|
|
# model specific inference part
|
|
return self._infer(**kwargs)
|
|
|
|
@abc.abstractmethod
|
|
def _infer(self, **kwargs) -> typing.Iterator[bytes]:
|
|
"""
|
|
Infer our payload through the model
|
|
:return: the response of the model
|
|
"""
|