






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



     """simple docstring"""



     def __init__(		self					:	int  ,  lowerCAmelCase__					:	int  ,  lowerCAmelCase__					:	Dict=1_3  ,  lowerCAmelCase__					:	Dict=7  ,  lowerCAmelCase__					:	Union[str, Any]=True  ,  lowerCAmelCase__					:	Any=True  ,  lowerCAmelCase__					:	Any=True  ,  lowerCAmelCase__					:	Tuple=True  ,  lowerCAmelCase__					:	Optional[int]=9_9  ,  lowerCAmelCase__					:	Tuple=2_4  ,  lowerCAmelCase__					:	Tuple=3_2  ,  lowerCAmelCase__					:	Dict=5  ,  lowerCAmelCase__					:	List[str]=4  ,  lowerCAmelCase__					:	Optional[Any]=3_7  ,  lowerCAmelCase__					:	Any="gelu"  ,  lowerCAmelCase__					:	int=0.1  ,  lowerCAmelCase__					:	List[Any]=0.1  ,  lowerCAmelCase__					:	Union[str, Any]=5_1_2  ,  lowerCAmelCase__					:	List[str]=1_6  ,  lowerCAmelCase__					:	Any=2  ,  lowerCAmelCase__					:	List[str]=0.02  ,  lowerCAmelCase__					:	Any=4  ,  )							->	Optional[int]:

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





     def     __UpperCamelCase						(		self					:	Tuple)							->	Dict:

      """simple docstring"""
      SCREAMING_SNAKE_CASE				    =		ids_tensor([self.batch_size, self.seq_length]  ,  self.vocab_size)

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

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

      SCREAMING_SNAKE_CASE				    =		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					:	Dict)							->	Optional[Any]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE				    =		self.prepare_config_and_inputs()
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		config_and_inputs
      SCREAMING_SNAKE_CASE				    =		{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': attention_mask}
      return config, inputs_dict

@require_flax
class 	a_(		lowercase__			,							unittest.TestCase	):



     """simple docstring"""



     __snake_case		:				str  					=True

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





     def     __UpperCamelCase						(		self					:	Any)							->	int:

      """simple docstring"""
      SCREAMING_SNAKE_CASE				    =		FlaxElectraModelTester(self)





     @slow
     def     __UpperCamelCase						(		self					:	Dict)							->	Optional[int]:

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








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,
)


__UpperCAmelCase      			=  {
    """configuration_switch_transformers""": [
        """SWITCH_TRANSFORMERS_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """SwitchTransformersConfig""",
        """SwitchTransformersOnnxConfig""",
    ]
}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     __UpperCAmelCase      			=  [
         """SWITCH_TRANSFORMERS_PRETRAINED_MODEL_ARCHIVE_LIST""",
         """SwitchTransformersEncoderModel""",
         """SwitchTransformersForConditionalGeneration""",
         """SwitchTransformersModel""",
         """SwitchTransformersPreTrainedModel""",
         """SwitchTransformersTop1Router""",
         """SwitchTransformersSparseMLP""",
     ]


if TYPE_CHECKING:
     from .configuration_switch_transformers import (
         SWITCH_TRANSFORMERS_PRETRAINED_CONFIG_ARCHIVE_MAP,
         SwitchTransformersConfig,
         SwitchTransformersOnnxConfig,
     )

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_switch_transformers import (
              SWITCH_TRANSFORMERS_PRETRAINED_MODEL_ARCHIVE_LIST,
              SwitchTransformersEncoderModel,
              SwitchTransformersForConditionalGeneration,
              SwitchTransformersModel,
              SwitchTransformersPreTrainedModel,
              SwitchTransformersSparseMLP,
              SwitchTransformersTopaRouter,
          )


else:
     import sys

     __UpperCAmelCase      			=  _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

