add: metacat can predict on spans in arbitrary spangroups (#391)

* add: ability to predict on other spangroups

* add: pr comments and better error

* fix: typo

* fix: linting

Author: jkgenser

medcat/config_meta_cat.py

@@ -37,6 +37,9 @@ class General(MixingConfig, BaseModel):
     a deployment."""
     pipe_batch_size_in_chars: int = 20000000
     """How many characters are piped at once into the meta_cat class"""
+    span_group: Optional[str] = None
+    """If set, the spacy span group that the metacat model will assign annotations. 
+    Otherwise defaults to doc._.ents or doc.ents per the annotate_overlapping settings"""
 
     class Config:
         extra = Extra.allow

medcat/meta_cat.py

@@ -5,7 +5,7 @@
 import numpy
 from multiprocessing import Lock
 from torch import nn, Tensor
-from spacy.tokens import Doc
+from spacy.tokens import Doc, Span
 from datetime import datetime
 from typing import Iterable, Iterator, Optional, Dict, List, Tuple, cast, Union
 from medcat.utils.hasher import Hasher
@@ -357,6 +357,20 @@ def load(cls, save_dir_path: str, config_dict: Optional[Dict] = None) -> "MetaCA
 
         return meta_cat
 
+    def get_ents(self, doc: Doc) -> Iterable[Span]:
+        spangroup_name = self.config.general.span_group
+        if spangroup_name:
+            try:
+                return doc.spans[spangroup_name]
+            except KeyError:
+                raise Exception(f"Configuration error MetaCAT was configured to set meta_anns on {spangroup_name} but this spangroup was not set on the doc.")
+
+        # Should we annotate overlapping entities
+        if self.config.general['annotate_overlapping']:
+            return doc._.ents
+
+        return doc.ents
+
     def prepare_document(self, doc: Doc, input_ids: List, offset_mapping: List, lowercase: bool) -> Tuple:
         """Prepares document.
 
@@ -381,11 +395,7 @@ def prepare_document(self, doc: Doc, input_ids: List, offset_mapping: List, lowe
         cntx_right = config.general['cntx_right']
         replace_center = config.general['replace_center']
 
-        # Should we annotate overlapping entities
-        if config.general['annotate_overlapping']:
-            ents = doc._.ents
-        else:
-            ents = doc.ents
+        ents = self.get_ents(doc)
 
         samples = []
         last_ind = 0
@@ -522,10 +532,7 @@ def _set_meta_anns(self,
 
                     predictions = all_predictions[start_ind:end_ind]
                     confidences = all_confidences[start_ind:end_ind]
-                    if config.general['annotate_overlapping']:
-                        ents = doc._.ents
-                    else:
-                        ents = doc.ents
+                    ents = self.get_ents(doc)
 
                     for ent in ents:
                         ent_ind = ent_id2ind[ent._.id]

tests/test_meta_cat.py

@@ -7,7 +7,8 @@
 from medcat.meta_cat import MetaCAT
 from medcat.config_meta_cat import ConfigMetaCAT
 from medcat.tokenizers.meta_cat_tokenizers import TokenizerWrapperBERT
-
+import spacy
+from spacy.tokens import Span
 
 class MetaCATTests(unittest.TestCase):
 
@@ -19,7 +20,7 @@ def setUpClass(cls) -> None:
         config.train['nepochs'] = 1
         config.model['input_size'] = 100
 
-        cls.meta_cat = MetaCAT(tokenizer=tokenizer, embeddings=None, config=config)
+        cls.meta_cat: MetaCAT = MetaCAT(tokenizer=tokenizer, embeddings=None, config=config)
 
         cls.tmp_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "tmp")
         os.makedirs(cls.tmp_dir, exist_ok=True)
@@ -44,6 +45,50 @@ def test_save_load(self):
 
         self.assertEqual(f1, n_f1)
 
+    def _prepare_doc_w_spangroup(self, spangroup_name: str):
+        """
+        Create spans under an arbitrary spangroup key
+        """
+        Span.set_extension('id', default=0, force=True)
+        Span.set_extension('meta_anns', default=None, force=True)
+        nlp = spacy.blank("en")
+        doc = nlp("Pt has diabetes and copd.")
+        span_0 = doc.char_span(7,15, label="diabetes")
+        assert span_0.text == 'diabetes'
+
+        span_1 = doc.char_span(20,24, label="copd")
+        assert span_1.text == 'copd'
+        doc.spans[spangroup_name] = [span_0, span_1]
+        return doc
+
+    def test_predict_spangroup(self):
+        json_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'resources', 'mct_export_for_meta_cat_test.json')
+        self.meta_cat.train(json_path, save_dir_path=self.tmp_dir)
+        self.meta_cat.save(self.tmp_dir)
+        n_meta_cat = MetaCAT.load(self.tmp_dir)
+
+        spangroup_name = "mock_span_group"
+        n_meta_cat.config.general.span_group = spangroup_name
+
+        doc = self._prepare_doc_w_spangroup(spangroup_name)
+        doc = n_meta_cat(doc)
+        spans = doc.spans[spangroup_name]
+        self.assertEqual(len(spans), 2)
+
+        # All spans are annotate
+        for span in spans:
+            self.assertEqual(span._.meta_anns['Status']['value'], "Affirmed")
+
+        # Informative error if spangroup is not set
+        doc = self._prepare_doc_w_spangroup("foo")
+        n_meta_cat.config.general.span_group = "bar"
+        try:
+            doc = n_meta_cat(doc)
+        except Exception as error:
+            self.assertIn("Configuration error", str(error))
+
+        n_meta_cat.config.general.span_group = None
+
 
 if __name__ == '__main__':
     unittest.main()