mirror of
https://git.isriupjv.fr/ISRI/ai-server
synced 2025-04-24 18:18:11 +02:00
52 lines
2 KiB
Python
52 lines
2 KiB
Python
import json
|
|
import os
|
|
import typing
|
|
import warnings
|
|
from pathlib import Path
|
|
|
|
from source import model
|
|
|
|
|
|
class ModelManager:
|
|
def __init__(self, model_library: os.PathLike | str):
|
|
self.model_library: Path = Path(model_library)
|
|
|
|
# the model types
|
|
self.model_types: dict[str, typing.Type[model.base.BaseModel]] = {}
|
|
# the models
|
|
self.models: dict[str, model.base.BaseModel] = {}
|
|
|
|
# the currently loaded model
|
|
# TODO(Faraphel): load more than one model at a time ? require a way more complex manager to handle memory issue
|
|
self.current_loaded_model: typing.Optional[model.base.BaseModel] = None
|
|
|
|
def register_model_type(self, name: str, model_type: typing.Type[model.base.BaseModel]):
|
|
self.model_types[name] = model_type
|
|
|
|
def reload(self):
|
|
for model_path in self.model_library.iterdir():
|
|
model_name: str = model_path.name
|
|
model_configuration_path: Path = model_path / "config.json"
|
|
|
|
# check if the configuration file exists
|
|
if not model_configuration_path.exists():
|
|
warnings.warn(f"Model {model_name!r} is missing a config.json file.")
|
|
continue
|
|
|
|
# load the configuration file
|
|
model_configuration = json.loads(model_configuration_path.read_text())
|
|
|
|
# get the model type for this model
|
|
model_type_name: str = model_configuration.get("type")
|
|
if model_type_name not in self.model_types:
|
|
warnings.warn("Field 'type' missing from the model configuration file.")
|
|
continue
|
|
|
|
# get the class of this model type
|
|
model_type = self.model_types.get(model_type_name)
|
|
if model_type is None:
|
|
warnings.warn(f"Model type {model_type_name!r} does not exists. Has it been registered ?")
|
|
continue
|
|
|
|
# load the model
|
|
self.models[model_name] = model_type(self, model_configuration, model_path)
|