



"""simple docstring"""







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


lowercase__		:				int          =	logging.get_logger(__name__)

lowercase__		:				str          =	{
    """Helsinki-NLP/opus-mt-en-de""": """https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/config.json""",
    # See all Marian models at https://huggingface.co/models?filter=marian
}







class       UpperCamelCase__   (	lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     """marian"""
      _SCREAMING_SNAKE_CASE  				=     ["""past_key_values"""]
      _SCREAMING_SNAKE_CASE  				=     {"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}



      def __init__(      self				:      Union[str, Any]    ,			SCREAMING_SNAKE_CASE_				:      Any=5_8_1_0_1    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]=None    ,			SCREAMING_SNAKE_CASE_				:      str=1_0_2_4    ,			SCREAMING_SNAKE_CASE_				:      List[str]=1_2    ,			SCREAMING_SNAKE_CASE_				:      Tuple=4_0_9_6    ,			SCREAMING_SNAKE_CASE_				:      Optional[Any]=1_6    ,			SCREAMING_SNAKE_CASE_				:      Tuple=1_2    ,			SCREAMING_SNAKE_CASE_				:      List[Any]=4_0_9_6    ,			SCREAMING_SNAKE_CASE_				:      Dict=1_6    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]=0.0    ,			SCREAMING_SNAKE_CASE_				:      List[Any]=0.0    ,			SCREAMING_SNAKE_CASE_				:      Dict=True    ,			SCREAMING_SNAKE_CASE_				:      List[str]=True    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]="gelu"    ,			SCREAMING_SNAKE_CASE_				:      Any=1_0_2_4    ,			SCREAMING_SNAKE_CASE_				:      Optional[Any]=0.1    ,			SCREAMING_SNAKE_CASE_				:      str=0.0    ,			SCREAMING_SNAKE_CASE_				:      List[Any]=0.0    ,			SCREAMING_SNAKE_CASE_				:      Any=0.02    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]=5_8_1_0_0    ,			SCREAMING_SNAKE_CASE_				:      List[Any]=False    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]=5_8_1_0_0    ,			SCREAMING_SNAKE_CASE_				:      List[str]=0    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]=0    ,			SCREAMING_SNAKE_CASE_				:      List[Any]=True    ,			**SCREAMING_SNAKE_CASE_				:      Union[str, Any]    ,			):
        lowerCAmelCase_   :				List[Any]								= vocab_size
        lowerCAmelCase_   :				Union[str, Any]								= decoder_vocab_size or vocab_size
        lowerCAmelCase_   :				Any								= max_position_embeddings
        lowerCAmelCase_   :				Any								= d_model
        lowerCAmelCase_   :				Optional[int]								= encoder_ffn_dim
        lowerCAmelCase_   :				Optional[int]								= encoder_layers
        lowerCAmelCase_   :				Union[str, Any]								= encoder_attention_heads
        lowerCAmelCase_   :				Optional[int]								= decoder_ffn_dim
        lowerCAmelCase_   :				int								= decoder_layers
        lowerCAmelCase_   :				Dict								= decoder_attention_heads
        lowerCAmelCase_   :				int								= dropout
        lowerCAmelCase_   :				int								= attention_dropout
        lowerCAmelCase_   :				Any								= activation_dropout
        lowerCAmelCase_   :				Optional[int]								= activation_function
        lowerCAmelCase_   :				List[str]								= init_std
        lowerCAmelCase_   :				Optional[Any]								= encoder_layerdrop
        lowerCAmelCase_   :				Dict								= decoder_layerdrop
        lowerCAmelCase_   :				str								= use_cache
        lowerCAmelCase_   :				List[str]								= encoder_layers
        lowerCAmelCase_   :				Union[str, Any]								= scale_embedding  # scale factor will be sqrt(d_model) if True
        lowerCAmelCase_   :				int								= share_encoder_decoder_embeddings
        super().__init__(
            pad_token_id=SCREAMING_SNAKE_CASE_    ,			eos_token_id=SCREAMING_SNAKE_CASE_    ,			is_encoder_decoder=SCREAMING_SNAKE_CASE_    ,			decoder_start_token_id=SCREAMING_SNAKE_CASE_    ,			forced_eos_token_id=SCREAMING_SNAKE_CASE_    ,			**SCREAMING_SNAKE_CASE_    ,			)







class       UpperCamelCase__   (	lowercase_  ):




      """simple docstring"""




      @property
      # Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.inputs
      def      SCREAMING_SNAKE_CASE__ (      self				:      Optional[Any]       ):
        if self.task in ["default", "seq2seq-lm"]:
          lowerCAmelCase_   :				Optional[int]								= OrderedDict(
              [
                  ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
                  ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
              ]       )

          if self.use_past:
            lowerCAmelCase_   :				List[Any]								= {0: 'batch'}
            lowerCAmelCase_   :				str								= {0: 'batch', 1: 'past_decoder_sequence + sequence'}
          else:
            lowerCAmelCase_   :				Any								= {0: 'batch', 1: 'decoder_sequence'}
            lowerCAmelCase_   :				Dict								= {0: 'batch', 1: 'decoder_sequence'}

          if self.use_past:
            self.fill_with_past_key_values_(SCREAMING_SNAKE_CASE_    ,			direction='inputs'       )
        elif self.task == "causal-lm":
          # TODO: figure this case out.
          lowerCAmelCase_   :				Optional[int]								= OrderedDict(
              [
                  ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
                  ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
              ]       )
          if self.use_past:
            lowerCAmelCase_						,lowerCAmelCase_   :				Union[str, Any]								= self.num_layers
            for i in range(SCREAMING_SNAKE_CASE_       ):
              lowerCAmelCase_   :				List[Any]								= {0: 'batch', 2: 'past_sequence + sequence'}
              lowerCAmelCase_   :				str								= {0: 'batch', 2: 'past_sequence + sequence'}
        else:
          lowerCAmelCase_   :				Union[str, Any]								= 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
      # Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.outputs
      def      SCREAMING_SNAKE_CASE__ (      self				:      int       ):
        if self.task in ["default", "seq2seq-lm"]:
          lowerCAmelCase_   :				Optional[int]								= super().outputs
        else:
          lowerCAmelCase_   :				Any								= super(SCREAMING_SNAKE_CASE_    ,			self       ).outputs
          if self.use_past:
            lowerCAmelCase_						,lowerCAmelCase_   :				List[str]								= self.num_layers
            for i in range(SCREAMING_SNAKE_CASE_       ):
              lowerCAmelCase_   :				Optional[Any]								= {0: 'batch', 2: 'past_sequence + sequence'}
              lowerCAmelCase_   :				List[str]								= {0: 'batch', 2: 'past_sequence + sequence'}
        return common_outputs



      def      SCREAMING_SNAKE_CASE__ (      self				:      List[Any]    ,			SCREAMING_SNAKE_CASE_				:      PreTrainedTokenizer    ,			SCREAMING_SNAKE_CASE_				:      int = -1    ,			SCREAMING_SNAKE_CASE_				:      int = -1    ,			SCREAMING_SNAKE_CASE_				:      bool = False    ,			SCREAMING_SNAKE_CASE_				:      Optional[TensorType] = None    ,			):
        lowerCAmelCase_   :				List[Any]								= self._generate_dummy_inputs_for_encoder_and_decoder(
            SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )

        # Generate decoder inputs
        lowerCAmelCase_   :				Union[str, Any]								= seq_length if not self.use_past else 1
        lowerCAmelCase_   :				Union[str, Any]								= self._generate_dummy_inputs_for_encoder_and_decoder(
            SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				str								= {F"decoder_{name}": tensor for name, tensor in decoder_inputs.items()}
        lowerCAmelCase_   :				List[str]								= dict(**SCREAMING_SNAKE_CASE_    ,			**SCREAMING_SNAKE_CASE_       )

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

          lowerCAmelCase_   :				Union[str, Any]								= torch.cat(
              [common_inputs['decoder_attention_mask'], torch.ones(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )]    ,			dim=1       )

          lowerCAmelCase_   :				Union[str, Any]								= []
          # If the number of encoder and decoder layers are present in the model configuration, both are considered
          lowerCAmelCase_						,lowerCAmelCase_   :				str								= self.num_layers
          lowerCAmelCase_   :				List[Any]								= min(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )
          lowerCAmelCase_   :				Optional[Any]								= max(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       ) - min_num_layers
          lowerCAmelCase_   :				str								= 'encoder' if num_encoder_layers > num_decoder_layers else 'decoder'

          for _ in range(SCREAMING_SNAKE_CASE_       ):
            common_inputs["past_key_values"].append(
                (
                    torch.zeros(SCREAMING_SNAKE_CASE_       ),
                    torch.zeros(SCREAMING_SNAKE_CASE_       ),
                    torch.zeros(SCREAMING_SNAKE_CASE_       ),
                    torch.zeros(SCREAMING_SNAKE_CASE_       ),
                )       )
          # TODO: test this.
          lowerCAmelCase_   :				Optional[int]								= encoder_shape if remaining_side_name == 'encoder' else decoder_shape
          for _ in range(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       ):
            common_inputs["past_key_values"].append((torch.zeros(SCREAMING_SNAKE_CASE_       ), torch.zeros(SCREAMING_SNAKE_CASE_       ))       )
        return common_inputs



      def      SCREAMING_SNAKE_CASE__ (      self				:      List[Any]    ,			SCREAMING_SNAKE_CASE_				:      PreTrainedTokenizer    ,			SCREAMING_SNAKE_CASE_				:      int = -1    ,			SCREAMING_SNAKE_CASE_				:      int = -1    ,			SCREAMING_SNAKE_CASE_				:      bool = False    ,			SCREAMING_SNAKE_CASE_				:      Optional[TensorType] = None    ,			):
        lowerCAmelCase_   :				str								= self._generate_dummy_inputs_for_encoder_and_decoder(
            SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )

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

          lowerCAmelCase_   :				int								= common_inputs['attention_mask'].dtype
          lowerCAmelCase_   :				List[str]								= torch.cat(
              [common_inputs['attention_mask'], torch.ones(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			dtype=SCREAMING_SNAKE_CASE_       )]    ,			dim=1       )
          lowerCAmelCase_   :				str								= [
              (torch.zeros(SCREAMING_SNAKE_CASE_       ), torch.zeros(SCREAMING_SNAKE_CASE_       )) for _ in range(SCREAMING_SNAKE_CASE_       )
          ]
        return common_inputs



      def      SCREAMING_SNAKE_CASE__ (      self				:      List[str]    ,			SCREAMING_SNAKE_CASE_				:      PreTrainedTokenizer    ,			SCREAMING_SNAKE_CASE_				:      int = -1    ,			SCREAMING_SNAKE_CASE_				:      int = -1    ,			SCREAMING_SNAKE_CASE_				:      bool = False    ,			SCREAMING_SNAKE_CASE_				:      Optional[TensorType] = None    ,			):
        # Copied from OnnxConfig.generate_dummy_inputs
        # Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
        # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
        lowerCAmelCase_   :				List[Any]								= compute_effective_axis_dimension(
            SCREAMING_SNAKE_CASE_    ,			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
        lowerCAmelCase_   :				List[Any]								= tokenizer.num_special_tokens_to_add(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				List[Any]								= compute_effective_axis_dimension(
            SCREAMING_SNAKE_CASE_    ,			fixed_dimension=OnnxConfig.default_fixed_sequence    ,			num_token_to_add=SCREAMING_SNAKE_CASE_       )

        # Generate dummy inputs according to compute batch and sequence
        lowerCAmelCase_   :				str								= [' '.join([tokenizer.unk_token]       ) * seq_length] * batch_size
        lowerCAmelCase_   :				List[str]								= dict(tokenizer(SCREAMING_SNAKE_CASE_    ,			return_tensors=SCREAMING_SNAKE_CASE_       )       )
        return common_inputs



      def      SCREAMING_SNAKE_CASE__ (      self				:      Dict    ,			SCREAMING_SNAKE_CASE_				:      PreTrainedTokenizer    ,			SCREAMING_SNAKE_CASE_				:      int = -1    ,			SCREAMING_SNAKE_CASE_				:      int = -1    ,			SCREAMING_SNAKE_CASE_				:      bool = False    ,			SCREAMING_SNAKE_CASE_				:      Optional[TensorType] = None    ,			):
        if self.task in ["default", "seq2seq-lm"]:
          lowerCAmelCase_   :				Dict								= self._generate_dummy_inputs_for_default_and_seqaseq_lm(
              SCREAMING_SNAKE_CASE_    ,			batch_size=SCREAMING_SNAKE_CASE_    ,			seq_length=SCREAMING_SNAKE_CASE_    ,			is_pair=SCREAMING_SNAKE_CASE_    ,			framework=SCREAMING_SNAKE_CASE_       )

        else:
          lowerCAmelCase_   :				Any								= self._generate_dummy_inputs_for_causal_lm(
              SCREAMING_SNAKE_CASE_    ,			batch_size=SCREAMING_SNAKE_CASE_    ,			seq_length=SCREAMING_SNAKE_CASE_    ,			is_pair=SCREAMING_SNAKE_CASE_    ,			framework=SCREAMING_SNAKE_CASE_       )

        return common_inputs



      def      SCREAMING_SNAKE_CASE__ (      self				:      Any    ,			SCREAMING_SNAKE_CASE_				:      Optional[Any]    ,			SCREAMING_SNAKE_CASE_				:      int    ,			SCREAMING_SNAKE_CASE_				:      Tuple    ,			SCREAMING_SNAKE_CASE_				:      Tuple       ):
        if self.task in ["default", "seq2seq-lm"]:
          lowerCAmelCase_   :				List[str]								= super()._flatten_past_key_values_(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )
        else:
          lowerCAmelCase_   :				Optional[Any]								= super(SCREAMING_SNAKE_CASE_    ,			self       )._flatten_past_key_values_(
              SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )





      @property
      def      SCREAMING_SNAKE_CASE__ (      self				:      List[str]       ):
        return 1E-4










"""simple docstring"""







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


lowercase__		:				int          =	{"""configuration_mbart""": ["""MBART_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MBartConfig""", """MBartOnnxConfig"""]}

try:
     if not is_sentencepiece_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     lowercase__		:				Optional[int]          =	["""MBartTokenizer"""]

try:
     if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     lowercase__		:				Dict          =	["""MBartTokenizerFast"""]

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     lowercase__		:				Any          =	[
         """MBART_PRETRAINED_MODEL_ARCHIVE_LIST""",
         """MBartForCausalLM""",
         """MBartForConditionalGeneration""",
         """MBartForQuestionAnswering""",
         """MBartForSequenceClassification""",
         """MBartModel""",
         """MBartPreTrainedModel""",
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     lowercase__		:				Tuple          =	[
         """TFMBartForConditionalGeneration""",
         """TFMBartModel""",
         """TFMBartPreTrainedModel""",
     ]

try:
     if not is_flax_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     lowercase__		:				Tuple          =	[
         """FlaxMBartForConditionalGeneration""",
         """FlaxMBartForQuestionAnswering""",
         """FlaxMBartForSequenceClassification""",
         """FlaxMBartModel""",
         """FlaxMBartPreTrainedModel""",
     ]


if TYPE_CHECKING:
     from .configuration_mbart import MBART_PRETRAINED_CONFIG_ARCHIVE_MAP, MBartConfig, MBartOnnxConfig

     try:
          if not is_sentencepiece_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_mbart import MBartTokenizer

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_mbart_fast import MBartTokenizerFast

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_mbart import (
              MBART_PRETRAINED_MODEL_ARCHIVE_LIST,
              MBartForCausalLM,
              MBartForConditionalGeneration,
              MBartForQuestionAnswering,
              MBartForSequenceClassification,
              MBartModel,
              MBartPreTrainedModel,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_mbart import TFMBartForConditionalGeneration, TFMBartModel, TFMBartPreTrainedModel

     try:
          if not is_flax_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_flax_mbart import (
              FlaxMBartForConditionalGeneration,
              FlaxMBartForQuestionAnswering,
              FlaxMBartForSequenceClassification,
              FlaxMBartModel,
              FlaxMBartPreTrainedModel,
          )

else:
     import sys

     lowercase__		:				int          =	_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)





