Create onnx_export.py

onnx_export.py

@@ -0,0 +1,108 @@
+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"""I am using the [ONNX Exporter tool 🤖🏎️]({SPACES_URL}) to convert this model to ONNX format. I, [{requesting_user}](https://huggingface.co/{requesting_user}), would like to add the converted ONNX model to this repository.
+                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,  # TODO
+                    create_pr=True,
+                )
+        finally:
+            shutil.rmtree(folder)
+        return "0", new_pr