Skip to content

Commit

Permalink
Refactor After Review
Browse files Browse the repository at this point in the history
  • Loading branch information
apaniukov committed Jan 18, 2024
1 parent 1b64b54 commit cd8aaf8
Show file tree
Hide file tree
Showing 3 changed files with 24 additions and 15 deletions.
27 changes: 13 additions & 14 deletions optimum/exporters/openvino/convert.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from transformers.utils import is_tf_available, is_torch_available

from openvino.runtime import PartialShape, save_model
from openvino.runtime.exceptions import OVTypeError
from openvino.runtime.utils.types import get_element_type
from openvino.tools.ovc import convert_model
from optimum.exporters.onnx.base import OnnxConfig
Expand Down Expand Up @@ -489,47 +490,45 @@ def export_models(
return outputs


OV_TOKENIZER_FILE_NAME = "openvino_tokenizer{}.xml"
OV_DETOKENIZER_FILE_NAME = "openvino_detokenizer{}.xml"
UNSUPPORTED_TOKENZIER_CLASSES = (
UNSUPPORTED_TOKENIZER_CLASSES = (
T5Tokenizer,
T5TokenizerFast,
)


def export_tokenizer(
tokenizer,
output_path: Union[str, Path],
output: Union[str, Path],
suffix: Optional[str] = "",
):
from openvino.runtime.exceptions import OVTypeError
from optimum.intel.openvino import OV_DETOKENIZER_NAME, OV_TOKENIZER_NAME # avoid circular imports

if isinstance(tokenizer, UNSUPPORTED_TOKENZIER_CLASSES):
logger.info("OpenVINO Tokenizer for this model is not supported.")
if isinstance(tokenizer, UNSUPPORTED_TOKENIZER_CLASSES):
logger.info(f"OpenVINO Tokenizer export for {type(tokenizer).__name__} is not supported.")
return

try:
from openvino_tokenizers import convert_tokenizer
except ModuleNotFoundError:
logger.info("Run `pip install openvino-tokenizers` to get OpenVINO tokenizer/detokenizer models.")

if not isinstance(output_path, Path):
output_path = Path(output_path)
if not isinstance(output, Path):
output = Path(output)

try:
converted = convert_tokenizer(tokenizer, with_detokenizer=True)
except NotImplementedError:
logger.info("Detokenizer is not supported, convert tokenizer only.")
logger.warning("Detokenizer is not supported, convert tokenizer only.")
converted = convert_tokenizer(tokenizer, with_detokenizer=False)
except OVTypeError:
logger.info("OpenVINO Tokenizer for this model is not supported.")
logger.warning(f"OpenVINO Tokenizer for {type(tokenizer).__name__} is not supported.")
return
except Exception as exception:
logger.warning(f"OpenVINO Tokenizer for this model is not supported. Exception: {exception}")
logger.warning(f"OpenVINO Tokenizer {type(tokenizer).__name__} is not supported. Exception: {exception}")
return

if not isinstance(converted, tuple):
converted = (converted,)

for model, file_name in zip(converted, (OV_TOKENIZER_FILE_NAME, OV_DETOKENIZER_FILE_NAME)):
save_model(model, output_path / file_name.format(suffix))
for model, file_name in zip(converted, (OV_TOKENIZER_NAME, OV_DETOKENIZER_NAME)):
save_model(model, output / file_name.format(suffix))
9 changes: 8 additions & 1 deletion optimum/intel/openvino/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,14 @@
import logging

from ..utils.import_utils import is_diffusers_available, is_nncf_available
from .utils import OV_DECODER_NAME, OV_DECODER_WITH_PAST_NAME, OV_ENCODER_NAME, OV_XML_FILE_NAME
from .utils import (
OV_DECODER_NAME,
OV_DECODER_WITH_PAST_NAME,
OV_DETOKENIZER_NAME,
OV_ENCODER_NAME,
OV_TOKENIZER_NAME,
OV_XML_FILE_NAME,
)


if is_nncf_available():
Expand Down
3 changes: 3 additions & 0 deletions optimum/intel/openvino/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,9 @@
OV_DECODER_NAME = "openvino_decoder_model.xml"
OV_DECODER_WITH_PAST_NAME = "openvino_decoder_with_past_model.xml"

OV_TOKENIZER_NAME = "openvino_tokenizer{}.xml"
OV_DETOKENIZER_NAME = "openvino_detokenizer{}.xml"

ONNX_WEIGHTS_NAME = "model.onnx"
ONNX_ENCODER_NAME = "encoder_model.onnx"
ONNX_DECODER_NAME = "decoder_model.onnx"
Expand Down

0 comments on commit cd8aaf8

Please sign in to comment.