Skip to content

Commit

Permalink
VoyageAI integration: correction (#64)
Browse files Browse the repository at this point in the history
Correcting the result type, encode_queries and encode_documents returns
with np.array for every embedding_type
  • Loading branch information
fzowl authored Jan 18, 2025
1 parent fa4021e commit c3e44ab
Showing 1 changed file with 12 additions and 7 deletions.
19 changes: 12 additions & 7 deletions src/pymilvus/model/dense/voyageai.py
Original file line number Diff line number Diff line change
Expand Up @@ -83,13 +83,18 @@ def _call_voyage_api(self, texts: List[str], input_type: Optional[str] = None):
output_dtype=self.embedding_type,
output_dimension=self.dim,
).embeddings
if self.embedding_type is None:

if self.embedding_type is None or self.embedding_type == "float":
results = [np.array(data, dtype=np.float32) for data in embeddings]
elif self.embedding_type == "binary":
results = [
np.unpackbits((np.array(result, dtype=np.int16) + 128).astype(np.uint8)).astype(bool)
for result in embeddings
]
elif self.embedding_type == "ubinary":
results = [np.unpackbits(np.array(result, dtype=np.uint8)).astype(bool) for result in embeddings]
else:
if self.embedding_type == "binary":
results = [struct.pack('b' * len(int8_vector), *int8_vector) for int8_vector in embeddings]
elif self.embedding_type == "ubinary":
results = [struct.pack('B' * len(uint8_vector), *uint8_vector) for uint8_vector in embeddings]
elif self.embedding_type == "float":
results = [np.array(result, dtype=np.float32) for result in embeddings]
raise ValueError(f"The provided embedding_type ({self.embedding_type}) is not supported by the selected model "
f"({self.model_name}). Leave this parameter empty for the default embedding_type (float). "
f"Please check the supported embedding_type values here: https://docs.voyageai.com/docs/embeddings")
return results

0 comments on commit c3e44ab

Please sign in to comment.