IMPROVEMENT: VoyageEmbeddings embed_general_texts (#13620)

- **Description:** add method embed_general_texts in VoyageEmebddings to
support input_type
  - **Issue:** 
  - **Dependencies:** 
  - **Tag maintainer:** 
  - **Twitter handle:** @Voyage_AI_
This commit is contained in:
Yujie Qian 2023-11-21 18:33:07 -08:00 committed by GitHub
parent 5064890fcf
commit ace9e64d62
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -115,10 +115,16 @@ class VoyageEmbeddings(BaseModel, Embeddings):
return params return params
def _get_embeddings( def _get_embeddings(
self, texts: List[str], batch_size: int, input_type: Optional[str] = None self,
texts: List[str],
batch_size: Optional[int] = None,
input_type: Optional[str] = None,
) -> List[List[float]]: ) -> List[List[float]]:
embeddings: List[List[float]] = [] embeddings: List[List[float]] = []
if batch_size is None:
batch_size = self.batch_size
if self.show_progress_bar: if self.show_progress_bar:
try: try:
from tqdm.auto import tqdm from tqdm.auto import tqdm
@ -168,9 +174,24 @@ class VoyageEmbeddings(BaseModel, Embeddings):
Args: Args:
text: The text to embed. text: The text to embed.
Returns:
Embedding for the text.
"""
return self._get_embeddings([text], input_type="query")[0]
def embed_general_texts(
self, texts: List[str], *, input_type: Optional[str] = None
) -> List[List[float]]:
"""Call out to Voyage Embedding endpoint for embedding general text.
Args:
texts: The list of texts to embed.
input_type: Type of the input text. Default to None, meaning the type is
unspecified. Other options: query, document.
Returns: Returns:
Embedding for the text. Embedding for the text.
""" """
return self._get_embeddings( return self._get_embeddings(
[text], batch_size=self.batch_size, input_type="query" texts, batch_size=self.batch_size, input_type=input_type
)[0] )