






"""simple docstring"""







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

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


lowerCAmelCase_		:			Tuple				   =					logging.get_logger(__name__)

lowerCAmelCase_		:			Tuple				   =					{
    '''facebook/blenderbot-3B''': '''https://huggingface.co/facebook/blenderbot-3B/resolve/main/config.json''',
    # See all Blenderbot models at https://huggingface.co/models?filter=blenderbot
}





class      UpperCamelCase_       (				a_     ):
  _A				:    Tuple =		'blenderbot'
  _A				:    Dict =		['past_key_values']
  _A				:    int =		{'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}



  def __init__(					self       ,						snake_case__=80_08       ,						snake_case__=1_28       ,						snake_case__=2       ,						snake_case__=1_02_40       ,						snake_case__=32       ,						snake_case__=24       ,						snake_case__=1_02_40       ,						snake_case__=32       ,						snake_case__=0.0       ,						snake_case__=0.0       ,						snake_case__=True       ,						snake_case__=True       ,						snake_case__="gelu"       ,						snake_case__=25_60       ,						snake_case__=0.1       ,						snake_case__=0.0       ,						snake_case__=0.0       ,						snake_case__=0.02       ,						snake_case__=1       ,						snake_case__=False       ,						snake_case__=0       ,						snake_case__=1       ,						snake_case__=2       ,						snake_case__=3       ,						snake_case__=2       ,						**snake_case__       ,						)	->	int:



         """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					=	use_cache
         UpperCAmelCase					=	encoder_layers
         UpperCAmelCase					=	scale_embedding  # scale factor will be sqrt(d_model) if True

         super().__init__(
             pad_token_id=snake_case__       ,						bos_token_id=snake_case__       ,						eos_token_id=snake_case__       ,						is_encoder_decoder=snake_case__       ,						decoder_start_token_id=snake_case__       ,						encoder_no_repeat_ngram_size=snake_case__       ,						forced_eos_token_id=snake_case__       ,						**snake_case__       ,						)





class      UpperCamelCase_       (				a_     ):



  @property
  def 						UpperCamelCase_ (					self )	->	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_(snake_case__       ,						direction="""inputs""" )
         elif self.task == "causal-lm":
                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(snake_case__ ):
                              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
  # Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.outputs
  def 						UpperCamelCase_ (					self )	->	Mapping[str, Mapping[int, str]]:



         """simple docstring"""





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



  def 						UpperCamelCase_ (					self       ,						snake_case__       ,						snake_case__ = -1       ,						snake_case__ = -1       ,						snake_case__ = False       ,						snake_case__ = None       ,						)	->	Mapping[str, Any]:



         """simple docstring"""





         UpperCAmelCase					=	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
             snake_case__       ,						snake_case__       ,						snake_case__       ,						snake_case__       ,						snake_case__ )
         # Generate decoder inputs
         UpperCAmelCase					=	seq_length if not self.use_past else 1
         UpperCAmelCase					=	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
             snake_case__       ,						snake_case__       ,						snake_case__       ,						snake_case__       ,						snake_case__ )
         UpperCAmelCase					=	{f'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
         UpperCAmelCase					=	dict(**snake_case__       ,						**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
                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
                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(snake_case__       ,						snake_case__ )]       ,						dim=1 )
                UpperCAmelCase					=	[]
                UpperCAmelCase			,				UpperCAmelCase					=	self.num_layers

                for _ in range(snake_case__ ):
                       common_inputs["past_key_values"].append(
                           (
                               torch.zeros(snake_case__ ),
                               torch.zeros(snake_case__ ),
                               torch.zeros(snake_case__ ),
                               torch.zeros(snake_case__ ),
                           ) )
         return common_inputs



  def 						UpperCamelCase_ (					self       ,						snake_case__       ,						snake_case__ = -1       ,						snake_case__ = -1       ,						snake_case__ = False       ,						snake_case__ = None       ,						)	->	Mapping[str, Any]:



         """simple docstring"""





         UpperCAmelCase					=	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
             snake_case__       ,						snake_case__       ,						snake_case__       ,						snake_case__       ,						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
                UpperCAmelCase			,				UpperCAmelCase					=	common_inputs["""input_ids"""].shape
                UpperCAmelCase					=	seqlen
                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(snake_case__       ,						snake_case__       ,						dtype=snake_case__ )]       ,						dim=1 )
                UpperCAmelCase					=	[
                    (torch.zeros(snake_case__ ), torch.zeros(snake_case__ )) for _ in range(snake_case__ )
                ]
         return common_inputs



  def 						UpperCamelCase_ (					self       ,						snake_case__       ,						snake_case__ = -1       ,						snake_case__ = -1       ,						snake_case__ = False       ,						snake_case__ = None       ,						)	->	Mapping[str, Any]:



         """simple docstring"""





         UpperCAmelCase					=	compute_effective_axis_dimension(
             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
         UpperCAmelCase					=	tokenizer.num_special_tokens_to_add(snake_case__ )
         UpperCAmelCase					=	compute_effective_axis_dimension(
             snake_case__       ,						fixed_dimension=OnnxConfig.default_fixed_sequence       ,						num_token_to_add=snake_case__ )

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



  def 						UpperCamelCase_ (					self       ,						snake_case__       ,						snake_case__ = -1       ,						snake_case__ = -1       ,						snake_case__ = False       ,						snake_case__ = None       ,						)	->	Mapping[str, Any]:



         """simple docstring"""





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

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

         return common_inputs



  def 						UpperCamelCase_ (					self       ,						snake_case__       ,						snake_case__       ,						snake_case__       ,						snake_case__ )	->	Any:



         """simple docstring"""





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



  def 						UpperCamelCase_ (					self       ,						snake_case__       ,						snake_case__ )	->	int:



         """simple docstring"""





         if direction not in ["inputs", "outputs"]:
                raise ValueError(f'''direction must either be "inputs" or "outputs", but {direction} was given''' )

         UpperCAmelCase					=	"""past_key_values""" if direction == """inputs""" else """present"""
         UpperCAmelCase			,				UpperCAmelCase					=	self.num_layers

         UpperCAmelCase					=	"""past_encoder_sequence"""
         UpperCAmelCase					=	"""past_decoder_sequence""" if direction == """inputs""" else """past_decoder_sequence + sequence"""

         for i in range(snake_case__ ):
                UpperCAmelCase					=	{0: """batch""", 2: decoder_sequence}
                UpperCAmelCase					=	{0: """batch""", 2: decoder_sequence}
                UpperCAmelCase					=	{0: """batch""", 2: encoder_sequence}
                UpperCAmelCase					=	{0: """batch""", 2: encoder_sequence}








"""simple docstring"""







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      UpperCamelCase_       (				unittest.TestCase     ):



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



         """simple docstring"""





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



  def 						UpperCamelCase_ (					self )	->	Optional[Any]:



         """simple docstring"""





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

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

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

         UpperCAmelCase					=	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[str]:



         """simple docstring"""





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





@require_flax
class      UpperCamelCase_       (				a_      ,    unittest.TestCase     ):
  _A				:    str =		True

  _A				:    List[Any] =		(
      (
          FlaxElectraModel,
          FlaxElectraForCausalLM,
          FlaxElectraForMaskedLM,
          FlaxElectraForPreTraining,
          FlaxElectraForTokenClassification,
          FlaxElectraForQuestionAnswering,
          FlaxElectraForMultipleChoice,
          FlaxElectraForSequenceClassification,
      )
      if is_flax_available()
      else ()
  )



  def 						UpperCamelCase_ (					self )	->	str:



         """simple docstring"""





         UpperCAmelCase					=	FlaxElectraModelTester(self )



  @slow
  def 						UpperCamelCase_ (					self )	->	Tuple:



         """simple docstring"""





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