

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 					lowercase		:



 '''simple docstring'''


 def __init__(self , __a , __a=13 , __a=7 , __a=True , __a=True , __a=True , __a=True , __a=99 , __a=32 , __a=5 , __a=4 , __a=37 , __a="gelu" , __a=0.1 , __a=0.1 , __a=512 , __a=16 , __a=2 , __a=0.02 , __a=3 , __a=4 , __a=None , )				->      Any:






  """simple docstring"""





  UpperCAmelCase__           =       parent
  UpperCAmelCase__           =       batch_size
  UpperCAmelCase__           =       seq_length
  UpperCAmelCase__           =       is_training
  UpperCAmelCase__           =       use_input_mask
  UpperCAmelCase__           =       use_token_type_ids
  UpperCAmelCase__           =       use_labels
  UpperCAmelCase__           =       vocab_size
  UpperCAmelCase__           =       hidden_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_labels
  UpperCAmelCase__           =       num_choices
  UpperCAmelCase__           =       scope




 def 						UpperCamelCase__							(self				)				->      Dict:






  """simple docstring"""





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

  UpperCAmelCase__           =       None
  if self.use_input_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__           =       None
  UpperCAmelCase__           =       None
  UpperCAmelCase__           =       None
  if self.use_labels:
   UpperCAmelCase__           =       ids_tensor([self.batch_size] , self.type_sequence_label_size				)
   UpperCAmelCase__           =       ids_tensor([self.batch_size, self.seq_length] , self.num_labels				)
   UpperCAmelCase__           =       ids_tensor([self.batch_size] , self.num_choices				)
   UpperCAmelCase__           =       ids_tensor([self.batch_size, self.seq_length] , 1				)

  UpperCAmelCase__           =       self.get_config()

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




 def 						UpperCamelCase__							(self				)				->      int:






  """simple docstring"""





  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 						UpperCamelCase__							(self				)				->      List[Any]:






  """simple docstring"""





  (
      (
      UpperCAmelCase__					
  )       ,					(
      UpperCAmelCase__					
  )       ,					(
      UpperCAmelCase__					
  )       ,					(
      UpperCAmelCase__					
  )       ,					(
      UpperCAmelCase__					
  )       ,					(
      UpperCAmelCase__					
  )       ,					(
      UpperCAmelCase__					
  )       ,					(
      UpperCAmelCase__					
  )       ,										
  )           =       self.prepare_config_and_inputs()

  UpperCAmelCase__           =       True
  UpperCAmelCase__           =       floats_tensor([self.batch_size, self.seq_length, self.hidden_size]				)
  UpperCAmelCase__           =       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 						UpperCamelCase__							(self , __a , __a , __a , __a , __a , __a , __a , __a , )				->      int:






  """simple docstring"""





  UpperCAmelCase__           =       ElectraModel(config=__a				)
  model.to(__a				)
  model.eval()
  UpperCAmelCase__           =       model(__a , attention_mask=__a , token_type_ids=__a				)
  UpperCAmelCase__           =       model(__a , token_type_ids=__a				)
  UpperCAmelCase__           =       model(__a				)
  self.parent.assertEqual(result.last_hidden_state.shape , (self.batch_size, self.seq_length, self.hidden_size)				)




 def 						UpperCamelCase__							(self , __a , __a , __a , __a , __a , __a , __a , __a , __a , )				->      Dict:






  """simple docstring"""





  UpperCAmelCase__           =       True
  UpperCAmelCase__           =       ElectraModel(__a				)
  model.to(__a				)
  model.eval()
  UpperCAmelCase__           =       model(
      __a , attention_mask=__a , token_type_ids=__a , encoder_hidden_states=__a , encoder_attention_mask=__a , )
  UpperCAmelCase__           =       model(
      __a , attention_mask=__a , token_type_ids=__a , encoder_hidden_states=__a , )
  UpperCAmelCase__           =       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 						UpperCamelCase__							(self , __a , __a , __a , __a , __a , __a , __a , __a , )				->      Optional[Any]:






  """simple docstring"""





  UpperCAmelCase__           =       ElectraForMaskedLM(config=__a				)
  model.to(__a				)
  model.eval()
  UpperCAmelCase__           =       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 						UpperCamelCase__							(self , __a , __a , __a , __a , __a , __a , __a , __a , __a , )				->      int:






  """simple docstring"""





  UpperCAmelCase__           =       ElectraForCausalLM(config=__a				)
  model.to(__a				)
  model.eval()
  UpperCAmelCase__           =       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 						UpperCamelCase__							(self , __a , __a , __a , __a , __a , __a , __a , __a , )				->      List[Any]:






  """simple docstring"""





  UpperCAmelCase__           =       self.num_labels
  UpperCAmelCase__           =       ElectraForTokenClassification(config=__a				)
  model.to(__a				)
  model.eval()
  UpperCAmelCase__           =       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 						UpperCamelCase__							(self , __a , __a , __a , __a , __a , __a , __a , __a , )				->      Any:






  """simple docstring"""





  UpperCAmelCase__           =       self.num_labels
  UpperCAmelCase__           =       ElectraForPreTraining(config=__a				)
  model.to(__a				)
  model.eval()
  UpperCAmelCase__           =       model(__a , attention_mask=__a , token_type_ids=__a , labels=__a				)
  self.parent.assertEqual(result.logits.shape , (self.batch_size, self.seq_length)				)




 def 						UpperCamelCase__							(self , __a , __a , __a , __a , __a , __a , __a , __a , )				->      str:






  """simple docstring"""





  UpperCAmelCase__           =       self.num_labels
  UpperCAmelCase__           =       ElectraForSequenceClassification(__a				)
  model.to(__a				)
  model.eval()
  UpperCAmelCase__           =       model(__a , attention_mask=__a , token_type_ids=__a , labels=__a				)
  self.parent.assertEqual(result.logits.shape , (self.batch_size, self.num_labels)				)




 def 						UpperCamelCase__							(self , __a , __a , __a , __a , __a , __a , __a , __a , )				->      List[str]:






  """simple docstring"""





  UpperCAmelCase__           =       ElectraForQuestionAnswering(config=__a				)
  model.to(__a				)
  model.eval()
  UpperCAmelCase__           =       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 						UpperCamelCase__							(self , __a , __a , __a , __a , __a , __a , __a , __a , )				->      List[Any]:






  """simple docstring"""





  UpperCAmelCase__           =       self.num_choices
  UpperCAmelCase__           =       ElectraForMultipleChoice(config=__a				)
  model.to(__a				)
  model.eval()
  UpperCAmelCase__           =       input_ids.unsqueeze(1				).expand(-1 , self.num_choices , -1				).contiguous()
  UpperCAmelCase__           =       token_type_ids.unsqueeze(1				).expand(-1 , self.num_choices , -1				).contiguous()
  UpperCAmelCase__           =       input_mask.unsqueeze(1				).expand(-1 , self.num_choices , -1				).contiguous()
  UpperCAmelCase__           =       model(
      __a , attention_mask=__a , token_type_ids=__a , labels=__a , )
  self.parent.assertEqual(result.logits.shape , (self.batch_size, self.num_choices)				)




 def 						UpperCamelCase__							(self				)				->      Dict:






  """simple docstring"""





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







@require_torch
class 					lowercase		( _UpperCamelCase				,       _UpperCamelCase				,       unittest.TestCase      ):



 '''simple docstring'''


 __SCREAMING_SNAKE_CASE      							=       (
     (
         ElectraModel,
         ElectraForPreTraining,
         ElectraForMaskedLM,
         ElectraForCausalLM,
         ElectraForMultipleChoice,
         ElectraForTokenClassification,
         ElectraForSequenceClassification,
         ElectraForQuestionAnswering,
     )
     if is_torch_available()
     else ()
 )
 __SCREAMING_SNAKE_CASE      							=       (
     {
         """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 {}
 )
 __SCREAMING_SNAKE_CASE      							=       True




 def 						UpperCamelCase__							(self , __a , __a , __a=False				)				->      Optional[Any]:






  """simple docstring"""





  UpperCAmelCase__           =       super()._prepare_for_class(__a , __a , return_labels=__a				)

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




 def 						UpperCamelCase__							(self				)				->      Union[str, Any]:






  """simple docstring"""





  UpperCAmelCase__           =       ElectraModelTester(self				)
  UpperCAmelCase__           =       ConfigTester(self , config_class=__a , hidden_size=37				)




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






  """simple docstring"""





  self.config_tester.run_common_tests()




 def 						UpperCamelCase__							(self				)				->      str:






  """simple docstring"""





  UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_electra_model(*__a				)




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






  """simple docstring"""





  UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs_for_decoder()
  self.model_tester.create_and_check_electra_model_as_decoder(*__a				)




 def 						UpperCamelCase__							(self				)				->      Union[str, Any]:






  """simple docstring"""





  UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs()
  for type in ["absolute", "relative_key", "relative_key_query"]:
   UpperCAmelCase__           =       type
   self.model_tester.create_and_check_electra_model(*__a				)




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






  """simple docstring"""





  UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_electra_for_masked_lm(*__a				)




 def 						UpperCamelCase__							(self				)				->      List[str]:






  """simple docstring"""





  UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_electra_for_token_classification(*__a				)




 def 						UpperCamelCase__							(self				)				->      Union[str, Any]:






  """simple docstring"""





  UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_electra_for_pretraining(*__a				)




 def 						UpperCamelCase__							(self				)				->      Tuple:






  """simple docstring"""





  UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_electra_for_sequence_classification(*__a				)




 def 						UpperCamelCase__							(self				)				->      Optional[int]:






  """simple docstring"""





  UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_electra_for_question_answering(*__a				)




 def 						UpperCamelCase__							(self				)				->      Optional[int]:






  """simple docstring"""





  UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_electra_for_multiple_choice(*__a				)




 @slow
 def 						UpperCamelCase__							(self				)				->      List[str]:






  """simple docstring"""





  for model_name in ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
   UpperCAmelCase__           =       ElectraModel.from_pretrained(__a				)
   self.assertIsNotNone(__a				)




 def 						UpperCamelCase__							(self				)				->      Any:






  """simple docstring"""





  UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs_for_decoder()
  self.model_tester.create_and_check_electra_for_causal_lm(*__a				)







@require_torch
class 					lowercase		( unittest.TestCase      ):



 '''simple docstring'''


 @slow
 def 						UpperCamelCase__							(self				)				->      Tuple:






  """simple docstring"""





  UpperCAmelCase__           =       ElectraModel.from_pretrained('google/electra-small-discriminator'				)
  UpperCAmelCase__           =       torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]				)
  UpperCAmelCase__           =       torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]				)
  UpperCAmelCase__           =       model(__a , attention_mask=__a				)[0]
  UpperCAmelCase__           =       torch.Size((1, 11, 256)				)
  self.assertEqual(output.shape , __a				)
  UpperCAmelCase__           =       torch.tensor(
      [[[0.44_71, 0.68_21, -0.32_65], [0.46_27, 0.52_55, -0.36_68], [0.45_32, 0.33_13, -0.43_44]]]				)

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






def 		UpperCamelCase_(    snake_case__:						int     ) ->				int:
 if not isinstance(snake_case__						,					snake_case__     ):
  UpperCAmelCase__           =       f"Input value of [number={number}] must be an integer"
  raise TypeError(snake_case__     )
 if number < 1:
  raise ValueError('Input must be a positive integer'     )
 return number * (2 * number - 1)


if __name__ == "__main__":
 import doctest

 doctest.testmod()



