




import unittest

from transformers import MraConfig, 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


if is_torch_available():
 import torch

 from transformers import (
     MraForMaskedLM,
     MraForMultipleChoice,
     MraForQuestionAnswering,
     MraForSequenceClassification,
     MraForTokenClassification,
     MraModel,
 )
 from transformers.models.mra.modeling_mra import MRA_PRETRAINED_MODEL_ARCHIVE_LIST







class   A:



 '''simple docstring'''


 def __init__(				self					:  str						,							A_					:  List[Any]						,							A_					:  List[str]=2						,							A_					:  str=8						,							A_					:  Optional[Any]=True						,							A_					:  Any=True						,							A_					:  List[str]=True						,							A_					:  Optional[int]=True						,							A_					:  int=99						,							A_					:  Dict=16						,							A_					:  int=5						,							A_					:  str=2						,							A_					:  int=36						,							A_					:  Optional[int]="gelu"						,							A_					:  Tuple=0.0						,							A_					:  List[Any]=0.0						,							A_					:  Optional[int]=512						,							A_					:  str=16						,							A_					:  int=2						,							A_					:  Tuple=0.02						,							A_					:  Optional[Any]=3						,							A_					:  Dict=4						,							A_					:  Optional[Any]=None						,							)   ->	List[Any]:

       """simple docstring"""






       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 				a__ (				self					:  int  )   ->	Optional[int]:

       """simple docstring"""






       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_		     =    self.get_config()

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

 def 				a__ (				self					:  Any  )   ->	Tuple:

       """simple docstring"""






       return MraConfig(
           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 				a__ (				self					:  Union[str, Any]  )   ->	Optional[int]:

       """simple docstring"""






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

 def 				a__ (				self					:  Dict  )   ->	str:

       """simple docstring"""






       (
           (
           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 				a__ (				self					:  Tuple						,							A_					:  List[str]						,							A_					:  Dict						,							A_					:  str						,							A_					:  Dict						,							A_					:  Union[str, Any]						,							A_					:  str						,							A_					:  Union[str, Any]  )   ->	Dict:

       """simple docstring"""






       lowerCamelCase_		     =    MraModel(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 				a__ (				self					:  List[str]						,							A_					:  int						,							A_					:  Optional[Any]						,							A_					:  str						,							A_					:  List[str]						,							A_					:  Any						,							A_					:  Dict						,							A_					:  Optional[Any]						,							A_					:  Dict						,							A_					:  str						,							)   ->	int:

       """simple docstring"""






       lowerCamelCase_		     =    True
       lowerCamelCase_		     =    MraModel(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 				a__ (				self					:  Dict						,							A_					:  Tuple						,							A_					:  Tuple						,							A_					:  Dict						,							A_					:  List[str]						,							A_					:  Optional[Any]						,							A_					:  Any						,							A_					:  int  )   ->	List[str]:

       """simple docstring"""






       lowerCamelCase_		     =    MraForMaskedLM(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 				a__ (				self					:  Any						,							A_					:  Union[str, Any]						,							A_					:  Optional[int]						,							A_					:  Any						,							A_					:  Any						,							A_					:  Union[str, Any]						,							A_					:  List[str]						,							A_					:  Union[str, Any]  )   ->	Union[str, Any]:

       """simple docstring"""






       lowerCamelCase_		     =    MraForQuestionAnswering(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 				a__ (				self					:  Optional[Any]						,							A_					:  int						,							A_					:  Optional[Any]						,							A_					:  List[Any]						,							A_					:  int						,							A_					:  Any						,							A_					:  Any						,							A_					:  Optional[Any]  )   ->	Any:

       """simple docstring"""






       lowerCamelCase_		     =    self.num_labels
       lowerCamelCase_		     =    MraForSequenceClassification(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 				a__ (				self					:  Optional[int]						,							A_					:  Dict						,							A_					:  int						,							A_					:  Optional[int]						,							A_					:  Dict						,							A_					:  str						,							A_					:  Dict						,							A_					:  List[str]  )   ->	Dict:

       """simple docstring"""






       lowerCamelCase_		     =    self.num_labels
       lowerCamelCase_		     =    MraForTokenClassification(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 				a__ (				self					:  Tuple						,							A_					:  Optional[Any]						,							A_					:  Any						,							A_					:  int						,							A_					:  Any						,							A_					:  List[Any]						,							A_					:  Dict						,							A_					:  List[Any]  )   ->	Optional[Any]:

       """simple docstring"""






       lowerCamelCase_		     =    self.num_choices
       lowerCamelCase_		     =    MraForMultipleChoice(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 				a__ (				self					:  Optional[int]  )   ->	str:

       """simple docstring"""






       lowerCamelCase_		     =    self.prepare_config_and_inputs()
       (
           (
           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(						UpperCamelCase							,      unittest.TestCase   ):



 '''simple docstring'''


 UpperCamelCase						   =						(
     (
         MraModel,
         MraForMaskedLM,
         MraForMultipleChoice,
         MraForQuestionAnswering,
         MraForSequenceClassification,
         MraForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 UpperCamelCase						   =						False
 UpperCamelCase						   =						False
 UpperCamelCase						   =						False
 UpperCamelCase						   =						False

 UpperCamelCase						   =						()

 def 				a__ (				self					:  str  )   ->	Dict:

       """simple docstring"""






       lowerCamelCase_		     =    MraModelTester(self  )
       lowerCamelCase_		     =    ConfigTester(self						,							config_class=A_						,							hidden_size=37  )

 def 				a__ (				self					:  Optional[Any]  )   ->	Tuple:

       """simple docstring"""






       self.config_tester.run_common_tests()

 def 				a__ (				self					:  Optional[int]  )   ->	str:

       """simple docstring"""






       lowerCamelCase_		     =    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*A_  )

 def 				a__ (				self					:  Dict  )   ->	str:

       """simple docstring"""






       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_model(*A_  )

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

       """simple docstring"""






       lowerCamelCase_		     =    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*A_  )

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

       """simple docstring"""






       lowerCamelCase_		     =    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_multiple_choice(*A_  )

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

       """simple docstring"""






       lowerCamelCase_		     =    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_question_answering(*A_  )

 def 				a__ (				self					:  Tuple  )   ->	int:

       """simple docstring"""






       lowerCamelCase_		     =    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_sequence_classification(*A_  )

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

       """simple docstring"""






       lowerCamelCase_		     =    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*A_  )

 @slow
 def 				a__ (				self					:  int  )   ->	str:

       """simple docstring"""






       for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             lowerCamelCase_		     =    MraModel.from_pretrained(A_  )
             self.assertIsNotNone(A_  )





 @unittest.skip(reason='MRA does not output attentions'  )
 def 				a__ (				self					:  List[Any]  )   ->	Tuple:

       """simple docstring"""






       return







@require_torch
class   A(						unittest.TestCase   ):



 '''simple docstring'''


 @slow
 def 				a__ (				self					:  int  )   ->	Any:

       """simple docstring"""






       lowerCamelCase_		     =    MraModel.from_pretrained('uw-madison/mra-base-512-4'  )
       lowerCamelCase_		     =    torch.arange(256  ).unsqueeze(0  )

       with torch.no_grad():
             lowerCamelCase_		     =    model(A_  )[0]

       lowerCamelCase_		     =    torch.Size((1, 256, 768)  )
       self.assertEqual(output.shape						,							A_  )

       lowerCamelCase_		     =    torch.tensor(
           [[[-0.0140, 0.0830, -0.0381], [0.1546, 0.1402, 0.0220], [0.1162, 0.0851, 0.0165]]]  )

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

 @slow
 def 				a__ (				self					:  int  )   ->	Any:

       """simple docstring"""






       lowerCamelCase_		     =    MraForMaskedLM.from_pretrained('uw-madison/mra-base-512-4'  )
       lowerCamelCase_		     =    torch.arange(256  ).unsqueeze(0  )

       with torch.no_grad():
             lowerCamelCase_		     =    model(A_  )[0]

       lowerCamelCase_		     =    50265

       lowerCamelCase_		     =    torch.Size((1, 256, vocab_size)  )
       self.assertEqual(output.shape						,							A_  )

       lowerCamelCase_		     =    torch.tensor(
           [[[9.2595, -3.6038, 11.8819], [9.3869, -3.2693, 11.0956], [11.8524, -3.4938, 13.1210]]]  )

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





 @slow
 def 				a__ (				self					:  Union[str, Any]  )   ->	Dict:

       """simple docstring"""






       lowerCamelCase_		     =    MraForMaskedLM.from_pretrained('uw-madison/mra-base-4096-8-d3'  )
       lowerCamelCase_		     =    torch.arange(4096  ).unsqueeze(0  )

       with torch.no_grad():
             lowerCamelCase_		     =    model(A_  )[0]

       lowerCamelCase_		     =    50265

       lowerCamelCase_		     =    torch.Size((1, 4096, vocab_size)  )
       self.assertEqual(output.shape						,							A_  )

       lowerCamelCase_		     =    torch.tensor(
           [[[5.4789, -2.3564, 7.5064], [7.9067, -1.3369, 9.9668], [9.0712, -1.8106, 7.0380]]]  )

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









from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


lowerCamelCase				:		str     =   {
    "configuration_nezha": ["NEZHA_PRETRAINED_CONFIG_ARCHIVE_MAP", "NezhaConfig"],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCamelCase				:		List[str]     =   [
     "NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST",
     "NezhaForNextSentencePrediction",
     "NezhaForMaskedLM",
     "NezhaForPreTraining",
     "NezhaForMultipleChoice",
     "NezhaForQuestionAnswering",
     "NezhaForSequenceClassification",
     "NezhaForTokenClassification",
     "NezhaModel",
     "NezhaPreTrainedModel",
 ]


if TYPE_CHECKING:
 from .configuration_nezha import NEZHA_PRETRAINED_CONFIG_ARCHIVE_MAP, NezhaConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_nezha import (
      NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST,
      NezhaForMaskedLM,
      NezhaForMultipleChoice,
      NezhaForNextSentencePrediction,
      NezhaForPreTraining,
      NezhaForQuestionAnswering,
      NezhaForSequenceClassification,
      NezhaForTokenClassification,
      NezhaModel,
      NezhaPreTrainedModel,
  )


else:
 import sys

 lowerCamelCase				:		str     =   _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



