Refactor linker

Author: cthoyt

scispacy/linking_utils.py

@@ -1,13 +1,65 @@
-from typing import List, Dict, NamedTuple, Optional, Set
+"""
+This submodule contains a data structure for storing lexical
+indexes over various biomedical vocabularies.
+
+There are several built-in vocabularies, which can be imported
+and instantiated like in:
+
+.. code-block:: python
+
+    from scispacy.linking_utils import UmlsKnowledgeBase
+
+    kb = UmlsKnowledgeBase()
+
+In general, new :class:`KnowledgeBase` objects can be constructed
+from a list of :class:`Entity` objects, or a path to a JSON or JSONL
+file containing dictionaries shaped the same way:
+
+.. code-block:: python
+
+    from scispacy.linking_utils import KnowledgeBase
+
+    # UMLS
+    kb = KnowledgeBase(
+        "https://ai2-s2-scispacy.s3-us-west-2.amazonaws.com/"
+        "data/kbs/2023-04-23/umls_mesh_2022.jsonl"
+    )
+
+"""
+
 import json
 from collections import defaultdict
+from contextlib import contextmanager
+from pathlib import Path
+from typing import (
+    List,
+    Dict,
+    NamedTuple,
+    Optional,
+    Set,
+    Union,
+    Iterable,
+    Tuple,
+    DefaultDict,
+    Generator,
+)
 
 from scispacy.file_cache import cached_path
 from scispacy.umls_semantic_type_tree import (
     UmlsSemanticTypeTree,
     construct_umls_tree_from_tsv,
 )
 
+__all__ = [
+    "Entity",
+    "KnowledgeBase",
+    "UmlsKnowledgeBase",
+    "Mesh",
+    "GeneOntology",
+    "HumanPhenotypeOntology",
+    "RxNorm",
+]
+
 
 class Entity(NamedTuple):
     concept_id: str
@@ -38,6 +90,53 @@ def __repr__(self):
 DEFAULT_UMLS_TYPES_PATH = "https://s3-us-west-2.amazonaws.com/ai2-s2-scispacy/data/umls_semantic_type_tree.tsv"
 
 
+@contextmanager
+def _iter_entities(
+    path_or_entities: Union[str, Path, Iterable[Entity]],
+) -> Generator[Iterable[Entity], None, None]:
+    """Iterate through entities from a JSON file, JSONL file, or pass through an existing iterable."""
+    if isinstance(path_or_entities, (str, Path)):
+        # normalize paths
+        path_or_entities = cached_path(path_or_entities)
+
+        # do the following inside a context manager to
+        # make sure the file gets closed properly
+        with open(path_or_entities) as file:
+            if path_or_entities.endswith("jsonl"):
+                yield (Entity(**json.loads(line)) for line in file)
+            else:
+                yield (Entity(**record) for record in json.load(file))
+    else:
+        yield path_or_entities
+
+
+def _index_entities(
+    entities: Iterable[Entity],
+) -> Tuple[Dict[str, Entity], Dict[str, Set[str]]]:
+    """Create indexes over entities for use in a :class:`KnowledgeBase`.
+
+    Parameters
+    ----------
+    entities :
+        An iterable (e.g., a list) of entity objects
+
+    Returns
+    -------
+    A pair of indexes for:
+
+    1. A mapping from local unique identifiers (e.g., CUIs for UMLS) to entity objects
+    2. A mapping from aliases (e.g., canonical names, aliases) to local unique identifiers
+    """
+    cui_to_entity: Dict[str, Entity] = {}
+    alias_to_cuis: DefaultDict[str, Set[str]] = defaultdict(set)
+    for entity in entities:
+        alias_to_cuis[entity.canonical_name].add(entity.concept_id)
+        for alias in entity.aliases:
+            alias_to_cuis[alias].add(entity.concept_id)
+        cui_to_entity[entity.concept_id] = entity
+    return cui_to_entity, dict(alias_to_cuis)
+
+
 class KnowledgeBase:
     """
     A class representing two commonly needed views of a Knowledge Base:
@@ -50,31 +149,20 @@ class KnowledgeBase:
         The file path to the json/jsonl representation of the KB to load.
     """
 
+    cui_to_entity: Dict[str, Entity]
+    alias_to_cuis: Dict[str, Set[str]]
+
     def __init__(
         self,
-        file_path: Optional[str] = None,
+        file_path: Union[None, str, Path, Iterable[Entity]] = None,
     ):
         if file_path is None:
             raise ValueError(
                 "Do not use the default arguments to KnowledgeBase. "
                 "Instead, use a subclass (e.g UmlsKnowledgeBase) or pass a path to a kb."
             )
-        if file_path.endswith("jsonl"):
-            raw = (json.loads(line) for line in open(cached_path(file_path)))
-        else:
-            raw = json.load(open(cached_path(file_path)))
-
-        alias_to_cuis: Dict[str, Set[str]] = defaultdict(set)
-        self.cui_to_entity: Dict[str, Entity] = {}
-
-        for concept in raw:
-            unique_aliases = set(concept["aliases"])
-            unique_aliases.add(concept["canonical_name"])
-            for alias in unique_aliases:
-                alias_to_cuis[alias].add(concept["concept_id"])
-            self.cui_to_entity[concept["concept_id"]] = Entity(**concept)
-
-        self.alias_to_cuis: Dict[str, Set[str]] = {**alias_to_cuis}
+        with _iter_entities(file_path) as entities:
+            self.cui_to_entity, self.alias_to_cuis = _index_entities(entities)
 
 
 class UmlsKnowledgeBase(KnowledgeBase):

tests/test_linking.py

@@ -5,6 +5,7 @@
 
 from scispacy.candidate_generation import CandidateGenerator, create_tfidf_ann_index
 from scispacy.linking import EntityLinker
+from scispacy.linking_utils import Entity, KnowledgeBase
 from scispacy.umls_utils import UmlsKnowledgeBase
 from scispacy.abbreviation import AbbreviationDetector
 from scispacy.util import scipy_supports_sparse_float16
@@ -27,32 +28,47 @@ def setUp(self):
         self.linker = EntityLinker(candidate_generator=candidate_generator, filter_for_definitions=False)
 
     def test_naive_entity_linking(self):
+        self._test_linker(self.linker)
+
+    def test_custom_loading(self):
+        entities = [
+            Entity(concept_id="C0000039", canonical_name="dipalmitoylphosphatidylcholine", types=["T109", "T121"])
+        ]
+        kb = KnowledgeBase(entities)
+        with tempfile.TemporaryDirectory() as dir_name:
+            concept_aliases, tfidf_vectorizer, ann_index = create_tfidf_ann_index(dir_name, kb)
+        candidate_generator = CandidateGenerator(ann_index, tfidf_vectorizer, concept_aliases, kb)
+        linker = EntityLinker(candidate_generator=candidate_generator, filter_for_definitions=False)
+
+        self._test_linker(linker)
+
+    def _test_linker(self, linker: EntityLinker) -> None:
         text = "There was a lot of Dipalmitoylphosphatidylcholine."
         doc = self.nlp(text)
 
         # Check that the linker returns nothing if we set the filter_for_definitions flag
         # and set the threshold very high for entities without definitions.
-        self.linker.filter_for_definitions = True
-        self.linker.no_definition_threshold = 3.0
-        doc = self.linker(doc)
+        linker.filter_for_definitions = True
+        linker.no_definition_threshold = 3.0
+        doc = linker(doc)
         assert doc.ents[0]._.kb_ents == []
 
         # Check that the linker returns only high confidence entities if we
         # set the threshold to something more reasonable.
-        self.linker.no_definition_threshold = 0.95
-        doc = self.linker(doc)
+        linker.no_definition_threshold = 0.95
+        doc = linker(doc)
         assert doc.ents[0]._.kb_ents == [("C0000039", 1.0)]
 
-        self.linker.filter_for_definitions = False
-        self.linker.threshold = 0.45
-        doc = self.linker(doc)
+        linker.filter_for_definitions = False
+        linker.threshold = 0.45
+        doc = linker(doc)
         # Without the filter_for_definitions filter, we get 2 entities for
         # the first mention.
         assert len(doc.ents[0]._.kb_ents) == 2
 
         id_with_score = doc.ents[0]._.kb_ents[0]
         assert id_with_score == ("C0000039", 1.0)
-        umls_entity = self.linker.kb.cui_to_entity[id_with_score[0]]
+        umls_entity = linker.kb.cui_to_entity[id_with_score[0]]
         assert umls_entity.concept_id == "C0000039"
         assert umls_entity.types == ["T109", "T121"]