From 25fc7578c45aa4227fb4fb102ee3cd8bc4a5d082 Mon Sep 17 00:00:00 2001 From: Ella Charlaix <80481427+echarlaix@users.noreply.github.com> Date: Thu, 28 Sep 2023 10:57:09 +0200 Subject: [PATCH] Add OpenVINO export CLI (#437) * add openvino export CLI * fix * add test * format --- optimum/commands/export/openvino.py | 106 ++++++++++++++++++ .../commands/register/register_openvino.py | 19 ++++ tests/openvino/test_exporters_cli.py | 59 ++++++++++ 3 files changed, 184 insertions(+) create mode 100644 optimum/commands/export/openvino.py create mode 100644 optimum/commands/register/register_openvino.py create mode 100644 tests/openvino/test_exporters_cli.py diff --git a/optimum/commands/export/openvino.py b/optimum/commands/export/openvino.py new file mode 100644 index 0000000000..b2d33e7647 --- /dev/null +++ b/optimum/commands/export/openvino.py @@ -0,0 +1,106 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Defines the command line for the export with OpenVINO.""" + +import sys +from pathlib import Path +from typing import TYPE_CHECKING, Optional + +from ...exporters import TasksManager +from ..base import BaseOptimumCLICommand, CommandInfo + + +if TYPE_CHECKING: + from argparse import ArgumentParser, Namespace, _SubParsersAction + + +def parse_args_openvino(parser: "ArgumentParser"): + required_group = parser.add_argument_group("Required arguments") + required_group.add_argument( + "-m", "--model", type=str, required=True, help="Model ID on huggingface.co or path on disk to load model from." + ) + required_group.add_argument( + "output", type=Path, help="Path indicating the directory where to store the generated OV model." + ) + optional_group = parser.add_argument_group("Optional arguments") + optional_group.add_argument( + "--task", + default="auto", + help=( + "The task to export the model for. If not specified, the task will be auto-inferred based on the model. Available tasks depend on the model, but are among:" + f" {str(TasksManager.get_all_tasks())}. For decoder models, use `xxx-with-past` to export the model using past key values in the decoder." + ), + ) + optional_group.add_argument("--cache_dir", type=str, default=None, help="Path indicating where to store cache.") + optional_group.add_argument( + "--framework", + type=str, + choices=["pt", "tf"], + default=None, + help=( + "The framework to use for the export. If not provided, will attempt to use the local checkpoint's original framework or what is available in the environment." + ), + ) + optional_group.add_argument( + "--trust-remote-code", + action="store_true", + help=( + "Allows to use custom code for the modeling hosted in the model repository. This option should only be set for repositories you trust and in which " + "you have read the code, as it will execute on your local machine arbitrary code present in the model repository." + ), + ) + optional_group.add_argument( + "--pad-token-id", + type=int, + default=None, + help=( + "This is needed by some models, for some tasks. If not provided, will attempt to use the tokenizer to guess it." + ), + ) + + +class OVExportCommand(BaseOptimumCLICommand): + COMMAND = CommandInfo(name="openvino", help="Export PyTorch models to OpenVINO IR.") + + def __init__( + self, + subparsers: "_SubParsersAction", + args: Optional["Namespace"] = None, + command: Optional["CommandInfo"] = None, + from_defaults_factory: bool = False, + parser: Optional["ArgumentParser"] = None, + ): + super().__init__( + subparsers, args=args, command=command, from_defaults_factory=from_defaults_factory, parser=parser + ) + self.args_string = " ".join(sys.argv[3:]) + + @staticmethod + def parse_args(parser: "ArgumentParser"): + return parse_args_openvino(parser) + + def run(self): + from ...exporters.openvino.__main__ import main_export + + # TODO : add input shapes + main_export( + model_name_or_path=self.args.model, + output=self.args.output, + task=self.args.task, + framework=self.args.framework, + cache_dir=self.args.cache_dir, + trust_remote_code=self.args.trust_remote_code, + pad_token_id=self.args.pad_token_id, + # **input_shapes, + ) diff --git a/optimum/commands/register/register_openvino.py b/optimum/commands/register/register_openvino.py new file mode 100644 index 0000000000..a1a74abaca --- /dev/null +++ b/optimum/commands/register/register_openvino.py @@ -0,0 +1,19 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ..export import ExportCommand +from ..export.openvino import OVExportCommand + + +REGISTER_COMMANDS = [(OVExportCommand, ExportCommand)] diff --git a/tests/openvino/test_exporters_cli.py b/tests/openvino/test_exporters_cli.py new file mode 100644 index 0000000000..6b1ba249ab --- /dev/null +++ b/tests/openvino/test_exporters_cli.py @@ -0,0 +1,59 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import subprocess +import unittest +from tempfile import TemporaryDirectory + +from parameterized import parameterized +from utils_tests import MODEL_NAMES + +from optimum.exporters.openvino.__main__ import main_export + + +class OVCLIExportTestCase(unittest.TestCase): + """ + Integration tests ensuring supported models are correctly exported. + """ + + SUPPORTED_ARCHITECTURES = ( + ["causal-lm", "gpt2"], + ["causal-lm-with-past", "gpt2"], + ["seq2seq-lm", "t5"], + ["seq2seq-lm-with-past", "t5"], + ["sequence-classification", "bert"], + ["question-answering", "distilbert"], + ["masked-lm", "bert"], + ["default", "blenderbot"], + ["default-with-past", "blenderbot"], + ["stable-diffusion", "stable-diffusion"], + ["stable-diffusion-xl", "stable-diffusion-xl"], + ["stable-diffusion-xl", "stable-diffusion-xl-refiner"], + ) + + def _openvino_export(self, model_name: str, task: str): + with TemporaryDirectory() as tmpdir: + main_export(model_name_or_path=model_name, output=tmpdir, task=task) + + @parameterized.expand(SUPPORTED_ARCHITECTURES) + def test_export(self, task: str, model_type: str): + self._openvino_export(MODEL_NAMES[model_type], task) + + @parameterized.expand(SUPPORTED_ARCHITECTURES) + def test_exporters_cli(self, task: str, model_type: str): + with TemporaryDirectory() as tmpdirname: + subprocess.run( + f"optimum-cli export openvino --model {MODEL_NAMES[model_type]} --task {task} {tmpdirname}", + shell=True, + check=True, + )