Improve some nlp utils (#178)

* only consider one model per lang

* cache lang to models func

* add func to process one text into spacy doc

* tests: add unit test for new nlp util func

* refactor: use single text process func in places

* tests: update mocks to use new func

Author: bdewilde

colandr/lib/extractors/locations.py

@@ -6,7 +6,7 @@
 
 from spacy.tokens import Span
 
-from ..nlp.utils import process_texts_into_docs
+from ..nlp.utils import process_text_into_doc
 from .metadata import Metadata
 
 
@@ -50,8 +50,7 @@ def extract_locations(self, record_id: int, text: str) -> list[Metadata]:
         if not text or not text.strip():
             return []
 
-        processed_docs_iter = process_texts_into_docs([text], max_len=None)
-        doc = next(iter(processed_docs_iter), None)
+        doc = process_text_into_doc(text, max_len=None)
         if doc is None:
             return []
 

colandr/lib/extractors/review_model.py

@@ -17,7 +17,7 @@
 from sklearn.preprocessing import MultiLabelBinarizer, StandardScaler
 from spacy.tokens import Doc, Span
 
-from ..nlp.utils import process_texts_into_docs
+from ..nlp.utils import process_text_into_doc, process_texts_into_docs
 from .metadata import Metadata
 
 
@@ -444,10 +444,7 @@ def _process_text(self, text_content: str) -> tuple[pd.DataFrame, list[dict]]:
             Tuple containing the feature DataFrame and original sentences list.
         """
         main_content, _ = self._split_references(text_content)
-        processed_docs_iter = process_texts_into_docs(
-            [main_content], max_len=None, exclude=("ner",)
-        )
-        doc = next(iter(processed_docs_iter), None)
+        doc = process_text_into_doc(main_content, max_len=None, exclude=("ner",))
         return self._extract_features_from_doc(doc)
 
     def _is_valid_sentence(self, sent: Optional[Span]) -> bool:

colandr/lib/nlp/utils.py

@@ -45,18 +45,19 @@ def detect_languages(texts: Iterable[str]) -> list[t.Optional[str]]:
     ]
 
 
-def get_lang_to_models() -> dict[str, list[str]]:
+@functools.cache
+def get_lang_to_models() -> dict[str, str]:
     """Get a mapping of ISO language code to installed spacy language models."""
-    lang_to_models = collections.defaultdict(list)
+    lang_to_models = {}
     models = spacy.util.get_installed_models()
     for model in models:
         if "_" in model:
             lang, _ = model.split("_", 1)
-            lang_to_models[lang].append(model)
+            lang_to_models[lang] = model
         else:
             LOGGER.warning("found unexpected spacy model name: %s", model)
 
-    return dict(lang_to_models)
+    return lang_to_models
 
 
 @functools.lru_cache(maxsize=10)
@@ -89,6 +90,41 @@ def load_spacy_lang(name: str, **kwargs) -> SpacyLang:
     return spacy_lang
 
 
+def process_text_into_doc(
+    text: str,
+    *,
+    max_len: t.Optional[int] = 1000,
+    fallback_lang: t.Optional[str] = "en",
+    **kwargs,
+) -> t.Optional[SpacyDoc]:
+    """
+    Args:
+        text
+        max_len: Maximum number of chars (code points) in text to include
+            when identifying its language and processing into a spacy document.
+        fallback_lang: Fallback language used in place of low-confidence predictions.
+        **kwargs: Passed as-is into :func:`load_spacy_lang()` .
+    """
+    # clean up whitespace, to make it easier on lang detector
+    text = text.strip().replace("\n", " ")
+    # truncate texts, optionally
+    if max_len is not None:
+        text = text[:max_len]
+    # identify most probable language (w/ optional fallback) for text
+    lang = detect_language(text) or fallback_lang
+    lang_models = get_lang_to_models()
+    if lang in lang_models:
+        spacy_lang: SpacyLang = load_spacy_lang(lang_models[lang], **kwargs)
+        spacy_doc = spacy_lang(text)
+        return spacy_doc
+    else:
+        LOGGER.info(
+            "unable to load spacy model for text with lang='%s'; doc set to null ...",
+            lang,
+        )
+        return None
+
+
 def process_texts_into_docs(
     texts: Iterable[str],
     *,
@@ -120,7 +156,7 @@ def process_texts_into_docs(
     lang_models = get_lang_to_models()
     for lang, tl_grp in itertools.groupby(text_langs, key=itemgetter(1)):
         if lang in lang_models:
-            spacy_lang = load_spacy_lang(lang_models[lang][0], **kwargs)
+            spacy_lang = load_spacy_lang(lang_models[lang], **kwargs)
             spacy_docs = spacy_lang.pipe((text for text, _ in tl_grp), n_process=1)
             for spacy_doc in spacy_docs:
                 yield spacy_doc

colandr/tasks.py

@@ -295,13 +295,12 @@ def get_fulltext_text_content_vector(fulltext_id: int):
         )
         return
 
-    docs = nlp_utils.process_texts_into_docs(
-        [fulltext["text_content"]],
+    doc = nlp_utils.process_text_into_doc(
+        fulltext["text_content"],
         max_len=3000,
         fallback_lang=None,
         exclude=("parser", "ner"),
     )
-    doc = next(iter(docs))
     text_content_vector_rep = doc.vector.tolist() if doc is not None else None
     if text_content_vector_rep is None:
         LOGGER.warning(
@@ -310,12 +309,12 @@ def get_fulltext_text_content_vector(fulltext_id: int):
         return
 
     fulltext["text_content_vector_rep"] = text_content_vector_rep
-    stmt = (
+    update_stmt = (
         sa.update(models.Study)
         .where(models.Study.id == fulltext_id)
         .values(fulltext=fulltext)
     )
-    db.session.execute(stmt)
+    db.session.execute(update_stmt)
     db.session.commit()
 
 

tests/lib/extractors/test_locations.py

@@ -48,8 +48,8 @@ def test_is_in_reference(self):
 
             assert extractor.is_in_reference(mock_ent) is False
 
-    @patch("colandr.lib.extractors.locations.process_texts_into_docs")
-    def test_extract_locations(self, mock_process_texts):
+    @patch("colandr.lib.extractors.locations.process_text_into_doc")
+    def test_extract_locations(self, mock_process_text):
         """Test extract_locations function."""
         extractor = LocationExtractor()
 
@@ -80,7 +80,7 @@ def test_extract_locations(self, mock_process_texts):
 
         mock_doc.ents = [mock_ent1, mock_ent2]
 
-        mock_process_texts.return_value = iter([mock_doc])
+        mock_process_text.return_value = mock_doc
 
         extractor.is_in_reference = MagicMock(return_value=False)
 

tests/lib/extractors/test_review_model.py

@@ -120,8 +120,8 @@ def test_compare_and_train_skips_retrain(self):
             assert retrained is False
             mock_train.assert_not_called()
 
-    @patch("colandr.lib.extractors.review_model.process_texts_into_docs")
-    def test_extract_metadata(self, mock_process_texts):
+    @patch("colandr.lib.extractors.review_model.process_text_into_doc")
+    def test_extract_metadata(self, mock_process_text):
         """Test the full metadata extraction integration."""
         model = ReviewModel()
 
@@ -141,11 +141,11 @@ def test_extract_metadata(self, mock_process_texts):
         mock_sent = self._create_mock_sentence(sent_text, has_verb=True)
         mock_doc = MagicMock()
         mock_doc.sents = [mock_sent]
-        mock_process_texts.return_value = iter([mock_doc])
+        mock_process_text.return_value = mock_doc
 
         results = model.extract_metadata(123, "some input text", threshold=0.5)
 
-        mock_process_texts.assert_called_once()
+        mock_process_text.assert_called_once()
         assert len(results) == 1
         result = results[0]
         assert result.record == 123

tests/lib/nlp/test_utils.py

@@ -78,6 +78,34 @@ def test_detect_languages(texts, exp_langs):
     assert obs_langs == exp_langs
 
 
+@pytest.mark.parametrize(
+    ["text", "max_len", "fallback_lang"],
+    [
+        (
+            "This is a short -- but not too short -- example English sentence.",
+            1000,
+            None,
+        ),
+        ("And this is another short example English sentence.", 100, "en"),
+        ("Esta es una frase corta de ejemplo en español.", None, None),
+    ],
+)
+def test_process_text_into_doc(text, max_len, fallback_lang, app):
+    doc = utils.process_text_into_doc(
+        text,
+        max_len=max_len,
+        fallback_lang=fallback_lang,
+        exclude=("parser", "ner"),
+    )
+    assert isinstance(doc, Doc) or doc is None
+    if doc.lang_ == "en":
+        spacy_lang = utils.load_spacy_lang(
+            utils.get_lang_to_models()["en"], exclude=("parser", "ner")
+        )
+        assert isinstance(spacy_lang, Language) and isinstance(doc, Doc)  # type guards
+        assert spacy_lang(text).to_bytes() == doc.to_bytes()
+
+
 @pytest.mark.parametrize(
     ["texts", "max_len", "fallback_lang"],
     [
@@ -124,7 +152,7 @@ def test_process_texts_into_docs(texts, max_len, fallback_lang, app):
     assert any(isinstance(doc, Doc) for doc in docs)
     # sanity-check vector value for first text only
     spacy_lang = utils.load_spacy_lang(
-        utils.get_lang_to_models()["en"][0], exclude=("parser", "ner")
+        utils.get_lang_to_models()["en"], exclude=("parser", "ner")
     )
     doc = docs[0]
     assert isinstance(spacy_lang, Language) and isinstance(doc, Doc)  # type guards