perf: minimize calls to load_embedder (#133)

Author: mplatzer

mostlyai/qa/_sampling.py

@@ -44,7 +44,7 @@
     ACCURACY_MAX_COLUMNS,
     ProgressCallbackWrapper,
 )
-from mostlyai.qa.assets import load_embedder, load_tokenizer
+from mostlyai.qa.assets import load_tokenizer
 
 
 _LOG = logging.getLogger(__name__)
@@ -290,11 +290,8 @@ def calculate_embeddings(
     progress: ProgressCallbackWrapper | None = None,
     progress_from: int | None = None,
     progress_to: int | None = None,
+    embedder: Any | None = None,
 ) -> np.ndarray:
-    # load embedder
-    t0 = time.time()
-    embedder = load_embedder()
-    _LOG.info(f"loaded load_embedder in {time.time() - t0:.2f}s")
     # split into buckets for calculating embeddings to avoid memory issues and report continuous progress
     steps = progress_to - progress_from if progress_to is not None and progress_from is not None else 1
     buckets = np.array_split(strings, steps)

mostlyai/qa/assets/__init__.py

@@ -43,6 +43,8 @@ def load_embedder():
     Load the embedder model.
     Can deal with read-only cache folder by attempting to download the model if it is not locally available.
     Users can set MOSTLY_HF_HOME environment variable to override the default cache folder.
+
+    Note that this method can take significant time to load the model. Thus, it is recommended to call this method once and reuse the returned object.
     """
     from sentence_transformers import SentenceTransformer
 

mostlyai/qa/reporting.py

@@ -45,6 +45,7 @@
     plot_store_distinct_categories_per_sequence,
     plot_store_sequences_per_distinct_category,
 )
+from mostlyai.qa.assets import load_embedder
 from mostlyai.qa.metrics import ModelMetrics, Accuracy, Similarity, Distances
 from mostlyai.qa._sampling import (
     calculate_embeddings,
@@ -288,6 +289,9 @@ def report(
             acc_cats_per_seq = acc_seqs_per_cat = pd.DataFrame({"column": [], "accuracy": [], "accuracy_max": []})
         progress.update(completed=25, total=100)
 
+        _LOG.info("load embedder")
+        embedder = load_embedder()
+
         _LOG.info("calculate embeddings for synthetic")
         syn_embeds = calculate_embeddings(
             strings=pull_data_for_embeddings(
@@ -300,6 +304,7 @@ def report(
             progress=progress,
             progress_from=25,
             progress_to=45,
+            embedder=embedder,
         )
         _LOG.info("calculate embeddings for training")
         trn_embeds = calculate_embeddings(
@@ -313,6 +318,7 @@ def report(
             progress=progress,
             progress_from=45,
             progress_to=65,
+            embedder=embedder,
         )
         if hol_tgt_data is not None:
             _LOG.info("calculate embeddings for holdout")
@@ -327,6 +333,7 @@ def report(
                 progress=progress,
                 progress_from=65,
                 progress_to=85,
+                embedder=embedder,
             )
         else:
             hol_embeds = None

mostlyai/qa/reporting_from_statistics.py

@@ -36,6 +36,7 @@
     ProgressCallbackWrapper,
 )
 from mostlyai.qa._filesystem import Statistics, TemporaryWorkspace
+from mostlyai.qa.assets import load_embedder
 
 _LOG = logging.getLogger(__name__)
 
@@ -162,6 +163,9 @@ def report_from_statistics(
             acc_cats_per_seq = acc_seqs_per_cat = pd.DataFrame({"column": [], "accuracy": [], "accuracy_max": []})
         progress.update(completed=40, total=100)
 
+        _LOG.info("load embedder")
+        embedder = load_embedder()
+
         _LOG.info("calculate embeddings for synthetic")
         syn_embeds = calculate_embeddings(
             strings=pull_data_for_embeddings(
@@ -174,6 +178,7 @@ def report_from_statistics(
             progress=progress,
             progress_from=40,
             progress_to=60,
+            embedder=embedder,
         )
 
         _LOG.info("report similarity")

tests/unit/test_distances.py

@@ -21,6 +21,7 @@
     plot_store_distances,
 )
 from mostlyai.qa._sampling import calculate_embeddings
+from mostlyai.qa.assets import load_embedder
 
 
 @pytest.fixture()
@@ -42,9 +43,10 @@ def cat_with_rare_and_none():
 
 def test_calculate_distances():
     n = 10
-    syn_embeds = calculate_embeddings(["a 0 1.0"] * n)
-    trn_embeds = calculate_embeddings(["a 0 0.0"] * n)
-    hol_embeds = calculate_embeddings(["a 0 1.0"] * n)
+    embedder = load_embedder()
+    syn_embeds = calculate_embeddings(["a 0 1.0"] * n, embedder=embedder)
+    trn_embeds = calculate_embeddings(["a 0 0.0"] * n, embedder=embedder)
+    hol_embeds = calculate_embeddings(["a 0 1.0"] * n, embedder=embedder)
     dcr_syn_trn, dcr_syn_hol, dcr_trn_hol = calculate_distances(
         syn_embeds=syn_embeds, trn_embeds=trn_embeds, hol_embeds=hol_embeds
     )
@@ -55,9 +57,9 @@ def test_calculate_distances():
     assert dcr_syn_hol.max() == 0
 
     # test specifically that near matches do not report a distance of 0 due to rounding
-    syn_embeds = calculate_embeddings(["a 0.0002"] * n)
-    trn_embeds = calculate_embeddings(["a 0.0001"] * n)
-    hol_embeds = calculate_embeddings(["a 0.0001"] * n)
+    syn_embeds = calculate_embeddings(["a 0.0002"] * n, embedder=embedder)
+    trn_embeds = calculate_embeddings(["a 0.0001"] * n, embedder=embedder)
+    hol_embeds = calculate_embeddings(["a 0.0001"] * n, embedder=embedder)
     dcr_syn_trn, dcr_syn_hol, dcr_trn_hol = calculate_distances(
         syn_embeds=syn_embeds, trn_embeds=trn_embeds, hol_embeds=hol_embeds
     )
@@ -66,7 +68,8 @@ def test_calculate_distances():
 
 
 def test_plot_store_dcr(workspace):
-    embeds = calculate_embeddings(["a 0.0002"] * 100)
+    embedder = load_embedder()
+    embeds = calculate_embeddings(["a 0.0002"] * 100, embedder=embedder)
     dcr_syn_trn, dcr_syn_hol, dcr_trn_hol = calculate_distances(syn_embeds=embeds, trn_embeds=embeds, hol_embeds=embeds)
     plot_store_distances(dcr_syn_trn, dcr_syn_hol, dcr_trn_hol, workspace)
     output_dir = workspace.workspace_dir / "figures"

tests/unit/test_html_report.py

@@ -14,6 +14,7 @@
 
 from mostlyai.qa import _accuracy, _html_report, _distances, _similarity
 from mostlyai.qa._common import CTX_COLUMN_PREFIX, TGT_COLUMN_PREFIX
+from mostlyai.qa.assets import load_embedder
 from mostlyai.qa.reporting import _calculate_metrics
 from mostlyai.qa._sampling import calculate_embeddings, pull_data_for_embeddings
 import pandas as pd
@@ -33,9 +34,11 @@ def test_generate_store_report(tmp_path, cols, workspace):
     acc_cats_per_seq = pd.DataFrame({"column": acc_uni["column"], "accuracy": 0.5, "accuracy_max": 0.5})
     acc_seqs_per_cat = pd.DataFrame({"column": acc_uni["column"], "accuracy": 0.5, "accuracy_max": 0.5})
     corr_trn = _accuracy.calculate_correlations(acc_trn)
-    syn_embeds = calculate_embeddings(pull_data_for_embeddings(df_tgt=syn))
-    trn_embeds = calculate_embeddings(pull_data_for_embeddings(df_tgt=trn))
-    hol_embeds = calculate_embeddings(pull_data_for_embeddings(df_tgt=hol))
+    embedder = load_embedder()
+    syn_embeds = calculate_embeddings(pull_data_for_embeddings(df_tgt=syn), embedder=embedder)
+    trn_embeds = calculate_embeddings(pull_data_for_embeddings(df_tgt=trn), embedder=embedder)
+    hol_embeds = calculate_embeddings(pull_data_for_embeddings(df_tgt=hol), embedder=embedder)
+
     sim_cosine_trn_hol, sim_cosine_trn_syn = _similarity.calculate_cosine_similarities(
         syn_embeds=syn_embeds,
         trn_embeds=trn_embeds,

tests/unit/test_similarity.py

@@ -16,6 +16,7 @@
 
 from mostlyai.qa._similarity import calculate_cosine_similarities, calculate_discriminator_auc
 from mostlyai.qa._sampling import calculate_embeddings
+from mostlyai.qa.assets import load_embedder
 
 
 def test_calculate_embeddings():
@@ -25,9 +26,10 @@ def test_calculate_embeddings():
     # semantically distant synthetic data
     syn_distant = ["quantum physics theory", "deep space exploration"]
 
-    trn_embeds = calculate_embeddings(trn)
-    syn_close_embeds = calculate_embeddings(syn_close)
-    syn_distant_embeds = calculate_embeddings(syn_distant)
+    embedder = load_embedder()
+    trn_embeds = calculate_embeddings(trn, embedder=embedder)
+    syn_close_embeds = calculate_embeddings(syn_close, embedder=embedder)
+    syn_distant_embeds = calculate_embeddings(syn_distant, embedder=embedder)
     assert np.all(trn_embeds[0] == trn_embeds[2])  # check that we retain row order
 
     # check that syn_close is closer to trn than syn_distant