style

app.py

@@ -1,12 +1,12 @@
 import csv
-import datetime
 import os
+from datetime import datetime
 from typing import Optional
-import gradio as gr
 
-from onnx_export import convert
+import gradio as gr
 from huggingface_hub import HfApi, Repository
 
+from onnx_export import convert
 
 DATASET_REPO_URL = "https://huggingface.co/datasets/optimum/exporters"
 DATA_FILENAME = "data.csv"
@@ -20,6 +20,7 @@ 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) -> str:
     if token == "" or model_id == "":
         return """
@@ -33,7 +34,7 @@ def onnx_export(token: str, model_id: str, task: str) -> str:
         error, commit_info = convert(api=api, model_id=model_id, task=task)
         if error != "0":
             return error
-        
+
         print("[commit_info]", commit_info)
 
         # save in a private dataset
@@ -57,6 +58,7 @@ def onnx_export(token: str, model_id: str, task: str) -> str:
     except Exception as e:
         return f"#### Error: {e}"
 
+
 TTILE_IMAGE = """
 <div
     style="
@@ -111,14 +113,23 @@ with gr.Blocks() as demo:
 
     with gr.Column():
         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)")
+        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)',
+        )
 
         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)
+
+    btn.click(
+        fn=onnx_export, inputs=[input_token, input_model, input_task], outputs=output
+    )
 
 """
 demo = gr.Interface(
@@ -136,4 +147,4 @@ demo = gr.Interface(
 )
 """
 
-demo.launch()
+demo.launch()

onnx_export.py

@@ -1,33 +1,35 @@
-from optimum.exporters.tasks import TasksManager
-
-from optimum.exporters.onnx import OnnxConfigWithPast, export, validate_model_outputs
-
-from tempfile import TemporaryDirectory
-
-from transformers import AutoConfig, AutoTokenizer, is_torch_available
-
-from pathlib import Path
-
+import argparse
 import os
 import shutil
-import argparse
-
-from typing import Optional, Tuple, List
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import List, Optional, Tuple
 
-from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download, get_repo_discussions
+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.tasks import TasksManager
+from transformers import AutoConfig, AutoTokenizer, is_torch_available
 
 SPACES_URL = "https://huggingface.co/spaces/optimum/exporters"
 
+
 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:
+        if (
+            discussion.status == "open"
+            and discussion.is_pull_request
+            and discussion.title == pr_title
+        ):
             return discussion
 
+
 def convert_onnx(model_id: str, task: str, folder: str) -> List:
 
     # Allocate the model
@@ -46,7 +48,7 @@ def convert_onnx(model_id: str, task: str, folder: str) -> List:
         and task in ["sequence_classification"]
     )
     if needs_pad_token_id:
-        #if args.pad_token_id is not None:
+        # if args.pad_token_id is not None:
         #    model.config.pad_token_id = args.pad_token_id
         try:
             tok = AutoTokenizer.from_pretrained(model_id)
@@ -76,18 +78,37 @@ def convert_onnx(model_id: str, task: str, folder: str) -> List:
         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([name for name in os.listdir(folder) if os.path.isfile(os.path.join(folder, name)) and not name.startswith(".")])
-    
+
+    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)]
+        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)]
-    
+        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) -> Tuple[int, "CommitInfo"]:
+def convert(
+    api: "HfApi", model_id: str, task: str, force: bool = False
+) -> 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)
@@ -98,7 +119,10 @@ def convert(api: "HfApi", model_id: str, task: str, force: bool = False) -> Tupl
         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
+            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"))
@@ -111,7 +135,9 @@ def convert(api: "HfApi", model_id: str, task: str, force: bool = False) -> Tupl
             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}")
+                raise Exception(
+                    f"Model {model_id} already has an open PR check out {url}"
+                )
             else:
                 operations = convert_onnx(model_id, task, folder)
 
@@ -159,4 +185,4 @@ if __name__ == "__main__":
     )
     args = parser.parse_args()
     api = HfApi()
-    convert(api, args.model_id, task=args.task, force=args.force)
+    convert(api, args.model_id, task=args.task, force=args.force)