diff --git a/app/machine_learning/models/__init__.py b/app/machine_learning/models/__init__.py index 41f56bc6..7bc5ce7d 100644 --- a/app/machine_learning/models/__init__.py +++ b/app/machine_learning/models/__init__.py @@ -40,6 +40,15 @@ def __init__(self, folder: str, tags: str, *, object_saver: ObjectSaver = None, else: self.object_saver = ObjectSaver({CONFIG_KEY: 'filesystem', 'filesystemDefaultPath': folder}) + def load_model(self, model_files: list[str]) -> list[Any]: + result = [] + for file in model_files: + model = self.object_saver.get_project_object(os.path.join(self.folder, file), using_json=False) + if model is None: + raise ValueError(f'Unable to load model "{file}".') + result.append(model) + return result + def get_model_info(self): folder_name = os.path.basename(self.folder.strip("/").strip("\\")).strip() tags = self.tags @@ -47,10 +56,6 @@ def get_model_info(self): tags = [folder_name] + self.tags return tags - @abstractmethod - def load_model(self): - raise NotImplementedError('"load_model" method is not implemented!') - @abstractmethod def save_model(self): raise NotImplementedError('"save_model" method is not implemented!') diff --git a/app/machine_learning/models/custom_defect_type_model.py b/app/machine_learning/models/custom_defect_type_model.py index 1051396c..bdcb1c91 100644 --- a/app/machine_learning/models/custom_defect_type_model.py +++ b/app/machine_learning/models/custom_defect_type_model.py @@ -20,24 +20,11 @@ class CustomDefectTypeModel(DefectTypeModel): - project_id: int | str def __init__(self, folder: str, app_config: dict[str, Any], project_id: int | str): - super().__init__(folder, 'custom defect type model') - self.project_id = project_id - self.object_saver = ObjectSaver(app_config) - - def load_model(self): - self.count_vectorizer_models = self.object_saver.get_project_object( - os.path.join(self.folder, "count_vectorizer_models"), self.project_id, using_json=False) - assert len(self.count_vectorizer_models) > 0 - self.models = self.object_saver.get_project_object(os.path.join(self.folder, "models"), self.project_id, - using_json=False) - assert len(self.models) > 0 + super().__init__(folder, 'custom defect type model', object_saver=ObjectSaver(app_config, project_id)) def save_model(self): self.object_saver.put_project_object(self.count_vectorizer_models, - os.path.join(self.folder, "count_vectorizer_models"), self.project_id, - using_json=False) - self.object_saver.put_project_object(self.models, os.path.join(self.folder, "models"), self.project_id, - using_json=False) + os.path.join(self.folder, "count_vectorizer_models"), using_json=False) + self.object_saver.put_project_object(self.models, os.path.join(self.folder, "models"), using_json=False) diff --git a/app/machine_learning/models/defect_type_model.py b/app/machine_learning/models/defect_type_model.py index 5b05958b..b98742eb 100644 --- a/app/machine_learning/models/defect_type_model.py +++ b/app/machine_learning/models/defect_type_model.py @@ -25,7 +25,7 @@ from app.commons.object_saving.object_saver import ObjectSaver from app.machine_learning.models import MlModel -from app.utils import text_processing, utils +from app.utils import text_processing class DefectTypeModel(MlModel): @@ -33,23 +33,8 @@ class DefectTypeModel(MlModel): def __init__(self, folder: str, tags: str = 'global defect type model', object_saver: ObjectSaver = None, app_config: dict[str, Any] = None) -> None: super().__init__(folder, tags, object_saver=object_saver, app_config=app_config) - self.count_vectorizer_models = {} - self.models = {} - self.load_model() - - def load_model(self): - if not utils.validate_folder(self.folder): - raise ValueError(f'Invalid model folder path: {self.folder}') - - count_vectorizer_models_file = os.path.join(self.folder, "count_vectorizer_models.pickle") - models_file = os.path.join(self.folder, "models.pickle") - if not utils.validate_file(count_vectorizer_models_file) or not utils.validate_file(models_file): - raise ValueError(f'Model folder path does not contains necessary files: {self.folder}') - - with open(count_vectorizer_models_file, "rb") as f: - self.count_vectorizer_models = pickle.load(f) - with open(models_file, "rb") as f: - self.models = pickle.load(f) + self.count_vectorizer_models, self.models = self.load_model( + ['count_vectorizer_models.pickle', 'models.pickle']) def save_model(self): os.makedirs(self.folder, exist_ok=True)