Skip to content

Commit

Permalink
community[patch]: Add Progress bar to HuggingFaceEmbeddings (#16758)
Browse files Browse the repository at this point in the history
- **Description:** Adds a function parameter to HuggingFaceEmbeddings
called `show_progress` that enables a `tqdm` progress bar if enabled.
Does not function if `multi_process = True`.
  - **Issue:** n/a
  - **Dependencies:** n/a
tylertitsworth authored Feb 5, 2024
1 parent ae33979 commit 304f3f5
Showing 1 changed file with 5 additions and 1 deletion.
6 changes: 5 additions & 1 deletion libs/community/langchain_community/embeddings/huggingface.py
Original file line number Diff line number Diff line change
@@ -49,6 +49,8 @@ class HuggingFaceEmbeddings(BaseModel, Embeddings):
"""Keyword arguments to pass when calling the `encode` method of the model."""
multi_process: bool = False
"""Run encode() on multiple GPUs."""
show_progress: bool = False
"""Whether to show a progress bar."""

def __init__(self, **kwargs: Any):
"""Initialize the sentence_transformer."""
@@ -88,7 +90,9 @@ def embed_documents(self, texts: List[str]) -> List[List[float]]:
embeddings = self.client.encode_multi_process(texts, pool)
sentence_transformers.SentenceTransformer.stop_multi_process_pool(pool)
else:
embeddings = self.client.encode(texts, **self.encode_kwargs)
embeddings = self.client.encode(
texts, show_progress_bar=self.show_progress, **self.encode_kwargs
)

return embeddings.tolist()

0 comments on commit 304f3f5

Please sign in to comment.