Add openvino export

app.py

@@ -0,0 +1,132 @@
+import csv
+import os
+from datetime import datetime
+from typing import Optional, Union
+import gradio as gr
+from huggingface_hub import HfApi, Repository
+from export import convert
+
+
+DATASET_REPO_URL = "https://huggingface.co/datasets/optimum/exporters"
+DATA_FILENAME = "data.csv"
+DATA_FILE = os.path.join("openvino", DATA_FILENAME)
+HF_TOKEN = os.environ.get("HF_WRITE_TOKEN")
+DATA_DIR = "exporters_data"
+
+repo = None
+if HF_TOKEN:
+    repo = Repository(local_dir=DATA_DIR, clone_from=DATASET_REPO_URL, token=HF_TOKEN)
+
+
+def 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, force=False)
+        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(DATA_DIR, 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 exported 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/echarlaix/openvino-export/resolve/main/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 your Transformers and Diffusers model to OpenVINO with 🤗 Optimum Intel (experimental)
+</h1>
+</div>
+"""
+
+DESCRIPTION = """
+This Space allows you to automatically export to the OpenVINO format various 🤗 Transformers and Diffusers PyTorch models hosted on the Hugging Face Hub.
+
+Once exported, you will be able to load the resulting model using the [🤗 Optimum Intel](https://huggingface.co/docs/optimum/intel/inference).
+
+To export your model, the steps are as 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: [distilbert-base-uncased-finetuned-sst-2-english](https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english))
+- Click "Export"
+- 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 🔥
+"""
+
+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="distilbert-base-uncased-finetuned-sst-2-english",
+            )
+            input_task = gr.Textbox(
+                value="auto",
+                max_lines=1,
+                label='Task (can be left to "auto", will be automatically inferred)',
+            )
+
+            btn = gr.Button("Export")
+            output = gr.Markdown(label="Output")
+
+    btn.click(
+        fn=export,
+        inputs=[input_token, input_model, input_task],
+        outputs=output,
+    )
+
+
+demo.launch()

export.py

@@ -0,0 +1,163 @@
+import argparse
+import os
+import shutil
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import List, Optional, Tuple
+import torch
+
+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 validate_model_outputs
+from optimum.exporters.tasks import TasksManager
+from transformers import AutoConfig, AutoTokenizer, is_torch_available
+from optimum.intel.openvino import (
+    _HEAD_TO_AUTOMODELS,
+    _TASK_ALIASES,
+    OVModelForAudioClassification,
+    OVModelForCausalLM,
+    OVModelForFeatureExtraction,
+    OVModelForImageClassification,
+    OVModelForMaskedLM,
+    OVModelForQuestionAnswering,
+    OVModelForSeq2SeqLM,
+    OVModelForSequenceClassification,
+    OVModelForTokenClassification,
+    OVStableDiffusionPipeline,
+)
+
+
+SPACES_URL = "https://huggingface.co/spaces/echarlaix/openvino-export"
+
+
+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
+        ):
+            return discussion
+
+
+def convert_openvino(model_id: str, task: str, folder: str) -> List:
+    task = _TASK_ALIASES.get(task, task)
+    if task not in _HEAD_TO_AUTOMODELS:
+        raise ValueError(f"The task '{task}' is not supported, only {_HEAD_TO_AUTOMODELS.keys()} tasks are supported")
+    auto_model_class = eval(_HEAD_TO_AUTOMODELS[task])
+    ov_model = auto_model_class.from_pretrained(model_id, export=True)
+    ov_model.save_pretrained(folder)
+    if not isinstance(ov_model, OVStableDiffusionPipeline):
+        try:
+            model = TasksManager.get_model_from_task(task, model_id)
+            onnx_config_class = TasksManager.get_exporter_config_constructor(
+                exporter="onnx",
+                model=model,
+                task=task,
+                model_name=model_id,
+                model_type=model.config.model_type.replace("_", "-"),
+            )
+            onnx_config = onnx_config_class(model.config)
+            inputs = onnx_config.generate_dummy_inputs(framework="pt")
+
+            if isinstance(ov_model, (OVModelForCausalLM, OVModelForSeq2SeqLM)):
+                ov_outputs = ov_model.generate(**inputs)
+                outputs = model.generate(**inputs)
+            else:
+                ov_outputs = ov_model(**inputs)
+                outputs = model(**inputs)
+
+            if isinstance(outputs, torch.Tensor):
+                outputs = {"logits": outputs}
+                ov_outputs = {"logits": ov_outputs}
+            for output_name in outputs:
+                if not torch.allclose(outputs[output_name], ov_outputs[output_name], atol=1e-3):
+                    raise ValueError(
+                        "The exported model does not have the same outputs as the original model. Export interrupted."
+                    )
+        except Exception as e:
+            raise
+
+    file_names = {elem for elem in os.listdir(folder) if os.path.isfile(os.path.join(folder, elem))}
+
+    operations = [
+        CommitOperationAdd(
+            path_in_repo=file_name, path_or_fileobj=os.path.join(folder, file_name)
+        )
+        for file_name in file_names if "openvino" in file_name
+    ]
+
+    dir_names = set(os.listdir(folder)) - file_names
+
+    for dir_name in dir_names.intersection({"vae_encoder", "vae_decoder", "text_encoder", "unet"}):
+        operations += [
+            CommitOperationAdd(
+                path_in_repo=os.path.join(dir_name, file_name),
+                path_or_fileobj=os.path.join(folder, dir_name, file_name),
+            )
+            for file_name in os.listdir(os.path.join(folder, dir_name)) if "openvino" in file_name
+        ]
+
+    return operations
+
+
+def convert(
+    api: "HfApi",
+    model_id: str,
+    task: str,
+    force: bool = False,
+) -> Tuple[int, "CommitInfo"]:
+    pr_title = "Adding OpenVINO file of this model"
+    info = api.model_info(model_id)
+    filenames = set(s.rfilename for s in info.siblings)
+
+    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)
+        new_pr = None
+        try:
+            pr = previous_pr(api, model_id, pr_title)
+            if "openvino_model.xml" in filenames and not force:
+                raise Exception(f"Model {model_id} is already converted, skipping..")
+            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}]({url})"
+                )
+            else:
+                operations = convert_openvino(model_id, task, folder)
+
+                commit_description = f"""
+                Beep boop I am the [OpenVINO exporter bot 🤖]({SPACES_URL}). On behalf of [{requesting_user}](https://huggingface.co/{requesting_user}), I would like to add to this repository the exported OpenVINO model.
+                """
+                new_pr = api.create_commit(
+                    repo_id=model_id,
+                    operations=operations,
+                    commit_message=pr_title,
+                    commit_description=commit_description,
+                    create_pr=True,
+                )
+        finally:
+            shutil.rmtree(folder)
+        return "0", new_pr

requirements.txt

@@ -1,3 +1,4 @@
 transformers
 diffusers
-git+https://github.com/huggingface/optimum-intel.git#egg=optimum[openvino]
+optimum==1.7.1
+git+https://github.com/huggingface/optimum-intel.git#egg=optimum-intel[openvino]