uploader.py

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