


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

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






     """simple docstring"""





     lowerCamelCase			: Union[str, Any]   =    parent
     lowerCamelCase			: Any   =    batch_size
     lowerCamelCase			: int   =    seq_length
     lowerCamelCase			: Optional[Any]   =    is_training
     lowerCamelCase			: str   =    use_input_mask
     lowerCamelCase			: Dict   =    use_token_type_ids
     lowerCamelCase			: Tuple   =    use_labels
     lowerCamelCase			: Union[str, Any]   =    vocab_size
     lowerCamelCase			: Dict   =    hidden_size
     lowerCamelCase			: int   =    num_hidden_layers
     lowerCamelCase			: Union[str, Any]   =    num_attention_heads
     lowerCamelCase			: List[str]   =    intermediate_size
     lowerCamelCase			: Any   =    hidden_act
     lowerCamelCase			: Any   =    hidden_dropout_prob
     lowerCamelCase			: Any   =    attention_probs_dropout_prob
     lowerCamelCase			: List[str]   =    max_position_embeddings
     lowerCamelCase			: Optional[Any]   =    type_vocab_size
     lowerCamelCase			: Dict   =    type_sequence_label_size
     lowerCamelCase			: Tuple   =    initializer_range
     lowerCamelCase			: Optional[Any]   =    num_labels
     lowerCamelCase			: int   =    num_choices
     lowerCamelCase			: str   =    scope






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Union[str, Any]   =    ids_tensor([self.batch_size, self.seq_length]						,    self.vocab_size			)

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

     lowerCamelCase			: List[str]   =    None
     if self.use_token_type_ids:
       lowerCamelCase			: Dict   =    ids_tensor([self.batch_size, self.seq_length]						,    self.type_vocab_size			)

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

     lowerCamelCase			: Tuple   =    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			):






     """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 _snake_case   (		self			):






     """simple docstring"""





     (
         (
         lowerCamelCase
     )   ,     (
         lowerCamelCase
     )   ,     (
         lowerCamelCase
     )   ,     (
         lowerCamelCase
     )   ,     (
         lowerCamelCase
     )   ,     (
         lowerCamelCase
     )   ,     (
         lowerCamelCase
     )   ,     (
         lowerCamelCase
     )   ,     
     )			: Optional[Any]   =    self.prepare_config_and_inputs()

     lowerCamelCase			: Any   =    True
     lowerCamelCase			: Dict   =    floats_tensor([self.batch_size, self.seq_length, self.hidden_size]			)
     lowerCamelCase			: Union[str, 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 _snake_case   (		self						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    ):






     """simple docstring"""





     lowerCamelCase			: Tuple   =    ElectraModel(config=__A			)
     model.to(__A			)
     model.eval()
     lowerCamelCase			: str   =    model(__A						,    attention_mask=__A						,    token_type_ids=__A			)
     lowerCamelCase			: List[Any]   =    model(__A						,    token_type_ids=__A			)
     lowerCamelCase			: Optional[Any]   =    model(__A			)
     self.parent.assertEqual(result.last_hidden_state.shape						,    (self.batch_size, self.seq_length, self.hidden_size)			)






   def _snake_case   (		self						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    ):






     """simple docstring"""





     lowerCamelCase			: Optional[int]   =    True
     lowerCamelCase			: Dict   =    ElectraModel(__A			)
     model.to(__A			)
     model.eval()
     lowerCamelCase			: Optional[Any]   =    model(
         __A						,    attention_mask=__A						,    token_type_ids=__A						,    encoder_hidden_states=__A						,    encoder_attention_mask=__A						,    )
     lowerCamelCase			: Optional[Any]   =    model(
         __A						,    attention_mask=__A						,    token_type_ids=__A						,    encoder_hidden_states=__A						,    )
     lowerCamelCase			: List[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 _snake_case   (		self						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    ):






     """simple docstring"""





     lowerCamelCase			: Tuple   =    ElectraForMaskedLM(config=__A			)
     model.to(__A			)
     model.eval()
     lowerCamelCase			: Dict   =    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						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    ):






     """simple docstring"""





     lowerCamelCase			: Dict   =    ElectraForCausalLM(config=__A			)
     model.to(__A			)
     model.eval()
     lowerCamelCase			: 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 _snake_case   (		self						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    ):






     """simple docstring"""





     lowerCamelCase			: Optional[Any]   =    self.num_labels
     lowerCamelCase			: Dict   =    ElectraForTokenClassification(config=__A			)
     model.to(__A			)
     model.eval()
     lowerCamelCase			: Union[str, Any]   =    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						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    ):






     """simple docstring"""





     lowerCamelCase			: Any   =    self.num_labels
     lowerCamelCase			: List[str]   =    ElectraForPreTraining(config=__A			)
     model.to(__A			)
     model.eval()
     lowerCamelCase			: int   =    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						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    ):






     """simple docstring"""





     lowerCamelCase			: List[str]   =    self.num_labels
     lowerCamelCase			: Tuple   =    ElectraForSequenceClassification(__A			)
     model.to(__A			)
     model.eval()
     lowerCamelCase			: Dict   =    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						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    ):






     """simple docstring"""





     lowerCamelCase			: Optional[Any]   =    ElectraForQuestionAnswering(config=__A			)
     model.to(__A			)
     model.eval()
     lowerCamelCase			: str   =    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						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    __A						,    ):






     """simple docstring"""





     lowerCamelCase			: str   =    self.num_choices
     lowerCamelCase			: List[Any]   =    ElectraForMultipleChoice(config=__A			)
     model.to(__A			)
     model.eval()
     lowerCamelCase			: str   =    input_ids.unsqueeze(1			).expand(-1						,    self.num_choices						,    -1			).contiguous()
     lowerCamelCase			: List[str]   =    token_type_ids.unsqueeze(1			).expand(-1						,    self.num_choices						,    -1			).contiguous()
     lowerCamelCase			: Dict   =    input_mask.unsqueeze(1			).expand(-1						,    self.num_choices						,    -1			).contiguous()
     lowerCamelCase			: 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 _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Any   =    self.prepare_config_and_inputs()
     (
         (
         lowerCamelCase
     )   ,     (
         lowerCamelCase
     )   ,     (
         lowerCamelCase
     )   ,     (
         lowerCamelCase
     )   ,     (
         lowerCamelCase
     )   ,     (
         lowerCamelCase
     )   ,     (
         lowerCamelCase
     )   ,     (
         lowerCamelCase
     )   ,     
     )			: Union[str, Any]   =    config_and_inputs
     lowerCamelCase			: Optional[int]   =    {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
     return config, inputs_dict



@require_torch
class      UpperCAmelCase_ ( UpperCamelCase		,			UpperCamelCase		,			unittest.TestCase   ):

   '''simple docstring'''




   __A			:				Optional[Any]							=		(
       (
           ElectraModel,
           ElectraForPreTraining,
           ElectraForMaskedLM,
           ElectraForCausalLM,
           ElectraForMultipleChoice,
           ElectraForTokenClassification,
           ElectraForSequenceClassification,
           ElectraForQuestionAnswering,
       )
       if is_torch_available()
       else ()
   )
   __A			:				str							=		(
       {
           "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 {}
   )
   __A			:				int							=		True






   def _snake_case   (		self						,    __A						,    __A						,    __A=False			):






     """simple docstring"""





     lowerCamelCase			: List[Any]   =    super()._prepare_for_class(__A						,    __A						,    return_labels=__A			)

     if return_labels:
       if model_class in get_values(__A			):
         lowerCamelCase			: List[str]   =    torch.zeros(
             (self.model_tester.batch_size, self.model_tester.seq_length)						,    dtype=torch.long						,    device=__A			)
     return inputs_dict






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: List[str]   =    ElectraModelTester(self			)
     lowerCamelCase			: int   =    ConfigTester(self						,    config_class=__A						,    hidden_size=37			)






   def _snake_case   (		self			):






     """simple docstring"""





     self.config_tester.run_common_tests()






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Optional[int]   =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_electra_model(*__A			)






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: List[Any]   =    self.model_tester.prepare_config_and_inputs_for_decoder()
     self.model_tester.create_and_check_electra_model_as_decoder(*__A			)






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Tuple   =    self.model_tester.prepare_config_and_inputs()
     for type in ["absolute", "relative_key", "relative_key_query"]:
       lowerCamelCase			: Union[str, Any]   =    type
       self.model_tester.create_and_check_electra_model(*__A			)






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: List[Any]   =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_electra_for_masked_lm(*__A			)






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Tuple   =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_electra_for_token_classification(*__A			)






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Optional[Any]   =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_electra_for_pretraining(*__A			)






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: int   =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_electra_for_sequence_classification(*__A			)






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Tuple   =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_electra_for_question_answering(*__A			)






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Dict   =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_electra_for_multiple_choice(*__A			)






   @slow
   def _snake_case   (		self			):






     """simple docstring"""





     for model_name in ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       lowerCamelCase			: str   =    ElectraModel.from_pretrained(__A			)
       self.assertIsNotNone(__A			)






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Optional[Any]   =    self.model_tester.prepare_config_and_inputs_for_decoder()
     self.model_tester.create_and_check_electra_for_causal_lm(*__A			)



@require_torch
class      UpperCAmelCase_ ( unittest.TestCase   ):

   '''simple docstring'''




   @slow
   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Union[str, Any]   =    ElectraModel.from_pretrained("google/electra-small-discriminator"			)
     lowerCamelCase			: Optional[Any]   =    torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]			)
     lowerCamelCase			: str   =    torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]			)
     lowerCamelCase			: Optional[int]   =    model(__A						,    attention_mask=__A			)[0]
     lowerCamelCase			: Optional[Any]   =    torch.Size((1, 11, 256)			)
     self.assertEqual(output.shape						,    __A			)
     lowerCamelCase			: int   =    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 importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp






def 		lowercase_(  ):


  '''simple docstring'''

  lowerCamelCase			: str   =    ArgumentParser(
      description=(
          "PyTorch TPU distributed training launch helper utility that will spawn up multiple distributed processes"
      )     )

  # Optional arguments for the launch helper
  parser.add_argument("--num_cores"				,						type=SCREAMING_SNAKE_CASE_				,						default=1				,						help="Number of TPU cores to use (1 or 8)."     )

  # positional
  parser.add_argument(
      "training_script"				,						type=SCREAMING_SNAKE_CASE_				,						help=(
          "The full path to the single TPU training "
          "program/script to be launched in parallel, "
          "followed by all the arguments for the "
          "training script"
      )				,						)

  # rest from the training program
  parser.add_argument("training_script_args"				,						nargs=SCREAMING_SNAKE_CASE_     )

  return parser.parse_args()






def 		lowercase_(  ):


  '''simple docstring'''

  lowerCamelCase			: Optional[Any]   =    parse_args()

  # Import training_script as a module.
  lowerCamelCase			: Tuple   =    Path(args.training_script     )
  sys.path.append(str(script_fpath.parent.resolve()     )     )
  lowerCamelCase			: int   =    script_fpath.stem
  lowerCamelCase			: Tuple   =    importlib.import_module(SCREAMING_SNAKE_CASE_     )

  # Patch sys.argv
  lowerCamelCase			: Union[str, Any]   =    [args.training_script] + args.training_script_args + ["--tpu_num_cores", str(args.num_cores     )]

  xmp.spawn(mod._mp_fn				,						args=()				,						nprocs=args.num_cores     )


if __name__ == "__main__":
   main()



