| from __future__ import annotations | |
| import os | |
| import pathlib | |
| import shlex | |
| import subprocess | |
| import slugify | |
| from huggingface_hub import HfApi | |
| from constants import (MODEL_LIBRARY_ORG_NAME, URL_TO_JOIN_MODEL_LIBRARY_ORG, | |
| UploadTarget) | |
| def join_model_library_org(hf_token: str) -> None: | |
| subprocess.run( | |
| shlex.split( | |
| f'curl -X POST -H "Authorization: Bearer {hf_token}" -H "Content-Type: application/json" {URL_TO_JOIN_MODEL_LIBRARY_ORG}' | |
| )) | |
| def upload(local_folder_path: str, | |
| target_repo_name: str, | |
| upload_to: str, | |
| private: bool = True, | |
| delete_existing_repo: bool = False, | |
| hf_token: str = '') -> str: | |
| hf_token = os.getenv('HF_TOKEN') or hf_token | |
| if not hf_token: | |
| raise ValueError | |
| api = HfApi(token=hf_token) | |
| if not local_folder_path: | |
| raise ValueError | |
| if not target_repo_name: | |
| target_repo_name = pathlib.Path(local_folder_path).name | |
| target_repo_name = slugify.slugify(target_repo_name) | |
| if upload_to == UploadTarget.PERSONAL_PROFILE.value: | |
| organization = api.whoami()['name'] | |
| elif upload_to == UploadTarget.MODEL_LIBRARY.value: | |
| organization = MODEL_LIBRARY_ORG_NAME | |
| join_model_library_org(hf_token) | |
| else: | |
| raise ValueError | |
| repo_id = f'{organization}/{target_repo_name}' | |
| if delete_existing_repo: | |
| try: | |
| api.delete_repo(repo_id, repo_type='model') | |
| except Exception: | |
| pass | |
| try: | |
| api.create_repo(repo_id, repo_type='model', private=private) | |
| api.upload_folder(repo_id=repo_id, | |
| folder_path=local_folder_path, | |
| path_in_repo='.', | |
| repo_type='model') | |
| url = f'https://huggingface.co/{repo_id}' | |
| message = f'Your model was successfully uploaded to {url}.' | |
| except Exception as e: | |
| message = str(e) | |
| return message | |