github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-08-04 18:53:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

community[patch]: Add Progress bar to HuggingFaceEmbeddings (#16758)

Browse Source 
- **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
This commit is contained in:
Tyler Titsworth 2024-02-05 14:33:34 -08:00 committed by GitHub
parent ae33979813
commit 304f3f5fc1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
libs/community/langchain_community/embeddings/huggingface.py
View File

@ -49,6 +49,8 @@ class HuggingFaceEmbeddings(BaseModel, Embeddings):
"""Keyword arguments to pass when calling the `encode` method of the model.""" """Keyword arguments to pass when calling the `encode` method of the model."""
multi_process: bool = False multi_process: bool = False
"""Run encode() on multiple GPUs.""" """Run encode() on multiple GPUs."""
show_progress: bool = False
"""Whether to show a progress bar."""
def __init__(self, **kwargs: Any): def __init__(self, **kwargs: Any):
"""Initialize the sentence_transformer.""" """Initialize the sentence_transformer."""
@ -88,7 +90,9 @@ class HuggingFaceEmbeddings(BaseModel, Embeddings):
embeddings = self.client.encode_multi_process(texts, pool) embeddings = self.client.encode_multi_process(texts, pool)
sentence_transformers.SentenceTransformer.stop_multi_process_pool(pool) sentence_transformers.SentenceTransformer.stop_multi_process_pool(pool)
else: 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() return embeddings.tolist()