feat(medcat): Allow loading with fewer addons

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,

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']

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)

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=[])