diff --git a/optimum/commands/export/openvino.py b/optimum/commands/export/openvino.py index 0eaddddef4..ade989857d 100644 --- a/optimum/commands/export/openvino.py +++ b/optimum/commands/export/openvino.py @@ -43,7 +43,32 @@ def parse_args_openvino(parser: "ArgumentParser"): ), ) 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.") @@ -68,15 +93,14 @@ def parse_args(parser: "ArgumentParser"): 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, - # fp16=self.args.fp16, - # no_post_process=self.args.no_post_process, - # framework=self.args.framework, + 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, + trust_remote_code=self.args.trust_remote_code, + pad_token_id=self.args.pad_token_id, # **input_shapes, - ) + ) \ No newline at end of file diff --git a/tests/openvino/test_exporters_cli.py b/tests/openvino/test_exporters_cli.py new file mode 100644 index 0000000000..e100c2ffe4 --- /dev/null +++ b/tests/openvino/test_exporters_cli.py @@ -0,0 +1,61 @@ +# 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 typing import Dict + +from parameterized import parameterized +from utils_tests import MODEL_NAMES + +from optimum.exporters.openvino.__main__ import main_export +from optimum.exporters.tasks import TasksManager + + +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, + )