Refactor/optimize embedding module

datastew/embedding/base.py

@@ -12,6 +12,8 @@
 _GLOBAL_LOCKS = {}
 _INIT_LOCK = Lock()
 
+_WHITESPACE_RE = re.compile(r"\s+")
+
 
 class EmbeddingModel(ABC):
     def __init__(self, model_name: str, cache: bool = False, cache_size: int = 10000):
@@ -37,25 +39,93 @@ def __init__(self, model_name: str, cache: bool = False, cache_size: int = 10000
             self._cache = None
             self._cache_lock = None
 
-    @abstractmethod
     def get_embedding(self, text: str) -> Sequence[float]:
         """Retrieve the embedding vector for a single text input.
 
         :param text: The input text to embed.
+        :raises TypeError: If the input is not a string.
+        :raises ValueError: If the input text is empty or only whitespace.
         :return: A sequence of floats representing the embedding.
         """
-        pass
+        if not isinstance(text, str):
+            raise TypeError(f"Expected string for 'text', got {type(text).__name__}.")
+
+        if not text.strip():
+            raise ValueError("Input 'text' cannot be empty or solely whitespace.")
+
+        text = self._sanitize(text)
+
+        if self._cache is not None:
+            cached = self._get_from_cache(text)
+            if cached:
+                return cached
+
+        embedding = self._generate_embedding(text)
+        self._add_to_cache(text, embedding)
+        return embedding
 
-    @abstractmethod
     def get_embeddings(self, messages: List[str]) -> Sequence[Sequence[float]]:
         """Retrieve embeddings for a list of text messages
 
         :param messages: A list of text messages to embed.
+        :raises TypeError: If 'messages' is not a list.
+        :raises ValueError: If 'messages' is empty.
         :return: A sequence of embedding vectors.
         """
+        if not isinstance(messages, list):
+            raise TypeError(f"Expected a list for 'messages', got {type(messages).__name__}")
+
+        if not messages:
+            raise ValueError("The 'messages' list cannot be empty.")
+
+        sanitized_messages = [self._sanitize(msg) for msg in messages if msg and isinstance(msg, str)]
+
+        if self._cache is not None:
+            embeddings, uncached_indices, uncached_messages = self._get_cached_embeddings(sanitized_messages)
+
+            if uncached_messages:
+                new_embeddings = self._generate_embeddings(uncached_messages)
+                self._add_batch_to_cache(uncached_messages, new_embeddings)
+                for idx, embedding in zip(uncached_indices, new_embeddings):
+                    embeddings[idx] = embedding
+
+            return [emb for emb in embeddings if emb is not None]
+
+        return self._generate_embeddings(sanitized_messages)
+
+    @abstractmethod
+    def _generate_embedding(self, text: str) -> Sequence[float]:
+        """Generate an embedding vector for a single sanitized text input
+
+        This is an abstract method that must be implemented by subclasses to interface
+        with the specific model's API or inference engine.
+
+        :param text: The sanitized input text to embed.
+        :return: A sequence of floats respresenting the generated embedding.
+        """
+        pass
+
+    @abstractmethod
+    def _generate_embeddings(self, messages: list[str]) -> Sequence[Sequence[float]]:
+        """Generate embedding vectors for a batch of sanitized text messages.
+
+        This is an abstract method that must be implemented by subclasses to interface
+        with the specific model's API or inference engine.
+
+        :param messages: A list of sanitized text messages to embed.
+        :return: A sequence of embedding vectors corresponding to the input messages.
+        """
         pass
 
-    def add_to_cache(self, text: str, embedding: Sequence[float]):
+    def _sanitize(self, message: str) -> str:
+        """Clean up the input text by trimming and converting to lowercase.
+
+        :param message: The input text message.
+        :return: Sanitized text.
+        """
+        return _WHITESPACE_RE.sub(" ", message).strip().lower()
+
+    def _add_to_cache(self, text: str, embedding: Sequence[float]):
         """Add a text-embedding pair to cache.
 
         :param text: The input text to be cached.
@@ -65,14 +135,20 @@ def add_to_cache(self, text: str, embedding: Sequence[float]):
             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."""
+    def _add_batch_to_cache(self, texts: Sequence[str], embeddings: Sequence[Sequence[float]]):
+        """Add a batch of text-embedding pairs to the cache.
+
+        Acquires the thread lock once for the entire batch update to minimize overhead.
+
+        :param texts: A sequence of input texts to be cached.
+        :param embeddings: A sequence of corresponding embedding vectors.
+        """
         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]]:
+    def _get_from_cache(self, text: str) -> Optional[Sequence[float]]:
         """Retrieve an embedding from the cache.
 
         :param text: Cached input text.
@@ -83,7 +159,7 @@ def get_from_cache(self, text: str) -> Optional[Sequence[float]]:
                 return self._cache.get(text, None)
         return None
 
-    def get_cached_embeddings(
+    def _get_cached_embeddings(
         self, messages: List[str]
     ) -> Tuple[List[Optional[Sequence[float]]], List[int], List[str]]:
         """Retrieve cached embeddings and identify uncached messages.
@@ -95,8 +171,7 @@ def get_cached_embeddings(
             - A list of 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
+            return [None] * len(messages), list(range(len(messages))), messages
 
         embeddings: List[Optional[Sequence[float]]] = []
         uncached_indices: List[int] = []
@@ -111,11 +186,3 @@ def get_cached_embeddings(
                     uncached_messages.append(msg)
 
         return embeddings, uncached_indices, uncached_messages
-
-    def sanitize(self, message: str) -> str:
-        """Clean up the input text by trimming and converting to lowercase.
-
-        :param message: The input text message.
-        :return: Sanitized text.
-        """
-        return re.sub(r"\s+", " ", message).strip().lower()

datastew/embedding/hugging_face.py

@@ -1,68 +1,33 @@
 import logging
-from typing import List, Sequence
+from typing import Sequence
 
+from cachetools import LRUCache
 from sentence_transformers import SentenceTransformer
 
 from datastew.embedding.base import EmbeddingModel
 
 
 class HuggingFaceAdapter(EmbeddingModel):
-    _model_cache = {}
-    _load_count = 0
+    _model_cache = LRUCache(maxsize=3)
 
     def __init__(self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2", cache: bool = False):
         super().__init__(model_name, cache)
 
         if model_name not in self._model_cache:
-            HuggingFaceAdapter._load_count += 1
             self._model_cache[model_name] = SentenceTransformer(model_name)
 
         self.model = self._model_cache[model_name]
 
-    def get_embedding(self, text: str) -> Sequence[float]:
-        if not text or not isinstance(text, str):
-            logging.warning("Empty or invalid text passed to get_embedding")
-            return []
-        text = self.sanitize(text)
-
-        if self._cache is not None:
-            cached = self.get_from_cache(text)
-            if cached:
-                return cached
-
+    def _generate_embedding(self, text: str) -> Sequence[float]:
         try:
-            embedding = self.model.encode(text)
-            embedding = [float(x) for x in embedding]
-            self.add_to_cache(text, embedding)
-            return embedding
+            return self.model.encode(text).tolist()
         except Exception as e:
             logging.error(f"Error getting embedding for {text}: {e}")
             raise
 
-    def get_embeddings(self, messages: List[str]) -> Sequence[Sequence[float]]:
-        sanitized_messages = [self.sanitize(msg) for msg in messages]
-        if self._cache is not None:
-            embeddings, uncached_indices, uncached_messages = self.get_cached_embeddings(sanitized_messages)
-
-            if uncached_messages:
-                try:
-                    new_embeddings = self.model.encode(uncached_messages, show_progress_bar=True)
-                    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):
-                        embeddings[idx] = embedding
-                except Exception as e:
-                    logging.error(f"Failed processing messages: {e}")
-                    raise
-
-            return [emb for emb in embeddings if emb is not None]
-
+    def _generate_embeddings(self, messages: list[str]) -> Sequence[Sequence[float]]:
         try:
-            embeddings = self.model.encode(sanitized_messages, show_progress_bar=True)
-            flattened_embeddings = [[float(element) for element in row] for row in embeddings]
-            return flattened_embeddings
+            return self.model.encode(messages, show_progress_bar=True).tolist()
         except Exception as e:
             logging.error(f"Failed processing messages: {e}")
             raise

datastew/embedding/ollama.py

@@ -1,5 +1,5 @@
 import logging
-from typing import List, Sequence
+from typing import Sequence
 
 from ollama import Client
 
@@ -13,44 +13,23 @@ def __init__(
         super().__init__(model_name, cache)
         self.client = Client(host)
 
-    def get_embedding(self, text: str) -> Sequence[float]:
-        if not text:
-            logging.warning("Empty text passed to get_embedding")
-            return []
-        text = self.sanitize(text)
-
-        if self._cache is not None:
-            cached = self.get_from_cache(text)
-            if cached:
-                return cached
+    def _generate_embedding(self, text: str) -> Sequence[float]:
         try:
-            embedding = self.client.embed(self.model_name, text).get("embeddings")[0]
-            self.add_to_cache(text, embedding)
-            return embedding
+            return self.client.embed(self.model_name, text).get("embeddings")[0]
         except Exception as e:
             logging.error(f"Error getting embedding for {text}: {e}")
             raise
 
-    def get_embeddings(self, messages: List[str]) -> Sequence[Sequence[float]]:
-        sanitized_messages = [self.sanitize(msg) for msg in messages]
-
-        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):
-                        embeddings[idx] = embedding
-                except Exception as e:
-                    logging.error(f"Failed processing messages: {e}")
-                    raise
-
-            return [emb for emb in embeddings if emb is not None]
+    def _generate_embeddings(self, messages: list[str]) -> Sequence[Sequence[float]]:
+        chunk_size = 500
+        all_embeddings = []
 
         try:
-            return self.client.embed(self.model_name, sanitized_messages).get("embeddings")
+            for i in range(0, len(messages), chunk_size):
+                chunk = messages[i : i + chunk_size]
+                response = self.client.embed(self.model_name, chunk)
+                all_embeddings.extend(response.get("embeddings", []))
+            return all_embeddings
         except Exception as e:
             logging.error(f"Failed processing messages: {e}")
             raise

datastew/embedding/openai.py

@@ -1,71 +1,40 @@
 import logging
-from typing import List, Sequence
+from typing import Sequence
 
-import openai
+from openai import OpenAI
 
 from datastew.embedding.base import EmbeddingModel
 
 
 class GPT4Adapter(EmbeddingModel):
-    def __init__(
-        self,
-        api_key: str,
-        model_name: str = "text-embedding-ada-002",
-        cache: bool = False,
-    ):
+    def __init__(self, api_key: str, model_name: str = "text-embedding-ada-002", cache: bool = False):
         """Initialize the GPT-4 adapter with OpenAI API key and model name.
 
         :param api_key: The API key for accessing OpenAI services.
         :param model_name: The specific embedding model to use, defaults to text-embedding-ada-002.
         :param cache: Enable or disable caching, defaults to False.
         """
         super().__init__(model_name, cache)
-        self.api_key = api_key
-        openai.api_key = api_key
-
-    def get_embedding(self, text: str) -> Sequence[float]:
-        if not text:
-            logging.warning("Empty text passed to get_embedding")
-            return []
-        text = self.sanitize(text)
-
-        if self._cache is not None:
-            cached = self.get_from_cache(text)
-            if cached:
-                return cached
+        self.client = OpenAI(api_key=api_key)
 
+    def _generate_embedding(self, text: str) -> Sequence[float]:
         try:
-            response = openai.embeddings.create(input=[text], model=self.model_name)
-            embedding = response.data[0].embedding
-            self.add_to_cache(text, embedding)
-            return embedding
+            response = self.client.embeddings.create(input=[text], model=self.model_name)
+            return response.data[0].embedding
         except Exception as e:
             logging.error(f"Error getting embedding for {text}: {e}")
             raise
 
-    def get_embeddings(self, messages: List[str]) -> Sequence[Sequence[float]]:
-        sanitized_messages = [self.sanitize(msg) for msg in messages]
-
-        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):
-                        embeddings[idx] = embedding
-                except Exception as e:
-                    logging.error(f"Error in processing chunk: {e}")
-                    raise
-
-            return [emb for emb in embeddings if emb is not None]
+    def _generate_embeddings(self, messages: list[str]) -> Sequence[Sequence[float]]:
+        chunk_size = 1000
+        all_embeddings = []
 
         try:
-            response = openai.embeddings.create(model=self.model_name, input=sanitized_messages)
-            embeddings = [item.embedding for item in response.data]
-            return embeddings
+            for i in range(0, len(messages), chunk_size):
+                chunk = messages[i : i + chunk_size]
+                response = self.client.embeddings.create(model=self.model_name, input=chunk)
+                all_embeddings.extend([item.embedding for item in response.data])
+            return all_embeddings
         except Exception as e:
             logging.error(f"Failed processing messages: {e}")
             raise