Fix dtype parsing from vectorizer kwargs (#237)

Fixing a bug in how we handled kwargs within the vectorizer classes
after a recent change was introduced.

Author: tylerhutcherson

redisvl/utils/vectorize/base.py

@@ -81,11 +81,13 @@ def batchify(self, seq: list, size: int, preprocess: Optional[Callable] = None):
             else:
                 yield seq[pos : pos + size]
 
-    def _process_embedding(self, embedding: List[float], as_buffer: bool, **kwargs):
+    def _process_embedding(
+        self, embedding: List[float], as_buffer: bool, dtype: Optional[str]
+    ):
         if as_buffer:
-            if "dtype" not in kwargs:
+            if not dtype:
                 raise RuntimeError(
                     "dtype is required if converting from float to byte string."
                 )
-            return array_to_buffer(embedding, kwargs["dtype"])
+            return array_to_buffer(embedding, dtype)
         return embedding

redisvl/utils/vectorize/text/azureopenai.py

@@ -190,11 +190,13 @@ def embed_many(
         if len(texts) > 0 and not isinstance(texts[0], str):
             raise TypeError("Must pass in a list of str values to embed.")
 
+        dtype = kwargs.pop("dtype", None)
+
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
             response = self._client.embeddings.create(input=batch, model=self.model)
             embeddings += [
-                self._process_embedding(r.embedding, as_buffer, **kwargs)
+                self._process_embedding(r.embedding, as_buffer, dtype)
                 for r in response.data
             ]
         return embeddings
@@ -231,8 +233,11 @@ def embed(
 
         if preprocess:
             text = preprocess(text)
+
+        dtype = kwargs.pop("dtype", None)
+
         result = self._client.embeddings.create(input=[text], model=self.model)
-        return self._process_embedding(result.data[0].embedding, as_buffer, **kwargs)
+        return self._process_embedding(result.data[0].embedding, as_buffer, dtype)
 
     @retry(
         wait=wait_random_exponential(min=1, max=60),
@@ -269,13 +274,15 @@ async def aembed_many(
         if len(texts) > 0 and not isinstance(texts[0], str):
             raise TypeError("Must pass in a list of str values to embed.")
 
+        dtype = kwargs.pop("dtype", None)
+
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
             response = await self._aclient.embeddings.create(
                 input=batch, model=self.model
             )
             embeddings += [
-                self._process_embedding(r.embedding, as_buffer, **kwargs)
+                self._process_embedding(r.embedding, as_buffer, dtype)
                 for r in response.data
             ]
         return embeddings
@@ -312,8 +319,11 @@ async def aembed(
 
         if preprocess:
             text = preprocess(text)
+
+        dtype = kwargs.pop("dtype", None)
+
         result = await self._aclient.embeddings.create(input=[text], model=self.model)
-        return self._process_embedding(result.data[0].embedding, as_buffer, **kwargs)
+        return self._process_embedding(result.data[0].embedding, as_buffer, dtype)
 
     @property
     def type(self) -> str:

redisvl/utils/vectorize/text/cohere.py

@@ -155,12 +155,16 @@ def embed(
                 "Must pass in a str value for cohere embedding input_type. \
                     See https://docs.cohere.com/reference/embed."
             )
+
         if preprocess:
             text = preprocess(text)
+
+        dtype = kwargs.pop("dtype", None)
+
         embedding = self._client.embed(
             texts=[text], model=self.model, input_type=input_type
         ).embeddings[0]
-        return self._process_embedding(embedding, as_buffer, **kwargs)
+        return self._process_embedding(embedding, as_buffer, dtype)
 
     @retry(
         wait=wait_random_exponential(min=1, max=60),
@@ -224,13 +228,15 @@ def embed_many(
                     See https://docs.cohere.com/reference/embed."
             )
 
+        dtype = kwargs.pop("dtype", None)
+
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
             response = self._client.embed(
                 texts=batch, model=self.model, input_type=input_type
             )
             embeddings += [
-                self._process_embedding(embedding, as_buffer, **kwargs)
+                self._process_embedding(embedding, as_buffer, dtype)
                 for embedding in response.embeddings
             ]
         return embeddings

redisvl/utils/vectorize/text/custom.py

@@ -172,9 +172,11 @@ def embed(
 
         if preprocess:
             text = preprocess(text)
-        else:
-            result = self._embed_func(text, **kwargs)
-        return self._process_embedding(result, as_buffer, **kwargs)
+
+        dtype = kwargs.pop("dtype", None)
+
+        result = self._embed_func(text, **kwargs)
+        return self._process_embedding(result, as_buffer, dtype)
 
     def embed_many(
         self,
@@ -210,11 +212,13 @@ def embed_many(
         if not self._embed_many_func:
             raise NotImplementedError
 
+        dtype = kwargs.pop("dtype", None)
+
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
             results = self._embed_many_func(batch, **kwargs)
             embeddings += [
-                self._process_embedding(r, as_buffer, **kwargs) for r in results
+                self._process_embedding(r, as_buffer, dtype) for r in results
             ]
         return embeddings
 
@@ -249,9 +253,11 @@ async def aembed(
 
         if preprocess:
             text = preprocess(text)
-        else:
-            result = await self._aembed_func(text, **kwargs)
-        return self._process_embedding(result, as_buffer, **kwargs)
+
+        dtype = kwargs.pop("dtype", None)
+
+        result = await self._aembed_func(text, **kwargs)
+        return self._process_embedding(result, as_buffer, dtype)
 
     async def aembed_many(
         self,
@@ -287,11 +293,13 @@ async def aembed_many(
         if not self._aembed_many_func:
             raise NotImplementedError
 
+        dtype = kwargs.pop("dtype", None)
+
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
             results = await self._aembed_many_func(batch, **kwargs)
             embeddings += [
-                self._process_embedding(r, as_buffer, **kwargs) for r in results
+                self._process_embedding(r, as_buffer, dtype) for r in results
             ]
         return embeddings
 

redisvl/utils/vectorize/text/huggingface.py

@@ -99,8 +99,11 @@ def embed(
 
         if preprocess:
             text = preprocess(text)
+
+        dtype = kwargs.pop("dtype", None)
+
         embedding = self._client.encode([text], **kwargs)[0]
-        return self._process_embedding(embedding.tolist(), as_buffer, **kwargs)
+        return self._process_embedding(embedding.tolist(), as_buffer, dtype)
 
     def embed_many(
         self,
@@ -133,12 +136,14 @@ def embed_many(
         if len(texts) > 0 and not isinstance(texts[0], str):
             raise TypeError("Must pass in a list of str values to embed.")
 
+        dtype = kwargs.pop("dtype", None)
+
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
             batch_embeddings = self._client.encode(batch, **kwargs)
             embeddings.extend(
                 [
-                    self._process_embedding(embedding.tolist(), as_buffer, **kwargs)
+                    self._process_embedding(embedding.tolist(), as_buffer, dtype)
                     for embedding in batch_embeddings
                 ]
             )

redisvl/utils/vectorize/text/mistral.py

@@ -140,11 +140,13 @@ def embed_many(
         if len(texts) > 0 and not isinstance(texts[0], str):
             raise TypeError("Must pass in a list of str values to embed.")
 
+        dtype = kwargs.pop("dtype", None)
+
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
             response = self._client.embeddings(model=self.model, input=batch)
             embeddings += [
-                self._process_embedding(r.embedding, as_buffer, **kwargs)
+                self._process_embedding(r.embedding, as_buffer, dtype)
                 for r in response.data
             ]
         return embeddings
@@ -181,8 +183,11 @@ def embed(
 
         if preprocess:
             text = preprocess(text)
+
+        dtype = kwargs.pop("dtype", None)
+
         result = self._client.embeddings(model=self.model, input=[text])
-        return self._process_embedding(result.data[0].embedding, as_buffer, **kwargs)
+        return self._process_embedding(result.data[0].embedding, as_buffer, dtype)
 
     @retry(
         wait=wait_random_exponential(min=1, max=60),
@@ -219,11 +224,13 @@ async def aembed_many(
         if len(texts) > 0 and not isinstance(texts[0], str):
             raise TypeError("Must pass in a list of str values to embed.")
 
+        dtype = kwargs.pop("dtype", None)
+
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
             response = await self._aclient.embeddings(model=self.model, input=batch)
             embeddings += [
-                self._process_embedding(r.embedding, as_buffer, **kwargs)
+                self._process_embedding(r.embedding, as_buffer, dtype)
                 for r in response.data
             ]
         return embeddings
@@ -260,8 +267,11 @@ async def aembed(
 
         if preprocess:
             text = preprocess(text)
+
+        dtype = kwargs.pop("dtype", None)
+
         result = await self._aclient.embeddings(model=self.model, input=[text])
-        return self._process_embedding(result.data[0].embedding, as_buffer, **kwargs)
+        return self._process_embedding(result.data[0].embedding, as_buffer, dtype)
 
     @property
     def type(self) -> str:

redisvl/utils/vectorize/text/openai.py

@@ -144,11 +144,13 @@ def embed_many(
         if len(texts) > 0 and not isinstance(texts[0], str):
             raise TypeError("Must pass in a list of str values to embed.")
 
+        dtype = kwargs.pop("dtype", None)
+
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
             response = self._client.embeddings.create(input=batch, model=self.model)
             embeddings += [
-                self._process_embedding(r.embedding, as_buffer, **kwargs)
+                self._process_embedding(r.embedding, as_buffer, dtype)
                 for r in response.data
             ]
         return embeddings
@@ -185,8 +187,11 @@ def embed(
 
         if preprocess:
             text = preprocess(text)
+
+        dtype = kwargs.pop("dtype", None)
+
         result = self._client.embeddings.create(input=[text], model=self.model)
-        return self._process_embedding(result.data[0].embedding, as_buffer, **kwargs)
+        return self._process_embedding(result.data[0].embedding, as_buffer, dtype)
 
     @retry(
         wait=wait_random_exponential(min=1, max=60),
@@ -223,13 +228,15 @@ async def aembed_many(
         if len(texts) > 0 and not isinstance(texts[0], str):
             raise TypeError("Must pass in a list of str values to embed.")
 
+        dtype = kwargs.pop("dtype", None)
+
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
             response = await self._aclient.embeddings.create(
                 input=batch, model=self.model
             )
             embeddings += [
-                self._process_embedding(r.embedding, as_buffer, **kwargs)
+                self._process_embedding(r.embedding, as_buffer, dtype)
                 for r in response.data
             ]
         return embeddings
@@ -266,8 +273,11 @@ async def aembed(
 
         if preprocess:
             text = preprocess(text)
+
+        dtype = kwargs.pop("dtype", None)
+
         result = await self._aclient.embeddings.create(input=[text], model=self.model)
-        return self._process_embedding(result.data[0].embedding, as_buffer, **kwargs)
+        return self._process_embedding(result.data[0].embedding, as_buffer, dtype)
 
     @property
     def type(self) -> str:

redisvl/utils/vectorize/text/vertexai.py

@@ -151,11 +151,13 @@ def embed_many(
         if len(texts) > 0 and not isinstance(texts[0], str):
             raise TypeError("Must pass in a list of str values to embed.")
 
+        dtype = kwargs.pop("dtype", None)
+
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
             response = self._client.get_embeddings(batch)
             embeddings += [
-                self._process_embedding(r.values, as_buffer, **kwargs) for r in response
+                self._process_embedding(r.values, as_buffer, dtype) for r in response
             ]
         return embeddings
 
@@ -191,8 +193,11 @@ def embed(
 
         if preprocess:
             text = preprocess(text)
+
+        dtype = kwargs.pop("dtype", None)
+
         result = self._client.get_embeddings([text])
-        return self._process_embedding(result[0].values, as_buffer, **kwargs)
+        return self._process_embedding(result[0].values, as_buffer, dtype)
 
     @property
     def type(self) -> str: