From 96d2d760427cdb6eebd15f522cf042558603d9c0 Mon Sep 17 00:00:00 2001 From: Mart Ratas Date: Wed, 24 Jun 2026 15:08:14 +0100 Subject: [PATCH 1/5] Add initial changes for optional addon filtering at load time --- medcat-v2/medcat/cat.py | 29 ++++++++++++++++++++++----- medcat-v2/medcat/pipeline/__init__.py | 4 ++-- medcat-v2/medcat/pipeline/pipeline.py | 16 +++++++++++++++ 3 files changed, 42 insertions(+), 7 deletions(-) diff --git a/medcat-v2/medcat/cat.py b/medcat-v2/medcat/cat.py index efe969ebe..4e1ef0087 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,20 @@ 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__ + if addon_config_dict is None: + addon_config_dict = {} + addons_to_keep = { + _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..5945baef2 100644 --- a/medcat-v2/medcat/pipeline/pipeline.py +++ b/medcat-v2/medcat/pipeline/pipeline.py @@ -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,8 @@ 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(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) From 5cdf351fdb00fe62b3a61f82ad8107714d68fea5 Mon Sep 17 00:00:00 2001 From: Mart Ratas Date: Wed, 24 Jun 2026 15:08:57 +0100 Subject: [PATCH 2/5] Add a bit of a comment on the code --- medcat-v2/medcat/cat.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/medcat-v2/medcat/cat.py b/medcat-v2/medcat/cat.py index 4e1ef0087..f06fc5dfd 100644 --- a/medcat-v2/medcat/cat.py +++ b/medcat-v2/medcat/cat.py @@ -863,6 +863,8 @@ def load_model_pack( 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 = { From 37ebd0a88efa6f7328ce4d3ec53b93cac212914e Mon Sep 17 00:00:00 2001 From: Mart Ratas Date: Wed, 24 Jun 2026 15:32:05 +0100 Subject: [PATCH 3/5] Add a few new tests for model load with addons --- medcat-v2/tests/test_cat.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/medcat-v2/tests/test_cat.py b/medcat-v2/tests/test_cat.py index db5d75c48..3b88aa9b3 100644 --- a/medcat-v2/tests/test_cat.py +++ b/medcat-v2/tests/test_cat.py @@ -545,11 +545,20 @@ 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()))) + def test_can_filter_addons_empty(self): # NONE -> empty addons = cat.CAT.load_addons(self.mpp, addon_types=[]) From 0e1a7804fbb9de00491fbdf27ae76ce154557769 Mon Sep 17 00:00:00 2001 From: Mart Ratas Date: Wed, 24 Jun 2026 15:40:11 +0100 Subject: [PATCH 4/5] Fix some linting issue --- medcat-v2/medcat/cat.py | 2 +- medcat-v2/medcat/pipeline/pipeline.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/medcat-v2/medcat/cat.py b/medcat-v2/medcat/cat.py index f06fc5dfd..cf38c8016 100644 --- a/medcat-v2/medcat/cat.py +++ b/medcat-v2/medcat/cat.py @@ -867,7 +867,7 @@ def load_model_pack( # out so they don't need to be loaded if addon_config_dict is None: addon_config_dict = {} - addons_to_keep = { + addons_to_keep: dict = { _ENABLED_ADDONS_PATH: [ addon_type.addon_type for addon_type in keep_addons_of_types] diff --git a/medcat-v2/medcat/pipeline/pipeline.py b/medcat-v2/medcat/pipeline/pipeline.py index 5945baef2..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 @@ -255,7 +255,8 @@ def _init_components(self, model_load_path: Optional[str], 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(addon_config_dict[_ENABLED_ADDONS_PATH]) + 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) From 83e40008225cfc961590e45dde8195643d1bf012 Mon Sep 17 00:00:00 2001 From: Mart Ratas Date: Thu, 25 Jun 2026 10:27:57 +0100 Subject: [PATCH 5/5] Add further assert to make sure configs remain equal upon load --- medcat-v2/tests/test_cat.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/medcat-v2/tests/test_cat.py b/medcat-v2/tests/test_cat.py index 3b88aa9b3..baf406593 100644 --- a/medcat-v2/tests/test_cat.py +++ b/medcat-v2/tests/test_cat.py @@ -558,6 +558,8 @@ def test_filter_non_existant_loads_no_addons(self): 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