Vectorizer improvements (#44)

This PR introduces a few improvements to the vectorizers module:
- Raises `TypeError` when the wrong input type is not passed in.
- Updates the methodology for fetching the output dimensions of the
embedding model to be dynamic based on the given model.

Author: tylerhutcherson

redisvl/__init__.py

@@ -1,5 +1,3 @@
-
-
 from redisvl.version import __version__
 
-all = ["__version__"]
+all = ["__version__"]

redisvl/cli/index.py

@@ -11,6 +11,7 @@
 
 logger = get_logger("[RedisVL]")
 
+
 class Index:
     usage = "\n".join(
         [

redisvl/cli/log.py

@@ -5,9 +5,8 @@
 
 # constants for logging
 coloredlogs.DEFAULT_DATE_FORMAT = "%H:%M:%S"
-coloredlogs.DEFAULT_LOG_FORMAT = (
-    "%(asctime)s %(name)s %(levelname)s   %(message)s"
-)
+coloredlogs.DEFAULT_LOG_FORMAT = "%(asctime)s %(name)s %(levelname)s   %(message)s"
+
 
 def get_logger(name, log_level="info", fmt=None):
     """Return a logger instance"""

redisvl/cli/main.py

@@ -2,8 +2,8 @@
 import sys
 
 from redisvl.cli.index import Index
-from redisvl.cli.version import Version
 from redisvl.cli.log import get_logger
+from redisvl.cli.version import Version
 
 logger = get_logger(__name__)
 
@@ -42,4 +42,4 @@ def index(self):
 
     def version(self):
         Version()
-        exit(0)
+        exit(0)

redisvl/cli/version.py

@@ -1,9 +1,10 @@
-import sys
 import argparse
+import sys
 from argparse import Namespace
 
 from redisvl import __version__
 from redisvl.cli.log import get_logger
+
 logger = get_logger("[RedisVL]")
 
 
@@ -28,4 +29,4 @@ def version(self, args: Namespace):
         if args.short:
             print(__version__)
         else:
-            logger.info(f"RedisVL version {__version__}")
+            logger.info(f"RedisVL version {__version__}")

redisvl/vectorize/base.py

@@ -1,11 +1,12 @@
-from typing import Callable, Dict, List, Optional
+from typing import Callable, List, Optional
 
 from redisvl.utils.utils import array_to_buffer
 
 
 class BaseVectorizer:
-    def __init__(self, model: str, dims: int, api_config: Optional[Dict] = None):
-        self._dims = dims
+    _dims = None
+
+    def __init__(self, model: str):
         self._model = model
 
     @property

redisvl/vectorize/text/huggingface.py

@@ -6,17 +6,24 @@
 class HFTextVectorizer(BaseVectorizer):
     # TODO - add docstring
     def __init__(self, model: str, api_config: Optional[Dict] = None):
-        # TODO set dims based on model
-        dims = 768
-        super().__init__(model, dims, api_config)
+        super().__init__(model)
         try:
             from sentence_transformers import SentenceTransformer
         except ImportError:
             raise ImportError(
                 "HFTextVectorizer requires sentence-transformers library. Please install with pip install sentence-transformers"
             )
 
-        self._model_client = SentenceTransformer(model)
+        self._model_client = SentenceTransformer(self._model)
+
+        try:
+            self._dims = self._set_model_dims()
+        except:
+            raise ValueError("Error setting embedding model dimensions")
+
+    def _set_model_dims(self):
+        embedding = self._model_client.encode(["dimension check"])[0]
+        return len(embedding)
 
     def embed(
         self,
@@ -35,7 +42,13 @@ def embed(
 
         Returns:
             List[float]: Embedding.
+
+        Raises:
+            TypeError: If the wrong input type is passed in for the text.
         """
+        if not isinstance(text, str):
+            raise TypeError("Must pass in a str value to embed.")
+
         if preprocess:
             text = preprocess(text)
         embedding = self._model_client.encode([text])[0]
@@ -62,7 +75,15 @@ def embed_many(
 
         Returns:
             List[List[float]]: List of embeddings.
+
+        Raises:
+            TypeError: If the wrong input type is passed in for the test.
         """
+        if not isinstance(texts, list):
+                raise TypeError("Must pass in a list of str values to embed.")
+        if  len(texts) > 0 and not isinstance(texts[0], str):
+                raise TypeError("Must pass in a list of str values to embed.")
+
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
             batch_embeddings = self._model_client.encode(batch)

redisvl/vectorize/text/openai.py

@@ -1,19 +1,15 @@
 from typing import Callable, Dict, List, Optional
 
-from tenacity import (  # for exponential backoff
-    retry,
-    stop_after_attempt,
-    wait_random_exponential,
-)
+from tenacity import retry, stop_after_attempt, wait_random_exponential
+from tenacity.retry import retry_if_not_exception_type
 
 from redisvl.vectorize.base import BaseVectorizer
 
 
 class OpenAITextVectorizer(BaseVectorizer):
     # TODO - add docstring
     def __init__(self, model: str, api_config: Optional[Dict] = None):
-        dims = 1536
-        super().__init__(model, dims, api_config)
+        super().__init__(model)
         if not api_config:
             raise ValueError("OpenAI API key is required in api_config")
         try:
@@ -25,7 +21,23 @@ def __init__(self, model: str, api_config: Optional[Dict] = None):
         openai.api_key = api_config.get("api_key", None)
         self._model_client = openai.Embedding
 
-    @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
+        try:
+            self._dims = self._set_model_dims()
+        except:
+            raise ValueError("Error setting embedding model dimensions")
+
+    def _set_model_dims(self):
+        embedding = self._model_client.create(
+            input=["dimension test"],
+            engine=self._model
+        )["data"][0]["embedding"]
+        return len(embedding)
+
+    @retry(
+        wait=wait_random_exponential(min=1, max=60),
+        stop=stop_after_attempt(6),
+        retry=retry_if_not_exception_type(TypeError),
+    )
     def embed_many(
         self,
         texts: List[str],
@@ -46,7 +58,15 @@ def embed_many(
 
         Returns:
             List[List[float]]: List of embeddings.
+
+        Raises:
+            TypeError: If the wrong input type is passed in for the test.
         """
+        if not isinstance(texts, list):
+                raise TypeError("Must pass in a list of str values to embed.")
+        if  len(texts) > 0 and not isinstance(texts[0], str):
+                raise TypeError("Must pass in a list of str values to embed.")
+
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
             response = self._model_client.create(input=batch, engine=self._model)
@@ -56,7 +76,11 @@ def embed_many(
             ]
         return embeddings
 
-    @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
+    @retry(
+        wait=wait_random_exponential(min=1, max=60),
+        stop=stop_after_attempt(6),
+        retry=retry_if_not_exception_type(TypeError),
+    )
     def embed(
         self,
         text: str,
@@ -74,13 +98,23 @@ def embed(
 
         Returns:
             List[float]: Embedding.
+
+        Raises:
+            TypeError: If the wrong input type is passed in for the test.
         """
+        if not isinstance(text, str):
+            raise TypeError("Must pass in a str value to embed.")
+
         if preprocess:
             text = preprocess(text)
         result = self._model_client.create(input=[text], engine=self._model)
         return self._process_embedding(result["data"][0]["embedding"], as_buffer)
 
-    @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
+    @retry(
+        wait=wait_random_exponential(min=1, max=60),
+        stop=stop_after_attempt(6),
+        retry=retry_if_not_exception_type(TypeError),
+    )
     async def aembed_many(
         self,
         texts: List[str],
@@ -101,7 +135,15 @@ async def aembed_many(
 
         Returns:
             List[List[float]]: List of embeddings.
+
+        Raises:
+            TypeError: If the wrong input type is passed in for the test.
         """
+        if not isinstance(texts, list):
+                raise TypeError("Must pass in a list of str values to embed.")
+        if  len(texts) > 0 and not isinstance(texts[0], str):
+                raise TypeError("Must pass in a list of str values to embed.")
+
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
             response = await self._model_client.acreate(input=batch, engine=self._model)
@@ -111,7 +153,11 @@ async def aembed_many(
             ]
         return embeddings
 
-    @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
+    @retry(
+        wait=wait_random_exponential(min=1, max=60),
+        stop=stop_after_attempt(6),
+        retry=retry_if_not_exception_type(TypeError),
+    )
     async def aembed(
         self,
         text: str,
@@ -129,7 +175,13 @@ async def aembed(
 
         Returns:
             List[float]: Embedding.
+
+        Raises:
+            TypeError: If the wrong input type is passed in for the test.
         """
+        if not isinstance(text, str):
+            raise TypeError("Must pass in a str value to embed.")
+
         if preprocess:
             text = preprocess(text)
         result = await self._model_client.acreate(input=[text], engine=self._model)

tests/integration/test_vectorizers.py

@@ -35,6 +35,17 @@ def test_vectorizer_embed_many(vectorizer):
     )
 
 
+def test_vectorizer_bad_input(vectorizer):
+    with pytest.raises(TypeError):
+        vectorizer.embed(1)
+
+    with pytest.raises(TypeError):
+        vectorizer.embed({"foo": "bar"})
+
+    with pytest.raises(TypeError):
+        vectorizer.embed_many(42)
+
+
 @pytest.fixture(params=[OpenAITextVectorizer])
 def avectorizer(request, openai_key):
     # Here we use actual models for integration test
@@ -63,3 +74,15 @@ async def test_vectorizer_aembed_many(avectorizer):
     assert all(
         isinstance(emb, list) and len(emb) == avectorizer.dims for emb in embeddings
     )
+
+
+@pytest.mark.asyncio
+async def test_avectorizer_bad_input(avectorizer):
+    with pytest.raises(TypeError):
+        avectorizer.embed(1)
+
+    with pytest.raises(TypeError):
+        avectorizer.embed({"foo": "bar"})
+
+    with pytest.raises(TypeError):
+        avectorizer.embed_many(42)