Skip to content

Commit

Permalink
Add default user-agent to huggingface-cli (#1664)
Browse files Browse the repository at this point in the history
  • Loading branch information
Wauplin committed Sep 18, 2023
1 parent 2d20084 commit fc1eeaa
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 22 deletions.
16 changes: 7 additions & 9 deletions src/huggingface_hub/commands/upload.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@
from huggingface_hub import logging
from huggingface_hub._commit_scheduler import CommitScheduler
from huggingface_hub.commands import BaseHuggingfaceCLICommand
from huggingface_hub.hf_api import create_repo, upload_file, upload_folder
from huggingface_hub.hf_api import HfApi
from huggingface_hub.utils import disable_progress_bars, enable_progress_bars


Expand Down Expand Up @@ -134,7 +134,7 @@ def __init__(self, args: Namespace) -> None:
self.commit_message: Optional[str] = args.commit_message
self.commit_description: Optional[str] = args.commit_description
self.create_pr: bool = args.create_pr
self.token: Optional[str] = args.token
self.api: HfApi = HfApi(token=args.token, library_name="huggingface-cli")
self.quiet: bool = args.quiet # disable warnings and progress bars

# Check `--every` is valid
Expand Down Expand Up @@ -222,7 +222,7 @@ def _upload(self) -> str:
path_in_repo=path_in_repo,
private=self.private,
every=self.every,
token=self.token,
hf_api=self.api,
)
print(f"Scheduling commits every {self.every} minutes to {scheduler.repo_id}.")
try: # Block main thread until KeyboardInterrupt
Expand All @@ -235,33 +235,31 @@ def _upload(self) -> str:
# Otherwise, create repo and proceed with the upload
if not os.path.isfile(self.local_path) and not os.path.isdir(self.local_path):
raise FileNotFoundError(f"No such file or directory: '{self.local_path}'.")
repo_id = create_repo(
repo_id=self.repo_id, repo_type=self.repo_type, exist_ok=True, private=self.private, token=self.token
repo_id = self.api.create_repo(
repo_id=self.repo_id, repo_type=self.repo_type, exist_ok=True, private=self.private
).repo_id

# File-based upload
if os.path.isfile(self.local_path):
return upload_file(
return self.api.upload_file(
path_or_fileobj=self.local_path,
path_in_repo=self.path_in_repo,
repo_id=repo_id,
repo_type=self.repo_type,
revision=self.revision,
token=self.token,
commit_message=self.commit_message,
commit_description=self.commit_description,
create_pr=self.create_pr,
)

# Folder-based upload
else:
return upload_folder(
return self.api.upload_folder(
folder_path=self.local_path,
path_in_repo=self.path_in_repo,
repo_id=repo_id,
repo_type=self.repo_type,
revision=self.revision,
token=self.token,
commit_message=self.commit_message,
commit_description=self.commit_description,
create_pr=self.create_pr,
Expand Down
22 changes: 9 additions & 13 deletions tests/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ def test_upload_basic(self) -> None:
self.assertEqual(cmd.commit_description, None)
self.assertEqual(cmd.create_pr, False)
self.assertEqual(cmd.every, None)
self.assertEqual(cmd.token, None)
self.assertEqual(cmd.api.token, None)
self.assertEqual(cmd.quiet, False)

def test_upload_with_all_options(self) -> None:
Expand Down Expand Up @@ -135,7 +135,7 @@ def test_upload_with_all_options(self) -> None:
self.assertEqual(cmd.commit_description, "My commit description")
self.assertEqual(cmd.create_pr, True)
self.assertEqual(cmd.every, 5)
self.assertEqual(cmd.token, "my-token")
self.assertEqual(cmd.api.token, "my-token")
self.assertEqual(cmd.quiet, True)

def test_upload_implicit_local_path_when_folder_exists(self) -> None:
Expand Down Expand Up @@ -211,8 +211,8 @@ def test_every_as_float(self) -> None:
cmd = UploadCommand(self.parser.parse_args(["upload", DUMMY_MODEL_ID, ".", "--every", "0.5"]))
self.assertEqual(cmd.every, 0.5)

@patch("huggingface_hub.commands.upload.upload_folder")
@patch("huggingface_hub.commands.upload.create_repo")
@patch("huggingface_hub.commands.upload.HfApi.upload_folder")
@patch("huggingface_hub.commands.upload.HfApi.create_repo")
def test_upload_folder_mock(self, create_mock: Mock, upload_mock: Mock) -> None:
with SoftTemporaryDirectory() as cache_dir:
cmd = UploadCommand(
Expand All @@ -222,16 +222,13 @@ def test_upload_folder_mock(self, create_mock: Mock, upload_mock: Mock) -> None:
)
cmd.run()

create_mock.assert_called_once_with(
repo_id="my-model", repo_type="model", exist_ok=True, private=True, token=None
)
create_mock.assert_called_once_with(repo_id="my-model", repo_type="model", exist_ok=True, private=True)
upload_mock.assert_called_once_with(
folder_path=cache_dir,
path_in_repo=".",
repo_id=create_mock.return_value.repo_id,
repo_type="model",
revision=None,
token=None,
commit_message=None,
commit_description=None,
create_pr=False,
Expand All @@ -240,8 +237,8 @@ def test_upload_folder_mock(self, create_mock: Mock, upload_mock: Mock) -> None:
delete_patterns=["*.json"],
)

@patch("huggingface_hub.commands.upload.upload_file")
@patch("huggingface_hub.commands.upload.create_repo")
@patch("huggingface_hub.commands.upload.HfApi.upload_file")
@patch("huggingface_hub.commands.upload.HfApi.create_repo")
def test_upload_file_mock(self, create_mock: Mock, upload_mock: Mock) -> None:
with SoftTemporaryDirectory() as cache_dir:
file_path = Path(cache_dir) / "file.txt"
Expand All @@ -254,21 +251,20 @@ def test_upload_file_mock(self, create_mock: Mock, upload_mock: Mock) -> None:
cmd.run()

create_mock.assert_called_once_with(
repo_id="my-dataset", repo_type="dataset", exist_ok=True, private=False, token=None
repo_id="my-dataset", repo_type="dataset", exist_ok=True, private=False
)
upload_mock.assert_called_once_with(
path_or_fileobj=str(file_path),
path_in_repo="logs/file.txt",
repo_id=create_mock.return_value.repo_id,
repo_type="dataset",
revision=None,
token=None,
commit_message=None,
commit_description=None,
create_pr=True,
)

@patch("huggingface_hub.commands.upload.create_repo")
@patch("huggingface_hub.commands.upload.HfApi.create_repo")
def test_upload_missing_path(self, create_mock: Mock) -> None:
cmd = UploadCommand(self.parser.parse_args(["upload", "my-model", "/path/to/missing_file", "logs/file.txt"]))
with self.assertRaises(FileNotFoundError):
Expand Down

0 comments on commit fc1eeaa

Please sign in to comment.