diff --git a/datastew/embedding.py b/datastew/embedding.py index 7f5c90d..f303875 100644 --- a/datastew/embedding.py +++ b/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,23 +298,23 @@ 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] @@ -294,7 +322,7 @@ def get_embeddings(self, messages: List[str]) -> Sequence[Sequence[float]]: 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: diff --git a/tests/test_embedding.py b/tests/test_embedding.py index 7cc29c7..490180d 100644 --- a/tests/test_embedding.py +++ b/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)