diff --git a/optimum/commands/export/openvino.py b/optimum/commands/export/openvino.py new file mode 100644 index 0000000000..37bcb77c4f --- /dev/null +++ b/optimum/commands/export/openvino.py @@ -0,0 +1,84 @@ +# 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 subprocess +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.") + + +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 + + 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, + 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..7ae12fa75b --- /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 ..openvino.base import OVCommand +from ..export.openvino import OVExportCommand + + +REGISTER_COMMANDS = [OVExportCommand]