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
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]]:
embeddings: List[List[float]] = []
if batch_size is None:
batch_size = self.batch_size
if self.show_progress_bar:
try:
from tqdm.auto import tqdm
@ -168,9 +174,24 @@ class VoyageEmbeddings(BaseModel, Embeddings):
Args:
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:
Embedding for the text.
"""
return self._get_embeddings(
[text], batch_size=self.batch_size, input_type="query"
)[0]
texts, batch_size=self.batch_size, input_type=input_type
)