




import unittest

from transformers import ElectraConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import (
          MODEL_FOR_PRETRAINING_MAPPING,
          ElectraForCausalLM,
          ElectraForMaskedLM,
          ElectraForMultipleChoice,
          ElectraForPreTraining,
          ElectraForQuestionAnswering,
          ElectraForSequenceClassification,
          ElectraForTokenClassification,
          ElectraModel,
      )
      from transformers.models.electra.modeling_electra import ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST






class      A     :
     """simple docstring"""

     def __init__(      self					:				Dict		,					A__					:				Optional[Any]		,					A__					:				List[str]=1_3		,					A__					:				List[str]=7		,					A__					:				str=True		,					A__					:				List[Any]=True		,					A__					:				Union[str, Any]=True		,					A__					:				Union[str, Any]=True		,					A__					:				Optional[Any]=9_9		,					A__					:				Union[str, Any]=3_2		,					A__					:				int=5		,					A__					:				List[Any]=4		,					A__					:				List[str]=3_7		,					A__					:				str="gelu"		,					A__					:				Optional[int]=0.1		,					A__					:				List[Any]=0.1		,					A__					:				str=5_1_2		,					A__					:				List[Any]=1_6		,					A__					:				Any=2		,					A__					:				Optional[int]=0.02		,					A__					:				Optional[int]=3		,					A__					:				Dict=4		,					A__					:				Optional[Any]=None		,					)			->    str:
      lowerCamelCase__ 					=			parent
      lowerCamelCase__ 					=			batch_size
      lowerCamelCase__ 					=			seq_length
      lowerCamelCase__ 					=			is_training
      lowerCamelCase__ 					=			use_input_mask
      lowerCamelCase__ 					=			use_token_type_ids
      lowerCamelCase__ 					=			use_labels
      lowerCamelCase__ 					=			vocab_size
      lowerCamelCase__ 					=			hidden_size
      lowerCamelCase__ 					=			num_hidden_layers
      lowerCamelCase__ 					=			num_attention_heads
      lowerCamelCase__ 					=			intermediate_size
      lowerCamelCase__ 					=			hidden_act
      lowerCamelCase__ 					=			hidden_dropout_prob
      lowerCamelCase__ 					=			attention_probs_dropout_prob
      lowerCamelCase__ 					=			max_position_embeddings
      lowerCamelCase__ 					=			type_vocab_size
      lowerCamelCase__ 					=			type_sequence_label_size
      lowerCamelCase__ 					=			initializer_range
      lowerCamelCase__ 					=			num_labels
      lowerCamelCase__ 					=			num_choices
      lowerCamelCase__ 					=			scope




     def 	__snake_case				(      self					:				Dict     )			->    Any:
      lowerCamelCase__ 					=			ids_tensor([self.batch_size, self.seq_length]		,					self.vocab_size     )

      lowerCamelCase__ 					=			None
      if self.use_input_mask:
       lowerCamelCase__ 					=			random_attention_mask([self.batch_size, self.seq_length]     )

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

      lowerCamelCase__ 					=			None
      lowerCamelCase__ 					=			None
      lowerCamelCase__ 					=			None
      if self.use_labels:
       lowerCamelCase__ 					=			ids_tensor([self.batch_size]		,					self.type_sequence_label_size     )
       lowerCamelCase__ 					=			ids_tensor([self.batch_size, self.seq_length]		,					self.num_labels     )
       lowerCamelCase__ 					=			ids_tensor([self.batch_size]		,					self.num_choices     )
       lowerCamelCase__ 					=			ids_tensor([self.batch_size, self.seq_length]		,					1     )

      lowerCamelCase__ 					=			self.get_config()

      return (
          config,
          input_ids,
          token_type_ids,
          input_mask,
          sequence_labels,
          token_labels,
          choice_labels,
          fake_token_labels,
      )




     def 	__snake_case				(      self					:				Any     )			->    Union[str, Any]:
      return ElectraConfig(
          vocab_size=self.vocab_size		,					hidden_size=self.hidden_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		,					is_decoder=A__		,					initializer_range=self.initializer_range		,					)




     def 	__snake_case				(      self					:				List[str]     )			->    int:
      (
          (
          lowerCamelCase__						
      )					,						(
          lowerCamelCase__						
      )					,						(
          lowerCamelCase__						
      )					,						(
          lowerCamelCase__						
      )					,						(
          lowerCamelCase__						
      )					,						(
          lowerCamelCase__						
      )					,						(
          lowerCamelCase__						
      )					,						(
          lowerCamelCase__						
      )					,												
      ) 					=			self.prepare_config_and_inputs()

      lowerCamelCase__ 					=			True
      lowerCamelCase__ 					=			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]     )
      lowerCamelCase__ 					=			ids_tensor([self.batch_size, self.seq_length]		,					vocab_size=2     )

      return (
          config,
          input_ids,
          token_type_ids,
          input_mask,
          sequence_labels,
          token_labels,
          choice_labels,
          encoder_hidden_states,
          encoder_attention_mask,
      )




     def 	__snake_case				(      self					:				int		,					A__					:				Optional[Any]		,					A__					:				List[Any]		,					A__					:				Tuple		,					A__					:				Dict		,					A__					:				Any		,					A__					:				Any		,					A__					:				Tuple		,					A__					:				int		,					)			->    List[Any]:
      lowerCamelCase__ 					=			ElectraModel(config=A__     )
      model.to(A__     )
      model.eval()
      lowerCamelCase__ 					=			model(A__		,					attention_mask=A__		,					token_type_ids=A__     )
      lowerCamelCase__ 					=			model(A__		,					token_type_ids=A__     )
      lowerCamelCase__ 					=			model(A__     )
      self.parent.assertEqual(result.last_hidden_state.shape		,					(self.batch_size, self.seq_length, self.hidden_size)     )




     def 	__snake_case				(      self					:				Optional[int]		,					A__					:				Any		,					A__					:				str		,					A__					:				Optional[Any]		,					A__					:				Tuple		,					A__					:				Optional[Any]		,					A__					:				Tuple		,					A__					:				int		,					A__					:				Tuple		,					A__					:				Optional[Any]		,					)			->    Union[str, Any]:
      lowerCamelCase__ 					=			True
      lowerCamelCase__ 					=			ElectraModel(A__     )
      model.to(A__     )
      model.eval()
      lowerCamelCase__ 					=			model(
          A__		,					attention_mask=A__		,					token_type_ids=A__		,					encoder_hidden_states=A__		,					encoder_attention_mask=A__		,					)
      lowerCamelCase__ 					=			model(
          A__		,					attention_mask=A__		,					token_type_ids=A__		,					encoder_hidden_states=A__		,					)
      lowerCamelCase__ 					=			model(A__		,					attention_mask=A__		,					token_type_ids=A__     )
      self.parent.assertEqual(result.last_hidden_state.shape		,					(self.batch_size, self.seq_length, self.hidden_size)     )




     def 	__snake_case				(      self					:				Optional[int]		,					A__					:				str		,					A__					:				List[Any]		,					A__					:				List[str]		,					A__					:				Optional[Any]		,					A__					:				Any		,					A__					:				int		,					A__					:				Tuple		,					A__					:				int		,					)			->    Tuple:
      lowerCamelCase__ 					=			ElectraForMaskedLM(config=A__     )
      model.to(A__     )
      model.eval()
      lowerCamelCase__ 					=			model(A__		,					attention_mask=A__		,					token_type_ids=A__		,					labels=A__     )
      self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.seq_length, self.vocab_size)     )




     def 	__snake_case				(      self					:				Optional[Any]		,					A__					:				List[Any]		,					A__					:				Dict		,					A__					:				Dict		,					A__					:				Dict		,					A__					:				Optional[Any]		,					A__					:				str		,					A__					:				Optional[Any]		,					A__					:				Any		,					A__					:				Dict		,					)			->    Optional[int]:
      lowerCamelCase__ 					=			ElectraForCausalLM(config=A__     )
      model.to(A__     )
      model.eval()
      lowerCamelCase__ 					=			model(A__		,					attention_mask=A__		,					token_type_ids=A__		,					labels=A__     )
      self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.seq_length, self.vocab_size)     )




     def 	__snake_case				(      self					:				Dict		,					A__					:				Any		,					A__					:				str		,					A__					:				str		,					A__					:				int		,					A__					:				str		,					A__					:				Union[str, Any]		,					A__					:				Optional[int]		,					A__					:				str		,					)			->    Dict:
      lowerCamelCase__ 					=			self.num_labels
      lowerCamelCase__ 					=			ElectraForTokenClassification(config=A__     )
      model.to(A__     )
      model.eval()
      lowerCamelCase__ 					=			model(A__		,					attention_mask=A__		,					token_type_ids=A__		,					labels=A__     )
      self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.seq_length, self.num_labels)     )




     def 	__snake_case				(      self					:				Union[str, Any]		,					A__					:				int		,					A__					:				List[str]		,					A__					:				List[str]		,					A__					:				Any		,					A__					:				int		,					A__					:				int		,					A__					:				Tuple		,					A__					:				Dict		,					)			->    List[str]:
      lowerCamelCase__ 					=			self.num_labels
      lowerCamelCase__ 					=			ElectraForPreTraining(config=A__     )
      model.to(A__     )
      model.eval()
      lowerCamelCase__ 					=			model(A__		,					attention_mask=A__		,					token_type_ids=A__		,					labels=A__     )
      self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.seq_length)     )




     def 	__snake_case				(      self					:				int		,					A__					:				Union[str, Any]		,					A__					:				Any		,					A__					:				List[Any]		,					A__					:				Union[str, Any]		,					A__					:				Union[str, Any]		,					A__					:				int		,					A__					:				List[str]		,					A__					:				Union[str, Any]		,					)			->    List[Any]:
      lowerCamelCase__ 					=			self.num_labels
      lowerCamelCase__ 					=			ElectraForSequenceClassification(A__     )
      model.to(A__     )
      model.eval()
      lowerCamelCase__ 					=			model(A__		,					attention_mask=A__		,					token_type_ids=A__		,					labels=A__     )
      self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.num_labels)     )




     def 	__snake_case				(      self					:				Any		,					A__					:				Optional[Any]		,					A__					:				str		,					A__					:				int		,					A__					:				Tuple		,					A__					:				Tuple		,					A__					:				Union[str, Any]		,					A__					:				Union[str, Any]		,					A__					:				List[Any]		,					)			->    Union[str, Any]:
      lowerCamelCase__ 					=			ElectraForQuestionAnswering(config=A__     )
      model.to(A__     )
      model.eval()
      lowerCamelCase__ 					=			model(
          A__		,					attention_mask=A__		,					token_type_ids=A__		,					start_positions=A__		,					end_positions=A__		,					)
      self.parent.assertEqual(result.start_logits.shape		,					(self.batch_size, self.seq_length)     )
      self.parent.assertEqual(result.end_logits.shape		,					(self.batch_size, self.seq_length)     )




     def 	__snake_case				(      self					:				Optional[Any]		,					A__					:				str		,					A__					:				List[str]		,					A__					:				List[str]		,					A__					:				Tuple		,					A__					:				Optional[int]		,					A__					:				str		,					A__					:				Any		,					A__					:				Dict		,					)			->    Tuple:
      lowerCamelCase__ 					=			self.num_choices
      lowerCamelCase__ 					=			ElectraForMultipleChoice(config=A__     )
      model.to(A__     )
      model.eval()
      lowerCamelCase__ 					=			input_ids.unsqueeze(1     ).expand(-1		,					self.num_choices		,					-1     ).contiguous()
      lowerCamelCase__ 					=			token_type_ids.unsqueeze(1     ).expand(-1		,					self.num_choices		,					-1     ).contiguous()
      lowerCamelCase__ 					=			input_mask.unsqueeze(1     ).expand(-1		,					self.num_choices		,					-1     ).contiguous()
      lowerCamelCase__ 					=			model(
          A__		,					attention_mask=A__		,					token_type_ids=A__		,					labels=A__		,					)
      self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.num_choices)     )




     def 	__snake_case				(      self					:				Any     )			->    List[Any]:
      lowerCamelCase__ 					=			self.prepare_config_and_inputs()
      (
          (
          lowerCamelCase__						
      )					,						(
          lowerCamelCase__						
      )					,						(
          lowerCamelCase__						
      )					,						(
          lowerCamelCase__						
      )					,						(
          lowerCamelCase__						
      )					,						(
          lowerCamelCase__						
      )					,						(
          lowerCamelCase__						
      )					,						(
          lowerCamelCase__						
      )					,												
      ) 					=			config_and_inputs
      lowerCamelCase__ 					=			{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
      return config, inputs_dict






@require_torch
class      A     (						_lowerCAmelCase     ,  _lowerCAmelCase     ,  unittest.TestCase						):
     """simple docstring"""

     _snake_case:		str       =(
         (
             ElectraModel,
             ElectraForPreTraining,
             ElectraForMaskedLM,
             ElectraForCausalLM,
             ElectraForMultipleChoice,
             ElectraForTokenClassification,
             ElectraForSequenceClassification,
             ElectraForQuestionAnswering,
         )
         if is_torch_available()
         else ()
     )
     _snake_case:		Optional[int]       =(
         {
             "feature-extraction": ElectraModel,
             "fill-mask": ElectraForMaskedLM,
             "question-answering": ElectraForQuestionAnswering,
             "text-classification": ElectraForSequenceClassification,
             "text-generation": ElectraForCausalLM,
             "token-classification": ElectraForTokenClassification,
             "zero-shot": ElectraForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     _snake_case:		str       =True




     def 	__snake_case				(      self					:				Optional[Any]		,					A__					:				List[Any]		,					A__					:				Tuple		,					A__					:				Union[str, Any]=False     )			->    Union[str, Any]:
      lowerCamelCase__ 					=			super()._prepare_for_class(A__		,					A__		,					return_labels=A__     )

      if return_labels:
       if model_class in get_values(A__     ):
        lowerCamelCase__ 					=			torch.zeros(
            (self.model_tester.batch_size, self.model_tester.seq_length)		,					dtype=torch.long		,					device=A__     )
      return inputs_dict




     def 	__snake_case				(      self					:				List[Any]     )			->    Union[str, Any]:
      lowerCamelCase__ 					=			ElectraModelTester(self     )
      lowerCamelCase__ 					=			ConfigTester(self		,					config_class=A__		,					hidden_size=3_7     )




     def 	__snake_case				(      self					:				Any     )			->    Any:
      self.config_tester.run_common_tests()




     def 	__snake_case				(      self					:				Union[str, Any]     )			->    Optional[Any]:
      lowerCamelCase__ 					=			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_electra_model(*A__     )




     def 	__snake_case				(      self					:				Optional[Any]     )			->    List[str]:
      lowerCamelCase__ 					=			self.model_tester.prepare_config_and_inputs_for_decoder()
      self.model_tester.create_and_check_electra_model_as_decoder(*A__     )




     def 	__snake_case				(      self					:				Optional[Any]     )			->    Optional[int]:
      lowerCamelCase__ 					=			self.model_tester.prepare_config_and_inputs()
      for type in ["absolute", "relative_key", "relative_key_query"]:
       lowerCamelCase__ 					=			type
       self.model_tester.create_and_check_electra_model(*A__     )




     def 	__snake_case				(      self					:				Any     )			->    str:
      lowerCamelCase__ 					=			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_electra_for_masked_lm(*A__     )




     def 	__snake_case				(      self					:				Any     )			->    Union[str, Any]:
      lowerCamelCase__ 					=			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_electra_for_token_classification(*A__     )




     def 	__snake_case				(      self					:				Tuple     )			->    Tuple:
      lowerCamelCase__ 					=			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_electra_for_pretraining(*A__     )




     def 	__snake_case				(      self					:				int     )			->    Optional[Any]:
      lowerCamelCase__ 					=			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_electra_for_sequence_classification(*A__     )




     def 	__snake_case				(      self					:				Dict     )			->    Optional[Any]:
      lowerCamelCase__ 					=			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_electra_for_question_answering(*A__     )




     def 	__snake_case				(      self					:				Any     )			->    Tuple:
      lowerCamelCase__ 					=			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_electra_for_multiple_choice(*A__     )




     @slow
     def 	__snake_case				(      self					:				Optional[int]     )			->    List[str]:
      for model_name in ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       lowerCamelCase__ 					=			ElectraModel.from_pretrained(A__     )
       self.assertIsNotNone(A__     )




     def 	__snake_case				(      self					:				Optional[int]     )			->    Any:
      lowerCamelCase__ 					=			self.model_tester.prepare_config_and_inputs_for_decoder()
      self.model_tester.create_and_check_electra_for_causal_lm(*A__     )






@require_torch
class      A     (						unittest.TestCase						):
     """simple docstring"""

     @slow
     def 	__snake_case				(      self					:				Tuple     )			->    Any:
      lowerCamelCase__ 					=			ElectraModel.from_pretrained('google/electra-small-discriminator'     )
      lowerCamelCase__ 					=			torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]     )
      lowerCamelCase__ 					=			torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]     )
      lowerCamelCase__ 					=			model(A__		,					attention_mask=A__     )[0]
      lowerCamelCase__ 					=			torch.Size((1, 1_1, 2_5_6)     )
      self.assertEqual(output.shape		,					A__     )
      lowerCamelCase__ 					=			torch.tensor(
          [[[0.4471, 0.6821, -0.3265], [0.4627, 0.5255, -0.3668], [0.4532, 0.3313, -0.4344]]]     )

      self.assertTrue(torch.allclose(output[:, 1:4, 1:4]		,					A__		,					atol=1E-4     )     )





from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__SCREAMING_SNAKE_CASE      :				int				      ={"""configuration_nat""": ["""NAT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """NatConfig"""]}


try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __SCREAMING_SNAKE_CASE      :				Optional[int]				      =[
          """NAT_PRETRAINED_MODEL_ARCHIVE_LIST""",
          """NatForImageClassification""",
          """NatModel""",
          """NatPreTrainedModel""",
          """NatBackbone""",
      ]

if TYPE_CHECKING:
      from .configuration_nat import NAT_PRETRAINED_CONFIG_ARCHIVE_MAP, NatConfig

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_nat import (
                NAT_PRETRAINED_MODEL_ARCHIVE_LIST,
                NatBackbone,
                NatForImageClassification,
                NatModel,
                NatPreTrainedModel,
            )

else:
      import sys

      __SCREAMING_SNAKE_CASE      :				Any				      =_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)