Create app.py

app.py

@@ -0,0 +1,160 @@
+import csv
+import os
+from datetime import datetime
+from typing import Optional, Union
+
+import gradio as gr
+from huggingface_hub import HfApi, Repository
+
+from onnx_export import convert
+
+from apscheduler.schedulers.background import BackgroundScheduler
+
+DATASET_REPO_URL = "https://huggingface.co/datasets/optimum/exporters"
+DATA_FILENAME = "data.csv"
+DATA_FILE = os.path.join("data", DATA_FILENAME)
+
+HF_TOKEN = os.environ.get("HF_WRITE_TOKEN")
+
+DATADIR = "exporters_data"
+
+repo: Optional[Repository] = None
+# if HF_TOKEN:
+#     repo = Repository(local_dir=DATADIR, clone_from=DATASET_REPO_URL, token=HF_TOKEN)
+
+
+def onnx_export(token: str, model_id: str, task: str, opset: Union[int, str]) -> str:
+    if token == "" or model_id == "":
+        return """
+        ### Invalid input 🐞
+
+        Please fill a token and model name.
+        """
+    try:
+        if opset == "":
+            opset = None
+        else:
+            opset = int(opset)
+
+        api = HfApi(token=token)
+
+        error, commit_info = convert(api=api, model_id=model_id, task=task, opset=opset)
+        if error != "0":
+            return error
+
+        print("[commit_info]", commit_info)
+
+        # save in a private dataset
+        if repo is not None:
+            repo.git_pull(rebase=True)
+            with open(os.path.join(DATADIR, DATA_FILE), "a") as csvfile:
+                writer = csv.DictWriter(
+                    csvfile, fieldnames=["model_id", "pr_url", "time"]
+                )
+                writer.writerow(
+                    {
+                        "model_id": model_id,
+                        "pr_url": commit_info.pr_url,
+                        "time": str(datetime.now()),
+                    }
+                )
+            commit_url = repo.push_to_hub()
+            print("[dataset]", commit_url)
+
+        pr_revision = commit_info.pr_revision.replace("/", "%2F")
+
+        return f"#### This model was successfully exported and a PR was open using your token, here: [{commit_info.pr_url}]({commit_info.pr_url}). If you would like to use the exported model without waiting for the PR to be approved, head to https://huggingface.co/{model_id}/tree/{pr_revision}"
+    except Exception as e:
+        return f"#### Error: {e}"
+
+
+TTILE_IMAGE = """
+<div
+    style="
+        display: block;
+        margin-left: auto;
+        margin-right: auto;
+        width: 50%;
+    "
+>
+<img src="https://i.ibb.co/m5VnjSsQ/Blue-and-White-Illustrative-Profile-Twitter-Header.png"/>
+</div>
+"""
+
+TITLE = """
+<div
+    style="
+        display: inline-flex;
+        align-items: center;
+        text-align: center;
+        max-width: 1400px;
+        gap: 0.8rem;
+        font-size: 2.2rem;
+    "
+>
+<h1 style="font-weight: 900; margin-bottom: 10px; margin-top: 10px;">
+    Export transformers model to ONNX with HF Optimum exporters.
+</h1>
+</div>
+"""
+
+# for some reason https://huggingface.co/settings/tokens is not showing as a link by default?
+DESCRIPTION = """
+This Space enables automatic export of Hugging Face transformers PyTorch models to [ONNX](https://onnx.ai/). It creates a pull request on the target model repository, allowing model owners to review and merge the ONNX export, making their models accessible across a wide range of devices and platforms.
+
+Once exported, the model can be seamlessly integrated with [HF Optimum](https://huggingface.co/docs/optimum/), maintaining compatibility with the transformers API. For detailed implementation, check out [this comprehensive guide](https://huggingface.co/docs/optimum/main/en/onnxruntime/usage_guides/models).
+
+Quick Start Guide:
+1. Obtain a read-access token from [https://huggingface.co/settings/tokens](https://huggingface.co/settings/tokens) (read access is sufficient for PR creation)
+2. Enter a model ID from the Hub (e.g., [textattack/distilbert-base-cased-CoLA](https://huggingface.co/textattack/distilbert-base-cased-CoLA))
+3. Click "Export to ONNX"
+4. Done! You'll receive feedback on the export status and, if successful, the URL of the created pull request
+
+Important Note: For models exceeding 2 GB, the ONNX export will be saved in an `onnx/` subfolder. When loading such models with Optimum, remember to include the `subfolder="onnx"` parameter."""
+
+with gr.Blocks() as demo:
+    gr.HTML(TTILE_IMAGE)
+    gr.HTML(TITLE)
+
+    with gr.Row():
+        with gr.Column(scale=50):
+            gr.Markdown(DESCRIPTION)
+
+        with gr.Column(scale=50):
+            input_token = gr.Textbox(
+                max_lines=1,
+                label="Hugging Face token",
+            )
+            input_model = gr.Textbox(
+                max_lines=1,
+                label="Model name",
+                placeholder="textattack/distilbert-base-cased-CoLA",
+            )
+            input_task = gr.Textbox(
+                value="auto",
+                max_lines=1,
+                label='Task (can be left to "auto", will be automatically inferred)',
+            )
+            onnx_opset = gr.Textbox(
+                placeholder="for example 14, can be left blank",
+                max_lines=1,
+                label="ONNX opset (optional, can be left blank)",
+            )
+
+            btn = gr.Button("Export to ONNX")
+            output = gr.Markdown(label="Output")
+
+    btn.click(
+        fn=onnx_export,
+        inputs=[input_token, input_model, input_task, onnx_opset],
+        outputs=output,
+    )
+
+def restart_space():
+    HfApi().restart_space(repo_id="onnx/export", token=HF_TOKEN, factory_reboot=True)
+
+scheduler = BackgroundScheduler()
+scheduler.add_job(restart_space, "interval", seconds=21600)
+scheduler.start()
+
+demo.launch()