Add auth to snapshot download (#340)

* Add API token to snapshot download

* Fix style

* Replace token with use_auth_token

* Add truthiness check for token

* Assign string value for token

* Add case for None token

* Refactor unit tests

* Add test cases

* Style

* Add unit test for private model info

* Revert "Add unit test for private model info"

This reverts commit eca1e8b.

* Make repo public after test

Author: lewtun

src/huggingface_hub/snapshot_download.py

@@ -4,7 +4,7 @@
 
 from .constants import HUGGINGFACE_HUB_CACHE
 from .file_download import cached_download, hf_hub_url
-from .hf_api import HfApi
+from .hf_api import HfApi, HfFolder
 
 
 REPO_ID_SEPARATOR = "__"
@@ -18,6 +18,7 @@ def snapshot_download(
     library_name: Optional[str] = None,
     library_version: Optional[str] = None,
     user_agent: Union[Dict, str, None] = None,
+    use_auth_token: Union[bool, str, None] = None,
 ) -> str:
     """
     Downloads a whole snapshot of a repo's files at the specified revision.
@@ -41,8 +42,19 @@ def snapshot_download(
     if isinstance(cache_dir, Path):
         cache_dir = str(cache_dir)
 
+    if isinstance(use_auth_token, str):
+        token = use_auth_token
+    elif use_auth_token:
+        token = HfFolder.get_token()
+        if token is None:
+            raise EnvironmentError(
+                "You specified use_auth_token=True, but a Hugging Face token was not found."
+            )
+    else:
+        token = None
+
     _api = HfApi()
-    model_info = _api.model_info(repo_id=repo_id, revision=revision)
+    model_info = _api.model_info(repo_id=repo_id, revision=revision, token=token)
 
     storage_folder = os.path.join(
         cache_dir, repo_id.replace("/", REPO_ID_SEPARATOR) + "." + model_info.sha
@@ -67,6 +79,7 @@ def snapshot_download(
             library_name=library_name,
             library_version=library_version,
             user_agent=user_agent,
+            use_auth_token=use_auth_token,
         )
 
         if os.path.exists(path + ".lock"):

tests/test_snapshot_download.py

@@ -3,7 +3,9 @@
 import time
 import unittest
 
+import requests
 from huggingface_hub import HfApi, Repository
+from huggingface_hub.hf_api import HfFolder
 from huggingface_hub.snapshot_download import snapshot_download
 from tests.testing_constants import ENDPOINT_STAGING, PASS, USER
 
@@ -88,3 +90,64 @@ def test_download_model(self):
 
             # folder name contains the revision's commit sha.
             self.assertTrue(self.first_commit_hash in storage_folder)
+
+    def test_download_private_model(self):
+        self._api.update_repo_visibility(self._token, REPO_NAME, private=True)
+
+        # Test download fails without token
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            with self.assertRaisesRegex(
+                requests.exceptions.HTTPError, "404 Client Error"
+            ):
+                _ = snapshot_download(
+                    f"{USER}/{REPO_NAME}", revision="main", cache_dir=tmpdirname
+                )
+
+        # Test we can download with token from cache
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            HfFolder.save_token(self._token)
+            storage_folder = snapshot_download(
+                f"{USER}/{REPO_NAME}",
+                revision="main",
+                cache_dir=tmpdirname,
+                use_auth_token=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 we can download with explicit token
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            storage_folder = snapshot_download(
+                f"{USER}/{REPO_NAME}",
+                revision="main",
+                cache_dir=tmpdirname,
+                use_auth_token=self._token,
+            )
+
+            # 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)
+
+        self._api.update_repo_visibility(self._token, REPO_NAME, private=False)