|
import os |
|
import shutil |
|
from tempfile import TemporaryDirectory |
|
from typing import List, Optional, Tuple |
|
|
|
from huggingface_hub import ( |
|
CommitOperationAdd, |
|
HfApi, |
|
) |
|
from huggingface_hub.file_download import repo_folder_name |
|
from optimum.exporters.onnx import main_export |
|
|
|
SPACES_URL = "https://huggingface.co/spaces/onnx/export" |
|
|
|
|
|
def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discussion"]: |
|
try: |
|
discussions = api.get_repo_discussions(repo_id=model_id) |
|
except Exception: |
|
return None |
|
for discussion in discussions: |
|
if ( |
|
discussion.status == "open" |
|
and discussion.is_pull_request |
|
and discussion.title == pr_title |
|
): |
|
return discussion |
|
|
|
|
|
def export_and_git_add(model_id: str, task: str, folder: str, opset: int) -> List: |
|
main_export( |
|
model_name_or_path=model_id, |
|
output=folder, |
|
task=task, |
|
opset=opset, |
|
) |
|
|
|
n_files = len( |
|
[ |
|
name |
|
for name in os.listdir(folder) |
|
if os.path.isfile(os.path.join(folder, name)) and not name.startswith(".") |
|
] |
|
) |
|
|
|
if n_files == 1: |
|
operations = [ |
|
CommitOperationAdd( |
|
path_in_repo=file_name, path_or_fileobj=os.path.join(folder, file_name) |
|
) |
|
for file_name in os.listdir(folder) |
|
] |
|
else: |
|
operations = [ |
|
CommitOperationAdd( |
|
path_in_repo=os.path.join("onnx", file_name), |
|
path_or_fileobj=os.path.join(folder, file_name), |
|
) |
|
for file_name in os.listdir(folder) |
|
] |
|
|
|
return operations |
|
|
|
|
|
def convert( |
|
api: "HfApi", |
|
model_id: str, |
|
task: str, |
|
force: bool = False, |
|
opset: int = None, |
|
) -> Tuple[int, "CommitInfo"]: |
|
pr_title = "Adding ONNX file of this model" |
|
info = api.model_info(model_id) |
|
filenames = set(s.rfilename for s in info.siblings) |
|
|
|
requesting_user = api.whoami()["name"] |
|
|
|
with TemporaryDirectory() as d: |
|
folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models")) |
|
os.makedirs(folder) |
|
new_pr = None |
|
try: |
|
pr = previous_pr(api, model_id, pr_title) |
|
if "model.onnx" in filenames and not force: |
|
raise Exception(f"Model {model_id} is already converted, skipping the export.") |
|
elif pr is not None and not force: |
|
url = f"https://huggingface.co/{model_id}/discussions/{pr.num}" |
|
new_pr = pr |
|
raise Exception( |
|
f"Model {model_id} already has an open PR check out [{url}]({url})" |
|
) |
|
else: |
|
operations = export_and_git_add(model_id, task, folder, opset) |
|
|
|
commit_description = f""" |
|
Beep boop I am the [ONNX export bot 🤖🏎️]({SPACES_URL}). On behalf of [{requesting_user}](https://huggingface.co/{requesting_user}), I would like to add to this repository the model converted to ONNX. |
|
|
|
What is ONNX? It stands for "Open Neural Network Exchange", and is the most commonly used open standard for machine learning interoperability. You can find out more at [onnx.ai](https://onnx.ai/)! |
|
|
|
The exported ONNX model can be then be consumed by various backends as TensorRT or TVM, or simply be used in a few lines with 🤗 Optimum through ONNX Runtime, check out how [here](https://huggingface.co/docs/optimum/main/en/onnxruntime/usage_guides/models)! |
|
""" |
|
new_pr = api.create_commit( |
|
repo_id=model_id, |
|
operations=operations, |
|
commit_message=pr_title, |
|
commit_description=commit_description, |
|
create_pr=True, |
|
) |
|
finally: |
|
shutil.rmtree(folder) |
|
return "0", new_pr |
|
|
