Merge pull request #182 from SCAI-BIO/refactor/embedding-cache-optimization

Refactor/embedding cache optimization

Author: mehmetcanay

datastew/embedding.py

@@ -10,6 +10,9 @@
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
+_GLOBAL_CACHES = {}
+_GLOBAL_LOCKS = {}
+
 
 class EmbeddingModel(ABC):
     def __init__(self, model_name: str, cache: bool = False, cache_size: int = 10000):
@@ -20,8 +23,19 @@ def __init__(self, model_name: str, cache: bool = False, cache_size: int = 10000
         :param cache_size: Maximum cache size when caching is enabled, defaults to 10,000.
         """
         self.model_name = model_name
-        self._cache = LRUCache(maxsize=cache_size) if cache else None
-        self._cache_lock = Lock() if cache else None
+        self.use_cache = cache
+
+        if self.use_cache:
+            if model_name not in _GLOBAL_CACHES:
+                _GLOBAL_CACHES[model_name] = LRUCache(maxsize=cache_size)
+                _GLOBAL_LOCKS[model_name] = Lock()
+
+            self._cache = _GLOBAL_CACHES[model_name]
+            self._cache_lock = _GLOBAL_LOCKS[model_name]
+
+        else:
+            self._cache = None
+            self._cache_lock = None
 
     @abstractmethod
     def get_embedding(self, text: str) -> Sequence[float]:
@@ -47,22 +61,31 @@ def add_to_cache(self, text: str, embedding: Sequence[float]):
         :param text: The input text to be cached.
         :param embedding: The embedding of the input text.
         """
-        if self._cache_lock and self._cache is not None:
+        if self._cache_lock is not None and self._cache is not None:
             with self._cache_lock:
                 self._cache[text] = embedding
 
+    def add_batch_to_cache(self, texts: Sequence[str], embeddings: Sequence[Sequence[float]]):
+        """Acquire lock once for the entire batch update."""
+        if self._cache_lock is not None and self._cache is not None:
+            with self._cache_lock:
+                for text, emb in zip(texts, embeddings):
+                    self._cache[text] = emb
+
     def get_from_cache(self, text: str) -> Optional[Sequence[float]]:
         """Retrieve an embedding from the cache.
 
         :param text: Cached input text.
         :return: Embedding of the cached input text or `None` if not present.
         """
-        if self._cache_lock and self._cache is not None:
+        if self._cache_lock is not None and self._cache is not None:
             with self._cache_lock:
                 return self._cache.get(text, None)
         return None
 
-    def get_cached_embeddings(self, messages: List[str]) -> Tuple[List[Sequence[float]], List[int], List[str]]:
+    def get_cached_embeddings(
+        self, messages: List[str]
+    ) -> Tuple[List[Optional[Sequence[float]]], List[int], List[str]]:
         """Retrieve cached embeddings and identify uncached messages.
 
         :param messages: A list of input text messages.
@@ -71,16 +94,22 @@ def get_cached_embeddings(self, messages: List[str]) -> Tuple[List[Sequence[floa
             - A list of indices for uncached messages.
             - A list of uncached messages.
         """
-        embeddings, uncached_indices, uncached_messages = [], [], []
+        if self._cache_lock is None or self._cache is None:
+            empty_embeddings: List[Optional[Sequence[float]]] = [None] * len(messages)
+            return empty_embeddings, list(range(len(messages))), messages
 
-        for i, msg in enumerate(messages):
-            cached = self.get_from_cache(msg)
-            if cached:
+        embeddings: List[Optional[Sequence[float]]] = []
+        uncached_indices: List[int] = []
+        uncached_messages: List[str] = []
+
+        with self._cache_lock:
+            for i, msg in enumerate(messages):
+                cached = self._cache.get(msg, None)
                 embeddings.append(cached)
-            else:
-                embeddings.append(None)
-                uncached_indices.append(i)
-                uncached_messages.append(msg)
+                if cached is None:
+                    uncached_indices.append(i)
+                    uncached_messages.append(msg)
+
         return embeddings, uncached_indices, uncached_messages
 
     def sanitize(self, message: str) -> str:
@@ -120,8 +149,7 @@ def get_embedding(self, text: str) -> Sequence[float]:
             return []
         text = self.sanitize(text)
 
-        if self._cache:
-            # Check cache
+        if self._cache is not None:
             cached = self.get_from_cache(text)
             if cached:
                 return cached
@@ -134,7 +162,7 @@ def get_embedding(self, text: str) -> Sequence[float]:
             return embedding
         except Exception as e:
             logging.error(f"Error getting embedding for {text}: {e}")
-            return []
+            raise
 
     def get_embeddings(self, messages: List[str]) -> Sequence[Sequence[float]]:
         """Retrieve embeddings for a list of text messages.
@@ -145,19 +173,19 @@ def get_embeddings(self, messages: List[str]) -> Sequence[Sequence[float]]:
 
         sanitized_messages = [self.sanitize(msg) for msg in messages]
 
-        if self._cache:
+        if self._cache is not None:
             embeddings, uncached_indices, uncached_messages = self.get_cached_embeddings(sanitized_messages)
 
             if uncached_messages:
                 try:
                     response = openai.embeddings.create(model=self.model_name, input=uncached_messages)
                     new_embeddings = [item.embedding for item in response.data]
+                    self.add_batch_to_cache(uncached_messages, new_embeddings)
                     for idx, embedding in zip(uncached_indices, new_embeddings):
-                        self.add_to_cache(sanitized_messages[idx], embedding)
                         embeddings[idx] = embedding
                 except Exception as e:
                     logging.error(f"Error in processing chunk: {e}")
-                    return []
+                    raise
 
             return [emb for emb in embeddings if emb is not None]
 
@@ -167,12 +195,12 @@ def get_embeddings(self, messages: List[str]) -> Sequence[Sequence[float]]:
             return embeddings
         except Exception as e:
             logging.error(f"Failed processing messages: {e}")
-            return []
+            raise
 
 
 class HuggingFaceAdapter(EmbeddingModel):
     _model_cache = {}
-    _load_count = 0 # For testing
+    _load_count = 0  # For testing
 
     def __init__(self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2", cache: bool = False):
         """Initialize the Hugging Face adapter with a specified model name.
@@ -200,7 +228,7 @@ def get_embedding(self, text: str) -> Sequence[float]:
         text = self.sanitize(text)
 
         # Check cache
-        if self._cache:
+        if self._cache is not None:
             cached = self.get_from_cache(text)
             if cached:
                 return cached
@@ -212,7 +240,7 @@ def get_embedding(self, text: str) -> Sequence[float]:
             return embedding
         except Exception as e:
             logging.error(f"Error getting embedding for {text}: {e}")
-            return []
+            raise
 
     def get_embeddings(self, messages: List[str]) -> Sequence[Sequence[float]]:
         """Retrieve embeddings for a list of text messages using MPNet.
@@ -221,7 +249,7 @@ def get_embeddings(self, messages: List[str]) -> Sequence[Sequence[float]]:
         :return: A sequence of embedding vectors.
         """
         sanitized_messages = [self.sanitize(msg) for msg in messages]
-        if self._cache:
+        if self._cache is not None:
             embeddings, uncached_indices, uncached_messages = self.get_cached_embeddings(sanitized_messages)
 
             if uncached_messages:
@@ -230,12 +258,12 @@ def get_embeddings(self, messages: List[str]) -> Sequence[Sequence[float]]:
                     flattened_embeddings = [
                         [float(element) for element in row] for row in new_embeddings if row is not None
                     ]
+                    self.add_batch_to_cache(uncached_messages, flattened_embeddings)
                     for idx, embedding in zip(uncached_indices, flattened_embeddings):
-                        self.add_to_cache(sanitized_messages[idx], embedding)
                         embeddings[idx] = embedding
                 except Exception as e:
                     logging.error(f"Failed processing messages: {e}")
-                    return []
+                    raise
 
             return [emb for emb in embeddings if emb is not None]
 
@@ -245,7 +273,7 @@ def get_embeddings(self, messages: List[str]) -> Sequence[Sequence[float]]:
             return flattened_embeddings
         except Exception as e:
             logging.error(f"Failed processing messages: {e}")
-            return []
+            raise
 
 
 class OllamaAdapter(EmbeddingModel):
@@ -260,7 +288,7 @@ def get_embedding(self, text: str) -> Sequence[float]:
             logging.warning("Empty text passed to get_embedding")
             return []
         text = self.sanitize(text)
-        if self._cache:
+        if self._cache is not None:
             cached = self.get_from_cache(text)
             if cached:
                 return cached
@@ -270,31 +298,31 @@ def get_embedding(self, text: str) -> Sequence[float]:
             return embedding
         except Exception as e:
             logging.error(f"Error getting embedding for {text}: {e}")
-            return []
+            raise
 
     def get_embeddings(self, messages: List[str]) -> Sequence[Sequence[float]]:
         sanitized_messages = [self.sanitize(msg) for msg in messages]
 
-        if self._cache:
+        if self._cache is not None:
             embeddings, uncached_indices, uncached_messages = self.get_cached_embeddings(sanitized_messages)
 
             if uncached_messages:
                 try:
                     new_embeddings = self.client.embed(self.model_name, uncached_messages).get("embeddings")
+                    self.add_batch_to_cache(uncached_messages, new_embeddings)
                     for idx, embedding in zip(uncached_indices, new_embeddings):
-                        self.add_to_cache(sanitized_messages[idx], embedding)
                         embeddings[idx] = embedding
                 except Exception as e:
                     logging.error(f"Failed processing messages: {e}")
-                    return []
+                    raise
 
             return [emb for emb in embeddings if emb is not None]
 
         try:
             return self.client.embed(self.model_name, sanitized_messages).get("embeddings")
         except Exception as e:
             logging.error(f"Failed processing messages: {e}")
-            return []
+            raise
 
 
 class Vectorizer:

tests/test_embedding.py

@@ -1,82 +1,84 @@
 import unittest
-from time import time
 from typing import Sequence
+from unittest.mock import patch
 
-from datastew.embedding import HuggingFaceAdapter
+from datastew.embedding import _GLOBAL_CACHES, HuggingFaceAdapter
 
 
 class TestEmbedding(unittest.TestCase):
 
     def setUp(self):
-        self.hugging_face_adapter = HuggingFaceAdapter(cache=True)
+        _GLOBAL_CACHES.clear()
+        self.adapter = HuggingFaceAdapter(cache=True)
 
     def test_hugging_face_adapter_get_embedding(self):
         text = "This is a test sentence."
-        embedding = self.hugging_face_adapter.get_embedding(text)
+        embedding = self.adapter.get_embedding(text)
         self.assertIsInstance(embedding, Sequence)
 
     def test_hugging_face_adapter_get_embeddings(self):
         messages = [f"This is message {i}." for i in range(20)]
-        embeddings = self.hugging_face_adapter.get_embeddings(messages)
+        embeddings = self.adapter.get_embeddings(messages)
         self.assertIsInstance(embeddings, Sequence)
         self.assertEqual(len(embeddings), len(messages))
 
     def test_sanitization(self):
         text1 = " Test"
         text2 = "test "
-        embedding1 = self.hugging_face_adapter.get_embedding(text1)
-        embedding2 = self.hugging_face_adapter.get_embedding(text2)
+        embedding1 = self.adapter.get_embedding(text1)
+        embedding2 = self.adapter.get_embedding(text2)
         self.assertSequenceEqual(embedding1, embedding2)
 
-    def test_caching_get_embedding(self):
-        text = "This is a test sentence."
-        if self.hugging_face_adapter._cache:
-            self.hugging_face_adapter._cache.clear()
+    def test_caching_get_embedding_deterministic(self):
+        text = "deterministic test sentence"
 
-        # Measure time for the first call
-        start_time = time()
-        embedding1 = self.hugging_face_adapter.get_embedding(text)
-        first_call_time = time() - start_time
+        with patch.object(self.adapter.model, "encode", wraps=self.adapter.model.encode) as spy_encode:
+            emb1 = self.adapter.get_embedding(text)
+            spy_encode.assert_called_once()
 
-        # Measure time for the second call
-        start_time = time()
-        embedding2 = self.hugging_face_adapter.get_embedding(text)
-        second_call_time = time() - start_time
+            emb2 = self.adapter.get_embedding(text)
+            self.assertEqual(spy_encode.call_count, 1)
+            self.assertSequenceEqual(emb1, emb2)
 
-        self.assertLess(second_call_time, first_call_time)
-        self.assertSequenceEqual(embedding1, embedding2)
+    def test_caching_get_embeddings_batch(self):
+        messages = [f"Batch message {i}." for i in range(5)]
 
-    def test_caching_get_embeddings(self):
-        messages = [f"This is message {i}." for i in range(20)]
-        if self.hugging_face_adapter._cache:
-            self.hugging_face_adapter._cache.clear()
+        with patch.object(self.adapter.model, "encode", wraps=self.adapter.model.encode) as spy_encode:
+            embeddings1 = self.adapter.get_embeddings(messages)
+            self.assertEqual(spy_encode.call_count, 1)
+            embeddings2 = self.adapter.get_embeddings(messages)
+            self.assertEqual(spy_encode.call_count, 1)
 
-        start_time = time()
-        embeddings1 = self.hugging_face_adapter.get_embeddings(messages)
-        first_call_time = time() - start_time
+            for emb1, emb2 in zip(embeddings1, embeddings2):
+                self.assertSequenceEqual(emb1, emb2)
 
-        start_time = time()
-        embeddings2 = self.hugging_face_adapter.get_embeddings(messages)
-        second_call_time = time() - start_time
+    def test_partial_cache_hit_batch(self):
+        messages = ["A", "B", "C"]
+        # Cache "A" manually
+        self.adapter.add_batch_to_cache([self.adapter.sanitize("A")], [[0.1, 0.2]])
 
-        self.assertLess(second_call_time, first_call_time)
+        with patch.object(self.adapter.model, "encode", wraps=self.adapter.model.encode) as spy_encode:
+            embeddings = self.adapter.get_embeddings(messages)
+            # encode should only be called once, and only for ["b", "c"]
+            spy_encode.assert_called_once_with(["b", "c"], show_progress_bar=True)
+            self.assertEqual(len(embeddings), 3)
+            self.assertEqual(embeddings[0], [0.1, 0.2])
 
-        for emb1, emb2 in zip(embeddings1, embeddings2):
-            self.assertSequenceEqual(emb1, emb2)
+    def test_exception_handling_aborts_silence(self):
+        messages = ["Fail 1", "Fail 2"]
+
+        with patch.object(self.adapter.model, "encode", side_effect=Exception("API Down")):
+            with self.assertRaises(Exception) as context:
+                self.adapter.get_embeddings(messages)
+
+            self.assertTrue("API Down" in str(context.exception))
 
     def test_huggingface_model_loaded_once(self):
-        # Reset class-level cache and counter before test
         HuggingFaceAdapter._model_cache.clear()
         HuggingFaceAdapter._load_count = 0
 
-        # Instantiate multiple adapters
         adapter1 = HuggingFaceAdapter(cache=True)
         adapter2 = HuggingFaceAdapter(cache=True)
-        adapter3 = HuggingFaceAdapter(cache=True)
 
-        # All adapters should use the same underlying model object
         self.assertIs(adapter1.model, adapter2.model)
-        self.assertIs(adapter2.model, adapter3.model)
-
-        # And the model should have been loaded only once
         self.assertEqual(HuggingFaceAdapter._load_count, 1)