Repository: huggingface/optimum

Issue:
Make  `ORTModelForConditionalGeneration` use exporters
### Feature request

Currently, `ORTModelForConditionalGeneration` use `transformers.onnx`to export models. Make it use `optimum.exporters.onnx`.

### Motivation

`transformers.onnx` is not going to be supported anymore.

### Your contribution


I will implement this.


Discussion:

Resolving patch:
diff --git a/optimum/onnxruntime/modeling_seq2seq.py b/optimum/onnxruntime/modeling_seq2seq.py
--- a/optimum/onnxruntime/modeling_seq2seq.py
+++ b/optimum/onnxruntime/modeling_seq2seq.py
@@ -29,14 +29,12 @@
 from transformers.file_utils import add_start_docstrings_to_model_forward, default_cache_path
 from transformers.generation_utils import GenerationMixin
 from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
-from transformers.onnx import FeaturesManager, export
 
 import onnxruntime as ort
 from huggingface_hub import hf_hub_download
 
-from ..exporters.onnx.model_configs import SpeechSeq2SeqDecoderOnnxConfig, SpeechSeq2SeqEncoderOnnxConfig
-from ..onnx.configuration import DecoderOnnxConfig, EncoderOnnxConfig
-from ..onnx.modeling_seq2seq import _DecoderWithLMhead
+from ..exporters.onnx.convert import export_encoder_decoder_model as export
+from ..exporters.tasks import TasksManager
 from .io_binding import TypeHelper
 from .modeling_ort import ORTModel
 from .utils import (
@@ -519,94 +517,35 @@ def _from_transformers(
         save_dir = Path(save_dir).joinpath(model_id, subfolder)
         save_dir.mkdir(parents=True, exist_ok=True)
 
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-
-        framework = FeaturesManager.determine_framework(os.path.join(model_id, subfolder))
-        model_class = FeaturesManager.get_model_class_for_feature(cls.export_feature, framework)
-        model = model_class.from_pretrained(
+        model = TasksManager.get_model_from_task(
+            cls.export_feature,
             model_id,
             subfolder=subfolder,
-            config=config,
-            cache_dir=cache_dir,
             revision=revision,
+            cache_dir=cache_dir,
+            config=config,
             use_auth_token=use_auth_token,
             local_files_only=local_files_only,
+            force_download=force_download,
         )
 
-        _, model_onnx_config = FeaturesManager.check_supported_model_or_raise(model, feature=cls.export_feature)
-        onnx_config = model_onnx_config(model.config)
-        onnx_opset = onnx_config.default_onnx_opset
-
-        # Extract the encoder and decoder for ONNX export
-        encoder = model.get_encoder()
-        decoder = model.get_decoder()
+        model_type = model.config.model_type.replace("_", "-")
+        model_name = getattr(model, "name", None)
 
-        # Concatenate the decoder with the language model head for ONNX export
-        decoder_with_lm_head = _DecoderWithLMhead(model)
-
-        # Get the encoder and decoder ONNX configs
-        onnx_config_encoder = cls.get_encoder_onnx_config(encoder.config)
-        onnx_config_decoder = cls.get_decoder_onnx_config(decoder.config, cls.export_feature, use_past=False)
-        if use_cache:
-            onnx_config_decoder_with_past = cls.get_decoder_onnx_config(
-                decoder.config, cls.export_feature, use_past=True
-            )
-
-        if config.model_type == "whisper":
-            from ..exporters.onnx.convert import export as export_optimum
-
-            # Export the encoder
-            export_optimum(
-                encoder,
-                onnx_config_encoder,
-                onnx_opset,
-                save_dir.joinpath(ONNX_ENCODER_NAME),
-            )
-
-            # Export the decoder without the past key values
-            export_optimum(
-                model,
-                onnx_config_decoder,
-                onnx_opset,
-                save_dir.joinpath(ONNX_DECODER_NAME),
-            )
-
-            # Export the decoder with the past key values
-            if use_cache:
-                export_optimum(
-                    model,
-                    onnx_config_decoder_with_past,
-                    onnx_opset,
-                    save_dir.joinpath(ONNX_DECODER_WITH_PAST_NAME),
-                )
-        else:
-            # Export the encoder
-            export(
-                preprocessor=tokenizer,
-                model=encoder,
-                config=onnx_config_encoder,
-                opset=onnx_opset,
-                output=save_dir.joinpath(ONNX_ENCODER_NAME),
-            )
-
-            # Export the decoder without the past key values
-            export(
-                preprocessor=tokenizer,
-                model=decoder_with_lm_head,
-                config=onnx_config_decoder,
-                opset=onnx_opset,
-                output=save_dir.joinpath(ONNX_DECODER_NAME),
-            )
-
-            # Export the decoder with the past key values
-            if use_cache:
-                export(
-                    preprocessor=tokenizer,
-                    model=decoder_with_lm_head,
-                    config=onnx_config_decoder_with_past,
-                    opset=onnx_opset,
-                    output=save_dir.joinpath(ONNX_DECODER_WITH_PAST_NAME),
-                )
+        onnx_config_constructor = TasksManager.get_exporter_config_constructor(
+            model_type, "onnx", task=cls.export_feature, model_name=model_name
+        )
+        onnx_config = onnx_config_constructor(model.config, use_past=use_cache)
+        onnx_opset = onnx_config.DEFAULT_ONNX_OPSET
+
+        export(
+            model,
+            onnx_config,
+            onnx_opset,
+            save_dir.joinpath(ONNX_ENCODER_NAME),
+            save_dir.joinpath(ONNX_DECODER_NAME),
+            save_dir.joinpath(ONNX_DECODER_WITH_PAST_NAME),
+        )
 
         return cls._from_pretrained(
             save_dir,
@@ -867,12 +806,8 @@ def __init__(
         self.session_outputs = {output_key.name: idx for idx, output_key in enumerate(self.session.get_outputs())}
         self.session_input_names = list(self.session_inputs.keys())
         self.session_output_names = list(self.session_outputs.keys())
-        self.key_value_input_names = [
-            key for key in self.session_input_names if ("key_values" in key or ".key" in key or ".value" in key)
-        ]
-        self.key_value_output_names = [
-            key for key in self.session_output_names if ("key_values" in key or ".key" in key or ".value" in key)
-        ]
+        self.key_value_input_names = [key for key in self.session_input_names if (".key" in key or ".value" in key)]
+        self.key_value_output_names = [key for key in self.session_output_names if (".key" in key or ".value" in key)]
         self.name_to_np_type = TypeHelper.get_io_numpy_type_map(self.session) if self.use_io_binding else None
 
     def prepare_output_buffer(
@@ -895,7 +830,7 @@ def prepare_output_buffer(
         elif output_name == "logits":
             output_shape = (batch_size, sequence_length, self.normalized_config.vocab_size)
             output_buffer = torch.empty(np.prod(output_shape), dtype=torch_type, device=self._device).contiguous()
-        elif "key_values" in output_name or ".key" in output_name or ".value" in output_name:
+        elif ".key" in output_name or ".value" in output_name:
             num_attention_heads = self.normalized_config.num_attention_heads
             hidden_size = self.normalized_config.hidden_size
             embed_size_per_head = hidden_size // num_attention_heads
@@ -1175,14 +1110,6 @@ def _initialize_encoder(
             main_input_name=self.main_input_name,
         )
 
-    def get_encoder_onnx_config(encoder_config: "PretrainedConfig") -> EncoderOnnxConfig:
-        return EncoderOnnxConfig(encoder_config, task="default")
-
-    def get_decoder_onnx_config(
-        decoder_config: "PretrainedConfig", export_feature: str, use_past: bool = False
-    ) -> DecoderOnnxConfig:
-        return DecoderOnnxConfig(decoder_config, export_feature, use_past=use_past)
-
     @add_start_docstrings_to_model_forward(
         SEQ2SEQ_ONNX_MODEL_DOCSTRING.format("batch_size, sequence_length")
         + TRANSLATION_EXAMPLE.format(
@@ -1302,14 +1229,6 @@ def _initialize_encoder(
             main_input_name=self.main_input_name,
         )
 
-    def get_encoder_onnx_config(encoder_config: "PretrainedConfig") -> SpeechSeq2SeqEncoderOnnxConfig:
-        return SpeechSeq2SeqEncoderOnnxConfig(encoder_config, task="default")
-
-    def get_decoder_onnx_config(
-        decoder_config: "PretrainedConfig", export_feature: str, use_past: bool = False
-    ) -> SpeechSeq2SeqDecoderOnnxConfig:
-        return SpeechSeq2SeqDecoderOnnxConfig(decoder_config, export_feature, use_past=use_past)
-
     @add_start_docstrings_to_model_forward(
         SPEECH_SEQ2SEQ_ONNX_MODEL_DOCSTRING.format("batch_size, feature_size, sequence_length")
         + AUTOMATIC_SPEECH_RECOGNITION_EXAMPLE.format(


Test patch: