feat(medcat-trainer): Avoid loading meta cat upon model pack upload

medcat-trainer/webapp/api/api/models.py

@@ -13,7 +13,8 @@
 from medcat.cat import CAT
 from medcat.cdb import CDB
 from medcat.vocab import Vocab
-from medcat.components.addons.meta_cat.meta_cat import MetaCAT, MetaCATAddon
+from medcat.components.addons.meta_cat.meta_cat import MetaCAT
+from medcat.config.config_meta_cat import ConfigMetaCAT
 from polymorphic.models import PolymorphicModel
 
 from core.settings import MEDIA_ROOT
@@ -111,22 +112,19 @@ def save(self, *args, skip_load=False, **kwargs):
             # DeID model packs do not have a vocab.dat file
             logger.warn('Error loading the Vocab from this model pack - '
                         f'if this is a DeID model pack, this is expected: {vocab_path_abs}')
+        # load dynamically to avoid circular imports
+        from .utils import load_meta_cat_info_from_model_folder
 
         # load MetaCATs
         try:
+            meta_cat_addons = load_meta_cat_info_from_model_folder(unpacked_model_pack_path)
             metaCATmodels = []
-            # should raise an error if there already is a MetaCAT model with this definition
-            addons = CAT.load_addons(unpacked_model_pack_path)
-            meta_cat_addons = [
-                (addon_path, addon) for addon_path, addon in addons
-                if isinstance(addon, MetaCATAddon)]
-            for meta_cat_dir, meta_cat_addon in meta_cat_addons:
-                meta_cat = meta_cat_addon.mc
+            for meta_cat_dir, meta_cat_cnf in meta_cat_addons:
                 mc_model = MetaCATModel()
                 mc_model.meta_cat_dir = meta_cat_dir.replace(f'{MEDIA_ROOT}/', '')
-                mc_model.name = f'{meta_cat.config.general.category_name} - {meta_cat.config.model.model_name}'
+                mc_model.name = f'{meta_cat_cnf.general.category_name} - {meta_cat_cnf.model.model_name}'
                 mc_model.save(unpack_load_meta_cat_dir=False)
-                mc_model.get_or_create_meta_tasks_and_values(meta_cat)
+                mc_model.get_or_create_meta_tasks_and_values(meta_cat_cnf)
                 metaCATmodels.append(mc_model)
             self.meta_cats.set(metaCATmodels)  # Use set() instead of add() for atomic operation
         except Exception as exc:
@@ -211,8 +209,8 @@ class MetaCATModel(models.Model):
                                                   'is set via a model pack .zip upload',
                                         allow_folders=True, editable=False)
 
-    def get_or_create_meta_tasks_and_values(self, meta_cat: MetaCAT):
-        task = meta_cat.config.general.category_name
+    def get_or_create_meta_tasks_and_values(self, mc_config: ConfigMetaCAT):
+        task = mc_config.general.category_name
         mt = MetaTask.objects.filter(name=task).first()
         if not mt:
             mt = MetaTask()
@@ -224,7 +222,7 @@ def get_or_create_meta_tasks_and_values(self, meta_cat: MetaCAT):
             mt.save()
 
         mt_vs = []
-        for meta_task_value in meta_cat.config.general.category_value2id.keys():
+        for meta_task_value in mc_config.general.category_value2id.keys():
             mt_v = MetaTaskValue.objects.filter(name=meta_task_value).first()
             if not mt_v:
                 mt_v = MetaTaskValue()

medcat-trainer/webapp/api/api/tests/test_model_pack_addons.py

@@ -23,24 +23,20 @@
 from django.core.files.base import ContentFile
 from django.test import TestCase, override_settings
 
-from medcat.components.addons.meta_cat.meta_cat import MetaCATAddon
-from medcat.components.addons.relation_extraction.rel_cat import RelCATAddon
-
 from ..models import ModelPack
 
 
-def _make_meta_cat_addon(category_name="Status", model_name="bert"):
-    addon = MagicMock(spec=MetaCATAddon)
-    meta_cat = MagicMock()
-    meta_cat.config.general.category_name = category_name
-    meta_cat.config.model.model_name = model_name
-    meta_cat.config.general.category_value2id = {"True": 0, "False": 1}
-    addon.mc = meta_cat
-    return addon
+def _make_meta_cat_addon_cnf(category_name="Status", model_name="bert"):
+    meta_cat_cnf = MagicMock()
+    meta_cat_cnf.comp_name = "meta_cat"
+    meta_cat_cnf.general.category_name = category_name
+    meta_cat_cnf.model.model_name = model_name
+    meta_cat_cnf.general.category_value2id = {"True": 0, "False": 1}
+    return meta_cat_cnf
 
 
-def _make_rel_cat_addon():
-    return MagicMock(spec=RelCATAddon)
+def _make_rel_cat_addon_cnf():
+    return MagicMock()
 
 
 @override_settings(MEDIA_ROOT=tempfile.mkdtemp())
@@ -56,21 +52,19 @@ def _prepare_model_pack(self, name="addon-pack-test"):
         return model_pack, unpacked
 
     @contextmanager
-    def _register_model_pack(self, model_pack, addons):
+    def _register_model_pack(self, model_pack, addon_cnfs):
         with patch("api.models.CAT.attempt_unpack"), \
                 patch("api.models.CDB.load"), \
                 patch("api.models.Vocab.load"), \
-                patch("api.models.CAT.load_addons", return_value=addons) as load_addons:
+                patch("api.utils._load_global_cnf_addon_cnfs", return_value=addon_cnfs) as load_addons:
             model_pack.save()
             yield load_addons
 
     def test_register_model_pack_with_meta_cat_only(self):
         model_pack, unpacked = self._prepare_model_pack(name="meta-cat-pack")
-        comps = os.path.join(unpacked, "saved_components")
-        meta_cat_path = os.path.join(comps, "addon_meta_cat.Status")
-        addons = [(meta_cat_path, _make_meta_cat_addon())]
+        addon_cnfs = [_make_meta_cat_addon_cnf()]
 
-        with self._register_model_pack(model_pack, addons) as load_addons:
+        with self._register_model_pack(model_pack, addon_cnfs) as load_addons:
             load_addons.assert_called_once_with(unpacked)
 
         self.assertIsNotNone(model_pack.concept_db)
@@ -82,11 +76,9 @@ def test_register_model_pack_with_meta_cat_only(self):
 
     def test_register_model_pack_with_rel_cat_only(self):
         model_pack, unpacked = self._prepare_model_pack(name="rel-cat-pack")
-        comps = os.path.join(unpacked, "saved_components")
-        rel_cat_path = os.path.join(comps, "addon_rel_cat.rel_cat")
-        addons = [(rel_cat_path, _make_rel_cat_addon())]
+        addon_cnfs = [_make_rel_cat_addon_cnf()]
 
-        with self._register_model_pack(model_pack, addons) as load_addons:
+        with self._register_model_pack(model_pack, addon_cnfs) as load_addons:
             load_addons.assert_called_once_with(unpacked)
 
         self.assertIsNotNone(model_pack.concept_db)
@@ -95,15 +87,12 @@ def test_register_model_pack_with_rel_cat_only(self):
 
     def test_register_model_pack_registers_multiple_meta_cats(self):
         model_pack, unpacked = self._prepare_model_pack(name="multi-meta-cat-pack")
-        comps = os.path.join(unpacked, "saved_components")
-        addons = [
-            (os.path.join(comps, "addon_meta_cat.Status"),
-             _make_meta_cat_addon(category_name="Status", model_name="bert")),
-            (os.path.join(comps, "addon_meta_cat.Experiencer"),
-             _make_meta_cat_addon(category_name="Experiencer", model_name="roberta")),
+        addon_cnfs = [
+             _make_meta_cat_addon_cnf(category_name="Status", model_name="bert"),
+             _make_meta_cat_addon_cnf(category_name="Experiencer", model_name="roberta"),
         ]
 
-        with self._register_model_pack(model_pack, addons):
+        with self._register_model_pack(model_pack, addon_cnfs):
             pass
 
         self.assertEqual(model_pack.meta_cats.count(), 2)
@@ -114,15 +103,9 @@ def test_register_model_pack_registers_multiple_meta_cats(self):
 
     def test_register_model_pack_with_meta_cat_and_rel_cat(self):
         model_pack, unpacked = self._prepare_model_pack(name="mixed-addon-pack")
-        comps = os.path.join(unpacked, "saved_components")
-        meta_cat_path = os.path.join(comps, "addon_meta_cat.Status")
-        rel_cat_path = os.path.join(comps, "addon_rel_cat.rel_cat")
-        addons = [
-            (meta_cat_path, _make_meta_cat_addon()),
-            (rel_cat_path, _make_rel_cat_addon()),
-        ]
+        addon_cnfs = [_make_meta_cat_addon_cnf(), _make_rel_cat_addon_cnf(),]
 
-        with self._register_model_pack(model_pack, addons) as load_addons:
+        with self._register_model_pack(model_pack, addon_cnfs) as load_addons:
             load_addons.assert_called_once_with(unpacked)
 
         self.assertIsNotNone(model_pack.concept_db)
@@ -133,6 +116,7 @@ def test_register_model_pack_with_meta_cat_and_rel_cat(self):
         self.assertEqual(meta_cat_model.name, "Status - bert")
         self.assertTrue(meta_cat_model.meta_cat_dir.endswith("addon_meta_cat.Status"))
 
+
     def test_register_model_pack_without_addons(self):
         model_pack, unpacked = self._prepare_model_pack(name="no-addon-pack")
 

medcat-trainer/webapp/api/api/tests/test_model_inference.py → medcat-trainer/webapp/api/api/tests/test_with_dummy_model.py

@@ -10,8 +10,10 @@
 
 import medcat
 from medcat.cat import CAT
+from medcat.components.addons.meta_cat import MetaCATAddon
 
 from api.models import Document, ProjectAnnotateEntities, ModelPack
+from api.model_import import import_model_pack
 
 
 RAW_MODEL_PATH = os.path.join(
@@ -25,14 +27,15 @@
     "..", "..", "media"
 )
 MODEL_PATH = os.path.join(
-    MEDIA_PATH, "fake_model_pack.zip"
+    MEDIA_PATH, "dummy_model_pack.zip"
 )
+MODEL_PATH_UNPACKED = MODEL_PATH.removesuffix(".zip")
 
 
 HAS_KNOWN_FAILURE = medcat.__version__ in ("2.8.0", "2.8.1", "2.8.2", "2.8.3")
 
 
-class ModelInferenceTests(TestCase):
+class BaseRealModelTests(TestCase):
 
     @classmethod
     def setUpClass(cls):
@@ -55,14 +58,17 @@ def _copy_model(cls):
             MODEL_PATH
         )
 
+
+class ModelInferenceTests(BaseRealModelTests):
+
     def setUp(self):
         # A real user — the view reads request.user
         self.user = create_user(username="testuser", password="password", is_staff=True)
         self.client = APIClient()
         self.client.force_authenticate(self.user)
 
         self.model_pack = ModelPack.objects.create(
-            name='fake-model',
+            name='test-model',
             model_pack=MODEL_PATH,
         )
 
@@ -108,3 +114,74 @@ def test_can_use_model_for_inference(self):
         # The document should now be in prepared_documents
         self.assertTrue(self.project.prepared_documents.all())
         self.assertEqual(len(self.project.prepared_documents.all()), len(doc_ids))
+
+
+class ModelImportTests(BaseRealModelTests):
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls._create_meta_cat_for_model()
+
+    # NOTE: if MetaCATAddons are going to be loaded,
+    #       it would fail with this, but with previous
+    #       code an attempt would have been made.
+    @classmethod
+    def _create_meta_cat_for_model(cls):
+        # first, create the folder
+        from api.utils import _PATH_TO_META_CAT_START, _PATH_TO_CNF_ON_DISK
+        mc_path = os.path.join(MODEL_PATH_UNPACKED, _PATH_TO_META_CAT_START + "Subject")
+        os.mkdir(os.path.dirname(mc_path))
+        os.mkdir(mc_path)
+        # then add .serialised_by
+        sb_path = os.path.join(mc_path, ".serialised_by")
+        with open(sb_path, 'w') as f:
+            f.write("dill")
+        # finally, add config
+        from medcat.config.config_meta_cat import ConfigMetaCAT
+        # create config
+        cnf = ConfigMetaCAT()
+        cnf.general.category_name = "Subject"
+        # load off disk
+        import dill
+        gcnf_path = os.path.join(MODEL_PATH_UNPACKED, _PATH_TO_CNF_ON_DISK)
+        with open(gcnf_path, 'rb') as f:
+            data = dill.load(f)
+        data['addons'] = [cnf]
+        # save to disk
+        with open(gcnf_path, 'wb') as f:
+            dill.dump(data, f)
+
+    def setUp(self):
+        # A real user — the view reads request.user
+        self.user = create_user(username="testuser", password="password", is_staff=True)
+
+    # NOTE: mocking dispatch just so it won't complain about missing sender
+    @patch('api.model_import.dispatch')
+    def test_model_import_does_not_load_addons(self, mock_dispatch):
+        with patch.object(MetaCATAddon, "__init__") as mock_init:
+            import_model_pack(
+                MODEL_PATH,
+                name='test-model',
+                user=self.user,
+                description='Fake model!',
+                source_uri='https://some/address',
+            )
+        mock_init.assert_not_called()
+
+
+# class RealModelTests(TestCase):
+#     MODEL_PATH = "../../../medcat-v2/.temp/20230227__kch_gstt_trained_model_f76d2121b77c3e9a/"
+
+#     def test_can_read_from_real_model(self):
+#         from api.utils import load_meta_cat_info_from_model_folder
+#         infos = load_meta_cat_info_from_model_folder(self.MODEL_PATH)
+#         self.assertEqual(len(infos), 3)
+#         categories = {
+#             cnf.general.category_name
+#             for _, cnf in infos
+#         }
+#         self.assertEqual(
+#             categories,
+#             {"Presence", "Subject", "Time"}
+#         )

medcat-trainer/webapp/api/api/utils.py

@@ -2,6 +2,7 @@
 import logging
 import os
 from typing import List
+import dill
 
 import requests
 from background_task import background
@@ -11,6 +12,7 @@
 from django.dispatch import receiver
 from medcat.cat import CAT
 from medcat.cdb import CDB
+from medcat.config.config_meta_cat import ConfigMetaCAT
 from medcat.components.ner.trf.deid import DeIdModel
 from medcat.tokenizing.tokens import UnregisteredDataPathException
 from opentelemetry import trace
@@ -499,3 +501,46 @@ def env_str_to_bool(var: str, default: bool):
         elif val.lower() in ('0', 'false', 'f', 'n'):
             return False
     return val
+
+
+# NOTE: meta cats are saved at
+#       saved_components/addon_meta_cat.<category>
+#       however, their configs are (uniquely!) saved
+#       along with the config for the entire model pack
+#       at cdb/config/components/raw_dict.dat at 'addons'
+#       so we will load that directly
+
+_PATH_TO_CNF_ON_DISK = os.path.join(
+    "cdb", "config", "components", "raw_dict.dat"
+)
+_PATH_TO_META_CAT_START = os.path.join(
+    "saved_components", "addon_meta_cat."
+)
+
+def _get_meta_cat_path(model_folder: str, cnf: ConfigMetaCAT) -> str:
+    return os.path.join(
+        model_folder, _PATH_TO_META_CAT_START) + cnf.general.category_name
+
+
+def _load_global_cnf_addon_cnfs(model_pack_path: str) -> list:
+    global_cnf_path = os.path.join(model_pack_path, _PATH_TO_CNF_ON_DISK)
+    with open(os.path.join(model_pack_path, '.serialised_by')) as f:
+        contents = f.read()
+    if contents.strip().lower() != "dill":
+        raise ValueError(
+            "Unable to load MetaCAT info without loading the meta cats from this model "
+            "because of its serialisation type (%s) - we only have implementation for 'dill'")
+    with open(global_cnf_path, 'rb') as f:
+        cnf = dill.load(f)
+        return cnf['addons']
+
+
+def load_meta_cat_info_from_model_folder(
+    model_pack_path: str
+) -> list[tuple[str, ConfigMetaCAT]]:
+    addon_cnfs = _load_global_cnf_addon_cnfs(model_pack_path)
+    return [
+        (_get_meta_cat_path(model_pack_path, mc_cnf), mc_cnf)
+        for mc_cnf in addon_cnfs
+        if mc_cnf.comp_name == "meta_cat"
+    ]