diff --git a/medcat-trainer/webapp/api/api/models.py b/medcat-trainer/webapp/api/api/models.py index ec1c673da..eab137b41 100644 --- a/medcat-trainer/webapp/api/api/models.py +++ b/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() diff --git a/medcat-trainer/webapp/api/api/tests/test_model_pack_addons.py b/medcat-trainer/webapp/api/api/tests/test_model_pack_addons.py index 9b2797e63..49fc8cde3 100644 --- a/medcat-trainer/webapp/api/api/tests/test_model_pack_addons.py +++ b/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") diff --git a/medcat-trainer/webapp/api/api/tests/test_model_inference.py b/medcat-trainer/webapp/api/api/tests/test_with_dummy_model.py similarity index 52% rename from medcat-trainer/webapp/api/api/tests/test_model_inference.py rename to medcat-trainer/webapp/api/api/tests/test_with_dummy_model.py index f2dbc395a..1d760b0d9 100644 --- a/medcat-trainer/webapp/api/api/tests/test_model_inference.py +++ b/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,6 +58,9 @@ 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) @@ -62,7 +68,7 @@ def setUp(self): 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"} +# ) diff --git a/medcat-trainer/webapp/api/api/utils.py b/medcat-trainer/webapp/api/api/utils.py index 20c236276..c786fa3b8 100644 --- a/medcat-trainer/webapp/api/api/utils.py +++ b/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. +# 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" + ]