more seq2seq models

huggingface · Jun 25, 2024 · 5ab4358 · 5ab4358
1 parent c618f35
commit 5ab4358
Show file tree

Hide file tree

Showing 5 changed files with 400 additions and 70 deletions.
diff --git a/optimum/exporters/openvino/__main__.py b/optimum/exporters/openvino/__main__.py
@@ -23,9 +23,8 @@
 
 from optimum.exporters import TasksManager
 from optimum.exporters.onnx.base import OnnxConfig
-from optimum.exporters.onnx.constants import SDPA_ARCHS_ONNX_EXPORT_NOT_SUPPORTED
 from optimum.exporters.openvino.convert import export_from_model
-from optimum.intel.utils.import_utils import is_openvino_tokenizers_available, is_transformers_version
+from optimum.intel.utils.import_utils import is_openvino_tokenizers_available
 from optimum.utils.save_utils import maybe_load_preprocessors
 
 

diff --git a/optimum/exporters/openvino/convert.py b/optimum/exporters/openvino/convert.py
@@ -32,7 +32,13 @@
 from optimum.exporters.onnx.convert import check_dummy_inputs_are_allowed
 from optimum.exporters.onnx.convert import export_pytorch as export_pytorch_to_onnx
 from optimum.exporters.onnx.convert import export_tensorflow as export_tensorflow_onnx
-from optimum.exporters.utils import _get_submodels_and_export_configs, ENCODER_NAME, DECODER_NAME, DECODER_WITH_PAST_NAME, _get_submodels_for_export_encoder_decoder
+from optimum.exporters.utils import (
+    DECODER_NAME,
+    DECODER_WITH_PAST_NAME,
+    ENCODER_NAME,
+    _get_submodels_and_export_configs,
+    _get_submodels_for_export_encoder_decoder,
+)
 from optimum.intel.utils.import_utils import (
     _nncf_version,
     _optimum_intel_version,
@@ -576,13 +582,13 @@ def export_from_model(
 
     logging.disable(logging.INFO)
 
-    if (model.config.is_encoder_decoder and task.startswith(TasksManager._ENCODER_DECODER_TASKS)) and stateful and not custom_architecture:
+    if (
+        (model.config.is_encoder_decoder and task.startswith(TasksManager._ENCODER_DECODER_TASKS))
+        and stateful
+        and not custom_architecture
+    ):
         export_config, models_and_export_configs = _get_encoder_decoder_stateful_models_for_export(
-            model=model,
-            task=task,
-            preprocessors=preprocessors,
-            library_name=library_name,
-            _variant="default"
+            model=model, task=task, preprocessors=preprocessors, library_name=library_name, _variant="default"
         )
         stateful = [False, True]
     else:
@@ -773,23 +779,22 @@ def _get_encoder_decoder_stateful_models_for_export(
     )
 
     export_config.variant = _variant
-    all_variants = "\n".join(
-        [f"    - {name}: {description}" for name, description in export_config.VARIANTS.items()]
-    )
+    all_variants = "\n".join([f"    - {name}: {description}" for name, description in export_config.VARIANTS.items()])
     logger.info(f"Using the export variant {export_config.variant}. Available variants are:\n{all_variants}")
 
     models_for_export = _get_submodels_for_export_encoder_decoder(model, use_past=True)
 
     encoder_export_config = export_config.with_behavior("encoder")
     models_for_export[ENCODER_NAME] = (models_for_export[ENCODER_NAME], encoder_export_config)
 
-    decoder_export_config_with_past = export_config.with_behavior("decoder", use_past=True, use_past_in_inputs=True, stateful=True)
+    decoder_export_config_with_past = export_config.with_behavior(
+        "decoder", use_past=True, use_past_in_inputs=True, stateful=True
+    )
     decoder_with_past_model = models_for_export.pop(DECODER_WITH_PAST_NAME)
     models_for_export[DECODER_NAME] = (
         decoder_with_past_model,
         decoder_export_config_with_past,
     )
     logger.info(models_for_export.keys())
 
-
     return None, models_for_export