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