From fdbfa6b2c8adf85c1e1a1a757ffd01ec02ceb8ab Mon Sep 17 00:00:00 2001 From: Unai Garay Maestre Date: Wed, 24 Jan 2024 19:16:16 +0100 Subject: [PATCH] Adds progress bar to VertexAIEmbeddings (#14542) - **Description:** Adds progress bar to VertexAIEmbeddings - **Issue:** related issue https://github.com/langchain-ai/langchain/issues/13637 Signed-off-by: ugm2 --------- Signed-off-by: ugm2 --- .../langchain_community/embeddings/vertexai.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/libs/community/langchain_community/embeddings/vertexai.py b/libs/community/langchain_community/embeddings/vertexai.py index e0f0834587f..3b6677d664d 100644 --- a/libs/community/langchain_community/embeddings/vertexai.py +++ b/libs/community/langchain_community/embeddings/vertexai.py @@ -30,6 +30,8 @@ class VertexAIEmbeddings(_VertexAICommon, Embeddings): # Instance context instance: Dict[str, Any] = {} #: :meta private: + show_progress_bar: bool = False + """Whether to show a tqdm progress bar. Must have `tqdm` installed.""" @root_validator() def validate_environment(cls, values: Dict) -> Dict: @@ -302,7 +304,20 @@ class VertexAIEmbeddings(_VertexAICommon, Embeddings): # In such case, batches have texts that were not processed yet. embeddings.extend(first_batch_result) tasks = [] - for batch in batches: + if self.show_progress_bar: + try: + from tqdm import tqdm + + iter_ = tqdm(batches, desc="VertexAIEmbeddings") + except ImportError: + logger.warning( + "Unable to show progress bar because tqdm could not be imported. " + "Please install with `pip install tqdm`." + ) + iter_ = batches + else: + iter_ = batches + for batch in iter_: tasks.append( self.instance["task_executor"].submit( self._get_embeddings_with_retry,