





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





       '''simple docstring'''







       def __init__(    self					:       Union[str, Any]				,							A					:       List[Any]				,							A					:       Tuple=13				,							A					:       Optional[int]=7				,							A					:       List[str]=True				,							A					:       int=True				,							A					:       str=True				,							A					:       Dict=True				,							A					:       Optional[int]=99				,							A					:       Optional[Any]=32				,							A					:       Optional[Any]=5				,							A					:       Tuple=4				,							A					:       Tuple=37				,							A					:       List[str]="gelu"				,							A					:       Any=0.1				,							A					:       str=0.1				,							A					:       Optional[int]=512				,							A					:       Union[str, Any]=16				,							A					:       Optional[int]=2				,							A					:       str=0.02				,							A					:       Optional[int]=3				,							A					:       List[Any]=4				,							A					:       str=None				,							):
         __snake_case:		Tuple					    =		parent
         __snake_case:		int					    =		batch_size
         __snake_case:		Tuple					    =		seq_length
         __snake_case:		str					    =		is_training
         __snake_case:		int					    =		use_input_mask
         __snake_case:		Tuple					    =		use_token_type_ids
         __snake_case:		List[str]					    =		use_labels
         __snake_case:		Tuple					    =		vocab_size
         __snake_case:		Optional[Any]					    =		hidden_size
         __snake_case:		str					    =		num_hidden_layers
         __snake_case:		List[Any]					    =		num_attention_heads
         __snake_case:		Union[str, Any]					    =		intermediate_size
         __snake_case:		Optional[int]					    =		hidden_act
         __snake_case:		Dict					    =		hidden_dropout_prob
         __snake_case:		List[str]					    =		attention_probs_dropout_prob
         __snake_case:		Tuple					    =		max_position_embeddings
         __snake_case:		Tuple					    =		type_vocab_size
         __snake_case:		List[str]					    =		type_sequence_label_size
         __snake_case:		Any					    =		initializer_range
         __snake_case:		Tuple					    =		num_labels
         __snake_case:		Any					    =		num_choices
         __snake_case:		Tuple					    =		scope




       def 				UpperCAmelCase__					(    self					:       Any						):
         __snake_case:		Optional[Any]					    =		ids_tensor([self.batch_size, self.seq_length]				,							self.vocab_size						)

         __snake_case:		Tuple					    =		None
         if self.use_input_mask:
           __snake_case:		str					    =		random_attention_mask([self.batch_size, self.seq_length]						)

         __snake_case:		str					    =		None
         if self.use_token_type_ids:
           __snake_case:		Optional[Any]					    =		ids_tensor([self.batch_size, self.seq_length]				,							self.type_vocab_size						)

         __snake_case:		str					    =		None
         __snake_case:		List[str]					    =		None
         __snake_case:		Optional[Any]					    =		None
         if self.use_labels:
           __snake_case:		Optional[int]					    =		ids_tensor([self.batch_size]				,							self.type_sequence_label_size						)
           __snake_case:		Union[str, Any]					    =		ids_tensor([self.batch_size, self.seq_length]				,							self.num_labels						)
           __snake_case:		str					    =		ids_tensor([self.batch_size]				,							self.num_choices						)
           __snake_case:		int					    =		ids_tensor([self.batch_size, self.seq_length]				,							1						)

         __snake_case:		List[str]					    =		self.get_config()

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




       def 				UpperCAmelCase__					(    self					:       List[str]						):
         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					:       Union[str, Any]						):
         (
             (
             __snake_case
         )    ,  (
             __snake_case
         )    ,  (
             __snake_case
         )    ,  (
             __snake_case
         )    ,  (
             __snake_case
         )    ,  (
             __snake_case
         )    ,  (
             __snake_case
         )    ,  (
             __snake_case
         )    ,  
         ):		Optional[int]					    =		self.prepare_config_and_inputs()

         __snake_case:		List[Any]					    =		True
         __snake_case:		Any					    =		floats_tensor([self.batch_size, self.seq_length, self.hidden_size]						)
         __snake_case:		List[Any]					    =		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					:       Optional[Any]				,							A					:       Optional[int]				,							A					:       Tuple				,							A					:       Tuple				,							A					:       List[Any]				,							A					:       List[str]				,							A					:       Any				,							A					:       Optional[int]				,							A					:       List[Any]				,							):
         __snake_case:		Optional[Any]					    =		ElectraModel(config=A						)
         model.to(A						)
         model.eval()
         __snake_case:		int					    =		model(A				,							attention_mask=A				,							token_type_ids=A						)
         __snake_case:		Any					    =		model(A				,							token_type_ids=A						)
         __snake_case:		str					    =		model(A						)
         self.parent.assertEqual(result.last_hidden_state.shape				,							(self.batch_size, self.seq_length, self.hidden_size)						)




       def 				UpperCAmelCase__					(    self					:       Optional[Any]				,							A					:       List[Any]				,							A					:       Optional[Any]				,							A					:       Optional[Any]				,							A					:       List[Any]				,							A					:       Union[str, Any]				,							A					:       Optional[Any]				,							A					:       Any				,							A					:       Tuple				,							A					:       Tuple				,							):
         __snake_case:		Optional[int]					    =		True
         __snake_case:		Union[str, Any]					    =		ElectraModel(A						)
         model.to(A						)
         model.eval()
         __snake_case:		Optional[int]					    =		model(
             A				,							attention_mask=A				,							token_type_ids=A				,							encoder_hidden_states=A				,							encoder_attention_mask=A				,							)
         __snake_case:		Dict					    =		model(
             A				,							attention_mask=A				,							token_type_ids=A				,							encoder_hidden_states=A				,							)
         __snake_case:		Optional[Any]					    =		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					:       Any				,							A					:       Tuple				,							A					:       Optional[int]				,							A					:       Optional[Any]				,							A					:       Tuple				,							A					:       Optional[int]				,							A					:       Tuple				,							A					:       List[str]				,							A					:       List[str]				,							):
         __snake_case:		List[Any]					    =		ElectraForMaskedLM(config=A						)
         model.to(A						)
         model.eval()
         __snake_case:		Tuple					    =		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					:       Tuple				,							A					:       List[str]				,							A					:       Optional[Any]				,							A					:       List[str]				,							A					:       Optional[int]				,							A					:       Tuple				,							A					:       Union[str, Any]				,							A					:       Union[str, Any]				,							A					:       Dict				,							A					:       Dict				,							):
         __snake_case:		int					    =		ElectraForCausalLM(config=A						)
         model.to(A						)
         model.eval()
         __snake_case:		Optional[int]					    =		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					:       List[Any]				,							A					:       Tuple				,							A					:       Dict				,							A					:       Dict				,							A					:       Any				,							A					:       Union[str, Any]				,							A					:       Any				,							A					:       List[str]				,							A					:       Dict				,							):
         __snake_case:		Any					    =		self.num_labels
         __snake_case:		int					    =		ElectraForTokenClassification(config=A						)
         model.to(A						)
         model.eval()
         __snake_case:		Tuple					    =		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					:       List[Any]				,							A					:       Any				,							A					:       int				,							A					:       Union[str, Any]				,							A					:       Any				,							A					:       Optional[Any]				,							A					:       Union[str, Any]				,							A					:       List[str]				,							A					:       List[str]				,							):
         __snake_case:		List[Any]					    =		self.num_labels
         __snake_case:		Optional[int]					    =		ElectraForPreTraining(config=A						)
         model.to(A						)
         model.eval()
         __snake_case:		Optional[Any]					    =		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					:       List[Any]				,							A					:       Optional[int]				,							A					:       Optional[int]				,							A					:       Any				,							A					:       List[str]				,							A					:       Any				,							A					:       List[Any]				,							A					:       Tuple				,							A					:       Optional[Any]				,							):
         __snake_case:		int					    =		self.num_labels
         __snake_case:		Tuple					    =		ElectraForSequenceClassification(A						)
         model.to(A						)
         model.eval()
         __snake_case:		Tuple					    =		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					:       Dict				,							A					:       Any				,							A					:       Optional[int]				,							A					:       Any				,							A					:       int				,							A					:       List[Any]				,							A					:       List[Any]				,							A					:       List[str]				,							A					:       Dict				,							):
         __snake_case:		Union[str, Any]					    =		ElectraForQuestionAnswering(config=A						)
         model.to(A						)
         model.eval()
         __snake_case:		Optional[int]					    =		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					:       int				,							A					:       Union[str, Any]				,							A					:       Optional[int]				,							A					:       Union[str, Any]				,							A					:       Tuple				,							A					:       Any				,							A					:       Optional[Any]				,							A					:       List[str]				,							A					:       List[str]				,							):
         __snake_case:		str					    =		self.num_choices
         __snake_case:		Dict					    =		ElectraForMultipleChoice(config=A						)
         model.to(A						)
         model.eval()
         __snake_case:		Tuple					    =		input_ids.unsqueeze(1						).expand(-1				,							self.num_choices				,							-1						).contiguous()
         __snake_case:		int					    =		token_type_ids.unsqueeze(1						).expand(-1				,							self.num_choices				,							-1						).contiguous()
         __snake_case:		Dict					    =		input_mask.unsqueeze(1						).expand(-1				,							self.num_choices				,							-1						).contiguous()
         __snake_case:		Any					    =		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					:       Tuple						):
         __snake_case:		Any					    =		self.prepare_config_and_inputs()
         (
             (
             __snake_case
         )    ,  (
             __snake_case
         )    ,  (
             __snake_case
         )    ,  (
             __snake_case
         )    ,  (
             __snake_case
         )    ,  (
             __snake_case
         )    ,  (
             __snake_case
         )    ,  (
             __snake_case
         )    ,  
         ):		List[str]					    =		config_and_inputs
         __snake_case:		Any					    =		{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
         return config, inputs_dict







@require_torch
class        __snake_case			(				__lowerCamelCase  ,     __lowerCamelCase  ,     unittest.TestCase     ):





       '''simple docstring'''







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




       def 				UpperCAmelCase__					(    self					:       str				,							A					:       str				,							A					:       Optional[Any]				,							A					:       List[Any]=False						):
         __snake_case:		Union[str, Any]					    =		super()._prepare_for_class(A				,							A				,							return_labels=A						)

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




       def 				UpperCAmelCase__					(    self					:       Optional[Any]						):
         __snake_case:		int					    =		ElectraModelTester(self						)
         __snake_case:		str					    =		ConfigTester(self				,							config_class=A				,							hidden_size=37						)




       def 				UpperCAmelCase__					(    self					:       Optional[Any]						):
         self.config_tester.run_common_tests()




       def 				UpperCAmelCase__					(    self					:       Tuple						):
         __snake_case:		Optional[int]					    =		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_electra_model(*A						)




       def 				UpperCAmelCase__					(    self					:       str						):
         __snake_case:		Union[str, Any]					    =		self.model_tester.prepare_config_and_inputs_for_decoder()
         self.model_tester.create_and_check_electra_model_as_decoder(*A						)




       def 				UpperCAmelCase__					(    self					:       int						):
         __snake_case:		Tuple					    =		self.model_tester.prepare_config_and_inputs()
         for type in ["absolute", "relative_key", "relative_key_query"]:
           __snake_case:		Any					    =		type
           self.model_tester.create_and_check_electra_model(*A						)




       def 				UpperCAmelCase__					(    self					:       List[str]						):
         __snake_case:		List[Any]					    =		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_electra_for_masked_lm(*A						)




       def 				UpperCAmelCase__					(    self					:       Dict						):
         __snake_case:		Any					    =		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_electra_for_token_classification(*A						)




       def 				UpperCAmelCase__					(    self					:       Optional[Any]						):
         __snake_case:		List[str]					    =		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_electra_for_pretraining(*A						)




       def 				UpperCAmelCase__					(    self					:       Dict						):
         __snake_case:		List[str]					    =		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_electra_for_sequence_classification(*A						)




       def 				UpperCAmelCase__					(    self					:       Union[str, Any]						):
         __snake_case:		Optional[int]					    =		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_electra_for_question_answering(*A						)




       def 				UpperCAmelCase__					(    self					:       str						):
         __snake_case:		Any					    =		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_electra_for_multiple_choice(*A						)




       @slow
       def 				UpperCAmelCase__					(    self					:       Dict						):
         for model_name in ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           __snake_case:		Tuple					    =		ElectraModel.from_pretrained(A						)
           self.assertIsNotNone(A						)






       def 				UpperCAmelCase__					(    self					:       Any						):
         __snake_case:		Dict					    =		self.model_tester.prepare_config_and_inputs_for_decoder()
         self.model_tester.create_and_check_electra_for_causal_lm(*A						)







@require_torch
class        __snake_case			(				unittest.TestCase     ):





       '''simple docstring'''







       @slow
       def 				UpperCAmelCase__					(    self					:       Any						):
         __snake_case:		Dict					    =		ElectraModel.from_pretrained("""google/electra-small-discriminator"""						)
         __snake_case:		Tuple					    =		torch.tensor([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]						)
         __snake_case:		str					    =		torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]						)
         __snake_case:		Optional[Any]					    =		model(A				,							attention_mask=A						)[0]
         __snake_case:		Dict					    =		torch.Size((1, 11, 256)						)
         self.assertEqual(output.shape				,							A						)
         __snake_case:		Any					    =		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						)						)













import argparse

import torch

from transformers import NystromformerConfig, NystromformerForMaskedLM







def      A__				(       SCREAMING_SNAKE_CASE__)  ->       Optional[Any]:
  if "model" in orig_key:
    __snake_case:		str					    =		orig_key.replace("""model."""   ,   """""")
  if "norm1" in orig_key:
    __snake_case:		Tuple					    =		orig_key.replace("""norm1"""   ,   """attention.output.LayerNorm""")
  if "norm2" in orig_key:
    __snake_case:		Any					    =		orig_key.replace("""norm2"""   ,   """output.LayerNorm""")
  if "norm" in orig_key:
    __snake_case:		List[str]					    =		orig_key.replace("""norm"""   ,   """LayerNorm""")
  if "transformer" in orig_key:
    __snake_case:		str					    =		orig_key.split(""".""")[0].split("""_""")[-1]
    __snake_case:		int					    =		orig_key.replace(F'''transformer_{layer_num}'''   ,   F'''encoder.layer.{layer_num}''')
  if "mha.attn" in orig_key:
    __snake_case:		int					    =		orig_key.replace("""mha.attn"""   ,   """attention.self""")
  if "mha" in orig_key:
    __snake_case:		List[Any]					    =		orig_key.replace("""mha"""   ,   """attention""")
  if "W_q" in orig_key:
    __snake_case:		Dict					    =		orig_key.replace("""W_q"""   ,   """self.query""")
  if "W_k" in orig_key:
    __snake_case:		str					    =		orig_key.replace("""W_k"""   ,   """self.key""")
  if "W_v" in orig_key:
    __snake_case:		List[Any]					    =		orig_key.replace("""W_v"""   ,   """self.value""")
  if "ff1" in orig_key:
    __snake_case:		int					    =		orig_key.replace("""ff1"""   ,   """intermediate.dense""")
  if "ff2" in orig_key:
    __snake_case:		Dict					    =		orig_key.replace("""ff2"""   ,   """output.dense""")
  if "ff" in orig_key:
    __snake_case:		List[Any]					    =		orig_key.replace("""ff"""   ,   """output.dense""")
  if "mlm_class" in orig_key:
    __snake_case:		Union[str, Any]					    =		orig_key.replace("""mlm.mlm_class"""   ,   """cls.predictions.decoder""")
  if "mlm" in orig_key:
    __snake_case:		List[Any]					    =		orig_key.replace("""mlm"""   ,   """cls.predictions.transform""")
  if "cls" not in orig_key:
    __snake_case:		Dict					    =		"""nystromformer.""" + orig_key

  return orig_key







def      A__				(       SCREAMING_SNAKE_CASE__   ,   SCREAMING_SNAKE_CASE__)  ->       Optional[int]:
  for key in orig_state_dict.copy().keys():
    __snake_case:		str					    =		orig_state_dict.pop(SCREAMING_SNAKE_CASE__)

    if ("pooler" in key) or ("sen_class" in key) or ("conv.bias" in key):
      continue
    else:
      __snake_case:		Tuple					    =		val

  __snake_case:		str					    =		orig_state_dict["""cls.predictions.decoder.bias"""]
  __snake_case:		int					    =		(
      torch.arange(config.max_position_embeddings).expand((1, -1)) + 2
  )

  return orig_state_dict







def      A__				(       SCREAMING_SNAKE_CASE__   ,   SCREAMING_SNAKE_CASE__   ,   SCREAMING_SNAKE_CASE__)  ->       Optional[Any]:
  __snake_case:		int					    =		torch.load(SCREAMING_SNAKE_CASE__   ,   map_location="""cpu""")["""model_state_dict"""]
  __snake_case:		List[Any]					    =		NystromformerConfig.from_json_file(SCREAMING_SNAKE_CASE__)
  __snake_case:		Optional[int]					    =		NystromformerForMaskedLM(SCREAMING_SNAKE_CASE__)

  __snake_case:		Dict					    =		convert_checkpoint_helper(SCREAMING_SNAKE_CASE__   ,   SCREAMING_SNAKE_CASE__)

  model.load_state_dict(SCREAMING_SNAKE_CASE__)
  model.eval()
  model.save_pretrained(SCREAMING_SNAKE_CASE__)

  print(F'''Checkpoint successfuly converted. Model saved at {pytorch_dump_path}''')


if __name__ == "__main__":
      __UpperCAmelCase			:  str										=				argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          "--pytorch_model_path", default=None, type=str, required=True, help="Path to Nystromformer pytorch checkpoint."
      )
      parser.add_argument(
          "--config_file",
          default=None,
          type=str,
          required=True,
          help="The json file for Nystromformer model config.",
      )
      parser.add_argument(
          "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
      )
      __UpperCAmelCase			:  Optional[Any]										=				parser.parse_args()
      convert_nystromformer_checkpoint(args.pytorch_model_path, args.config_file, args.pytorch_dump_path)






