add opset

app.py

@@ -21,7 +21,7 @@ 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) -> str:
+def onnx_export(token: str, model_id: str, task: str, opset: int) -> str:
     if token == "" or model_id == "":
         return """
         ### Invalid input 🐞
@@ -31,7 +31,7 @@ def onnx_export(token: str, model_id: str, task: str) -> str:
     try:
         api = HfApi(token=token)
 
-        error, commit_info = convert(api=api, model_id=model_id, task=task)
+        error, commit_info = convert(api=api, model_id=model_id, task=task, opset=opset)
         if error != "0":
             return error
 
@@ -129,12 +129,19 @@ with gr.Blocks() as demo:
                 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("Convert to ONNX")
             output = gr.Markdown(label="Output")
 
     btn.click(
-        fn=onnx_export, inputs=[input_token, input_model, input_task], outputs=output
+        fn=onnx_export,
+        inputs=[input_token, input_model, input_task, onnx_opset],
+        outputs=output,
     )
 
 """

onnx_export.py

@@ -5,11 +5,14 @@ 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 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 OnnxConfigWithPast, export, validate_model_outputs
 from optimum.exporters.tasks import TasksManager
 from transformers import AutoConfig, AutoTokenizer, is_torch_available
 
@@ -30,7 +33,7 @@ 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) -> List:
+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")
@@ -58,8 +61,10 @@ def convert_onnx(model_id: str, task: str, folder: str) -> List:
                 "Could not infer the pad token id, which is needed in this case, please provide it with the --pad_token_id argument"
             )
 
+    print("opset:", opset)
     # Ensure the requested opset is sufficient
-    opset = onnx_config.DEFAULT_ONNX_OPSET
+    if opset == None:
+        opset = onnx_config.DEFAULT_ONNX_OPSET
 
     output = Path(folder).joinpath("model.onnx")
     onnx_inputs, onnx_outputs = export(
@@ -107,7 +112,11 @@ def convert_onnx(model_id: str, task: str, folder: str) -> List:
 
 
 def convert(
-    api: "HfApi", model_id: str, task: str, force: bool = False
+    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)
@@ -139,7 +148,7 @@ def convert(
                     f"Model {model_id} already has an open PR check out [{url}]({url})"
                 )
             else:
-                operations = convert_onnx(model_id, task, folder)
+                operations = convert_onnx(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.