




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


snake_case__ :						str     =			logging.get_logger(__name__)

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     snake_case  (	_snake_case							):


   '''simple docstring'''







   UpperCamelCase__	:       str     =      "mbart"
   UpperCamelCase__	:       Union[str, Any]     =      ["past_key_values"]
   UpperCamelCase__	:       Tuple     =      {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}




   def __init__(					self     :	Any						,    lowerCamelCase_     :	Union[str, Any]=5_0265						,    lowerCamelCase_     :	Optional[Any]=1024						,    lowerCamelCase_     :	Optional[int]=12						,    lowerCamelCase_     :	List[str]=4096						,    lowerCamelCase_     :	str=16						,    lowerCamelCase_     :	Union[str, Any]=12						,    lowerCamelCase_     :	Any=4096						,    lowerCamelCase_     :	Optional[int]=16						,    lowerCamelCase_     :	Dict=0.0						,    lowerCamelCase_     :	Optional[Any]=0.0						,    lowerCamelCase_     :	Dict=True						,    lowerCamelCase_     :	Optional[int]=True						,    lowerCamelCase_     :	int="gelu"						,    lowerCamelCase_     :	str=1024						,    lowerCamelCase_     :	List[Any]=0.1						,    lowerCamelCase_     :	str=0.0						,    lowerCamelCase_     :	int=0.0						,    lowerCamelCase_     :	str=0.02						,    lowerCamelCase_     :	int=0.0						,    lowerCamelCase_     :	List[str]=False						,    lowerCamelCase_     :	Union[str, Any]=1						,    lowerCamelCase_     :	Optional[Any]=0						,    lowerCamelCase_     :	List[Any]=2						,    lowerCamelCase_     :	Any=2						,    **lowerCamelCase_     :	List[Any]						,    )					->Optional[Any]:
     '''simple docstring'''







     UpperCAmelCase__    						=   vocab_size
     UpperCAmelCase__    						=   max_position_embeddings
     UpperCAmelCase__    						=   d_model
     UpperCAmelCase__    						=   encoder_ffn_dim
     UpperCAmelCase__    						=   encoder_layers
     UpperCAmelCase__    						=   encoder_attention_heads
     UpperCAmelCase__    						=   decoder_ffn_dim
     UpperCAmelCase__    						=   decoder_layers
     UpperCAmelCase__    						=   decoder_attention_heads
     UpperCAmelCase__    						=   dropout
     UpperCAmelCase__    						=   attention_dropout
     UpperCAmelCase__    						=   activation_dropout
     UpperCAmelCase__    						=   activation_function
     UpperCAmelCase__    						=   init_std
     UpperCAmelCase__    						=   encoder_layerdrop
     UpperCAmelCase__    						=   decoder_layerdrop
     UpperCAmelCase__    						=   classifier_dropout
     UpperCAmelCase__    						=   use_cache
     UpperCAmelCase__    						=   encoder_layers
     UpperCAmelCase__    						=   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     snake_case  (	_snake_case							):


   '''simple docstring'''







   @property
   def 	UpperCAmelCase       (					self     :	Tuple			)					->Mapping[str, Mapping[int, str]]:
     '''simple docstring'''







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

       if self.use_past:
         UpperCAmelCase__    						=   {0: """batch"""}
         UpperCAmelCase__    						=   {0: """batch""", 1: """past_decoder_sequence + sequence"""}
       else:
         UpperCAmelCase__    						=   {0: """batch""", 1: """decoder_sequence"""}
         UpperCAmelCase__    						=   {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.
       UpperCAmelCase__    						=   OrderedDict(
           [
               ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
               ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
           ]			)
       if self.use_past:
         UpperCAmelCase__						,					UpperCAmelCase__    						=   self.num_layers
         for i in range(lowerCamelCase_			):
           UpperCAmelCase__    						=   {0: """batch""", 2: """past_sequence + sequence"""}
           UpperCAmelCase__    						=   {0: """batch""", 2: """past_sequence + sequence"""}
     else:
       UpperCAmelCase__    						=   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 	UpperCAmelCase       (					self     :	Dict			)					->Mapping[str, Mapping[int, str]]:
     '''simple docstring'''







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




   def 	UpperCAmelCase       (					self     :	Optional[Any]						,    lowerCamelCase_     :	PreTrainedTokenizer						,    lowerCamelCase_     :	int = -1						,    lowerCamelCase_     :	int = -1						,    lowerCamelCase_     :	bool = False						,    lowerCamelCase_     :	Optional[TensorType] = None						,    )					->Mapping[str, Any]:
     '''simple docstring'''







     UpperCAmelCase__    						=   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
         lowerCamelCase_						,    lowerCamelCase_						,    lowerCamelCase_						,    lowerCamelCase_						,    lowerCamelCase_			)

     # Generate decoder inputs
     UpperCAmelCase__    						=   seq_length if not self.use_past else 1
     UpperCAmelCase__    						=   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
         lowerCamelCase_						,    lowerCamelCase_						,    lowerCamelCase_						,    lowerCamelCase_						,    lowerCamelCase_			)
     UpperCAmelCase__    						=   {f'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
     UpperCAmelCase__    						=   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
       UpperCAmelCase__						,					UpperCAmelCase__    						=   common_inputs["""input_ids"""].shape
       UpperCAmelCase__    						=   common_inputs["""decoder_input_ids"""].shape[1]
       UpperCAmelCase__						,					UpperCAmelCase__    						=   self.num_attention_heads
       UpperCAmelCase__    						=   (
           batch,
           num_encoder_attention_heads,
           encoder_seq_length,
           self._config.hidden_size // num_encoder_attention_heads,
       )
       UpperCAmelCase__    						=   decoder_seq_length + 3
       UpperCAmelCase__    						=   (
           batch,
           num_decoder_attention_heads,
           decoder_past_length,
           self._config.hidden_size // num_decoder_attention_heads,
       )

       UpperCAmelCase__    						=   torch.cat(
           [common_inputs["""decoder_attention_mask"""], torch.ones(lowerCamelCase_						,    lowerCamelCase_			)]						,    dim=1			)

       UpperCAmelCase__    						=   []
       # If the number of encoder and decoder layers are present in the model configuration, both are considered
       UpperCAmelCase__						,					UpperCAmelCase__    						=   self.num_layers
       UpperCAmelCase__    						=   min(lowerCamelCase_						,    lowerCamelCase_			)
       UpperCAmelCase__    						=   max(lowerCamelCase_						,    lowerCamelCase_			) - min_num_layers
       UpperCAmelCase__    						=   """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.
       UpperCAmelCase__    						=   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 	UpperCAmelCase       (					self     :	Union[str, Any]						,    lowerCamelCase_     :	PreTrainedTokenizer						,    lowerCamelCase_     :	int = -1						,    lowerCamelCase_     :	int = -1						,    lowerCamelCase_     :	bool = False						,    lowerCamelCase_     :	Optional[TensorType] = None						,    )					->Mapping[str, Any]:
     '''simple docstring'''







     UpperCAmelCase__    						=   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
       UpperCAmelCase__						,					UpperCAmelCase__    						=   common_inputs["""input_ids"""].shape
       # Not using the same length for past_key_values
       UpperCAmelCase__    						=   seqlen + 2
       UpperCAmelCase__						,					UpperCAmelCase__    						=   self.num_layers
       UpperCAmelCase__						,					UpperCAmelCase__    						=   self.num_attention_heads
       UpperCAmelCase__    						=   (
           batch,
           num_encoder_attention_heads,
           past_key_values_length,
           self._config.hidden_size // num_encoder_attention_heads,
       )

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




   def 	UpperCAmelCase       (					self     :	Any						,    lowerCamelCase_     :	PreTrainedTokenizer						,    lowerCamelCase_     :	int = -1						,    lowerCamelCase_     :	int = -1						,    lowerCamelCase_     :	bool = False						,    lowerCamelCase_     :	Optional[TensorType] = None						,    )					->Mapping[str, Any]:
     '''simple docstring'''







     UpperCAmelCase__    						=   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
     UpperCAmelCase__    						=   tokenizer.num_special_tokens_to_add(lowerCamelCase_			)
     UpperCAmelCase__    						=   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
     UpperCAmelCase__    						=   [""" """.join([tokenizer.unk_token]			) * seq_length] * batch_size
     UpperCAmelCase__    						=   dict(tokenizer(lowerCamelCase_						,    return_tensors=lowerCamelCase_			)			)
     return common_inputs




   def 	UpperCAmelCase       (					self     :	List[Any]						,    lowerCamelCase_     :	PreTrainedTokenizer						,    lowerCamelCase_     :	int = -1						,    lowerCamelCase_     :	int = -1						,    lowerCamelCase_     :	bool = False						,    lowerCamelCase_     :	Optional[TensorType] = None						,    )					->Mapping[str, Any]:
     '''simple docstring'''







     if self.task in ["default", "seq2seq-lm"]:
       UpperCAmelCase__    						=   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":
       UpperCAmelCase__    						=   self._generate_dummy_inputs_for_causal_lm(
           lowerCamelCase_						,    batch_size=lowerCamelCase_						,    seq_length=lowerCamelCase_						,    is_pair=lowerCamelCase_						,    framework=lowerCamelCase_			)
     else:
       UpperCAmelCase__    						=   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 	UpperCAmelCase       (					self     :	Any						,    lowerCamelCase_     :	str						,    lowerCamelCase_     :	Dict						,    lowerCamelCase_     :	Union[str, Any]						,    lowerCamelCase_     :	int			)					->Dict:
     '''simple docstring'''







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










import unittest

import numpy as np

from transformers import is_tf_available
from transformers.testing_utils import require_tf


if is_tf_available():
   import tensorflow as tf

   from transformers.activations_tf import get_tf_activation





@require_tf
class     snake_case  (	unittest.TestCase							):


   '''simple docstring'''







   def 	UpperCAmelCase       (					self     :	Tuple			)					->Union[str, Any]:
     '''simple docstring'''







     UpperCAmelCase__    						=   tf.constant([-100, -1.0, -0.1, 0, 0.1, 1.0, 100.0]			)
     UpperCAmelCase__    						=   get_tf_activation("""gelu"""			)
     UpperCAmelCase__    						=   get_tf_activation("""gelu_10"""			)

     UpperCAmelCase__    						=   gelu(lowerCamelCase_			)
     UpperCAmelCase__    						=   geluaa(lowerCamelCase_			)

     UpperCAmelCase__    						=   tf.where(y_gelu_aa < 10.0						,    1.0						,    0.0			)

     self.assertEqual(tf.math.reduce_max(lowerCamelCase_			).numpy().item()						,    10.0			)
     self.assertTrue(np.allclose(y_gelu * clipped_mask						,    y_gelu_aa * clipped_mask			)			)







   def 	UpperCAmelCase       (					self     :	List[Any]			)					->Any:
     '''simple docstring'''







     get_tf_activation("""gelu"""			)
     get_tf_activation("""gelu_10"""			)
     get_tf_activation("""gelu_fast"""			)
     get_tf_activation("""gelu_new"""			)
     get_tf_activation("""glu"""			)
     get_tf_activation("""mish"""			)
     get_tf_activation("""quick_gelu"""			)
     get_tf_activation("""relu"""			)
     get_tf_activation("""sigmoid"""			)
     get_tf_activation("""silu"""			)
     get_tf_activation("""swish"""			)
     get_tf_activation("""tanh"""			)
     with self.assertRaises(lowerCamelCase_			):
       get_tf_activation("""bogus"""			)
     with self.assertRaises(lowerCamelCase_			):
       get_tf_activation(lowerCamelCase_			)




