optimum/onnxruntime/modeling_decoder.py

#  Copyright 2022 The HuggingFace Team. All rights reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
"""Classes handling causal-lm related architectures in ONNX Runtime."""

import logging
import os
import re
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Tuple, Union

import onnx
import torch
from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE
from onnx.tools import update_model_dims
from transformers import AutoModelForCausalLM, GenerationConfig
from transformers.file_utils import add_end_docstrings, add_start_docstrings_to_model_forward
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.utils import cached_file

from onnxruntime import InferenceSession, SessionOptions

from ..exporters.onnx import MODEL_TYPES_REQUIRING_POSITION_IDS, main_export
from ..exporters.tasks import TasksManager
from ..onnx.utils import check_model_uses_external_data
from ..utils import NormalizedConfigManager, is_transformers_version
from ..utils.file_utils import find_files_matching_pattern
from ..utils.save_utils import maybe_save_preprocessors
from .constants import (
    DECODER_MERGED_ONNX_FILE_PATTERN,
    DECODER_ONNX_FILE_PATTERN,
    DECODER_WITH_PAST_ONNX_FILE_PATTERN,
    ONNX_FILE_PATTERN,
)
from .modeling_ort import ONNX_MODEL_END_DOCSTRING, ORTModel
from .utils import prepare_providers_and_provider_options


if TYPE_CHECKING:
    from transformers import PretrainedConfig

if is_transformers_version(">=", "4.25.0"):
    from transformers.generation import GenerationMixin
else:
    from transformers.generation_utils import GenerationMixin  # type: ignore # noqa: F401


logger = logging.getLogger(__name__)

DECODER_INPUTS_DOCSTRING = r"""
    Args:
        input_ids (`torch.LongTensor`):
            Indices of decoder input sequence tokens in the vocabulary of shape `(batch_size, sequence_length)`.
        attention_mask (`torch.LongTensor`, *optional*):
            Mask to avoid performing attention on padding token indices, of shape
            `(batch_size, sequence_length)`. Mask values selected in `[0, 1]`.
        past_key_values (`tuple(tuple(torch.FloatTensor), *optional*, defaults to `None`)`
            Contains the precomputed key and value hidden states of the attention blocks used to speed up decoding.
            The tuple is of length `config.n_layers` with each tuple having 2 tensors of shape
            `(batch_size, num_heads, sequence_length, embed_size_per_head)`.
"""

CAUSALLM_ONNX_MODEL_DOCSTRING = r"""
    Args:
        input_ids (`torch.LongTensor`):
            Indices of decoder input sequence tokens in the vocabulary of shape `(batch_size, sequence_length)`.
        attention_mask (`torch.LongTensor`):
            Mask to avoid performing attention on padding token indices, of shape
            `(batch_size, sequence_length)`. Mask values selected in `[0, 1]`.
        past_key_values (`tuple(tuple(torch.FloatTensor), *optional*, defaults to `None`)`
            Contains the precomputed key and value hidden states of the attention blocks used to speed up decoding.
            The tuple is of length `config.n_layers` with each tuple having 2 tensors of shape
            `(batch_size, num_heads, sequence_length, embed_size_per_head)`.
"""

_TOKENIZER_FOR_DOC = "AutoTokenizer"

TEXT_GENERATION_EXAMPLE = r"""
    Example of text generation:

    ```python
    >>> from transformers import {processor_class}
    >>> from optimum.onnxruntime import {model_class}
    >>> import torch

    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
    >>> model = {model_class}.from_pretrained("{checkpoint}")

    >>> inputs = tokenizer("My name is Arthur and I live in", return_tensors="pt")

    >>> gen_tokens = model.generate(**inputs,do_sample=True,temperature=0.9, min_length=20,max_length=20)
    >>> tokenizer.batch_decode(gen_tokens)  # doctest: +IGNORE_RESULT
    ```

    Example using `transformers.pipelines`:

    ```python
    >>> from transformers import {processor_class}, pipeline
    >>> from optimum.onnxruntime import {model_class}

    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
    >>> model = {model_class}.from_pretrained("{checkpoint}")
    >>> onnx_gen = pipeline("text-generation", model=model, tokenizer=tokenizer)

    >>> text = "My name is Arthur and I live in"
    >>> gen = onnx_gen(text)
    ```
"""


@add_end_docstrings(ONNX_MODEL_END_DOCSTRING)
class ORTModelForCausalLM(ORTModel, GenerationMixin):
    """
    ONNX model with a causal language modeling head for ONNX Runtime inference. This class officially supports bloom, codegen, falcon, gpt2, gpt-bigcode, gpt_neo, gpt_neox, gptj, llama.
    """

    auto_model_class = AutoModelForCausalLM
    main_input_name = "input_ids"
    _supports_cache_class = False

    def __init__(
        self,
        *args,
        config: "PretrainedConfig" = None,
        session: "InferenceSession" = None,
        use_io_binding: Optional[bool] = None,
        generation_config: Optional["GenerationConfig"] = None,
        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
        **kwargs,
    ):
        # DEPRECATED BEHAVIOR
        if args:
            logger.warning(
                "Instantiating an ORTModelForCausalLM with positional arguments is deprecated and will be removed in the next version. "
                "Please use the keywords arguments {config, session, use_io_binding, generation_config, model_save_dir, use_cache} instead."
            )
            # the old signature is ORTModelForCausalLM(model, config, use_io_binding, model_save_dir, preprocessors, generation_config, use_cache)
            session = args[0]
            if len(args) > 1:
                config = args[1]
            if len(args) > 2:
                use_io_binding = args[2]
            if len(args) > 3:
                model_save_dir = args[3]
            if len(args) > 4:
                _ = args[4]
            if len(args) > 5:
                generation_config = args[5]
            if len(args) > 6:
                _ = args[6]

        if kwargs.get("model", None) is not None:
            logger.warning(
                "Passing the inference session as `model` argument to an ORTModelForCausalLM is deprecated. Please use `session` instead."
            )
            session = kwargs.pop("model")
        if kwargs:
            logger.warning(
                f"Some keyword arguments were passed to the ORTModelForCausalLM constructor that are not part of its signature: {', '.join(kwargs.keys())}. "
                "These arguments will be ignored in the current version and will raise an error in the next version."
            )

        if config is None:
            raise ValueError(
                "The parameter config is required. Please pass a config or use the from_pretrained method."
            )
        if session is None:
            raise ValueError(
                "The parameter session is required. Please pass a session or use the from_pretrained method."
            )
        ## END OF DEPRECATED BEHAVIOR
        super().__init__(config=config, session=session, use_io_binding=use_io_binding, model_save_dir=model_save_dir)

        self.normalized_config = NormalizedConfigManager.get_normalized_config_class(config.model_type)(config)
        self.key_value_input_names = [key for key in self.input_names if (".key" in key) or (".value" in key)]
        self.key_value_output_names = [key for key in self.output_names if (".key" in key) or (".value" in key)]
        self.can_use_cache = len(self.key_value_input_names) > 0 and len(self.key_value_output_names) > 0
        self.is_merged = "use_cache_branch" in self.input_names
        self.generation_config = generation_config

        # Reference: https://github.com/huggingface/optimum/pull/1381
        model_type = self.config.model_type
        if model_type in MODEL_TYPES_REQUIRING_POSITION_IDS and "position_ids" not in self.input_names:
            logger.warning(
                f"ORTModelForCausalLM loaded a legacy ONNX model with no position_ids input, although the model type {model_type} "
                "requires it. for correct batched generation. We strongly encourage to re-export the model with "
                "a newer version of Optimum for better performance and more reliable generation. "
            )

        if not self.can_use_cache and self.generation_config.use_cache:
            logger.warning(
                "`model.generation_config.use_cache=True` but the loaded model does not support using the past key values cache."
                "Please re-export the original model once again with `use_cache=True` to be able to use it during generation. "
                "Or set `model.generation_config.use_cache=False` to avoid errors from attempting to use the cache. "
                "To re-export your model, simply set `export=True` as in `from_pretrained(..., export=True, use_cache=True)`."
            )

        if self.config.model_type == "gemma":
            self.embed_size_per_head = self.normalized_config.head_dim
        else:
            self.embed_size_per_head = self.normalized_config.hidden_size // self.normalized_config.num_attention_heads
        if self.config.model_type in {"gemma", "mistral", "llama", "qwen2", "qwen3", "qwen3_moe", "granite"}:
            self.num_key_value_heads = self.normalized_config.num_key_value_heads
        elif self.config.model_type == "falcon":
            self.num_key_value_heads = (
                self.config.num_kv_heads
                if (self.config.new_decoder_architecture or not self.config.multi_query)
                else 1
            )
        else:
            self.num_key_value_heads = self.normalized_config.num_attention_heads

    @property
    def use_cache(self):
        logger.warning(
            "The `ORTModelForCausalLM.use_cache` property is deprecated and will be removed in a future version. "
            "Please rather use `ORTModelForCausalLM.can_use_cache` to check if a model supports using cache during generation. "
            "And use `ORTModelForCausalLM.generation_config.use_cache` to check if the model is configured to use cache during generation."
        )
        return self.can_use_cache

    @property
    def use_merged(self):
        logger.warning(
            "The `ORTModelForCausalLM.use_merged` property is deprecated and will be removed in a future version. "
            "Please rather use `ORTModelForCausalLM.is_merged` to check if the underlying model is merged or not."
        )
        return self.is_merged

    @add_start_docstrings_to_model_forward(
        CAUSALLM_ONNX_MODEL_DOCSTRING.format("batch_size, sequence_length")
        + TEXT_GENERATION_EXAMPLE.format(
            processor_class=_TOKENIZER_FOR_DOC,
            model_class="ORTModelForCausalLM",
            checkpoint="optimum/gpt2",
        )
    )
    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
        position_ids: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        use_torch = isinstance(input_ids, torch.Tensor)
        self.raise_on_numpy_input_io_binding(use_torch)
        use_cache = use_cache if use_cache is not None else self.config.use_cache

        if use_cache and not self.can_use_cache:
            raise ValueError(
                f"`use_cache={use_cache}` was passed to the model but the loaded model only supports `use_cache={self.can_use_cache}`. "
                f"Please load your current model with `use_cache={self.can_use_cache}` or export the original model "
                f"once again with `use_cache={use_cache}` when calling the `from_pretrained` method. "
                "To re-export your model, simply set `export=True` in the `from_pretrained` method."
            )

        if past_key_values is not None and isinstance(past_key_values[0], tuple):
            # Flattens the past_key_values to a single tuple
            past_key_values = sum(past_key_values, ())

        if "position_ids" in self.input_names and position_ids is None:
            if attention_mask is not None:
                # Create position_ids from attention_mask
                position_ids = attention_mask.cumsum(-1) - 1
                position_ids.masked_fill_(attention_mask == 0, 1)
                if past_key_values is not None:
                    position_ids = position_ids[:, -1].unsqueeze(-1)
            else:
                raise ValueError(
                    "The model requires position_ids for batched generation but none were provided. "
                    "Please provide position_ids or attention_mask (from which position_ids can be inferred)."
                )

        use_cache_branch = None
        if self.is_merged:
            # Uses cache branch of merged decoders depending on whether real past key values are passed
            use_cache_branch = torch.full((1,), past_key_values is not None, dtype=torch.bool, device=self.device)

        if past_key_values is None and len(self.key_value_input_names) > 0:
            # Generates the input pkv for the first forward of the model (merged or with past)
            batch_size, seq_len = input_ids.shape
            if self.config.model_type == "gpt_bigcode":
                shape = (batch_size, 0, self.embed_size_per_head * 2)
            else:
                shape = (batch_size, self.num_key_value_heads, 0, self.embed_size_per_head)
            tensor = torch.empty(shape, dtype=self.dtype, device=self.device)
            past_key_values = tuple(tensor for _ in range(len(self.key_value_input_names)))

        model_inputs = {
            "input_ids": input_ids,
            "position_ids": position_ids,
            "attention_mask": attention_mask,
            "use_cache_branch": use_cache_branch,
        }
        if len(self.key_value_input_names) > 0:
            model_inputs.update(zip(self.key_value_input_names, past_key_values))

        known_output_shapes = None
        outputs_to_not_bind = None
        if use_cache:
            # Infers the shape of the output pkv
            batch_size, seq_len = input_ids.shape
            if self.config.model_type == "gpt_bigcode":
                pkv_seq_len, embed_size_per_head_2 = past_key_values[0].shape[1:]
                pkv_output_shape = (batch_size, pkv_seq_len + seq_len, embed_size_per_head_2)
            else:
                num_key_value_heads, pkv_seq_len, embed_size_per_head = past_key_values[0].shape[1:]
                pkv_output_shape = (batch_size, num_key_value_heads, pkv_seq_len + seq_len, embed_size_per_head)
            known_output_shapes = dict.fromkeys(self.key_value_output_names, pkv_output_shape)
        else:
            # Don't bind the output pkv if not used/returned
            outputs_to_not_bind = self.key_value_output_names

        if self.use_io_binding:
            output_shapes, output_buffers = self._prepare_io_binding(
                model_inputs,
                outputs_to_not_bind=outputs_to_not_bind,
                known_output_shapes=known_output_shapes,
            )

            if self.device.type == "cpu":
                self.session.run_with_iobinding(self._io_binding)
            else:
                self._io_binding.synchronize_inputs()
                self.session.run_with_iobinding(self._io_binding)
                self._io_binding.synchronize_outputs()

            loss = output_buffers.get("loss", None)
            logits = output_buffers["logits"].view(output_shapes["logits"])

            if use_cache:
                past_key_values = tuple(
                    output_buffers.pop(name).view(output_shapes[name]) for name in self.key_value_output_names
                )
        else:
            onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs)
            onnx_outputs = self.session.run(None, onnx_inputs)
            model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs)

            loss = model_outputs.pop("loss", None)
            logits = model_outputs.pop("logits")

            if use_cache:
                past_key_values = tuple(model_outputs.pop(name) for name in self.key_value_output_names)

        if use_cache and self.config.model_type != "gpt_bigcode":
            # Tuple of tuple of length `n_layers`, with each tuple of length equal to the number of self-attention and per decoder layer
            past_key_values = tuple(past_key_values[i : i + 2] for i in range(0, len(past_key_values), 2))

        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=past_key_values)

    def prepare_inputs_for_generation(self, *args, **kwargs):
        if is_transformers_version("<", "4.46.0"):
            return self._prepare_inputs_for_generation_legacy(*args, **kwargs)
        else:
            return super().prepare_inputs_for_generation(*args, **kwargs)

    # Adapted from transformers.models.gpt2.modeling_gpt2.GPT2LMHeadModel.prepare_inputs_for_generation
    def _prepare_inputs_for_generation_legacy(
        self,
        input_ids,
        attention_mask=None,
        past_key_values=None,
        token_type_ids=None,
        position_ids=None,
        use_cache=None,
        **kwargs,
    ):
        if past_key_values is not None:
            if self.config.model_type == "gpt_bigcode":
                if self.config.multi_query:
                    past_length = past_key_values[0].shape[1]
                else:
                    past_length = past_key_values[0].shape[2]
            else:
                past_length = past_key_values[0][0].shape[2]

            if input_ids.shape[1] > past_length:
                remove_prefix_length = past_length
            else:
                remove_prefix_length = input_ids.shape[1] - 1
            input_ids = input_ids[:, remove_prefix_length:]

        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "past_key_values": past_key_values,
            "token_type_ids": token_type_ids,
            "position_ids": position_ids,
            "use_cache": use_cache,
        }

    @staticmethod
    def _reorder_cache(
        past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor
    ) -> Tuple[Tuple[torch.Tensor]]:
        if isinstance(past_key_values, tuple) and isinstance(past_key_values[0], tuple):
            # GPT2 style
            return tuple(
                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
                for layer_past in past_key_values
            )
        elif isinstance(past_key_values, tuple) and isinstance(past_key_values[0], torch.Tensor):
            # GPT BigCode style
            return tuple(layer_past.index_select(0, beam_idx.to(layer_past.device)) for layer_past in past_key_values)
        else:
            raise ValueError(
                f"Unexpected past_key_values: {past_key_values}. "
                "Expected tuple of tuples (GPT2 style) or tuple of tensors (GPT BigCode style)."
            )

    @classmethod
    def _from_pretrained(
        cls,
        model_id: Union[str, Path],
        config: "PretrainedConfig",
        # hub options
        subfolder: str = "",
        revision: str = "main",
        force_download: bool = False,
        local_files_only: bool = False,
        trust_remote_code: bool = False,
        cache_dir: str = HUGGINGFACE_HUB_CACHE,
        token: Optional[Union[bool, str]] = None,
        # file options
        file_name: Optional[str] = None,
        # session options
        provider: str = "CPUExecutionProvider",
        providers: Optional[Sequence[str]] = None,
        provider_options: Optional[Union[Sequence[Dict[str, Any]], Dict[str, Any]]] = None,
        session_options: Optional[SessionOptions] = None,
        # inference options
        use_cache: bool = True,
        use_merged: Optional[bool] = None,
        use_io_binding: Optional[bool] = None,
        generation_config: Optional[GenerationConfig] = None,
        # other arguments
        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
    ) -> "ORTModelForCausalLM":
        onnx_files = find_files_matching_pattern(
            model_id,
            ONNX_FILE_PATTERN,
            glob_pattern="**/*.onnx",
            subfolder=subfolder,
            token=token,
            revision=revision,
        )

        if len(onnx_files) == 0:
            raise FileNotFoundError(f"Could not find any ONNX model file in {model_id}")

        if len(onnx_files) == 1:
            subfolder = onnx_files[0].parent
            _file_name = onnx_files[0].name
            if file_name and file_name != _file_name:
                raise FileNotFoundError(f"Trying to load {file_name} but only found {_file_name}")
            file_name = _file_name

        else:
            model_files = []
            # Check first for merged models and then for decoder / decoder_with_past models
            if use_merged is not False:
                model_files = [p for p in onnx_files if re.search(DECODER_MERGED_ONNX_FILE_PATTERN, str(p))]
                use_merged = len(model_files) != 0

            if use_merged is False:
                pattern = DECODER_WITH_PAST_ONNX_FILE_PATTERN if use_cache else DECODER_ONNX_FILE_PATTERN
                model_files = [p for p in onnx_files if re.search(pattern, str(p))]

            # if file_name is specified we don't filter legacy models
            if not model_files or file_name:
                model_files = onnx_files
            else:
                logger.warning(
                    f"Legacy models found in {model_files} will be loaded. "
                    "Legacy models will be deprecated in the next version of optimum, please re-export your model"
                )
            _file_name = model_files[0].name
            subfolder = model_files[0].parent

            defaut_file_name = file_name or "model.onnx"
            for file in model_files:
                if file.name == defaut_file_name:
                    _file_name = file.name
                    subfolder = file.parent
                    break

            file_name = _file_name

            if len(model_files) > 1:
                logger.warning(
                    f"Too many ONNX model files were found in {' ,'.join(map(str, model_files))}. "
                    "specify which one to load by using the `file_name` and/or the `subfolder` arguments. "
                    f"Loading the file {file_name} in the subfolder {subfolder}."
                )

        if os.path.isdir(model_id):
            model_id = subfolder
            subfolder = ""

        if isinstance(subfolder, Path):
            subfolder = subfolder.as_posix()

        model_cache_path = cached_file(
            model_id,
            filename=file_name,
            # hub options
            token=token,
            revision=revision,
            subfolder=subfolder,
            cache_dir=cache_dir,
            force_download=force_download,
            local_files_only=local_files_only,
        )

        # model_save_dir can be provided in kwargs as a TemporaryDirectory instance, in which case we want to keep it
        # instead of the path only.
        if model_save_dir is None:
            model_save_dir = Path(model_cache_path).parent

        try:
            cached_file(
                model_id,
                filename=file_name + "_data",
                # hub options
                token=token,
                revision=revision,
                subfolder=subfolder,
                cache_dir=cache_dir,
                force_download=force_download,
                local_files_only=local_files_only,
            )
        except EnvironmentError:
            # If the external data file is not found, we assume that the model is not using external data.
            pass

        # This should be removed at some point
        onnx_model = onnx.load(str(model_cache_path), load_external_data=False)
        model_uses_external_data = check_model_uses_external_data(onnx_model)
        if model_uses_external_data:
            onnx_model = onnx.load(str(model_cache_path), load_external_data=True)
        input_dims = {
            node.name: [dim.dim_value or dim.dim_param for dim in node.type.tensor_type.shape.dim]
            for node in onnx_model.graph.input
        }
        output_dims = {
            node.name: [dim.dim_value or dim.dim_param for dim in node.type.tensor_type.shape.dim]
            for node in onnx_model.graph.output
        }
        override_dims = False
        # Since v1.7.0 decoder with past models have fixed sequence length of 1
        # To keep these models compatible we set this dimension to dynamic
        if input_dims["input_ids"][1] == 1:
            input_dims["input_ids"][1] = "sequence_length"
            output_dims["logits"][1] = "sequence_length"
            override_dims = True
        # Since https://github.com/huggingface/optimum/pull/871/
        # changed axis notation/naming during export, we need to update the dims
        for input_name in input_dims.keys():
            if "past" in input_name and input_dims[input_name][2] == "past_sequence_length + sequence_length":
                input_dims[input_name][2] = "past_sequence_length"
                override_dims = True
        if override_dims:
            # this is kinda dangerous, warning the user is the least we can do
            logger.warning(
                "The ONNX model was probably exported with an older version of optimum. "
                "We are updating the input/output dimensions and overwriting the model file "
                "with new dimensions. This is necessary for the model to work correctly with "
                "the current version of optimum. If you encounter any issues, please re-export "
                "the model with the latest version of optimum for optimal performance."
            )
            onnx_model = update_model_dims.update_inputs_outputs_dims(onnx_model, input_dims, output_dims)
            onnx.save(
                onnx_model,
                str(model_cache_path),
                save_as_external_data=model_uses_external_data,
                location=Path(model_cache_path).name + "_data",
                all_tensors_to_one_file=True,
                convert_attribute=True,
                size_threshold=0,
            )
        del onnx_model

        # Important: for encoder-decoder models used with CausalLM, we need to set the is_decoder flag to True
        # and the is_encoder_decoder flag to False. This is needed for the model to work correctly with generation logic.
        if hasattr(config, "is_decoder"):
            config.is_decoder = True
        if hasattr(config, "is_encoder_decoder"):
            config.is_encoder_decoder = False

        if generation_config is None:
            try:
                generation_config = GenerationConfig.from_pretrained(
                    model_id,
                    token=token,
                    revision=revision,
                    subfolder=subfolder,
                    cache_dir=cache_dir,
                    force_download=force_download,
                    local_files_only=local_files_only,
                )
            except OSError:
                logger.info("Generation config file not found, creating a new one from model config.")
                generation_config = GenerationConfig.from_model_config(config)

        # TODO: not sure if setting config.use_cache is needed for older versions of transformers
        generation_config.use_cache = use_cache
        config.use_cache = use_cache

        if is_transformers_version(">=", "4.45.0"):
            misplaced_generation_parameters = config._get_non_default_generation_parameters()
            if len(misplaced_generation_parameters) > 0:
                logger.warning(
                    "Moving the following attributes in the config to the generation config: "
                    f"{misplaced_generation_parameters}. You are seeing this warning because you've set "
                    "generation parameters in the model config, as opposed to in the generation config.",
                )
                for param_name, param_value in misplaced_generation_parameters.items():
                    setattr(generation_config, param_name, param_value)
                    setattr(config, param_name, None)

        providers, provider_options = prepare_providers_and_provider_options(
            provider=provider, providers=providers, provider_options=provider_options
        )
        session = InferenceSession(
            model_cache_path,
            providers=providers,
            provider_options=provider_options,
            sess_options=session_options,
        )

        return cls(
            config=config,
            session=session,
            use_io_binding=use_io_binding,
            generation_config=generation_config,
            model_save_dir=model_save_dir,
        )

    @classmethod
    def _export(
        cls,
        model_id: Union[str, Path],
        config: "PretrainedConfig",
        # hub options
        subfolder: str = "",
        revision: str = "main",
        force_download: bool = False,
        local_files_only: bool = False,
        trust_remote_code: bool = False,
        cache_dir: str = HUGGINGFACE_HUB_CACHE,
        token: Optional[Union[bool, str]] = None,
        # inference options
        use_cache: bool = True,
        **kwargs,
    ) -> "ORTModelForCausalLM":
        # this is garanteed to work since we it uses a mapping from model classes to task names
        # instead of relying on the hub metadata or the model configuration
        task = TasksManager._infer_task_from_model_or_model_class(model_class=cls.auto_model_class)
        if use_cache:
            task += "-with-past"

        if kwargs.get("task", None) is not None:
            raise ValueError(
                f"The `task` argument is not needed when exporting a model with `{cls.__name__}`. "
                f"The `task` is automatically inferred from the class as `{task}`."
            )

        save_dir = TemporaryDirectory()
        save_dir_path = Path(save_dir.name)

        main_export(
            model_name_or_path=model_id,
            output=save_dir_path,
            task=task,
            do_validation=False,
            no_post_process=False,
            legacy=False,
            subfolder=subfolder,
            revision=revision,
            cache_dir=cache_dir,
            token=token,
            local_files_only=local_files_only,
            force_download=force_download,
            trust_remote_code=trust_remote_code,
        )
        maybe_save_preprocessors(model_id, save_dir_path, src_subfolder=subfolder)

        return cls._from_pretrained(
            save_dir_path,
            config,
            use_cache=use_cache,
            model_save_dir=save_dir,
            **kwargs,
        )

    def _save_config(self, save_directory):
        """
        Save the model and generation configs to the specified directory.

        Args:
            save_directory (`str` or `os.PathLike`):
                Directory where the model and generation configs will be saved.
        """
        self.config.save_pretrained(save_directory)
        self.generation_config.save_pretrained(save_directory)