

import unittest

from transformers import YosoConfig, is_torch_available
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 (
           YosoForMaskedLM,
           YosoForMultipleChoice,
           YosoForQuestionAnswering,
           YosoForSequenceClassification,
           YosoForTokenClassification,
           YosoModel,
       )
       from transformers.models.yoso.modeling_yoso import YOSO_PRETRAINED_MODEL_ARCHIVE_LIST







class 		__magic_name__     :

   '''simple docstring'''




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






     """simple docstring"""






     a__    =parent
     a__    =batch_size
     a__    =seq_length
     a__    =is_training
     a__    =use_input_mask
     a__    =use_token_type_ids
     a__    =use_labels
     a__    =vocab_size
     a__    =hidden_size
     a__    =num_hidden_layers
     a__    =num_attention_heads
     a__    =intermediate_size
     a__    =hidden_act
     a__    =hidden_dropout_prob
     a__    =attention_probs_dropout_prob
     a__    =max_position_embeddings
     a__    =type_vocab_size
     a__    =type_sequence_label_size
     a__    =initializer_range
     a__    =num_labels
     a__    =num_choices
     a__    =scope



   def 					_UpperCAmelCase					(							self			)				->       Tuple:






     """simple docstring"""






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

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

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

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

     a__    =self.get_config()

     return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels



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






     """simple docstring"""






     return YosoConfig(
         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=lowercase_,       initializer_range=self.initializer_range,       )



   def 					_UpperCAmelCase					(							self			)				->       int:






     """simple docstring"""






     a__    =self.get_config()
     a__    =300
     return config



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






     """simple docstring"""






     (
         (
         a__
     ), (
         a__
     ), (
         a__
     ), (
         a__
     ), (
         a__
     ), (
         a__
     ), (
         a__
     ), 
     )    =self.prepare_config_and_inputs()

     a__    =True
     a__    =floats_tensor([self.batch_size, self.seq_length, self.hidden_size]			)
     a__    =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,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_			)				->       Optional[Any]:






     """simple docstring"""






     a__    =YosoModel(config=lowercase_			)
     model.to(lowercase_			)
     model.eval()
     a__    =model(lowercase_,       attention_mask=lowercase_,       token_type_ids=lowercase_			)
     a__    =model(lowercase_,       token_type_ids=lowercase_			)
     a__    =model(lowercase_			)
     self.parent.assertEqual(result.last_hidden_state.shape,       (self.batch_size, self.seq_length, self.hidden_size)			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       )				->       Optional[Any]:






     """simple docstring"""






     a__    =True
     a__    =YosoModel(lowercase_			)
     model.to(lowercase_			)
     model.eval()
     a__    =model(
         lowercase_,       attention_mask=lowercase_,       token_type_ids=lowercase_,       encoder_hidden_states=lowercase_,       encoder_attention_mask=lowercase_,       )
     a__    =model(
         lowercase_,       attention_mask=lowercase_,       token_type_ids=lowercase_,       encoder_hidden_states=lowercase_,       )
     a__    =model(lowercase_,       attention_mask=lowercase_,       token_type_ids=lowercase_			)
     self.parent.assertEqual(result.last_hidden_state.shape,       (self.batch_size, self.seq_length, self.hidden_size)			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_			)				->       Any:






     """simple docstring"""






     a__    =YosoForMaskedLM(config=lowercase_			)
     model.to(lowercase_			)
     model.eval()
     a__    =model(lowercase_,       attention_mask=lowercase_,       token_type_ids=lowercase_,       labels=lowercase_			)
     self.parent.assertEqual(result.logits.shape,       (self.batch_size, self.seq_length, self.vocab_size)			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_			)				->       Any:






     """simple docstring"""






     a__    =YosoForQuestionAnswering(config=lowercase_			)
     model.to(lowercase_			)
     model.eval()
     a__    =model(
         lowercase_,       attention_mask=lowercase_,       token_type_ids=lowercase_,       start_positions=lowercase_,       end_positions=lowercase_,       )
     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,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_			)				->       Union[str, Any]:






     """simple docstring"""






     a__    =self.num_labels
     a__    =YosoForSequenceClassification(lowercase_			)
     model.to(lowercase_			)
     model.eval()
     a__    =model(lowercase_,       attention_mask=lowercase_,       token_type_ids=lowercase_,       labels=lowercase_			)
     self.parent.assertEqual(result.logits.shape,       (self.batch_size, self.num_labels)			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_			)				->       Optional[int]:






     """simple docstring"""






     a__    =self.num_labels
     a__    =YosoForTokenClassification(config=lowercase_			)
     model.to(lowercase_			)
     model.eval()
     a__    =model(lowercase_,       attention_mask=lowercase_,       token_type_ids=lowercase_,       labels=lowercase_			)
     self.parent.assertEqual(result.logits.shape,       (self.batch_size, self.seq_length, self.num_labels)			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_			)				->       List[str]:






     """simple docstring"""






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



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






     """simple docstring"""






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







@require_torch
class 		__magic_name__     (				lowerCamelCase__						,			lowerCamelCase__						,			unittest.TestCase				):

   '''simple docstring'''




   lowerCamelCase__					:      int        =  (
       (
           YosoModel,
           YosoForMaskedLM,
           YosoForMultipleChoice,
           YosoForQuestionAnswering,
           YosoForSequenceClassification,
           YosoForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   lowerCamelCase__					:      Optional[int]        =  False
   lowerCamelCase__					:      Optional[int]        =  False
   lowerCamelCase__					:      List[str]        =  False

   lowerCamelCase__					:      int        =  ()
   lowerCamelCase__					:      Any        =  (
       {
           'feature-extraction': YosoModel,
           'fill-mask': YosoForMaskedLM,
           'question-answering': YosoForQuestionAnswering,
           'text-classification': YosoForSequenceClassification,
           'token-classification': YosoForTokenClassification,
           'zero-shot': YosoForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )



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






     """simple docstring"""






     a__    =YosoModelTester(self			)
     a__    =ConfigTester(self,       config_class=lowercase_,       hidden_size=37			)



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






     """simple docstring"""






     self.config_tester.run_common_tests()



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






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowercase_			)



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






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     for type in ["absolute", "relative_key", "relative_key_query"]:
       a__    =type
       self.model_tester.create_and_check_model(*lowercase_			)



   def 					_UpperCAmelCase					(							self			)				->       Dict:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_lm(*lowercase_			)



   def 					_UpperCAmelCase					(							self			)				->       Any:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_multiple_choice(*lowercase_			)



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






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_question_answering(*lowercase_			)



   def 					_UpperCAmelCase					(							self			)				->       str:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_sequence_classification(*lowercase_			)



   def 					_UpperCAmelCase					(							self			)				->       str:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_token_classification(*lowercase_			)



   @slow
   def 					_UpperCAmelCase					(							self			)				->       Dict:






     """simple docstring"""






     for model_name in YOSO_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       a__    =YosoModel.from_pretrained(lowercase_			)
       self.assertIsNotNone(lowercase_			)



   def 					_UpperCAmelCase					(							self			)				->       Dict:






     """simple docstring"""






     return







@require_torch
class 		__magic_name__     (				unittest.TestCase				):

   '''simple docstring'''




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






     """simple docstring"""






     a__    =YosoModel.from_pretrained('''uw-madison/yoso-4096'''			)
     a__    =torch.tensor([[0, 1, 2, 3, 4, 5]]			)

     with torch.no_grad():
       a__    =model(lowercase_			)[0]

     a__    =torch.Size((1, 6, 768)			)
     self.assertEqual(output.shape,       lowercase_			)

     a__    =torch.tensor(
         [[[-0.0611, 0.1242, 0.0840], [0.0280, -0.0048, 0.1125], [0.0106, 0.0226, 0.0751]]]			)

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



   @slow
   def 					_UpperCAmelCase					(							self			)				->       str:






     """simple docstring"""






     a__    =YosoForMaskedLM.from_pretrained('''uw-madison/yoso-4096'''			)
     a__    =torch.tensor([[0, 1, 2, 3, 4, 5]]			)

     with torch.no_grad():
       a__    =model(lowercase_			)[0]

     a__    =50265

     a__    =torch.Size((1, 6, vocab_size)			)
     self.assertEqual(output.shape,       lowercase_			)

     a__    =torch.tensor(
         [[[-2.1313, -3.7285, -2.2407], [-2.7047, -3.3314, -2.6408], [0.0629, -2.5166, -0.3356]]]			)

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



   @slow
   def 					_UpperCAmelCase					(							self			)				->       str:






     """simple docstring"""






     a__    =YosoForMaskedLM.from_pretrained('''uw-madison/yoso-4096'''			)
     a__    =torch.arange(4096			).unsqueeze(0			)

     with torch.no_grad():
       a__    =model(lowercase_			)[0]

     a__    =50265

     a__    =torch.Size((1, 4096, vocab_size)			)
     self.assertEqual(output.shape,       lowercase_			)

     a__    =torch.tensor(
         [[[-2.3914, -4.3742, -5.0956], [-4.0988, -4.2384, -7.0406], [-3.1427, -3.7192, -6.6800]]]			)

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



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

   '''simple docstring'''




   def __init__(							self,       lowercase_,       lowercase_=13,       lowercase_=7,       lowercase_=True,       lowercase_=True,       lowercase_=True,       lowercase_=True,       lowercase_=99,       lowercase_=24,       lowercase_=32,       lowercase_=5,       lowercase_=4,       lowercase_=37,       lowercase_="gelu",       lowercase_=0.1,       lowercase_=0.1,       lowercase_=512,       lowercase_=16,       lowercase_=2,       lowercase_=0.02,       lowercase_=4,       )				->       int:






     """simple docstring"""






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



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






     """simple docstring"""






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

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

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

     a__    =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			)				->       List[Any]:






     """simple docstring"""






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







@require_flax
class 		__magic_name__     (				lowerCamelCase__						,			unittest.TestCase				):

   '''simple docstring'''




   lowerCamelCase__					:      List[str]        =  True

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



   def 					_UpperCAmelCase					(							self			)				->       str:






     """simple docstring"""






     a__    =FlaxElectraModelTester(self			)



   @slow
   def 					_UpperCAmelCase					(							self			)				->       str:






     """simple docstring"""






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