✏️ Add repo_type to list_repo_files

hezarai · Jun 11, 2024 · edb20e6 · edb20e6
1 parent a4d47ab
commit edb20e6
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/hezar/preprocessors/preprocessor.py b/hezar/preprocessors/preprocessor.py
@@ -81,7 +81,7 @@ def load(
         """
         subfolder = subfolder or cls.preprocessor_subfolder
         cache_dir = cache_dir or HEZAR_CACHE_DIR
-        preprocessor_files = list_repo_files(hub_or_local_path, subfolder=subfolder)
+        preprocessor_files = list_repo_files(hub_or_local_path, subfolder=subfolder, repo_type="model")
         preprocessors = PreprocessorsContainer()
         for f in preprocessor_files:
             if f.endswith(".yaml"):

diff --git a/hezar/utils/hub_utils.py b/hezar/utils/hub_utils.py
@@ -80,13 +80,14 @@ def clone_repo(repo_id: str, save_path: str, **kwargs):
     return repo.local_dir
 
 
-def list_repo_files(hub_or_local_path: str, subfolder: str = None):
+def list_repo_files(hub_or_local_path: str, subfolder: str = None, repo_type: str | RepoType = RepoType.MODEL):
     """
     List all files in a Hub or local model repo
 
     Args:
         hub_or_local_path: Path to hub or local repo
         subfolder: Optional subfolder path
+        repo_type: Repo type of either dataset or model
 
     Returns:
         A list of all file names
@@ -98,7 +99,7 @@ def list_repo_files(hub_or_local_path: str, subfolder: str = None):
             for file in files_:
                 files.append(os.path.relpath(os.path.join(root, file), hub_or_local_path))
     else:
-        files = HfApi().list_repo_files(hub_or_local_path, repo_type=str(RepoType.MODEL))
+        files = HfApi().list_repo_files(hub_or_local_path, repo_type=str(repo_type))
 
     if subfolder:
         files = [os.path.relpath(f, subfolder) for f in files if subfolder in f]