update

app.py

@@ -59,7 +59,9 @@ def onnx_export(token: str, model_id: str, task: str, opset: Union[int, str]) ->
             commit_url = repo.push_to_hub()
             print("[dataset]", commit_url)
 
-        return f"#### Success 🔥 Yay! This model was successfully exported and a PR was open using your token, here: [{commit_info.pr_url}]({commit_info.pr_url})"
+        pr_revision = commit_info.pr_revision.replace("/", "%2F")
+
+        return f"#### Success 🔥 Yay! 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 "
     except Exception as e:
         return f"#### Error: {e}"
 
@@ -89,7 +91,7 @@ TITLE = """
     "
 >
 <h1 style="font-weight: 900; margin-bottom: 10px; margin-top: 10px;">
-    Export transformers model to ONNX with 🤗 Optimum exporters 🏎️ (Beta)
+    Export transformers model to ONNX with 🤗 Optimum exporters 🏎️
 </h1>
 </div>
 """

onnx_export.py

@@ -1,20 +1,16 @@
 import argparse
 import os
 import shutil
-from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import List, Optional, Tuple
 
 from huggingface_hub import (
     CommitOperationAdd,
     HfApi,
-    get_repo_discussions,
-    hf_hub_download,
 )
 from huggingface_hub.file_download import repo_folder_name
-from optimum.exporters.onnx import OnnxConfigWithPast, export, validate_model_outputs
+from optimum.exporters.onnx import main_export
 from optimum.exporters.tasks import TasksManager
-from transformers import AutoConfig, AutoTokenizer, is_torch_available
 
 SPACES_URL = "https://huggingface.co/spaces/optimum/exporters"
 
@@ -33,54 +29,13 @@ def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discuss
             return discussion
 
 
-def convert_onnx(model_id: str, task: str, folder: str, opset: int) -> List:
-
-    # Allocate the model
-    model = TasksManager.get_model_from_task(task, model_id, framework="pt")
-    model_type = model.config.model_type.replace("_", "-")
-    model_name = getattr(model, "name", None)
-
-    onnx_config_constructor = TasksManager.get_exporter_config_constructor(
-        exporter="onnx",
-        model=model,
+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,
-        model_name=model_name,
-        model_type=model_type,
-    )
-    onnx_config = onnx_config_constructor(model.config)
-
-    needs_pad_token_id = (
-        isinstance(onnx_config, OnnxConfigWithPast)
-        and getattr(model.config, "pad_token_id", None) is None
-        and task in ["sequence_classification"]
+        opset=opset,
     )
-    if needs_pad_token_id:
-        # if args.pad_token_id is not None:
-        #    model.config.pad_token_id = args.pad_token_id
-        try:
-            tok = AutoTokenizer.from_pretrained(model_id)
-            model.config.pad_token_id = tok.pad_token_id
-        except Exception:
-            raise ValueError(
-                "Could not infer the pad token id, which is needed in this case, please provide it with the --pad_token_id argument"
-            )
-
-    # Ensure the requested opset is sufficient
-    if opset == None:
-        opset = onnx_config.DEFAULT_ONNX_OPSET
-
-    output = Path(folder).joinpath("model.onnx")
-    onnx_inputs, onnx_outputs = export(model, onnx_config, output, opset)
-
-    atol = onnx_config.ATOL_FOR_VALIDATION
-    if isinstance(atol, dict):
-        atol = atol[task.replace("-with-past", "")]
-
-    try:
-        validate_model_outputs(onnx_config, model, output, onnx_outputs, atol)
-        print(f"All good, model saved at: {output}")
-    except ValueError:
-        print(f"An error occured, but the model was saved at: {output.as_posix()}")
 
     n_files = len(
         [
@@ -122,15 +77,6 @@ def convert(
 
     requesting_user = api.whoami()["name"]
 
-    if task == "auto":
-        try:
-            task = TasksManager.infer_task_from_model(model_id)
-        except Exception as e:
-            return (
-                f"### Error: {e}. Please pass explicitely the task as it could not be infered.",
-                None,
-            )
-
     with TemporaryDirectory() as d:
         folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models"))
         os.makedirs(folder)
@@ -146,7 +92,7 @@ def convert(
                     f"Model {model_id} already has an open PR check out [{url}]({url})"
                 )
             else:
-                operations = convert_onnx(model_id, task, folder, opset)
+                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.