





"""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


_snake_case	:						Tuple					=					logging.get_logger(__name__)

_snake_case	:						Dict					=					{
    '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__					):
					UpperCamelCase						=				'mbart'
					UpperCamelCase						=				['past_key_values']
					UpperCamelCase						=				{'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}


					def __init__(   self     :int	,			__UpperCamelCase     :List[Any]=5_02_65	,			__UpperCamelCase     :Optional[Any]=10_24	,			__UpperCamelCase     :Optional[int]=12	,			__UpperCamelCase     :Tuple=40_96	,			__UpperCamelCase     :str=16	,			__UpperCamelCase     :Union[str, Any]=12	,			__UpperCamelCase     :Tuple=40_96	,			__UpperCamelCase     :str=16	,			__UpperCamelCase     :int=0.0	,			__UpperCamelCase     :Optional[Any]=0.0	,			__UpperCamelCase     :Optional[Any]=True	,			__UpperCamelCase     :Union[str, Any]=True	,			__UpperCamelCase     :Tuple="gelu"	,			__UpperCamelCase     :List[str]=10_24	,			__UpperCamelCase     :Union[str, Any]=0.1	,			__UpperCamelCase     :Dict=0.0	,			__UpperCamelCase     :List[str]=0.0	,			__UpperCamelCase     :Any=0.02	,			__UpperCamelCase     :Any=0.0	,			__UpperCamelCase     :int=False	,			__UpperCamelCase     :Optional[Any]=1	,			__UpperCamelCase     :Optional[int]=0	,			__UpperCamelCase     :Union[str, Any]=2	,			__UpperCamelCase     :str=2	,			**__UpperCamelCase     :Optional[Any]	,			):
												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=lowercase_	,			bos_token_id=lowercase_	,			eos_token_id=lowercase_	,			is_encoder_decoder=lowercase_	,			forced_eos_token_id=lowercase_	,			**lowercase_	,			)







class    _UpperCAmelCase  (      lowerCamelCase__					):


					@property
					def       lowerCamelCase			(   self     :Optional[int]						):






												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_(lowercase_	,			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(lowercase_						):
																																	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     :str						):






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


					def       lowerCamelCase			(   self     :int	,			__UpperCamelCase     :List[str]	,			__UpperCamelCase     :Union[str, Any] = -1	,			__UpperCamelCase     :Any = -1	,			__UpperCamelCase     :Optional[int] = False	,			__UpperCamelCase     :Any = None	,			):
												A						=    self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
												    lowercase_	,			lowercase_	,			lowercase_	,			lowercase_	,			lowercase_						)

												# Generate decoder inputs
												A						=    seq_length if not self.use_past else 1
												A						=    self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
												    lowercase_	,			lowercase_	,			lowercase_	,			lowercase_	,			lowercase_						)
												A						=    {f"decoder_{name}": tensor for name, tensor in decoder_inputs.items()}
												A						=    dict(**lowercase_	,			**lowercase_						)

												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(lowercase_	,			lowercase_						)]	,			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(lowercase_	,			lowercase_						)
																			A						=    max(lowercase_	,			lowercase_						) - min_num_layers
																			A						=    "encoder" if num_encoder_layers > num_decoder_layers else "decoder"

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


					def       lowerCamelCase			(   self     :Union[str, Any]	,			__UpperCamelCase     :Any	,			__UpperCamelCase     :List[Any] = -1	,			__UpperCamelCase     :Tuple = -1	,			__UpperCamelCase     :Union[str, Any] = False	,			__UpperCamelCase     :Optional[Any] = None	,			):
												A						=    self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
												    lowercase_	,			lowercase_	,			lowercase_	,			lowercase_	,			lowercase_						)

												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(lowercase_	,			lowercase_	,			dtype=lowercase_						)]	,			dim=1						)
																			A						=    [
																			    (torch.zeros(lowercase_						), torch.zeros(lowercase_						)) for _ in range(lowercase_						)
																			]
												return common_inputs


					def       lowerCamelCase			(   self     :Any	,			__UpperCamelCase     :int	,			__UpperCamelCase     :int = -1	,			__UpperCamelCase     :List[Any] = -1	,			__UpperCamelCase     :Optional[Any] = False	,			__UpperCamelCase     :int = None	,			):
												A						=    compute_effective_axis_dimension(
												    lowercase_	,			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(lowercase_						)
												A						=    compute_effective_axis_dimension(
												    lowercase_	,			fixed_dimension=OnnxConfig.default_fixed_sequence	,			num_token_to_add=lowercase_						)

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


					def       lowerCamelCase			(   self     :Optional[int]	,			__UpperCamelCase     :Any	,			__UpperCamelCase     :Dict = -1	,			__UpperCamelCase     :Dict = -1	,			__UpperCamelCase     :Union[str, Any] = False	,			__UpperCamelCase     :int = None	,			):






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

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

												return common_inputs


					def       lowerCamelCase			(   self     :str	,			__UpperCamelCase     :int	,			__UpperCamelCase     :List[Any]	,			__UpperCamelCase     :Union[str, Any]	,			__UpperCamelCase     :Optional[Any]						):






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





import unittest

import numpy as np

from transformers import ElectraConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
       from transformers.models.electra.modeling_flax_electra import (
           FlaxElectraForCausalLM,
           FlaxElectraForMaskedLM,
           FlaxElectraForMultipleChoice,
           FlaxElectraForPreTraining,
           FlaxElectraForQuestionAnswering,
           FlaxElectraForSequenceClassification,
           FlaxElectraForTokenClassification,
           FlaxElectraModel,
       )







class 		__magic_name__     (				unittest.TestCase				):

   '''simple docstring'''




   def __init__(							self,       lowercase_,       lowercase_=13,       lowercase_=7,       lowercase_=True,       lowercase_=True,       lowercase_=True,       lowercase_=True,       lowercase_=99,       lowercase_=24,       lowercase_=32,       lowercase_=5,       lowercase_=4,       lowercase_=37,       lowercase_="gelu",       lowercase_=0.1,       lowercase_=0.1,       lowercase_=512,       lowercase_=16,       lowercase_=2,       lowercase_=0.02,       lowercase_=4,       )				->       int:






     """simple docstring"""






     a__    =parent
     a__    =batch_size
     a__    =seq_length
     a__    =is_training
     a__    =use_attention_mask
     a__    =use_token_type_ids
     a__    =use_labels
     a__    =vocab_size
     a__    =hidden_size
     a__    =embedding_size
     a__    =num_hidden_layers
     a__    =num_attention_heads
     a__    =intermediate_size
     a__    =hidden_act
     a__    =hidden_dropout_prob
     a__    =attention_probs_dropout_prob
     a__    =max_position_embeddings
     a__    =type_vocab_size
     a__    =type_sequence_label_size
     a__    =initializer_range
     a__    =num_choices



   def 					_UpperCAmelCase					(							self			)				->       List[Any]:






     """simple docstring"""






     a__    =ids_tensor([self.batch_size, self.seq_length],       self.vocab_size			)

     a__    =None
     if self.use_attention_mask:
       a__    =random_attention_mask([self.batch_size, self.seq_length]			)

     a__    =None
     if self.use_token_type_ids:
       a__    =ids_tensor([self.batch_size, self.seq_length],       self.type_vocab_size			)

     a__    =ElectraConfig(
         vocab_size=self.vocab_size,       hidden_size=self.hidden_size,       embedding_size=self.embedding_size,       num_hidden_layers=self.num_hidden_layers,       num_attention_heads=self.num_attention_heads,       intermediate_size=self.intermediate_size,       hidden_act=self.hidden_act,       hidden_dropout_prob=self.hidden_dropout_prob,       attention_probs_dropout_prob=self.attention_probs_dropout_prob,       max_position_embeddings=self.max_position_embeddings,       type_vocab_size=self.type_vocab_size,       initializer_range=self.initializer_range,       )

     return config, input_ids, token_type_ids, attention_mask



   def 					_UpperCAmelCase					(							self			)				->       List[Any]:






     """simple docstring"""






     a__    =self.prepare_config_and_inputs()
     a__, a__, a__, a__    =config_and_inputs
     a__    ={'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
     return config, inputs_dict







@require_flax
class 		__magic_name__     (				lowerCamelCase__						,			unittest.TestCase				):

   '''simple docstring'''




   lowerCamelCase__					:      List[str]        =  True

   lowerCamelCase__					:      Tuple        =  (
       (
           FlaxElectraModel,
           FlaxElectraForCausalLM,
           FlaxElectraForMaskedLM,
           FlaxElectraForPreTraining,
           FlaxElectraForTokenClassification,
           FlaxElectraForQuestionAnswering,
           FlaxElectraForMultipleChoice,
           FlaxElectraForSequenceClassification,
       )
       if is_flax_available()
       else ()
   )



   def 					_UpperCAmelCase					(							self			)				->       str:






     """simple docstring"""






     a__    =FlaxElectraModelTester(self			)



   @slow
   def 					_UpperCAmelCase					(							self			)				->       str:






     """simple docstring"""






     for model_class_name in self.all_model_classes:
       if model_class_name == FlaxElectraForMaskedLM:
         a__    =model_class_name.from_pretrained('''google/electra-small-generator'''			)
       else:
         a__    =model_class_name.from_pretrained('''google/electra-small-discriminator'''			)
       a__    =model(np.ones((1, 1)			)			)
       self.assertIsNotNone(lowercase_			)
