


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


SCREAMING_SNAKE_CASE		:Any											=       logging.get_logger(__name__)

SCREAMING_SNAKE_CASE		:int											=       {
    '''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       __lowerCAmelCase       (	_a   ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE						      =    'mbart'
  _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]      ,			_lowerCAmelCase    :    Optional[Any]=5_0_2_6_5      ,			_lowerCAmelCase    :    str=1_0_2_4      ,			_lowerCAmelCase    :    str=1_2      ,			_lowerCAmelCase    :    Any=4_0_9_6      ,			_lowerCAmelCase    :    str=1_6      ,			_lowerCAmelCase    :    Optional[int]=1_2      ,			_lowerCAmelCase    :    Optional[Any]=4_0_9_6      ,			_lowerCAmelCase    :    int=1_6      ,			_lowerCAmelCase    :    Union[str, Any]=0.0      ,			_lowerCAmelCase    :    List[str]=0.0      ,			_lowerCAmelCase    :    int=True      ,			_lowerCAmelCase    :    List[Any]=True      ,			_lowerCAmelCase    :    Union[str, Any]="gelu"      ,			_lowerCAmelCase    :    List[Any]=1_0_2_4      ,			_lowerCAmelCase    :    Optional[Any]=0.1      ,			_lowerCAmelCase    :    List[str]=0.0      ,			_lowerCAmelCase    :    str=0.0      ,			_lowerCAmelCase    :    str=0.02      ,			_lowerCAmelCase    :    int=0.0      ,			_lowerCAmelCase    :    int=False      ,			_lowerCAmelCase    :    Dict=1      ,			_lowerCAmelCase    :    Tuple=0      ,			_lowerCAmelCase    :    List[str]=2      ,			_lowerCAmelCase    :    Union[str, Any]=2      ,			**_lowerCAmelCase    :    Any      ,			)				->	Optional[int]:




      """simple docstring"""
      snake_case_         =   vocab_size
      snake_case_         =   max_position_embeddings
      snake_case_         =   d_model
      snake_case_         =   encoder_ffn_dim
      snake_case_         =   encoder_layers
      snake_case_         =   encoder_attention_heads
      snake_case_         =   decoder_ffn_dim
      snake_case_         =   decoder_layers
      snake_case_         =   decoder_attention_heads
      snake_case_         =   dropout
      snake_case_         =   attention_dropout
      snake_case_         =   activation_dropout
      snake_case_         =   activation_function
      snake_case_         =   init_std
      snake_case_         =   encoder_layerdrop
      snake_case_         =   decoder_layerdrop
      snake_case_         =   classifier_dropout
      snake_case_         =   use_cache
      snake_case_         =   encoder_layers
      snake_case_         =   scale_embedding  # scale factor will be sqrt(d_model) if True
      super().__init__(
          pad_token_id=_lowerCAmelCase      ,			bos_token_id=_lowerCAmelCase      ,			eos_token_id=_lowerCAmelCase      ,			is_encoder_decoder=_lowerCAmelCase      ,			forced_eos_token_id=_lowerCAmelCase      ,			**_lowerCAmelCase      ,			)






class       __lowerCAmelCase       (	_a   ):






  """simple docstring"""





  @property
  def    lowerCAmelCase__					(       self    :    List[Any]					)				->	List[str]:




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

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

          if self.use_past:
              self.fill_with_past_key_values_(_lowerCAmelCase      ,			direction="inputs"					)
      elif self.task == "causal-lm":
          # TODO: figure this case out.
          snake_case_         =   OrderedDict(
              [
                  ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                  ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
              ]					)
          if self.use_past:
              snake_case_      ,      snake_case_         =   self.num_layers
              for i in range(_lowerCAmelCase					):
                  snake_case_         =   {0: "batch", 2: "past_sequence + sequence"}
                  snake_case_         =   {0: "batch", 2: "past_sequence + sequence"}
      else:
          snake_case_         =   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    :    Optional[Any]					)				->	Dict:




      """simple docstring"""
      if self.task in ["default", "seq2seq-lm"]:
          snake_case_         =   super().outputs
      else:
          snake_case_         =   super(_lowerCAmelCase      ,			self					).outputs
          if self.use_past:
              snake_case_      ,      snake_case_         =   self.num_layers
              for i in range(_lowerCAmelCase					):
                  snake_case_         =   {0: "batch", 2: "past_sequence + sequence"}
                  snake_case_         =   {0: "batch", 2: "past_sequence + sequence"}
      return common_outputs



  def    lowerCAmelCase__					(       self    :    Tuple      ,			_lowerCAmelCase    :    PreTrainedTokenizer      ,			_lowerCAmelCase    :    int = -1      ,			_lowerCAmelCase    :    int = -1      ,			_lowerCAmelCase    :    bool = False      ,			_lowerCAmelCase    :    Optional[TensorType] = None      ,			)				->	Any:




      """simple docstring"""
      snake_case_         =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
          _lowerCAmelCase      ,			_lowerCAmelCase      ,			_lowerCAmelCase      ,			_lowerCAmelCase      ,			_lowerCAmelCase					)

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

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

          snake_case_         =   torch.cat(
              [common_inputs["decoder_attention_mask"], torch.ones(_lowerCAmelCase      ,			_lowerCAmelCase					)]      ,			dim=1					)

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

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



  def    lowerCAmelCase__					(       self    :    Dict      ,			_lowerCAmelCase    :    PreTrainedTokenizer      ,			_lowerCAmelCase    :    int = -1      ,			_lowerCAmelCase    :    int = -1      ,			_lowerCAmelCase    :    bool = False      ,			_lowerCAmelCase    :    Optional[TensorType] = None      ,			)				->	Tuple:




      """simple docstring"""
      snake_case_         =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
          _lowerCAmelCase      ,			_lowerCAmelCase      ,			_lowerCAmelCase      ,			_lowerCAmelCase      ,			_lowerCAmelCase					)

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

          snake_case_         =   common_inputs["attention_mask"].dtype
          snake_case_         =   torch.cat(
              [common_inputs["attention_mask"], torch.ones(_lowerCAmelCase      ,			_lowerCAmelCase      ,			dtype=_lowerCAmelCase					)]      ,			dim=1					)
          snake_case_         =   [
              (torch.zeros(_lowerCAmelCase					), torch.zeros(_lowerCAmelCase					)) for _ in range(_lowerCAmelCase					)
          ]
      return common_inputs



  def    lowerCAmelCase__					(       self    :    str      ,			_lowerCAmelCase    :    PreTrainedTokenizer      ,			_lowerCAmelCase    :    int = -1      ,			_lowerCAmelCase    :    int = -1      ,			_lowerCAmelCase    :    bool = False      ,			_lowerCAmelCase    :    Optional[TensorType] = None      ,			)				->	Dict:




      """simple docstring"""
      # 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
      snake_case_         =   compute_effective_axis_dimension(
          _lowerCAmelCase      ,			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
      snake_case_         =   tokenizer.num_special_tokens_to_add(_lowerCAmelCase					)
      snake_case_         =   compute_effective_axis_dimension(
          _lowerCAmelCase      ,			fixed_dimension=OnnxConfig.default_fixed_sequence      ,			num_token_to_add=_lowerCAmelCase					)

      # Generate dummy inputs according to compute batch and sequence
      snake_case_         =   [" ".join([tokenizer.unk_token]					) * seq_length] * batch_size
      snake_case_         =   dict(tokenizer(_lowerCAmelCase      ,			return_tensors=_lowerCAmelCase					)					)
      return common_inputs



  def    lowerCAmelCase__					(       self    :    int      ,			_lowerCAmelCase    :    PreTrainedTokenizer      ,			_lowerCAmelCase    :    int = -1      ,			_lowerCAmelCase    :    int = -1      ,			_lowerCAmelCase    :    bool = False      ,			_lowerCAmelCase    :    Optional[TensorType] = None      ,			)				->	Tuple:




      """simple docstring"""
      if self.task in ["default", "seq2seq-lm"]:
          snake_case_         =   self._generate_dummy_inputs_for_default_and_seqaseq_lm(
              _lowerCAmelCase      ,			batch_size=_lowerCAmelCase      ,			seq_length=_lowerCAmelCase      ,			is_pair=_lowerCAmelCase      ,			framework=_lowerCAmelCase					)

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

      return common_inputs



  def    lowerCAmelCase__					(       self    :    List[str]      ,			_lowerCAmelCase    :    Optional[Any]      ,			_lowerCAmelCase    :    List[Any]      ,			_lowerCAmelCase    :    List[Any]      ,			_lowerCAmelCase    :    int					)				->	Dict:




      """simple docstring"""
      if self.task in ["default", "seq2seq-lm"]:
          snake_case_         =   super()._flatten_past_key_values_(_lowerCAmelCase      ,			_lowerCAmelCase      ,			_lowerCAmelCase      ,			_lowerCAmelCase					)
      else:
          snake_case_         =   super(_lowerCAmelCase      ,			self					)._flatten_past_key_values_(
              _lowerCAmelCase      ,			_lowerCAmelCase      ,			_lowerCAmelCase      ,			_lowerCAmelCase					)








import math
from collections import OrderedDict

import torch
from packaging import version
from torch import Tensor, nn

from .utils import logging


SCREAMING_SNAKE_CASE		:Tuple											=       logging.get_logger(__name__)


class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def __init__(       self    :    Any					)				->	List[str]:




     """simple docstring"""
     super().__init__()
     if version.parse(torch.__version__					) < version.parse("1.12.0"					):
        raise ImportError(
            F'''You are using torch=={torch.__version__}, but torch>=1.12.0 is required to use '''
            "PytorchGELUTanh. Please upgrade torch."					)



  def    lowerCAmelCase__					(       self    :    Optional[int]      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return nn.functional.gelu(_lowerCAmelCase      ,			approximate="tanh"					)






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def    lowerCAmelCase__					(       self    :    int      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi					) * (input + 0.044_715 * torch.pow(_lowerCAmelCase      ,			3.0					))					))






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def __init__(       self    :    int      ,			_lowerCAmelCase    :    bool = False					)				->	int:




     """simple docstring"""
     super().__init__()
     if use_gelu_python:
        snake_case_         =   self._gelu_python
     else:
        snake_case_         =   nn.functional.gelu



  def    lowerCAmelCase__					(       self    :    Optional[Any]      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return input * 0.5 * (1.0 + torch.erf(input / math.sqrt(2.0					)					))



  def    lowerCAmelCase__					(       self    :    Dict      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return self.act(_lowerCAmelCase					)






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def    lowerCAmelCase__					(       self    :    Optional[int]      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return 0.5 * input * (1.0 + torch.tanh(input * 0.7_978_845_608 * (1.0 + 0.044_715 * input * input)					))






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def    lowerCAmelCase__					(       self    :    int      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return input * torch.sigmoid(1.702 * input					)






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def __init__(       self    :    List[str]      ,			_lowerCAmelCase    :    float      ,			_lowerCAmelCase    :    float					)				->	Any:




     """simple docstring"""
     if min > max:
        raise ValueError(F'''min should be < max (got min: {min}, max: {max})'''					)

     super().__init__()
     snake_case_         =   min
     snake_case_         =   max



  def    lowerCAmelCase__					(       self    :    Union[str, Any]      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return torch.clip(gelu(_lowerCAmelCase					)      ,			self.min      ,			self.max					)






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def __init__(       self    :    Any					)				->	str:




     """simple docstring"""
     super().__init__()
     snake_case_         =   math.sqrt(2 / math.pi					)



  def    lowerCAmelCase__					(       self    :    Optional[Any]      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return 0.5 * input * (1 + torch.tanh(self.precomputed_constant * (input + 0.044_715 * torch.pow(_lowerCAmelCase      ,			3					))					))






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def    lowerCAmelCase__					(       self    :    str      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return nn.functional.silu(_lowerCAmelCase					)






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def __init__(       self    :    Optional[int]					)				->	List[Any]:




     """simple docstring"""
     super().__init__()
     if version.parse(torch.__version__					) < version.parse("1.9.0"					):
        snake_case_         =   self._mish_python
     else:
        snake_case_         =   nn.functional.mish



  def    lowerCAmelCase__					(       self    :    Tuple      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return input * torch.tanh(nn.functional.softplus(_lowerCAmelCase					)					)



  def    lowerCAmelCase__					(       self    :    Optional[Any]      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return self.act(_lowerCAmelCase					)






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def    lowerCAmelCase__					(       self    :    List[str]      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return input






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def    lowerCAmelCase__					(       self    :    List[Any]      ,			_lowerCAmelCase    :    List[Any]      ,			_lowerCAmelCase    :    Union[str, Any]=0.707_107      ,			_lowerCAmelCase    :    int=0.282_095					)				->	List[str]:




     """simple docstring"""
     snake_case_         =   (input - mu).div(sigma * math.sqrt(2.0					)					)
     return 0.5 * (1.0 + torch.erf(_lowerCAmelCase					))






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def    lowerCAmelCase__					(       self    :    Any      ,			_lowerCAmelCase    :    Union[str, Any]					)				->	Any:




     """simple docstring"""
     snake_case_         =   nn.functional.relu(_lowerCAmelCase					)
     snake_case_         =   torch.square(_lowerCAmelCase					)
     return squared






class       __lowerCAmelCase       (	a   ):






  """simple docstring"""





  def __getitem__(       self    :    Optional[Any]      ,			_lowerCAmelCase    :    str					)				->	Optional[int]:




     """simple docstring"""
     snake_case_         =   super().__getitem__(_lowerCAmelCase					)
     snake_case_      ,      snake_case_         =   content if isinstance(_lowerCAmelCase      ,			_lowerCAmelCase					) else (content, {})
     return cls(**_lowerCAmelCase					)






SCREAMING_SNAKE_CASE		:Optional[int]											=       {
    '''gelu''': GELUActivation,
    '''gelu_10''': (ClippedGELUActivation, {'''min''': -10, '''max''': 10}),
    '''gelu_fast''': FastGELUActivation,
    '''gelu_new''': NewGELUActivation,
    '''gelu_python''': (GELUActivation, {'''use_gelu_python''': True}),
    '''gelu_pytorch_tanh''': PytorchGELUTanh,
    '''gelu_accurate''': AccurateGELUActivation,
    '''laplace''': LaplaceActivation,
    '''linear''': LinearActivation,
    '''mish''': MishActivation,
    '''quick_gelu''': QuickGELUActivation,
    '''relu''': nn.ReLU,
    '''relu2''': ReLUSquaredActivation,
    '''relu6''': nn.ReLUa,
    '''sigmoid''': nn.Sigmoid,
    '''silu''': SiLUActivation,
    '''swish''': SiLUActivation,
    '''tanh''': nn.Tanh,
}
SCREAMING_SNAKE_CASE		:Optional[int]											=       ClassInstantier(ACTaCLS)



def      _lowerCAmelCase  (							lowerCAmelCase_							:Optional[Any]							)->Union[str, Any]:
   '''simple docstring'''






   if activation_string in ACTaFN:
      return ACTaFN[activation_string]
   else:
      raise KeyError(F'''function {activation_string} not found in ACT2FN mapping {list(ACTaFN.keys()							)}'''							)


# For backwards compatibility with: from activations import gelu_python
SCREAMING_SNAKE_CASE		:Tuple											=       get_activation('''gelu_python''')
SCREAMING_SNAKE_CASE		:int											=       get_activation('''gelu_new''')
SCREAMING_SNAKE_CASE		:List[Any]											=       get_activation('''gelu''')
SCREAMING_SNAKE_CASE		:int											=       get_activation('''gelu_fast''')
SCREAMING_SNAKE_CASE		:Tuple											=       get_activation('''quick_gelu''')
SCREAMING_SNAKE_CASE		:Tuple											=       get_activation('''silu''')
SCREAMING_SNAKE_CASE		:Any											=       get_activation('''mish''')



SCREAMING_SNAKE_CASE		:Optional[Any]											=       get_activation('''linear''')




