Skip to content

Commit

Permalink
improve: Turned off progress bar of huggingface models if cached
Browse files Browse the repository at this point in the history
  • Loading branch information
hh-space-invader committed Nov 17, 2024
1 parent 544d253 commit c8f2f0a
Showing 1 changed file with 13 additions and 1 deletion.
14 changes: 13 additions & 1 deletion fastembed/common/model_management.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,11 @@

import requests
from huggingface_hub import snapshot_download
from huggingface_hub.utils import RepositoryNotFoundError
from huggingface_hub.utils import (
RepositoryNotFoundError,
disable_progress_bars,
enable_progress_bars,
)
from loguru import logger
from tqdm import tqdm

Expand Down Expand Up @@ -119,6 +123,12 @@ def download_files_from_huggingface(
if extra_patterns is not None:
allow_patterns.extend(extra_patterns)

snapshot_dir = Path(cache_dir) / f"models--{hf_source_repo.replace('/', '--')}"
is_cached = snapshot_dir.exists()

if is_cached:
disable_progress_bars()

return snapshot_download(
repo_id=hf_source_repo,
allow_patterns=allow_patterns,
Expand Down Expand Up @@ -265,6 +275,8 @@ def download_model(
f"Could not download model from HuggingFace: {e} "
"Falling back to other sources."
)
finally:
enable_progress_bars()
if url_source or local_files_only:
try:
return cls.retrieve_model_gcs(
Expand Down

0 comments on commit c8f2f0a

Please sign in to comment.