Skip to content

Commit

Permalink
format
Browse files Browse the repository at this point in the history
  • Loading branch information
echarlaix committed Sep 27, 2023
1 parent 101c924 commit 51fff5c
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 16 deletions.
4 changes: 2 additions & 2 deletions optimum/commands/export/openvino.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,6 @@ def parse_args_openvino(parser: "ArgumentParser"):
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(
Expand All @@ -70,6 +69,7 @@ def parse_args_openvino(parser: "ArgumentParser"):
),
)


class OVExportCommand(BaseOptimumCLICommand):
COMMAND = CommandInfo(name="openvino", help="Export PyTorch models to OpenVINO IR.")

Expand Down Expand Up @@ -103,4 +103,4 @@ def run(self):
trust_remote_code=self.args.trust_remote_code,
pad_token_id=self.args.pad_token_id,
# **input_shapes,
)
)
26 changes: 12 additions & 14 deletions tests/openvino/test_exporters_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,11 @@
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):
Expand All @@ -29,18 +27,18 @@ class OVCLIExportTestCase(unittest.TestCase):
"""

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"],
["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):
Expand Down

0 comments on commit 51fff5c

Please sign in to comment.