CU-8699vkmu4: Allow load with merging config(s) (#53)

* CU-8699vkmu4: Add option to load MetaCATs with config

* CU-8699vkmu4: Add small test to make sure MetaCAT config is merged correctly

* CU-8699vkmu4: Allow model config when loading model pack

* CU-8699vkmu4: Add small tests for CAT model config addition

* CU-8699vkmu4: Allow merging in arbitrary addon configs upon model load if/when specified

* CU-8699vkmu4: Add small test for addon config merge upon CAT load

* CU-8699vkmu4: Make a method to a class method

* CU-8699vkmu4: Limit indentation

* CU-8699vkmu4: Generalis load_addons method to for all addons

Author: mart-r

medcat-v2/medcat/cat.py

@@ -49,6 +49,8 @@ def __init__(self,
                  vocab: Union[Vocab, None] = None,
                  config: Optional[Config] = None,
                  model_load_path: Optional[str] = None,
+                 config_dict: Optional[dict] = None,
+                 addon_config_dict: Optional[dict[str, dict]] = None,
                  ) -> None:
         self.cdb = cdb
         self.vocab = vocab
@@ -60,20 +62,24 @@ def __init__(self,
         elif config is not None:
             self.cdb.config = config
         self.config = config
+        if config_dict:
+            self.config.merge_config(config_dict)
 
         self._trainer: Optional[Trainer] = None
-        self._pipeline = self._recreate_pipe(model_load_path)
+        self._pipeline = self._recreate_pipe(model_load_path, addon_config_dict)
         self.usage_monitor = UsageMonitor(
             self._get_hash, self.config.general.usage_monitor)
 
-    def _recreate_pipe(self, model_load_path: Optional[str] = None
+    def _recreate_pipe(self, model_load_path: Optional[str] = None,
+                       addon_config_dict: Optional[dict[str, dict]] = None,
                        ) -> Pipeline:
         if hasattr(self, "_pipeline"):
             old_pipe = self._pipeline
         else:
             old_pipe = None
         self._pipeline = Pipeline(self.cdb, self.vocab, model_load_path,
-                                  old_pipe=old_pipe)
+                                  old_pipe=old_pipe,
+                                  addon_config_dict=addon_config_dict)
         return self._pipeline
 
     @classmethod
@@ -668,11 +674,21 @@ def attempt_unpack(cls, zip_path: str) -> str:
         return model_pack_path
 
     @classmethod
-    def load_model_pack(cls, model_pack_path: str) -> 'CAT':
+    def load_model_pack(cls, model_pack_path: str,
+                        config_dict: Optional[dict] = None,
+                        addon_config_dict: Optional[dict[str, dict]] = None
+                        ) -> 'CAT':
         """Load the model pack from file.
 
         Args:
             model_pack_path (str): The model pack path.
+            config_dict (Optional[dict]): The model config to
+                merge in before initialising the pipe. Defaults to None.
+            addon_config_dict (Optional[dict]): The Addon-specific
+                config dict to merge in before pipe initialisation.
+                If specified, it needs to have an addon dict per name.
+                For instance, `{"meta_cat.Subject": {}}` would apply
+                to the specific MetaCAT.
 
         Raises:
             ValueError: If the saved data does not represent a model pack.
@@ -703,7 +719,9 @@ def load_model_pack(cls, model_pack_path: str) -> 'CAT':
                             TOKENIZER_PREFIX,
                             # components will be loaded semi-manually
                             # within the creation of pipe
-                            COMPONENTS_FOLDER})
+                            COMPONENTS_FOLDER},
+                          config_dict=config_dict,
+                          addon_config_dict=addon_config_dict)
         # NOTE: deserialising of components that need serialised
         #       will be dealt with upon pipeline creation automatically
         if not isinstance(cat, CAT):
@@ -730,33 +748,40 @@ def load_cdb(cls, model_pack_path: str) -> CDB:
 
     @classmethod
     def load_addons(
-            cls, model_pack_path: str, meta_cat_config_dict: Optional[dict] = None
+            cls, model_pack_path: str,
+            addon_config_dict: Optional[dict[str, dict]] = None
             ) -> list[tuple[str, AddonComponent]]:
         """Load addons based on a model pack path.
 
         Args:
             model_pack_path (str): path to model pack, zip or dir.
-            meta_cat_config_dict (Optional[dict]):
-                A config dict that will overwrite existing configs in meta_cat.
-                e.g. meta_cat_config_dict = {'general': {'device': 'cpu'}}.
-                Defaults to None.
+            addon_config_dict (Optional[dict]): The Addon-specific
+                config dict to merge in before pipe initialisation.
+                If specified, it needs to have an addon dict per name.
+                For instance,
+                `{"meta_cat.Subject": {'general': {'device': 'cpu'}}}`
+                would apply to the specific MetaCAT.
 
         Returns:
             List[tuple(str, AddonComponent)]: list of pairs of adddon names the addons.
         """
         components_folder = os.path.join(model_pack_path, COMPONENTS_FOLDER)
         if not os.path.exists(components_folder):
             return []
-        addon_paths = [
-            folder_path
+        addon_paths_and_names = [
+            (folder_path, folder_name.removeprefix(AddonComponent.NAME_PREFIX))
             for folder_name in os.listdir(components_folder)
             if os.path.isdir(folder_path := os.path.join(
                 components_folder, folder_name))
             and folder_name.startswith(AddonComponent.NAME_PREFIX)
         ]
         loaded_addons = [
-            addon for addon_path in addon_paths
-            if isinstance(addon := deserialise(addon_path), AddonComponent)
+            addon for addon_path, addon_name in addon_paths_and_names
+            if isinstance(addon := (
+                deserialise(addon_path, model_config=addon_config_dict.get(addon_name))
+                if addon_config_dict else
+                deserialise(addon_path)
+                ), AddonComponent)
         ]
         return [(addon.full_name, addon) for addon in loaded_addons]
 

medcat-v2/medcat/components/addons/meta_cat/meta_cat.py

@@ -221,6 +221,8 @@ def deserialise_from(cls, folder_path: str, **init_kwargs
                 "Inferring config from file at '%s'", folder_path,
                 config_path)
             cnf = ConfigMetaCAT.load(config_path)
+        if 'model_config' in init_kwargs:
+            cnf.merge_config(init_kwargs['model_config'])
         if 'tokenizer' in init_kwargs:
             tokenizer = init_kwargs['tokenizer']
         else:

medcat-v2/medcat/pipeline/pipeline.py

@@ -72,7 +72,8 @@ class Pipeline:
     def __init__(self, cdb: CDB, vocab: Optional[Vocab],
                  model_load_path: Optional[str],
                  # NOTE: upon reload, old pipe can be useful
-                 old_pipe: Optional['Pipeline'] = None):
+                 old_pipe: Optional['Pipeline'] = None,
+                 addon_config_dict: Optional[dict[str, dict]] = None):
         self.cdb = cdb
         # NOTE: Vocab is None in case of DeID models and thats fine then,
         #       but it should be non-None otherwise
@@ -81,7 +82,7 @@ def __init__(self, cdb: CDB, vocab: Optional[Vocab],
         self._tokenizer = self._init_tokenizer()
         self._components: list[CoreComponent] = []
         self._addons: list[AddonComponent] = []
-        self._init_components(model_load_path, old_pipe)
+        self._init_components(model_load_path, old_pipe, addon_config_dict)
 
     @property
     def tokenizer(self) -> BaseTokenizer:
@@ -172,8 +173,45 @@ def _load_saved_core_component(self, cct_name: str, comp_folder_path: str
                 f"'{comp.get_type().name}' instead.")
         return comp
 
+    @classmethod
+    def _attempt_merge(
+            cls, addon_cnf: ComponentConfig,
+            addon_config_dict: dict[str, dict]) -> None:
+        for name, config_dict in addon_config_dict.items():
+            if not name.startswith(addon_cnf.comp_name):
+                continue
+            # TODO: is there an option to do this in a more general way?
+            #       right now it's an implementation-specific code smell
+            if isinstance(addon_cnf, ConfigMetaCAT):
+                full_name = f"{addon_cnf.comp_name}.{addon_cnf.general.category_name}"
+                if name == full_name:
+                    addon_cnf.merge_config(config_dict)
+                    return
+                continue
+            logger.warning(
+                "No implementation specified for defining if/when %s"
+                "should apply to addon config (e.g %s)",
+                type(addon_cnf).__name__, name)
+            # if only 1 of the type, then just merge
+            similars = [cd for oname, cd in addon_config_dict.items()
+                        if oname.startswith(addon_cnf.comp_name)]
+            if len(similars) == 1:
+                logger.warning(
+                    "Since there is only 1 config for this type (%s) specified "
+                    "we will just merge the configs (@%s).",
+                    addon_cnf.comp_name, name)
+                addon_cnf.merge_config(config_dict)
+                return
+            else:
+                logger.warning(
+                    "There are %d similar configs (%s) specified, so unable to "
+                    "merge the config since it's ambiguous (@%s)",
+                    len(similars), addon_cnf.comp_name, name)
+
     def _init_components(self, model_load_path: Optional[str],
-                         old_pipe: Optional['Pipeline']) -> None:
+                         old_pipe: Optional['Pipeline'],
+                         addon_config_dict: Optional[dict[str, dict]],
+                         ) -> None:
         (loaded_core_component_paths,
          loaded_addon_component_paths) = self._get_loaded_components_paths(
              model_load_path)
@@ -186,6 +224,8 @@ def _init_components(self, model_load_path: Optional[str],
                     CoreComponentType[cct_name], model_load_path)
             self._components.append(comp)
         for addon_cnf in self.config.components.addons:
+            if addon_config_dict:
+                self._attempt_merge(addon_cnf, addon_config_dict)
             addon = self._init_addon(
                 addon_cnf, loaded_addon_component_paths, old_pipe)
             # mark as not dirty at loat / init time

medcat-v2/tests/test_cat.py

@@ -71,6 +71,19 @@ def setUpClass(cls):
             cls.model.config.components.linking.train = False
 
 
+class ConfigMergeTests(unittest.TestCase):
+    spacy_model_name = 'en_core_web_lg'
+    model_dict = {
+        "general": {'nlp': {"modelname": spacy_model_name}}
+    }
+
+    def test_can_merge_config(self):
+        model = cat.CAT.load_model_pack(
+            EXAMPLE_MODEL_PACK_ZIP, config_dict=self.model_dict)
+        self.assertEqual(
+            model.config.general.nlp.modelname, self.spacy_model_name)
+
+
 class InferenceFromLoadedTests(TrainedModelTests):
 
     def test_can_load_model(self):
@@ -298,6 +311,22 @@ def test_can_load_meta_cat(self):
         _, addon = addons[0]
         self.assertIsInstance(addon, MetaCATAddon)
 
+    def test_can_load_meta_cat_with_addon_cnf(self, seed: int = -41):
+        mc: MetaCATAddon = cat.CAT.load_addons(
+            self.mpp, addon_config_dict={
+                "meta_cat.Status": {
+                    "general": {"seed": seed}}})[0][1]
+        self.assertEqual(mc.config.general.seed, seed)
+
+    def test_can_merge_cnf_upon_load(self, use_seed: int = -4):
+        loaded = cat.CAT.load_model_pack(
+            self.mpp,
+            addon_config_dict={
+                "meta_cat.Status": {"general": {"seed": use_seed}}
+            })
+        addon: MetaCATAddon = list(loaded._pipeline.iter_addons())[0]
+        self.assertEqual(addon.config.general.seed, use_seed)
+
 
 class CatWithChangesMetaCATTests(CatWithMetaCATTests):
     EXPECTED_HASH = "0b22401059a08380"