[Snapshot download] Allow to load local repo id with snapshot download (#505)

* [Snapshotdownload] Add more parameter names

* up

* fix style

* add revision to cache

* update tests

* make style

* improve snapshot download

* add tests

* delete dummy folders

* Remove tree

* finish cette mer**

Co-authored-by: Lysandre Debut <[email protected]>

Author: patrickvonplaten

src/huggingface_hub/constants.py

@@ -15,6 +15,8 @@
 CONFIG_NAME = "config.json"
 REPOCARD_NAME = "README.md"
 
+DEFAULT_REVISION = "main"
+
 HUGGINGFACE_CO_URL_HOME = "https://huggingface.co/"
 
 ENV_VARS_TRUE_VALUES = {"1", "ON", "YES", "TRUE"}

src/huggingface_hub/file_download.py

@@ -21,6 +21,7 @@
 
 from . import __version__
 from .constants import (
+    DEFAULT_REVISION,
     HUGGINGFACE_CO_URL_TEMPLATE,
     HUGGINGFACE_HUB_CACHE,
     REPO_TYPES,
@@ -110,7 +111,7 @@ def hf_hub_url(
         repo_id = REPO_TYPES_URL_PREFIXES[repo_type] + repo_id
 
     if revision is None:
-        revision = "main"
+        revision = DEFAULT_REVISION
     return HUGGINGFACE_CO_URL_TEMPLATE.format(
         repo_id=repo_id, revision=revision, filename=filename
     )

src/huggingface_hub/snapshot_download.py

@@ -1,24 +1,33 @@
 import os
+from glob import glob
 from pathlib import Path
 from typing import Dict, Optional, Union
 
-from .constants import HUGGINGFACE_HUB_CACHE
+from .constants import DEFAULT_REVISION, HUGGINGFACE_HUB_CACHE
 from .file_download import cached_download, hf_hub_url
 from .hf_api import HfApi, HfFolder
+from .utils import logging
 
 
 REPO_ID_SEPARATOR = "__"
 # ^ make sure this substring is not allowed in repo_ids on hf.co
 
 
+logger = logging.get_logger(__name__)
+
+
 def snapshot_download(
     repo_id: str,
     revision: Optional[str] = None,
     cache_dir: Union[str, Path, None] = None,
     library_name: Optional[str] = None,
     library_version: Optional[str] = None,
     user_agent: Union[Dict, str, None] = None,
+    proxies=None,
+    etag_timeout=10,
+    resume_download=False,
     use_auth_token: Union[bool, str, None] = None,
+    local_files_only=False,
 ) -> str:
     """
     Downloads a whole snapshot of a repo's files at the specified revision.
@@ -39,6 +48,8 @@ def snapshot_download(
     """
     if cache_dir is None:
         cache_dir = HUGGINGFACE_HUB_CACHE
+    if revision is None:
+        revision = DEFAULT_REVISION
     if isinstance(cache_dir, Path):
         cache_dir = str(cache_dir)
 
@@ -53,18 +64,96 @@ def snapshot_download(
     else:
         token = None
 
-    _api = HfApi()
-    model_info = _api.model_info(repo_id=repo_id, revision=revision, token=token)
+    # remove all `/` occurances to correctly convert repo to directory name
+    repo_id_flattened = repo_id.replace("/", REPO_ID_SEPARATOR)
+
+    # if we have no internet connection we will look for the
+    # last modified folder in the cache
+    if local_files_only:
+        # possible repos have <path/to/cache_dir>/<flatten_repo_id> prefix
+        repo_folders_prefix = os.path.join(cache_dir, repo_id_flattened)
+
+        # list all possible folders that can correspond to the repo_id
+        # and are of the format <flattened-repo-id>.<revision>.<commit-sha>
+        # now let's list all cached repos that have to be included in the revision.
+        # There are 3 cases that we have to consider.
+
+        # 1) cached repos of format <repo_id>.{revision}.<any-hash>
+        # -> in this case {revision} has to be a branch
+        repo_folders_branch = glob(repo_folders_prefix + "." + revision + ".*")
+
+        # 2) cached repos of format <repo_id>.{revision}
+        # -> in this case {revision} has to be a commit sha
+        repo_folders_commit_only = glob(repo_folders_prefix + "." + revision)
+
+        # 3) cached repos of format <repo_id>.<any-branch>.{revision}
+        # -> in this case {revision} also has to be a commit sha
+        repo_folders_branch_commit = glob(repo_folders_prefix + ".*." + revision)
+
+        # combine all possible fetched cached repos
+        repo_folders = (
+            repo_folders_branch + repo_folders_commit_only + repo_folders_branch_commit
+        )
+
+        if len(repo_folders) == 0:
+            raise ValueError(
+                "Cannot find the requested files in the cached path and outgoing traffic has been"
+                " disabled. To enable model look-ups and downloads online, set 'local_files_only'"
+                " to False."
+            )
 
-    storage_folder = os.path.join(
-        cache_dir, repo_id.replace("/", REPO_ID_SEPARATOR) + "." + model_info.sha
-    )
+        # check if repo id was previously cached from a commit sha revision
+        # and passed {revision} is not a commit sha
+        # in this case snapshotting repos locally might lead to unexpected
+        # behavior the user should be warned about
 
-    for model_file in model_info.siblings:
-        url = hf_hub_url(
-            repo_id, filename=model_file.rfilename, revision=model_info.sha
+        # get all folders that were cached with just a sha commit revision
+        all_repo_folders_from_sha = set(glob(repo_folders_prefix + ".*")) - set(
+            glob(repo_folders_prefix + ".*.*")
         )
-        relative_filepath = os.path.join(*model_file.rfilename.split("/"))
+        # 1) is there any repo id that was previously cached from a commit sha?
+        has_a_sha_revision_been_cached = len(all_repo_folders_from_sha) > 0
+        # 2) is the passed {revision} is a branch
+        is_revision_a_branch = (
+            len(repo_folders_commit_only + repo_folders_branch_commit) == 0
+        )
+
+        if has_a_sha_revision_been_cached and is_revision_a_branch:
+            # -> in this case let's warn the user
+            logger.warn(
+                f"The repo {repo_id} was previously downloaded from a commit hash revision "
+                f"and has created the following cached directories {all_repo_folders_from_sha}."
+                f" In this case, trying to load a repo from the branch {revision} in offline "
+                "mode might lead to unexpected behavior by not taking into account the latest "
+                "commits."
+            )
+
+        # find last modified folder
+        storage_folder = max(repo_folders, key=os.path.getmtime)
+
+        # get commit sha
+        repo_id_sha = storage_folder.split(".")[-1]
+        model_files = os.listdir(storage_folder)
+    else:
+        # if we have internet connection we retrieve the correct folder name from the huggingface api
+        _api = HfApi()
+        model_info = _api.model_info(repo_id=repo_id, revision=revision, token=token)
+
+        storage_folder = os.path.join(cache_dir, repo_id_flattened + "." + revision)
+
+        # if passed revision is not identical to the commit sha
+        # then revision has to be a branch name, e.g. "main"
+        # in this case make sure that the branch name is included
+        # cached storage folder name
+        if revision != model_info.sha:
+            storage_folder += f".{model_info.sha}"
+
+        repo_id_sha = model_info.sha
+        model_files = [f.rfilename for f in model_info.siblings]
+
+    for model_file in model_files:
+        url = hf_hub_url(repo_id, filename=model_file, revision=repo_id_sha)
+        relative_filepath = os.path.join(*model_file.split("/"))
 
         # Create potential nested dir
         nested_dirname = os.path.dirname(
@@ -79,7 +168,11 @@ def snapshot_download(
             library_name=library_name,
             library_version=library_version,
             user_agent=user_agent,
+            proxies=proxies,
+            etag_timeout=etag_timeout,
+            resume_download=resume_download,
             use_auth_token=use_auth_token,
+            local_files_only=local_files_only,
         )
 
         if os.path.exists(path + ".lock"):

tests/test_snapshot_download.py

@@ -1,4 +1,5 @@
 import os
+import shutil
 import tempfile
 import time
 import unittest
@@ -46,8 +47,15 @@ def setUp(self) -> None:
 
         self.second_commit_hash = repo.git_head_hash()
 
+        with repo.commit("Add file to other branch", branch="other"):
+            with open("dummy_file_2.txt", "w+") as f:
+                f.write("v4")
+
+        self.third_commit_hash = repo.git_head_hash()
+
     def tearDown(self) -> None:
         self._api.delete_repo(name=REPO_NAME, token=self._token)
+        shutil.rmtree(REPO_NAME)
 
     def test_download_model(self):
         # Test `main` branch
@@ -155,3 +163,148 @@ def test_download_private_model(self):
         self._api.update_repo_visibility(
             token=self._token, name=REPO_NAME, private=False
         )
+
+    def test_download_model_local_only(self):
+        # Test no branch specified
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            # first download folder to cache it
+            snapshot_download(f"{USER}/{REPO_NAME}", cache_dir=tmpdirname)
+
+            # now load from cache
+            storage_folder = snapshot_download(
+                f"{USER}/{REPO_NAME}",
+                cache_dir=tmpdirname,
+                local_files_only=True,
+            )
+
+            # folder contains the two files contributed and the .gitattributes
+            folder_contents = os.listdir(storage_folder)
+            self.assertEqual(len(folder_contents), 3)
+            self.assertTrue("dummy_file.txt" in folder_contents)
+            self.assertTrue("dummy_file_2.txt" in folder_contents)
+            self.assertTrue(".gitattributes" in folder_contents)
+
+            with open(os.path.join(storage_folder, "dummy_file.txt"), "r") as f:
+                contents = f.read()
+                self.assertEqual(contents, "v2")
+
+            # folder name contains the revision's commit sha.
+            self.assertTrue(self.second_commit_hash in storage_folder)
+
+        # Test with specific revision branch
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            # first download folder to cache it
+            snapshot_download(
+                f"{USER}/{REPO_NAME}",
+                revision="other",
+                cache_dir=tmpdirname,
+            )
+
+            # now load from cache
+            storage_folder = snapshot_download(
+                f"{USER}/{REPO_NAME}",
+                revision="other",
+                cache_dir=tmpdirname,
+                local_files_only=True,
+            )
+
+            # folder contains the two files contributed and the .gitattributes
+            folder_contents = os.listdir(storage_folder)
+            self.assertEqual(len(folder_contents), 3)
+            self.assertTrue("dummy_file.txt" in folder_contents)
+            self.assertTrue(".gitattributes" in folder_contents)
+
+            with open(os.path.join(storage_folder, "dummy_file.txt"), "r") as f:
+                contents = f.read()
+                self.assertEqual(contents, "v2")
+
+            # folder name contains the revision's commit sha.
+            self.assertTrue(self.third_commit_hash in storage_folder)
+
+        # Test with specific revision hash
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            # first download folder to cache it
+            snapshot_download(
+                f"{USER}/{REPO_NAME}",
+                revision=self.first_commit_hash,
+                cache_dir=tmpdirname,
+            )
+
+            # now load from cache
+            storage_folder = snapshot_download(
+                f"{USER}/{REPO_NAME}",
+                revision=self.first_commit_hash,
+                cache_dir=tmpdirname,
+                local_files_only=True,
+            )
+
+            # folder contains the two files contributed and the .gitattributes
+            folder_contents = os.listdir(storage_folder)
+            self.assertEqual(len(folder_contents), 2)
+            self.assertTrue("dummy_file.txt" in folder_contents)
+            self.assertTrue(".gitattributes" in folder_contents)
+
+            with open(os.path.join(storage_folder, "dummy_file.txt"), "r") as f:
+                contents = f.read()
+                self.assertEqual(contents, "v1")
+
+            # folder name contains the revision's commit sha.
+            self.assertTrue(self.first_commit_hash in storage_folder)
+
+    def test_download_model_local_only_multiple(self):
+        # Test `main` branch
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            # download both from branch and from commit
+            snapshot_download(
+                f"{USER}/{REPO_NAME}",
+                cache_dir=tmpdirname,
+            )
+
+            snapshot_download(
+                f"{USER}/{REPO_NAME}",
+                revision=self.first_commit_hash,
+                cache_dir=tmpdirname,
+            )
+
+            # now load from cache and make sure warning to be raised
+            with self.assertWarns(Warning):
+                snapshot_download(
+                    f"{USER}/{REPO_NAME}",
+                    cache_dir=tmpdirname,
+                    local_files_only=True,
+                )
+
+        # cache multiple commits and make sure correct commit is taken
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            # first download folder to cache it
+            snapshot_download(
+                f"{USER}/{REPO_NAME}",
+                cache_dir=tmpdirname,
+            )
+
+            # now load folder from another branch
+            snapshot_download(
+                f"{USER}/{REPO_NAME}",
+                revision="other",
+                cache_dir=tmpdirname,
+            )
+
+            # now make sure that loading "main" branch gives correct branch
+            storage_folder = snapshot_download(
+                f"{USER}/{REPO_NAME}",
+                cache_dir=tmpdirname,
+                local_files_only=True,
+            )
+
+            # folder contains the two files contributed and the .gitattributes
+            folder_contents = os.listdir(storage_folder)
+            self.assertEqual(len(folder_contents), 3)
+            self.assertTrue("dummy_file.txt" in folder_contents)
+            self.assertTrue(".gitattributes" in folder_contents)
+
+            with open(os.path.join(storage_folder, "dummy_file.txt"), "r") as f:
+                contents = f.read()
+                self.assertEqual(contents, "v2")
+
+            # folder name contains the 2nd commit sha and not the 3rd
+            self.assertTrue(self.second_commit_hash in storage_folder)