Skip to content

Commit

Permalink
feat(models): add HF upload/download class
Browse files Browse the repository at this point in the history
  • Loading branch information
leafspark committed Aug 31, 2024
1 parent 5d6b6fb commit fb9addb
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 1 deletion.
4 changes: 3 additions & 1 deletion requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -8,4 +8,6 @@ pynvml~=11.5.3
PySide6~=6.7.2
flask~=3.0.3
python-dotenv~=1.0.1
safetensors~=0.4.4
safetensors~=0.4.4
setuptools~=68.2.0
huggingface-hub~=0.24.6
30 changes: 30 additions & 0 deletions src/HFTransfer.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
from huggingface_hub import HfApi, snapshot_download, hf_hub_url


class HFTransfer:
def __init__(self, token) -> None:
self.api = HfApi(token=token)

def upload_file(self, path, remote_path, repo, repo_type) -> str:
status = self.api.upload_file(
path_or_fileobj=path,
path_in_repo=remote_path,
repo_id=repo,
repo_type=repo_type,
)
return status

def upload_folder(self, local_folder, remote_path, repo, repo_type) -> str:
status = self.api.upload_folder(
folder_path=local_folder,
path_in_repo=remote_path,
repo_id=repo,
repo_type=repo_type,
)
return status

def download_repo(self, repo) -> str:
return snapshot_download(repo_id=repo)

def get_download_link(self, repo, file) -> str:
return hf_hub_url(repo_id=repo, filename=file)

0 comments on commit fb9addb

Please sign in to comment.