Spaces:
Runtime error
Runtime error
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 os | |
import shutil | |
import argparse | |
from typing import Optional, Tuple, List | |
from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download, get_repo_discussions | |
from huggingface_hub.file_download import repo_folder_name | |
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_onnx(model_id: str, task: str, folder: str) -> 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( | |
model_type, "onnx", task=task, model_name=model_name | |
) | |
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"] | |
) | |
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 | |
opset = onnx_config.DEFAULT_ONNX_OPSET | |
output = Path(folder).joinpath("model.onnx") | |
onnx_inputs, onnx_outputs = export( | |
model, | |
onnx_config, | |
opset, | |
output, | |
) | |
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([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)] | |
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)] | |
return operations | |
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) | |
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 "model.onnx" 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}") | |
else: | |
operations = convert_onnx(model_id, task, folder) | |
new_pr = api.create_commit( | |
repo_id=model_id, | |
operations=operations, | |
commit_message=pr_title, | |
create_pr=True, | |
) | |
finally: | |
shutil.rmtree(folder) | |
return "0", new_pr | |
if __name__ == "__main__": | |
DESCRIPTION = """ | |
Simple utility tool to convert automatically a model on the hub to onnx format. | |
It is PyTorch exclusive for now. | |
It works by downloading the weights (PT), converting them locally, and uploading them back | |
as a PR on the hub. | |
""" | |
parser = argparse.ArgumentParser(description=DESCRIPTION) | |
parser.add_argument( | |
"--model_id", | |
type=str, | |
help="The name of the model on the hub to convert. E.g. `gpt2` or `facebook/wav2vec2-base-960h`", | |
) | |
parser.add_argument( | |
"--task", | |
type=str, | |
help="The task the model is performing", | |
) | |
parser.add_argument( | |
"--force", | |
action="store_true", | |
help="Create the PR even if it already exists of if the model was already converted.", | |
) | |
args = parser.parse_args() | |
api = HfApi() | |
convert(api, args.model_id, task=args.task, force=args.force) |