diff --git a/medcat-v2/medcat/cat.py b/medcat-v2/medcat/cat.py index efe969ebe..cf38c8016 100644 --- a/medcat-v2/medcat/cat.py +++ b/medcat-v2/medcat/cat.py @@ -23,7 +23,7 @@ from medcat.storage.mp_ents_save import BatchAnnotationSaver from medcat.utils.fileutils import ensure_folder_if_parent from medcat.utils.hasher import Hasher -from medcat.pipeline import Pipeline +from medcat.pipeline import Pipeline, _ENABLED_ADDONS_PATH from medcat.tokenizing.tokens import MutableDocument, MutableEntity from medcat.tokenizing.tokenizers import SaveableTokenizer, TOKENIZER_PREFIX from medcat.data.entities import Entity, Entities, OnlyCUIEntities @@ -816,10 +816,12 @@ def _get_missing_plugins(cls, model_pack_path: str) -> list[MissingPluginInfo]: return missing_plugins @classmethod - def load_model_pack(cls, model_pack_path: str, - config_dict: Optional[dict] = None, - addon_config_dict: Optional[dict[str, dict]] = None - ) -> 'CAT': + def load_model_pack( + cls, model_pack_path: str, + config_dict: Optional[dict] = None, + addon_config_dict: Optional[dict[str, dict]] = None, + keep_addons_of_types: Optional[list[Type[AddonComponent]]] = None, + ) -> 'CAT': """Load the model pack from file. Args: @@ -831,6 +833,9 @@ def load_model_pack(cls, model_pack_path: str, If specified, it needs to have an addon dict per name. For instance, `{"meta_cat.Subject": {}}` would apply to the specific MetaCAT. + keep_addons_of_types (Optional[list[Type[AddonComponent]]]): + Only load addons of specified types. If `None`, all addons + will be loaded. Defaults to `None`. Raises: ValueError: If the saved data does not represent a model pack. @@ -855,6 +860,22 @@ def load_model_pack(cls, model_pack_path: str, # Load model card to check for required plugins missing_plugins = cls._get_missing_plugins(model_pack_path) + if keep_addons_of_types: + # add these to addon_config_dict + # which will propgate to __init__ + # and then in the pipe the rest will be filtered + # out so they don't need to be loaded + if addon_config_dict is None: + addon_config_dict = {} + addons_to_keep: dict = { + _ENABLED_ADDONS_PATH: [ + addon_type.addon_type + for addon_type in keep_addons_of_types] + } + # avoid mutating incoming dict + addons_to_keep.update(addon_config_dict) + addon_config_dict = addons_to_keep + try: # NOTE: ignoring addons since they will be loaded later / separately cat = deserialise(model_pack_path, model_load_path=model_pack_path, diff --git a/medcat-v2/medcat/pipeline/__init__.py b/medcat-v2/medcat/pipeline/__init__.py index 08f1fdd85..5b6f90465 100644 --- a/medcat-v2/medcat/pipeline/__init__.py +++ b/medcat-v2/medcat/pipeline/__init__.py @@ -1,4 +1,4 @@ -from .pipeline import Pipeline +from .pipeline import Pipeline, _ENABLED_ADDONS_PATH -__all__ = ['Pipeline'] +__all__ = ['Pipeline', '_ENABLED_ADDONS_PATH'] diff --git a/medcat-v2/medcat/pipeline/pipeline.py b/medcat-v2/medcat/pipeline/pipeline.py index 43d512f98..5e24bc421 100644 --- a/medcat-v2/medcat/pipeline/pipeline.py +++ b/medcat-v2/medcat/pipeline/pipeline.py @@ -1,4 +1,4 @@ -from typing import Optional, Iterable, Union +from typing import Optional, Iterable, Union, cast import logging import os import warnings @@ -24,6 +24,9 @@ logger = logging.getLogger(__name__) +# NOTE: used only temporarily and internally +_ENABLED_ADDONS_PATH = "==ENABLED_ADDONS==" + class DelegatingTokenizer(BaseTokenizer): """A delegating tokenizer. @@ -225,6 +228,17 @@ def _attempt_merge( "merge the config since it's ambiguous (@%s)", len(similars), addon_cnf.comp_name, name) + def _filter_config_addons_at_load(self, addon_types: list[str]): + logger.info("Filtering addon types at model load time. Only allowing %s", + addon_types) + for cnf in list(self.config.components.addons): + if any(cnf.comp_name == addon_type for addon_type in addon_types): + continue + logger.debug( + "Removing addon %s because it doesn't match any expected types", + cnf.comp_name) + self.config.components.addons.remove(cnf) + def _init_components(self, model_load_path: Optional[str], old_pipe: Optional['Pipeline'], addon_config_dict: Optional[dict[str, dict]], @@ -240,6 +254,9 @@ def _init_components(self, model_load_path: Optional[str], comp = self._init_component( CoreComponentType[cct_name], model_load_path) self._components.append(comp) + if addon_config_dict and _ENABLED_ADDONS_PATH in addon_config_dict: + self._filter_config_addons_at_load( + cast(list[str], addon_config_dict[_ENABLED_ADDONS_PATH])) for addon_cnf in self.config.components.addons: if addon_config_dict: self._attempt_merge(addon_cnf, addon_config_dict) diff --git a/medcat-v2/tests/test_cat.py b/medcat-v2/tests/test_cat.py index db5d75c48..baf406593 100644 --- a/medcat-v2/tests/test_cat.py +++ b/medcat-v2/tests/test_cat.py @@ -545,11 +545,22 @@ def test_can_load_meta_cat(self): addons = cat.CAT.load_addons(self.mpp) self.assert_has_one_meta_cat(addons) - def assert_has_one_meta_cat(self, addons: list[AddonComponent]): + def assert_has_one_meta_cat(self, addons: list[tuple[str, AddonComponent]]): self.assertEqual(len(addons), 1) _, addon = addons[0] self.assertIsInstance(addon, MetaCATAddon) + def test_filter_non_existant_loads_no_addons(self): + from medcat.components.addons.relation_extraction.rel_cat import RelCATAddon + loaded = cat.CAT.load_model_pack(self.mpp, keep_addons_of_types=[RelCATAddon]) + self.assertFalse(list(enumerate(loaded.pipe.iter_addons()))) + + def test_filter_meta_cat_loads_meta_cat(self): + loaded = cat.CAT.load_model_pack(self.mpp, keep_addons_of_types=[MetaCATAddon]) + self.assert_has_one_meta_cat(list(enumerate(loaded.pipe.iter_addons()))) + mc = next(loaded.pipe.iter_addons()) + self.assertEqual(mc.config, self.addon.config) + def test_can_filter_addons_empty(self): # NONE -> empty addons = cat.CAT.load_addons(self.mpp, addon_types=[])