merge main

Author: echarlaix

.github/workflows/build_documentation.yml

@@ -25,7 +25,7 @@ jobs:
       - uses: actions/checkout@v4
       - uses: actions/setup-node@v4
         with:
-          node-version: "18"
+          node-version: "20"
           cache-dependency-path: "kit/package-lock.json"
 
       - name: Set up Python

.github/workflows/build_pr_documentation.yml

@@ -26,7 +26,7 @@ jobs:
       - uses: actions/checkout@v4
       - uses: actions/setup-node@v4
         with:
-          node-version: "18"
+          node-version: "20"
           cache-dependency-path: "kit/package-lock.json"
 
       - name: Set up Python

optimum/commands/export/openvino.py

@@ -21,15 +21,8 @@
 
 from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE
 
-from ...exporters import TasksManager
-from ...intel.utils.import_utils import (
-    DIFFUSERS_IMPORT_ERROR,
-    is_diffusers_available,
-    is_nncf_available,
-)
-from ...intel.utils.modeling_utils import _infer_library_from_model_name_or_path
-from ...utils.save_utils import maybe_load_preprocessors
-from ..base import BaseOptimumCLICommand, CommandInfo
+from optimum.commands.base import BaseOptimumCLICommand, CommandInfo
+from optimum.exporters.tasks import TasksManager
 
 
 logger = logging.getLogger(__name__)
@@ -348,9 +341,13 @@ def parse_args(parser: "ArgumentParser"):
         return parse_args_openvino(parser)
 
     def run(self):
+        from optimum.utils.save_utils import maybe_load_preprocessors
+
         from ...exporters.openvino.__main__ import infer_task, main_export, maybe_convert_tokenizers
         from ...exporters.openvino.utils import save_preprocessors
         from ...intel.openvino.configuration import _DEFAULT_4BIT_WQ_CONFIG, OVConfig, get_default_quantization_config
+        from ...intel.utils.import_utils import DIFFUSERS_IMPORT_ERROR, is_diffusers_available, is_nncf_available
+        from ...intel.utils.modeling_utils import _infer_library_from_model_name_or_path
 
         if self.args.library is None:
             # TODO: add revision, subfolder and token to args

optimum/commands/neural_compressor/base.py

@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .. import BaseOptimumCLICommand, CommandInfo
+from optimum.commands.base import BaseOptimumCLICommand, CommandInfo
+
 from .quantize import INCQuantizeCommand
 
 

optimum/commands/neural_compressor/quantize.py

@@ -16,8 +16,8 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Optional
 
-from ...exporters import TasksManager
-from ..base import BaseOptimumCLICommand, CommandInfo
+from optimum.commands.base import BaseOptimumCLICommand, CommandInfo
+from optimum.exporters.tasks import TasksManager
 
 
 if TYPE_CHECKING:

optimum/commands/register/register_inc.py

@@ -15,4 +15,6 @@
 from ..neural_compressor.base import INCCommand
 
 
-REGISTER_COMMANDS = [INCCommand]
+REGISTER_COMMANDS = [
+    INCCommand,
+]

optimum/commands/register/register_openvino.py

@@ -12,8 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..export import ExportCommand
+from optimum.commands.export.base import ExportCommand
+
 from ..export.openvino import OVExportCommand
 
 
-REGISTER_COMMANDS = [(OVExportCommand, ExportCommand)]
+REGISTER_COMMANDS = [
+    (OVExportCommand, ExportCommand),
+]

optimum/exporters/openvino/__main__.py

@@ -26,8 +26,8 @@
 from transformers.utils import is_torch_available
 
 from openvino import Core, Type, save_model
-from optimum.exporters import TasksManager
 from optimum.exporters.onnx.base import OnnxConfig
+from optimum.exporters.tasks import TasksManager
 from optimum.intel.utils.import_utils import (
     is_nncf_available,
     is_openvino_tokenizers_available,

optimum/exporters/openvino/convert.py

@@ -24,12 +24,12 @@
 from packaging.version import Version
 from transformers.generation import GenerationMixin
 from transformers.models.speecht5.modeling_speecht5 import SpeechT5HifiGan
-from transformers.utils import is_tf_available, is_torch_available
+from transformers.utils import is_torch_available
 
 from openvino import Model, save_model
 from openvino.exceptions import OVTypeError
 from openvino.tools.ovc import convert_model
-from optimum.exporters import TasksManager
+from optimum.exporters.tasks import TasksManager
 from optimum.exporters.utils import (
     DECODER_NAME,
     ENCODER_NAME,
@@ -57,7 +57,6 @@
 
 from ...intel.utils.import_utils import is_nncf_available
 from ...intel.utils.modeling_utils import _infer_library_from_model_or_model_class
-from .model_patcher import patch_model_with_bettertransformer
 from .stateful import (
     ensure_export_task_support_stateful,
     ensure_model_type_support_stateful,
@@ -87,9 +86,6 @@
 if is_diffusers_available():
     from diffusers import DiffusionPipeline, ModelMixin
 
-if is_tf_available():
-    from transformers.modeling_tf_utils import TFPreTrainedModel
-
 
 if TYPE_CHECKING:
     from optimum.exporters.onnx.base import OnnxConfig
@@ -99,7 +95,7 @@
 def _set_runtime_options(
     models_and_export_configs: Dict[
         str,
-        Tuple[Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin", "DiffusionPipeline"], "OnnxConfig"],
+        Tuple[Union["PreTrainedModel", "ModelMixin", "DiffusionPipeline"], "OnnxConfig"],
     ],
     task: str,
     library_name: str,
@@ -141,7 +137,7 @@ def _save_model(
 
 
 def export(
-    model: Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin", "DiffusionPipeline"],
+    model: Union["PreTrainedModel", "ModelMixin", "DiffusionPipeline"],
     config: "OnnxConfig",
     output: Path,
     opset: Optional[int] = None,
@@ -154,10 +150,10 @@ def export(
     library_name: Optional[str] = None,
 ) -> Tuple[List[str], List[str]]:
     """
-    Exports a Pytorch or TensorFlow model to an OpenVINO Intermediate Representation.
+    Exports a Pytorch model to an OpenVINO Intermediate Representation.
 
     Args:
-        model ([`PreTrainedModel`] or [`TFPreTrainedModel`]):
+        model ([`PreTrainedModel`]):
             The model to export.
         config ([`~exporters.onnx.config.OnnxConfig`]):
             The ONNX configuration associated with the exported model.
@@ -179,10 +175,10 @@ def export(
         `Tuple[List[str], List[str]]`: A tuple with an ordered list of the model's inputs, and the named inputs from
         the ONNX configuration.
     """
-    if not (is_torch_available() or is_tf_available()):
+    if not is_torch_available():
         raise ImportError(
-            "Cannot convert because neither PyTorch nor TensorFlow are installed. "
-            "Please install torch or tensorflow first."
+            "Cannot convert because torch is not installed. "
+            "Please install torch with `pip install torch` and try again."
         )
 
     if "diffusers" in str(model.__class__) and not is_diffusers_available():
@@ -227,63 +223,8 @@ def export(
             patch_16bit_model=patch_16bit_model,
             library_name=library_name,
         )
-
-    elif is_tf_available() and issubclass(type(model), TFPreTrainedModel):
-        output.parent.mkdir(parents=True, exist_ok=True)
-        if opset is None:
-            opset = config.DEFAULT_ONNX_OPSET
-        if device == "cuda":
-            raise RuntimeError("`tf2onnx` does not support export on CUDA device.")
-        if input_shapes is not None:
-            logger.info("`input_shapes` argument is not supported by the Tensorflow ONNX export and will be ignored.")
-        return export_tensorflow(model, config, opset, output, ov_config=ov_config, library_name=library_name)
-
     else:
-        raise RuntimeError(
-            "You either provided a PyTorch model with only TensorFlow installed, or a TensorFlow model with only PyTorch installed."
-        )
-
-
-def export_tensorflow(
-    model: Union["PreTrainedModel", "ModelMixin"],
-    config: "OnnxConfig",
-    opset: int,
-    output: Path,
-    ov_config: Optional["OVConfig"] = None,
-    library_name: Optional[str] = None,
-):
-    """
-    Export the TensorFlow model to OpenVINO format.
-
-    Args:
-        model (Union[): The model to export.
-        config (OnnxConfig): The configuration of the model.
-        opset (int): The ONNX opset version to use.
-        output (Path): The path to save the model.
-
-    Returns:
-        input_names: list of input names from ONNX configuration
-        output_names: list of output names from ONNX configuration
-        bool:  True if the model was exported successfully.
-    """
-    from optimum.exporters.onnx.convert import export_tensorflow as export_tensorflow_onnx
-
-    onnx_path = Path(output).with_suffix(".onnx")
-    input_names, output_names = export_tensorflow_onnx(model, config, opset, onnx_path)
-    ov_model = convert_model(str(onnx_path))
-
-    library_name = _infer_library_from_model_or_model_class(model=model, library_name=library_name)
-
-    _save_model(
-        ov_model,
-        output.parent / output,
-        ov_config=ov_config,
-        library_name=library_name,
-        config=config,
-    )
-    del ov_model
-    gc.collect()
-    return input_names, output_names, True
+        raise RuntimeError("You either provided a non-PyTorch model or the PyTorch library is not installed.")
 
 
 def export_pytorch_via_onnx(
@@ -406,15 +347,6 @@ def export_pytorch(
     # TODO: temporary solution but statefulness should be added to the export config earlier
     config.stateful = stateful
 
-    if stateful:
-        # Trigger bettertransformer together with stateful model because OpenVINO HW-dependent transformations expect
-        # both of them are applied to demonstrate the best performance.
-        # TODO: Consider applying bettertransformer regardless of stateful flag -- requires additional validation.
-        model = patch_model_with_bettertransformer(model)
-        # TODO: Consider unpatching model after export is done in the end of this function.
-        #       Now it is left as-is because the model is not expected to be used after call export_pytorch, and
-        #       this function is one of the _internal_ steps in a bigger model conversion pipeline.
-
     with torch.no_grad():
         if hasattr(model, "config"):
             model.config.torchscript = False
@@ -523,7 +455,7 @@ def ts_patched_forward(*args, **kwargs):
 
 def export_models(
     models_and_export_configs: Dict[
-        str, Tuple[Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin", "DiffusionPipeline"], "OnnxConfig"]
+        str, Tuple[Union["PreTrainedModel", "ModelMixin", "DiffusionPipeline"], "OnnxConfig"]
     ],
     output_dir: Path,
     opset: Optional[int] = None,
@@ -540,7 +472,7 @@ def export_models(
     Export the models to OpenVINO IR format
 
     Args:
-        models_and_export_configs (Dict[ str, Tuple[Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"], "OnnxConfig"]):
+        models_and_export_configs (Dict[ str, Tuple[Union["PreTrainedModel", "ModelMixin"], "OnnxConfig"]):
         output_dir (Path): output directory for saving models
         opset (Optional[int], optional, Default to None): ONNX export opset
         output_names (Optional[List[str]], optional, Defaults to None): model output names
@@ -596,7 +528,7 @@ def export_models(
 
 
 def export_from_model(
-    model: Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin", "DiffusionPipeline"],
+    model: Union["PreTrainedModel", "ModelMixin", "DiffusionPipeline"],
     output: Union[str, Path],
     task: Optional[str] = None,
     ov_config: Optional["OVConfig"] = None,
@@ -720,7 +652,6 @@ def export_from_model(
             library_name=library_name,
             model_kwargs=model_kwargs,
             _variant="default",
-            legacy=False,
             exporter="openvino",
             stateful=stateful,
         )
@@ -937,7 +868,7 @@ def _add_version_info_to_model(model: Model, library_name: Optional[str] = None)
 
 
 def _get_multi_modal_submodels_and_export_configs(
-    model: Union["PreTrainedModel", "TFPreTrainedModel"],
+    model: "PreTrainedModel",
     task: str,
     library_name: str,
     int_dtype: str,
@@ -991,7 +922,7 @@ def _get_multi_modal_submodels_and_export_configs(
 
 
 def _get_submodels_and_export_configs(
-    model: Union["PreTrainedModel", "TFPreTrainedModel", "DiffusionPipeline"],
+    model: Union["PreTrainedModel", "DiffusionPipeline"],
     task: str,
     monolith: bool,
     custom_export_configs: Dict,
@@ -1002,7 +933,6 @@ def _get_submodels_and_export_configs(
     float_dtype: str = "fp32",
     fn_get_submodels: Optional[Callable] = None,
     preprocessors: Optional[List[Any]] = None,
-    legacy: bool = False,
     model_kwargs: Optional[Dict] = None,
     exporter: str = "openvino",
     stateful: bool = False,
@@ -1032,7 +962,6 @@ def _get_submodels_and_export_configs(
         float_dtype,
         fn_get_submodels,
         preprocessors,
-        legacy,
         model_kwargs,
         exporter,
     )
@@ -1405,7 +1334,7 @@ def get_flux_models_for_export(pipeline, exporter, int_dtype, float_dtype):
 
 
 def _get_encoder_decoder_stateful_models_for_export(
-    model: Union["PreTrainedModel", "TFPreTrainedModel"],
+    model: "PreTrainedModel",
     task: str,
     _variant: str,
     library_name: str,
@@ -1421,7 +1350,6 @@ def _get_encoder_decoder_stateful_models_for_export(
         int_dtype=int_dtype,
         float_dtype=float_dtype,
         preprocessors=preprocessors,
-        legacy=False,
     )
 
     export_config.variant = _variant
@@ -1444,7 +1372,7 @@ def _get_encoder_decoder_stateful_models_for_export(
 
 
 def _get_speecht5_tss_model_for_export(
-    model: Union["PreTrainedModel", "TFPreTrainedModel"],
+    model: "PreTrainedModel",
     task: str,
     library_name: str,
     int_dtype: str,
@@ -1467,7 +1395,6 @@ def _get_speecht5_tss_model_for_export(
         int_dtype=int_dtype,
         float_dtype=float_dtype,
         preprocessors=preprocessors,
-        legacy=False,
     )
     export_config.variant = "default"