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