




from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


_lowerCamelCase							=      logging.get_logger(__name__)

_lowerCamelCase							=      {
    """facebook/mbart-large-cc25""": """https://huggingface.co/facebook/mbart-large-cc25/resolve/main/config.json""",
    # See all MBART models at https://huggingface.co/models?filter=mbart
}


class        UpperCAmelCase__				(							lowerCamelCase_				):



     '''simple docstring'''
     _SCREAMING_SNAKE_CASE      :					List[Any]     					=    '''mbart'''
     _SCREAMING_SNAKE_CASE      :					int     					=    ['''past_key_values''']
     _SCREAMING_SNAKE_CASE      :					List[str]     					=    {'''num_attention_heads''': '''encoder_attention_heads''', '''hidden_size''': '''d_model'''}
     def __init__(				self ,    _lowerCAmelCase=50_265 ,    _lowerCAmelCase=1_024 ,    _lowerCAmelCase=12 ,    _lowerCAmelCase=4_096 ,    _lowerCAmelCase=16 ,    _lowerCAmelCase=12 ,    _lowerCAmelCase=4_096 ,    _lowerCAmelCase=16 ,    _lowerCAmelCase=0.0 ,    _lowerCAmelCase=0.0 ,    _lowerCAmelCase=True ,    _lowerCAmelCase=True ,    _lowerCAmelCase="gelu" ,    _lowerCAmelCase=1_024 ,    _lowerCAmelCase=0.1 ,    _lowerCAmelCase=0.0 ,    _lowerCAmelCase=0.0 ,    _lowerCAmelCase=0.02 ,    _lowerCAmelCase=0.0 ,    _lowerCAmelCase=False ,    _lowerCAmelCase=1 ,    _lowerCAmelCase=0 ,    _lowerCAmelCase=2 ,    _lowerCAmelCase=2 ,    **_lowerCAmelCase ,    ):
            a			   =vocab_size
            a			   =max_position_embeddings
            a			   =d_model
            a			   =encoder_ffn_dim
            a			   =encoder_layers
            a			   =encoder_attention_heads
            a			   =decoder_ffn_dim
            a			   =decoder_layers
            a			   =decoder_attention_heads
            a			   =dropout
            a			   =attention_dropout
            a			   =activation_dropout
            a			   =activation_function
            a			   =init_std
            a			   =encoder_layerdrop
            a			   =decoder_layerdrop
            a			   =classifier_dropout
            a			   =use_cache
            a			   =encoder_layers
            a			   =scale_embedding  # scale factor will be sqrt(d_model) if True
            super().__init__(
                pad_token_id=__UpperCAmelCase ,    bos_token_id=__UpperCAmelCase ,    eos_token_id=__UpperCAmelCase ,    is_encoder_decoder=__UpperCAmelCase ,    forced_eos_token_id=__UpperCAmelCase ,    **__UpperCAmelCase ,    )



class        UpperCAmelCase__				(							lowerCamelCase_				):



     '''simple docstring'''
     @property
     def        lowerCAmelCase__  (				self	):



            if self.task in ["default", "seq2seq-lm"]:
                   a			   =OrderedDict(
                       [
                           ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
                           ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
                       ]	)

                   if self.use_past:
                          a			   ={0: """batch"""}
                          a			   ={0: """batch""", 1: """past_decoder_sequence + sequence"""}
                   else:
                          a			   ={0: """batch""", 1: """decoder_sequence"""}
                          a			   ={0: """batch""", 1: """decoder_sequence"""}

                   if self.use_past:
                          self.fill_with_past_key_values_(__UpperCAmelCase ,    direction="""inputs"""	)
            elif self.task == "causal-lm":
                   # TODO: figure this case out.
                   a			   =OrderedDict(
                       [
                           ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
                           ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
                       ]	)
                   if self.use_past:
                          a   ,      a			   =self.num_layers
                          for i in range(__UpperCAmelCase	):
                                 a			   ={0: """batch""", 2: """past_sequence + sequence"""}
                                 a			   ={0: """batch""", 2: """past_sequence + sequence"""}
            else:
                   a			   =OrderedDict(
                       [
                           ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
                           ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
                           ("""decoder_input_ids""", {0: """batch""", 1: """decoder_sequence"""}),
                           ("""decoder_attention_mask""", {0: """batch""", 1: """decoder_sequence"""}),
                       ]	)

            return common_inputs
     @property
     def        lowerCAmelCase__  (				self	):



            if self.task in ["default", "seq2seq-lm"]:
                   a			   =super().outputs
            else:
                   a			   =super(__UpperCAmelCase ,    self	).outputs
                   if self.use_past:
                          a   ,      a			   =self.num_layers
                          for i in range(__UpperCAmelCase	):
                                 a			   ={0: """batch""", 2: """past_sequence + sequence"""}
                                 a			   ={0: """batch""", 2: """past_sequence + sequence"""}
            return common_outputs
     def        lowerCAmelCase__  (				self ,    _lowerCAmelCase ,    _lowerCAmelCase = -1 ,    _lowerCAmelCase = -1 ,    _lowerCAmelCase = False ,    _lowerCAmelCase = None ,    ):
            a			   =self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
                __UpperCAmelCase ,    __UpperCAmelCase ,    __UpperCAmelCase ,    __UpperCAmelCase ,    __UpperCAmelCase	)

            # Generate decoder inputs
            a			   =seq_length if not self.use_past else 1
            a			   =self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
                __UpperCAmelCase ,    __UpperCAmelCase ,    __UpperCAmelCase ,    __UpperCAmelCase ,    __UpperCAmelCase	)
            a			   ={F'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
            a			   =dict(**__UpperCAmelCase ,    **__UpperCAmelCase	)

            if self.use_past:
                   if not is_torch_available():
                          raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""	)
                   else:
                          import torch
                   a   ,      a			   =common_inputs["""input_ids"""].shape
                   a			   =common_inputs["""decoder_input_ids"""].shape[1]
                   a   ,      a			   =self.num_attention_heads
                   a			   =(
                       batch,
                       num_encoder_attention_heads,
                       encoder_seq_length,
                       self._config.hidden_size // num_encoder_attention_heads,
                   )
                   a			   =decoder_seq_length + 3
                   a			   =(
                       batch,
                       num_decoder_attention_heads,
                       decoder_past_length,
                       self._config.hidden_size // num_decoder_attention_heads,
                   )

                   a			   =torch.cat(
                       [common_inputs["""decoder_attention_mask"""], torch.ones(__UpperCAmelCase ,    __UpperCAmelCase	)] ,    dim=1	)

                   a			   =[]
                   # If the number of encoder and decoder layers are present in the model configuration, both are considered
                   a   ,      a			   =self.num_layers
                   a			   =min(__UpperCAmelCase ,    __UpperCAmelCase	)
                   a			   =max(__UpperCAmelCase ,    __UpperCAmelCase	) - min_num_layers
                   a			   ="""encoder""" if num_encoder_layers > num_decoder_layers else """decoder"""

                   for _ in range(__UpperCAmelCase	):
                          common_inputs["past_key_values"].append(
                              (
                                  torch.zeros(__UpperCAmelCase	),
                                  torch.zeros(__UpperCAmelCase	),
                                  torch.zeros(__UpperCAmelCase	),
                                  torch.zeros(__UpperCAmelCase	),
                              )	)
                   # TODO: test this.
                   a			   =encoder_shape if remaining_side_name == """encoder""" else decoder_shape
                   for _ in range(__UpperCAmelCase ,    __UpperCAmelCase	):
                          common_inputs["past_key_values"].append((torch.zeros(__UpperCAmelCase	), torch.zeros(__UpperCAmelCase	))	)
            return common_inputs
     def        lowerCAmelCase__  (				self ,    _lowerCAmelCase ,    _lowerCAmelCase = -1 ,    _lowerCAmelCase = -1 ,    _lowerCAmelCase = False ,    _lowerCAmelCase = None ,    ):
            a			   =self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
                __UpperCAmelCase ,    __UpperCAmelCase ,    __UpperCAmelCase ,    __UpperCAmelCase ,    __UpperCAmelCase	)

            if self.use_past:
                   if not is_torch_available():
                          raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""	)
                   else:
                          import torch
                   a   ,      a			   =common_inputs["""input_ids"""].shape
                   # Not using the same length for past_key_values
                   a			   =seqlen + 2
                   a   ,      a			   =self.num_layers
                   a   ,      a			   =self.num_attention_heads
                   a			   =(
                       batch,
                       num_encoder_attention_heads,
                       past_key_values_length,
                       self._config.hidden_size // num_encoder_attention_heads,
                   )

                   a			   =common_inputs["""attention_mask"""].dtype
                   a			   =torch.cat(
                       [common_inputs["""attention_mask"""], torch.ones(__UpperCAmelCase ,    __UpperCAmelCase ,    dtype=__UpperCAmelCase	)] ,    dim=1	)
                   a			   =[
                       (torch.zeros(__UpperCAmelCase	), torch.zeros(__UpperCAmelCase	)) for _ in range(__UpperCAmelCase	)
                   ]
            return common_inputs
     def        lowerCAmelCase__  (				self ,    _lowerCAmelCase ,    _lowerCAmelCase = -1 ,    _lowerCAmelCase = -1 ,    _lowerCAmelCase = False ,    _lowerCAmelCase = None ,    ):
            a			   =compute_effective_axis_dimension(
                __UpperCAmelCase ,    fixed_dimension=OnnxConfig.default_fixed_batch ,    num_token_to_add=0	)

            # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
            a			   =tokenizer.num_special_tokens_to_add(__UpperCAmelCase	)
            a			   =compute_effective_axis_dimension(
                __UpperCAmelCase ,    fixed_dimension=OnnxConfig.default_fixed_sequence ,    num_token_to_add=__UpperCAmelCase	)

            # Generate dummy inputs according to compute batch and sequence
            a			   =[""" """.join([tokenizer.unk_token]	) * seq_length] * batch_size
            a			   =dict(tokenizer(__UpperCAmelCase ,    return_tensors=__UpperCAmelCase	)	)
            return common_inputs
     def        lowerCAmelCase__  (				self ,    _lowerCAmelCase ,    _lowerCAmelCase = -1 ,    _lowerCAmelCase = -1 ,    _lowerCAmelCase = False ,    _lowerCAmelCase = None ,    ):



            if self.task in ["default", "seq2seq-lm"]:
                   a			   =self._generate_dummy_inputs_for_default_and_seqaseq_lm(
                       __UpperCAmelCase ,    batch_size=__UpperCAmelCase ,    seq_length=__UpperCAmelCase ,    is_pair=__UpperCAmelCase ,    framework=__UpperCAmelCase	)

            elif self.task == "causal-lm":
                   a			   =self._generate_dummy_inputs_for_causal_lm(
                       __UpperCAmelCase ,    batch_size=__UpperCAmelCase ,    seq_length=__UpperCAmelCase ,    is_pair=__UpperCAmelCase ,    framework=__UpperCAmelCase	)
            else:
                   a			   =self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
                       __UpperCAmelCase ,    batch_size=__UpperCAmelCase ,    seq_length=__UpperCAmelCase ,    is_pair=__UpperCAmelCase ,    framework=__UpperCAmelCase	)

            return common_inputs






     def        lowerCAmelCase__  (				self ,    _lowerCAmelCase ,    _lowerCAmelCase ,    _lowerCAmelCase ,    _lowerCAmelCase	):



            if self.task in ["default", "seq2seq-lm"]:
                   a			   =super()._flatten_past_key_values_(__UpperCAmelCase ,    __UpperCAmelCase ,    __UpperCAmelCase ,    __UpperCAmelCase	)
            else:
                   a			   =super(__UpperCAmelCase ,    self	)._flatten_past_key_values_(
                       __UpperCAmelCase ,    __UpperCAmelCase ,    __UpperCAmelCase ,    __UpperCAmelCase	)






from typing import TYPE_CHECKING

from ..utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_tf_available, is_torch_available


__a			:	List[Any]							     =		{
    """configuration_utils""": ["""GenerationConfig"""],
    """streamers""": ["""TextIteratorStreamer""", """TextStreamer"""],
}

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__a			:	Dict							     =		[
				    """Constraint""",
				    """ConstraintListState""",
				    """DisjunctiveConstraint""",
				    """PhrasalConstraint""",
				]
				__a			:	int							     =		[
				    """BeamHypotheses""",
				    """BeamScorer""",
				    """BeamSearchScorer""",
				    """ConstrainedBeamSearchScorer""",
				]
				__a			:	List[Any]							     =		[
				    """EpsilonLogitsWarper""",
				    """EtaLogitsWarper""",
				    """ForcedBOSTokenLogitsProcessor""",
				    """ForcedEOSTokenLogitsProcessor""",
				    """HammingDiversityLogitsProcessor""",
				    """InfNanRemoveLogitsProcessor""",
				    """LogitsProcessor""",
				    """LogitsProcessorList""",
				    """LogitsWarper""",
				    """MinLengthLogitsProcessor""",
				    """MinNewTokensLengthLogitsProcessor""",
				    """NoBadWordsLogitsProcessor""",
				    """NoRepeatNGramLogitsProcessor""",
				    """PrefixConstrainedLogitsProcessor""",
				    """RepetitionPenaltyLogitsProcessor""",
				    """SequenceBiasLogitsProcessor""",
				    """EncoderRepetitionPenaltyLogitsProcessor""",
				    """TemperatureLogitsWarper""",
				    """TopKLogitsWarper""",
				    """TopPLogitsWarper""",
				    """TypicalLogitsWarper""",
				    """EncoderNoRepeatNGramLogitsProcessor""",
				    """ExponentialDecayLengthPenalty""",
				    """LogitNormalization""",
				]
				__a			:	Any							     =		[
				    """MaxNewTokensCriteria""",
				    """MaxLengthCriteria""",
				    """MaxTimeCriteria""",
				    """StoppingCriteria""",
				    """StoppingCriteriaList""",
				    """validate_stopping_criteria""",
				]
				__a			:	str							     =		[
				    """GenerationMixin""",
				    """top_k_top_p_filtering""",
				    """GreedySearchEncoderDecoderOutput""",
				    """GreedySearchDecoderOnlyOutput""",
				    """SampleEncoderDecoderOutput""",
				    """SampleDecoderOnlyOutput""",
				    """BeamSearchEncoderDecoderOutput""",
				    """BeamSearchDecoderOnlyOutput""",
				    """BeamSampleEncoderDecoderOutput""",
				    """BeamSampleDecoderOnlyOutput""",
				    """ContrastiveSearchEncoderDecoderOutput""",
				    """ContrastiveSearchDecoderOnlyOutput""",
				]

try:
				if not is_tf_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__a			:	List[Any]							     =		[
				    """TFForcedBOSTokenLogitsProcessor""",
				    """TFForcedEOSTokenLogitsProcessor""",
				    """TFLogitsProcessor""",
				    """TFLogitsProcessorList""",
				    """TFLogitsWarper""",
				    """TFMinLengthLogitsProcessor""",
				    """TFNoBadWordsLogitsProcessor""",
				    """TFNoRepeatNGramLogitsProcessor""",
				    """TFRepetitionPenaltyLogitsProcessor""",
				    """TFTemperatureLogitsWarper""",
				    """TFTopKLogitsWarper""",
				    """TFTopPLogitsWarper""",
				    """TFForceTokensLogitsProcessor""",
				    """TFSuppressTokensAtBeginLogitsProcessor""",
				    """TFSuppressTokensLogitsProcessor""",
				]
				__a			:	Any							     =		[
				    """TFGenerationMixin""",
				    """tf_top_k_top_p_filtering""",
				    """TFGreedySearchDecoderOnlyOutput""",
				    """TFGreedySearchEncoderDecoderOutput""",
				    """TFSampleEncoderDecoderOutput""",
				    """TFSampleDecoderOnlyOutput""",
				    """TFBeamSearchEncoderDecoderOutput""",
				    """TFBeamSearchDecoderOnlyOutput""",
				    """TFBeamSampleEncoderDecoderOutput""",
				    """TFBeamSampleDecoderOnlyOutput""",
				    """TFContrastiveSearchEncoderDecoderOutput""",
				    """TFContrastiveSearchDecoderOnlyOutput""",
				]

try:
				if not is_flax_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__a			:	str							     =		[
				    """FlaxForcedBOSTokenLogitsProcessor""",
				    """FlaxForcedEOSTokenLogitsProcessor""",
				    """FlaxLogitsProcessor""",
				    """FlaxLogitsProcessorList""",
				    """FlaxLogitsWarper""",
				    """FlaxMinLengthLogitsProcessor""",
				    """FlaxTemperatureLogitsWarper""",
				    """FlaxTopKLogitsWarper""",
				    """FlaxTopPLogitsWarper""",
				]
				__a			:	Tuple							     =		[
				    """FlaxGenerationMixin""",
				    """FlaxGreedySearchOutput""",
				    """FlaxSampleOutput""",
				    """FlaxBeamSearchOutput""",
				]

if TYPE_CHECKING:
				from .configuration_utils import GenerationConfig
				from .streamers import TextIteratorStreamer, TextStreamer

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .beam_constraints import Constraint, ConstraintListState, DisjunctiveConstraint, PhrasalConstraint
								from .beam_search import BeamHypotheses, BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer
								from .logits_process import (
								    EncoderNoRepeatNGramLogitsProcessor,
								    EncoderRepetitionPenaltyLogitsProcessor,
								    EpsilonLogitsWarper,
								    EtaLogitsWarper,
								    ExponentialDecayLengthPenalty,
								    ForcedBOSTokenLogitsProcessor,
								    ForcedEOSTokenLogitsProcessor,
								    HammingDiversityLogitsProcessor,
								    InfNanRemoveLogitsProcessor,
								    LogitNormalization,
								    LogitsProcessor,
								    LogitsProcessorList,
								    LogitsWarper,
								    MinLengthLogitsProcessor,
								    MinNewTokensLengthLogitsProcessor,
								    NoBadWordsLogitsProcessor,
								    NoRepeatNGramLogitsProcessor,
								    PrefixConstrainedLogitsProcessor,
								    RepetitionPenaltyLogitsProcessor,
								    SequenceBiasLogitsProcessor,
								    TemperatureLogitsWarper,
								    TopKLogitsWarper,
								    TopPLogitsWarper,
								    TypicalLogitsWarper,
								)
								from .stopping_criteria import (
								    MaxLengthCriteria,
								    MaxNewTokensCriteria,
								    MaxTimeCriteria,
								    StoppingCriteria,
								    StoppingCriteriaList,
								    validate_stopping_criteria,
								)
								from .utils import (
								    BeamSampleDecoderOnlyOutput,
								    BeamSampleEncoderDecoderOutput,
								    BeamSearchDecoderOnlyOutput,
								    BeamSearchEncoderDecoderOutput,
								    ContrastiveSearchDecoderOnlyOutput,
								    ContrastiveSearchEncoderDecoderOutput,
								    GenerationMixin,
								    GreedySearchDecoderOnlyOutput,
								    GreedySearchEncoderDecoderOutput,
								    SampleDecoderOnlyOutput,
								    SampleEncoderDecoderOutput,
								    top_k_top_p_filtering,
								)

				try:
								if not is_tf_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .tf_logits_process import (
								    TFForcedBOSTokenLogitsProcessor,
								    TFForcedEOSTokenLogitsProcessor,
								    TFForceTokensLogitsProcessor,
								    TFLogitsProcessor,
								    TFLogitsProcessorList,
								    TFLogitsWarper,
								    TFMinLengthLogitsProcessor,
								    TFNoBadWordsLogitsProcessor,
								    TFNoRepeatNGramLogitsProcessor,
								    TFRepetitionPenaltyLogitsProcessor,
								    TFSuppressTokensAtBeginLogitsProcessor,
								    TFSuppressTokensLogitsProcessor,
								    TFTemperatureLogitsWarper,
								    TFTopKLogitsWarper,
								    TFTopPLogitsWarper,
								)
								from .tf_utils import (
								    TFBeamSampleDecoderOnlyOutput,
								    TFBeamSampleEncoderDecoderOutput,
								    TFBeamSearchDecoderOnlyOutput,
								    TFBeamSearchEncoderDecoderOutput,
								    TFContrastiveSearchDecoderOnlyOutput,
								    TFContrastiveSearchEncoderDecoderOutput,
								    TFGenerationMixin,
								    TFGreedySearchDecoderOnlyOutput,
								    TFGreedySearchEncoderDecoderOutput,
								    TFSampleDecoderOnlyOutput,
								    TFSampleEncoderDecoderOutput,
								    tf_top_k_top_p_filtering,
								)

				try:
								if not is_flax_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .flax_logits_process import (
								    FlaxForcedBOSTokenLogitsProcessor,
								    FlaxForcedEOSTokenLogitsProcessor,
								    FlaxLogitsProcessor,
								    FlaxLogitsProcessorList,
								    FlaxLogitsWarper,
								    FlaxMinLengthLogitsProcessor,
								    FlaxTemperatureLogitsWarper,
								    FlaxTopKLogitsWarper,
								    FlaxTopPLogitsWarper,
								)
								from .flax_utils import FlaxBeamSearchOutput, FlaxGenerationMixin, FlaxGreedySearchOutput, FlaxSampleOutput
else:
				import sys

				__a			:	List[Any]							     =		_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
