Add caching to embedding (#443)

Author: cthoyt

src/pyobo/api/embedding.py

@@ -7,8 +7,14 @@
 import curies
 import numpy as np
 import pandas as pd
+from tqdm import tqdm
+from typing_extensions import Unpack
 
-from pyobo.api.names import get_definition, get_name, get_references
+from pyobo.api.names import get_definition, get_id_name_mapping, get_name
+from pyobo.api.utils import get_version_from_kwargs
+from pyobo.constants import GetOntologyKwargs, check_should_force
+from pyobo.identifier_utils import wrap_norm_prefix
+from pyobo.utils.path import CacheArtifact, get_cache_path
 
 if TYPE_CHECKING:
     import sentence_transformers
@@ -31,38 +37,61 @@ def get_text_embedding_model() -> sentence_transformers.SentenceTransformer:
 
 def _get_text(
     reference: str | curies.Reference | curies.ReferenceTuple,
+    /,
+    *,
+    name: str | None = None,
+    **kwargs: Unpack[GetOntologyKwargs],
 ) -> str | None:
-    name = get_name(reference)
+    if name is None:
+        name = get_name(reference, **kwargs)
     if name is None:
         return None
-    description = get_definition(reference)
+    description = get_definition(reference, **kwargs)
     if description:
         name += " " + description
     return name
 
 
+@wrap_norm_prefix
 def get_text_embeddings_df(
     prefix: str,
     *,
     model: sentence_transformers.SentenceTransformer | None = None,
+    **kwargs: Unpack[GetOntologyKwargs],
 ) -> pd.DataFrame:
     """Get embeddings for all entities in the resource.
 
     :param prefix: A reference, either as a string or Reference object
     :param model: A sentence transformer model. Defaults to ``all-MiniLM-L6-v2`` if not
         given.
+    :param kwargs: The keyword arguments to forward to ontology getter functions for
+        names, definitions, and version
+
+    :returns: A pandas dataframe with an index representing local unique identifiers and
+        columns for the values of the model returned vectors
     """
+    path = get_cache_path(
+        prefix, CacheArtifact.embeddings, version=get_version_from_kwargs(prefix, kwargs)
+    )
+    if path.is_file() and not check_should_force(kwargs):
+        df = pd.read_csv(path, sep="\t").set_index(0)
+        return df
+
+    id_to_name = get_id_name_mapping(prefix, **kwargs)
+
     luids, texts = [], []
-    for reference in get_references(prefix):
-        text = _get_text(reference)
+    for identifier, name in tqdm(id_to_name.items(), desc=f"[{prefix}] constructing text"):
+        text = _get_text(curies.ReferenceTuple(prefix, identifier), name=name, **kwargs)
         if text is None:
             continue
-        luids.append(reference.identifier)
+        luids.append(identifier)
         texts.append(text)
     if model is None:
         model = get_text_embedding_model()
-    res = model.encode(texts)
-    return pd.DataFrame(res, index=luids)
+    res = model.encode(texts, show_progress_bar=True)
+    df = pd.DataFrame(res, index=luids)
+    df.to_csv(path, sep="\t")  # index is important here!
+    return df
 
 
 def get_text_embedding(

src/pyobo/sources/bigg/bigg_compartment.py

@@ -55,7 +55,7 @@ def get_compartments(*, force: bool = False, version: str | None = None) -> dict
     """Get a dictionary of BiGG compartments."""
     rv = {}
     soup = get_soup(DATA_URL)
-    table = soup.find(**{"class": "myTable"})  # type:ignore[arg-type]
+    table = soup.find(class_="myTable")
     if table is None:
         raise ValueError
     for row in table.find_all("tr"):  # type:ignore[attr-defined]

src/pyobo/sources/omim_ps.py

@@ -2,6 +2,7 @@
 
 import logging
 from collections.abc import Iterable
+from typing import cast
 
 from bioversions.utils import get_soup
 
@@ -34,9 +35,14 @@ def iter_terms(self, force: bool = False) -> Iterable[Term]:
         if tbody is None:
             raise ValueError("omim.ps failed - scraper could not find table body in HTML")
         for row in tbody.find_all("tr"):
-            anchor = row.find("td").find("a")
+            td = row.find("td")
+            if td is None:
+                continue
+            anchor = td.find("a")
+            if anchor is None or anchor.text is None:
+                continue
             name = anchor.text.strip()
-            identifier = anchor.attrs["href"][len("/phenotypicSeries/") :]
+            identifier = cast(str, anchor.attrs["href"])[len("/phenotypicSeries/") :]
             yield Term.from_triple(PREFIX, identifier, name)
 
 

src/pyobo/utils/path.py

@@ -151,6 +151,8 @@ class CacheArtifact(enum.Enum):
     prefixes = "prefixes.json"
     metadata = "metadata.json"
 
+    embeddings = "embeddings.tsv.gz"
+
 
 def get_cache_path(
     ontology: str,

tests/test_struct/test_obo/test_struct_obo.py

@@ -362,6 +362,7 @@ def test_ror_metadata(self) -> None:
             property_value: foaf:homepage "https\://ror.org" xsd:anyURI
             property_value: doap:repository "https\://github.com/ror-community" xsd:anyURI
             property_value: foaf:logo "https\://ror.org/img/ror-logo.svg" xsd:anyURI
+            property_value: doap:mailing-list "support@ror.org" xsd:string
             property_value: doap:maintainer orcid:0000-0002-2916-3423
             """,  # add Maria Gould
             ontology,
@@ -386,6 +387,7 @@ def test_ror_metadata(self) -> None:
             Annotation(foaf:homepage "https://ror.org"^^xsd:anyURI)
             Annotation(doap:repository "https://github.com/ror-community"^^xsd:anyURI)
             Annotation(foaf:logo "https://ror.org/img/ror-logo.svg"^^xsd:anyURI)
+            Annotation(doap:mailing-list "support@ror.org"^^xsd:string)
             Annotation(doap:maintainer orcid:0000-0002-2916-3423)
             )
             """,
@@ -411,6 +413,7 @@ def test_ror_metadata(self) -> None:
             operated by California Digital Library, Crossref, and Datacite.</dcterms:description>
                     <dcterms:license>CC0-1.0</dcterms:license>
                     <dcterms:title>Research Organization Registry</dcterms:title>
+                    <doap:mailing-list>support@ror.org</doap:mailing-list>
                     <doap:maintainer rdf:resource="https://orcid.org/0000-0002-2916-3423"/>
                     <doap:repository rdf:datatype="http://www.w3.org/2001/XMLSchema#anyURI">https://github.com/ror-community</doap:repository>
                     <foaf:homepage rdf:datatype="http://www.w3.org/2001/XMLSchema#anyURI">https://ror.org</foaf:homepage>
@@ -429,6 +432,8 @@ def test_ror_metadata(self) -> None:
                 <owl:AnnotationProperty rdf:about="http://purl.org/dc/terms/license"/>
                 <!-- http://purl.org/dc/terms/title -->
                 <owl:AnnotationProperty rdf:about="http://purl.org/dc/terms/title"/>
+                <!-- http://usefulinc.com/ns/doap#mailing-list -->
+                <owl:AnnotationProperty rdf:about="http://usefulinc.com/ns/doap#mailing-list"/>
                 <!-- http://usefulinc.com/ns/doap#maintainer -->
                 <owl:AnnotationProperty rdf:about="http://usefulinc.com/ns/doap#maintainer"/>
                 <!-- http://usefulinc.com/ns/doap#repository -->