hfutils 0.1.1__tar.gz → 0.1.2__tar.gz

{hfutils-0.1.1/hfutils.egg-info → hfutils-0.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: hfutils
-Version: 0.1.1
+Version: 0.1.2
 Summary: Useful utilities for huggingface
 Home-page: https://github.com/deepghs/hfutils
 Author: narugo1992

{hfutils-0.1.1 → hfutils-0.1.2}/hfutils/config/meta.py

@@ -7,7 +7,7 @@ Overview:
 __TITLE__ = 'hfutils'
 
 #: Version of this project.
-__VERSION__ = '0.1.1'
+__VERSION__ = '0.1.2'
 
 #: Short description of the project, will be included in ``setup.py``.
 __DESCRIPTION__ = 'Useful utilities for huggingface'

{hfutils-0.1.1 → hfutils-0.1.2}/hfutils/entry/download.py

@@ -1,3 +1,4 @@
+import os
 import warnings
 from typing import Optional
 
@@ -46,10 +47,12 @@ def _add_download_subcommand(cli: click.Group) -> click.Group:
                   help='Max threads to download.', show_default=True)
     @click.option('-p', '--password', 'password', type=str, default=None,
                   help='Password for the archive file. Only applied when -a is used.', show_default=True)
+    @click.option('--tmpdir', 'tmpdir', type=str, default=None,
+                  help='Use custom temporary Directory.', show_default=True)
     @command_wrap()
     def download(repo_id: str, repo_type: RepoTypeTyping,
                  file_in_repo: Optional[str], archive_in_repo: Optional[str], dir_in_repo: Optional[str],
-                 output_path: str, revision: str, max_workers: int, password: Optional[str]):
+                 output_path: str, revision: str, max_workers: int, password: Optional[str], tmpdir: Optional[str]):
         """
         Download data from HuggingFace repositories.
 
@@ -71,7 +74,12 @@ def _add_download_subcommand(cli: click.Group) -> click.Group:
         :type max_workers: int
         :param password: Password for the archive file. Only applied when -a is used.
         :type password: str, optional
+        :param tmpdir: Use custom temporary Directory.
+        :type tmpdir: str, optional
         """
+        if tmpdir:
+            os.environ['TMPDIR'] = tmpdir
+
         if not file_in_repo and not archive_in_repo and not dir_in_repo:
             raise NoRemotePathAssignedWithDownload('No remote path in repository assigned.\n'
                                                    'One of the -f, -a, or -d option is required.')

{hfutils-0.1.1 → hfutils-0.1.2}/hfutils/operate/download.py

@@ -4,6 +4,8 @@ import shutil
 from concurrent.futures import ThreadPoolExecutor
 from typing import List, Optional
 
+import requests.exceptions
+
 from .base import RepoTypeTyping, list_files_in_repository, _IGNORE_PATTERN_UNSET, get_hf_client
 from .validate import is_local_file_ready
 from ..archive import archive_unpack
@@ -83,7 +85,7 @@ def download_archive_as_directory(local_directory: str, repo_id: str, file_in_re
 def download_directory_as_directory(local_directory: str, repo_id: str, dir_in_repo: str = '.',
                                     repo_type: RepoTypeTyping = 'dataset', revision: str = 'main',
                                     silent: bool = False, ignore_patterns: List[str] = _IGNORE_PATTERN_UNSET,
-                                    resume_download: bool = True, max_workers: int = 8,
+                                    resume_download: bool = True, max_workers: int = 8, max_retries: int = 5,
                                     hf_token: Optional[str] = None):
     """
     Download all files in a directory from a Hugging Face repository to a local directory.
@@ -104,15 +106,18 @@ def download_directory_as_directory(local_directory: str, repo_id: str, dir_in_r
     :type ignore_patterns: List[str]
     :param max_workers: Max workers when downloading. Default is ``8``.
     :type max_workers: int
+    :param max_retries: Max retry times when downloading. Default is ``5``.
+    :type max_retries: int
     :param resume_download: Resume the existing download.
     :type resume_download: bool
     :param hf_token: Huggingface token for API client, use ``HF_TOKEN`` variable if not assigned.
     :type hf_token: str, optional
     """
-    files = list_files_in_repository(repo_id, repo_type, dir_in_repo, revision, ignore_patterns)
+    files = list_files_in_repository(repo_id, repo_type, dir_in_repo, revision, ignore_patterns, hf_token=hf_token)
     progress = tqdm(files, silent=silent, desc=f'Downloading {dir_in_repo!r} ...')
 
     def _download_one_file(rel_file):
+        current_resume_download = resume_download
         try:
             dst_file = os.path.join(local_directory, rel_file)
             if os.path.exists(dst_file) and is_local_file_ready(
@@ -125,15 +130,28 @@ def download_directory_as_directory(local_directory: str, repo_id: str, dir_in_r
             ):
                 logging.info(f'Local file {rel_file} is ready, download skipped.')
             else:
-                download_file_to_file(
-                    local_file=dst_file,
-                    repo_id=repo_id,
-                    file_in_repo=f'{dir_in_repo}/{rel_file}',
-                    repo_type=repo_type,
-                    revision=revision,
-                    resume_download=resume_download,
-                    hf_token=hf_token,
-                )
+                tries = 0
+                while True:
+                    try:
+                        download_file_to_file(
+                            local_file=dst_file,
+                            repo_id=repo_id,
+                            file_in_repo=f'{dir_in_repo}/{rel_file}',
+                            repo_type=repo_type,
+                            revision=revision,
+                            resume_download=current_resume_download,
+                            hf_token=hf_token,
+                        )
+                    except requests.exceptions.RequestException as err:
+                        if tries < max_retries:
+                            tries += 1
+                            logging.warning(f'Download {rel_file!r} failed, retry ({tries}/{max_retries}) - {err!r}.')
+                            current_resume_download = True
+                        else:
+                            raise
+                    else:
+                        break
+
             progress.update()
         except Exception as err:
             logging.error(f'Unexpected error when downloading {rel_file!r} - {err!r}')

{hfutils-0.1.1 → hfutils-0.1.2}/hfutils/operate/upload.py

@@ -110,7 +110,7 @@ def upload_directory_as_directory(local_directory, repo_id: str, path_in_repo: s
     if clear:
         pre_exist_files = {
             tuple(file.split('/')) for file in
-            list_files_in_repository(repo_id, repo_type, path_in_repo, revision, ignore_patterns)
+            list_files_in_repository(repo_id, repo_type, path_in_repo, revision, ignore_patterns, hf_token=hf_token)
         }
     else:
         pre_exist_files = set()

{hfutils-0.1.1 → hfutils-0.1.2/hfutils.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: hfutils
-Version: 0.1.1
+Version: 0.1.2
 Summary: Useful utilities for huggingface
 Home-page: https://github.com/deepghs/hfutils
 Author: narugo1992