From f7fc4710d7c0b57712e313867d5b18d8585e56d3 Mon Sep 17 00:00:00 2001
From: Michael Platzer <michael.platzer@gmail.com>
Date: Thu, 27 Mar 2025 23:53:36 +0100
Subject: [PATCH] chore: log timings for loading embedder model

---
 mostlyai/qa/_sampling.py       | 4 +++-
 mostlyai/qa/assets/__init__.py | 4 ++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/mostlyai/qa/_sampling.py b/mostlyai/qa/_sampling.py
index 2992462..9abdf96 100644
--- a/mostlyai/qa/_sampling.py
+++ b/mostlyai/qa/_sampling.py
@@ -291,14 +291,16 @@ def calculate_embeddings(
     progress_from: int | None = None,
     progress_to: int | None = None,
 ) -> np.ndarray:
-    t0 = time.time()
     # 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)
     buckets = [b for b in buckets if len(b) > 0]
     # calculate embeddings for each bucket
+    t0 = time.time()
     embeds = []
     for i, bucket in enumerate(buckets, 1):
         embeds += [embedder.encode(bucket.tolist(), show_progress_bar=False)]
diff --git a/mostlyai/qa/assets/__init__.py b/mostlyai/qa/assets/__init__.py
index 6f53569..55f83a6 100644
--- a/mostlyai/qa/assets/__init__.py
+++ b/mostlyai/qa/assets/__init__.py
@@ -40,14 +40,14 @@ def load_tokenizer():
 
 def load_embedder():
     """
-    Load the embedder model. 
+    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.
     """
     from sentence_transformers import SentenceTransformer
 
     model_name = "sentence-transformers/all-MiniLM-L6-v2"
-    cache_folder=os.getenv("MOSTLY_HF_HOME")
+    cache_folder = os.getenv("MOSTLY_HF_HOME")
     try:
         # First try loading from local cache
         return SentenceTransformer(model_name_or_path=model_name, cache_folder=cache_folder, local_files_only=True)