From fed963f360bebdeba7280917a912370c07984c1a Mon Sep 17 00:00:00 2001 From: fodizoltan Date: Thu, 15 Feb 2024 17:06:20 +0100 Subject: [PATCH 1/2] - Change model to a required argument - Change the default batch size to 72 for model "voyage-2"/"voyage-02", and to 7 for all other models. (However, still allow the user to specify batch size) --- libs/community/langchain_community/embeddings/voyageai.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/libs/community/langchain_community/embeddings/voyageai.py b/libs/community/langchain_community/embeddings/voyageai.py index 93109d45c65b6..bff4ee3419767 100644 --- a/libs/community/langchain_community/embeddings/voyageai.py +++ b/libs/community/langchain_community/embeddings/voyageai.py @@ -115,13 +115,21 @@ def _invocation_params( def _get_embeddings( self, texts: List[str], + model: str, batch_size: Optional[int] = None, input_type: Optional[str] = None, ) -> List[List[float]]: embeddings: List[List[float]] = [] + self.model = model + if batch_size is None: batch_size = self.batch_size + else: + if self.model in ["voyage-2", "voyage-02"]: + batch_size = 72 + else: + batch_size = 7 if self.show_progress_bar: try: From 0404dd6d9ac40499c7c1ca9183a9d27dedd2be89 Mon Sep 17 00:00:00 2001 From: fodizoltan Date: Thu, 15 Feb 2024 20:31:05 +0100 Subject: [PATCH 2/2] - Change model to a required argument - Change the default batch size to 72 for model "voyage-2"/"voyage-02", and to 7 for all other models. (However, still allow the user to specify batch size) --- .../langchain_community/embeddings/voyageai.py | 14 +++----------- .../integration_tests/embeddings/test_voyageai.py | 2 +- 2 files changed, 4 insertions(+), 12 deletions(-) diff --git a/libs/community/langchain_community/embeddings/voyageai.py b/libs/community/langchain_community/embeddings/voyageai.py index bff4ee3419767..e0bc53852c17b 100644 --- a/libs/community/langchain_community/embeddings/voyageai.py +++ b/libs/community/langchain_community/embeddings/voyageai.py @@ -69,12 +69,12 @@ class VoyageEmbeddings(BaseModel, Embeddings): from langchain_community.embeddings import VoyageEmbeddings - voyage = VoyageEmbeddings(voyage_api_key="your-api-key") + voyage = VoyageEmbeddings(voyage_api_key="your-api-key", model="voyage-2") text = "This is a test query." query_result = voyage.embed_query(text) """ - model: str = "voyage-01" + model: str voyage_api_base: str = "https://api.voyageai.com/v1/embeddings" voyage_api_key: Optional[SecretStr] = None batch_size: int = 8 @@ -115,21 +115,13 @@ def _invocation_params( def _get_embeddings( self, texts: List[str], - model: str, batch_size: Optional[int] = None, input_type: Optional[str] = None, ) -> List[List[float]]: embeddings: List[List[float]] = [] - self.model = model - if batch_size is None: - batch_size = self.batch_size - else: - if self.model in ["voyage-2", "voyage-02"]: - batch_size = 72 - else: - batch_size = 7 + batch_size = 72 if self.model in ["voyage-2", "voyage-02"] else 7 if self.show_progress_bar: try: diff --git a/libs/community/tests/integration_tests/embeddings/test_voyageai.py b/libs/community/tests/integration_tests/embeddings/test_voyageai.py index b23dbd7f538b3..abdf08e94d5b0 100644 --- a/libs/community/tests/integration_tests/embeddings/test_voyageai.py +++ b/libs/community/tests/integration_tests/embeddings/test_voyageai.py @@ -2,7 +2,7 @@ from langchain_community.embeddings.voyageai import VoyageEmbeddings # Please set VOYAGE_API_KEY in the environment variables -MODEL = "voyage-01" +MODEL = "voyage-2" def test_voyagi_embedding_documents() -> None: