Spaces:
Runtime error
Runtime error
import csv | |
import os | |
from datetime import datetime | |
from typing import Optional | |
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" | |
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) -> str: | |
if token == "" or model_id == "": | |
return """ | |
### Invalid input π | |
Please fill a token and model name. | |
""" | |
try: | |
api = HfApi(token=token) | |
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 | |
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) | |
return f"#### Success π₯ Yay! This model was successfully converted and a PR was open using your token, here: [{commit_info.pr_url}]({commit_info.pr_url})" | |
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://huggingface.co/spaces/optimum/exporters/resolve/main/clean_hf_onnx.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;"> | |
Convert transformers model to ONNX with π€ Optimum exporters ποΈ (Beta) | |
</h1> | |
</div> | |
""" | |
# for some reason https://huggingface.co/settings/tokens is not showing as a link by default? | |
DESCRIPTION = """ | |
This Space allows to automatically convert to ONNX π€ transformers PyTorch models hosted on the Hugging Face Hub. It opens a PR on the target model, and it is up to the owner of the original model | |
to merge the PR to allow people to leverage the ONNX standard to share and use the model on a wide range of devices! | |
Once converted, the model can for example be used in the [π€ Optimum](https://huggingface.co/docs/optimum/) library following closely the transormers API. | |
Check out [this guide](https://huggingface.co/docs/optimum/main/en/onnxruntime/usage_guides/models) to see how! | |
The steps are the following: | |
- Paste a read-access token from [https://huggingface.co/settings/tokens](https://huggingface.co/settings/tokens). Read access is enough given that we will open a PR against the source repo. | |
- Input a model id from the Hub (for example: [textattack/distilbert-base-cased-CoLA](https://huggingface.co/textattack/distilbert-base-cased-CoLA)) | |
- Click "Convert to ONNX" | |
- That's it! You'll get feedback if it works or not, and if it worked, you'll get the URL of the opened PR! | |
Note: in case the model to convert is larger than 2 GB, it will be saved in a subfolder called `onnx/`. To load it from Optimum, the argument `subfolder="onnx"` should be provided. | |
""" | |
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", | |
visible=False, | |
) | |
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 | |
) | |
""" | |
demo = gr.Interface( | |
title="", | |
description=DESCRIPTION, | |
allow_flagging="never", | |
article="Check out the [π€ Optimum repoository on GitHub](https://github.com/huggingface/optimum) as well!", | |
inputs=[ | |
gr.Text(max_lines=1, label="Hugging Face token"), | |
gr.Text(max_lines=1, label="Model name", placeholder="textattack/distilbert-base-cased-CoLA"), | |
gr.Text(value="auto", max_lines=1, label="Task (can be left blank, will be automatically inferred)") | |
], | |
outputs=[gr.Markdown(label="output")], | |
fn=onnx_export, | |
) | |
""" | |
demo.launch() | |