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 <unaigaraymaestre@gmail.com>

---------

Signed-off-by: ugm2 <unaigaraymaestre@gmail.com>
This commit is contained in:
Unai Garay Maestre 2024-01-24 19:16:16 +01:00 committed by GitHub
parent 643fb3ab50
commit fdbfa6b2c8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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,