


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


   '''simple docstring'''







   def __init__(     self   ,							UpperCamelCase_   ,							UpperCamelCase_=2   ,							UpperCamelCase_=8   ,							UpperCamelCase_=True   ,							UpperCamelCase_=True   ,							UpperCamelCase_=True   ,							UpperCamelCase_=True   ,							UpperCamelCase_=99   ,							UpperCamelCase_=16   ,							UpperCamelCase_=5   ,							UpperCamelCase_=2   ,							UpperCamelCase_=36   ,							UpperCamelCase_="gelu"   ,							UpperCamelCase_=0.0   ,							UpperCamelCase_=0.0   ,							UpperCamelCase_=512   ,							UpperCamelCase_=16   ,							UpperCamelCase_=2   ,							UpperCamelCase_=0.02   ,							UpperCamelCase_=3   ,							UpperCamelCase_=4   ,							UpperCamelCase_=None   ,							):
        lowercase_       :Any     =       parent
        lowercase_       :Union[str, Any]     =       batch_size
        lowercase_       :Tuple     =       seq_length
        lowercase_       :int     =       is_training
        lowercase_       :Any     =       use_input_mask
        lowercase_       :List[str]     =       use_token_type_ids
        lowercase_       :Any     =       use_labels
        lowercase_       :Optional[int]     =       vocab_size
        lowercase_       :Dict     =       hidden_size
        lowercase_       :Any     =       num_hidden_layers
        lowercase_       :Union[str, Any]     =       num_attention_heads
        lowercase_       :Optional[int]     =       intermediate_size
        lowercase_       :Tuple     =       hidden_act
        lowercase_       :Union[str, Any]     =       hidden_dropout_prob
        lowercase_       :Dict     =       attention_probs_dropout_prob
        lowercase_       :List[Any]     =       max_position_embeddings
        lowercase_       :Tuple     =       type_vocab_size
        lowercase_       :Optional[Any]     =       type_sequence_label_size
        lowercase_       :str     =       initializer_range
        lowercase_       :List[str]     =       num_labels
        lowercase_       :Optional[int]     =       num_choices
        lowercase_       :Optional[Any]     =       scope
   def 		UpperCamelCase  (     self		):
        lowercase_       :Optional[Any]     =       ids_tensor([self.batch_size, self.seq_length]   ,							self.vocab_size		)

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

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

        lowercase_       :Union[str, Any]     =       None
        lowercase_       :Union[str, Any]     =       None
        lowercase_       :str     =       None
        if self.use_labels:
             lowercase_       :Tuple     =       ids_tensor([self.batch_size]   ,							self.type_sequence_label_size		)
             lowercase_       :Dict     =       ids_tensor([self.batch_size, self.seq_length]   ,							self.num_labels		)
             lowercase_       :Optional[int]     =       ids_tensor([self.batch_size]   ,							self.num_choices		)

        lowercase_       :List[Any]     =       self.get_config()

        return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
   def 		UpperCamelCase  (     self		):
        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=UpperCamelCase_   ,							initializer_range=self.initializer_range   ,							)
   def 		UpperCamelCase  (     self		):
        lowercase_       :Optional[int]     =       self.get_config()
        lowercase_       :List[Any]     =       300
        return config
   def 		UpperCamelCase  (     self		):
        (
            (
            lowercase_
        )    ,      (
            lowercase_
        )    ,      (
            lowercase_
        )    ,      (
            lowercase_
        )    ,      (
            lowercase_
        )    ,      (
            lowercase_
        )    ,      (
            lowercase_
        )    ,      
        )       :Dict     =       self.prepare_config_and_inputs()

        lowercase_       :Optional[Any]     =       True
        lowercase_       :str     =       floats_tensor([self.batch_size, self.seq_length, self.hidden_size]		)
        lowercase_       :Dict     =       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   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_		):
        lowercase_       :Tuple     =       MraModel(config=UpperCamelCase_		)
        model.to(UpperCamelCase_		)
        model.eval()
        lowercase_       :Optional[int]     =       model(UpperCamelCase_   ,							attention_mask=UpperCamelCase_   ,							token_type_ids=UpperCamelCase_		)
        lowercase_       :Any     =       model(UpperCamelCase_   ,							token_type_ids=UpperCamelCase_		)
        lowercase_       :Dict     =       model(UpperCamelCase_		)
        self.parent.assertEqual(result.last_hidden_state.shape   ,							(self.batch_size, self.seq_length, self.hidden_size)		)
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							):
        lowercase_       :Optional[int]     =       True
        lowercase_       :Optional[Any]     =       MraModel(UpperCamelCase_		)
        model.to(UpperCamelCase_		)
        model.eval()
        lowercase_       :Dict     =       model(
            UpperCamelCase_   ,							attention_mask=UpperCamelCase_   ,							token_type_ids=UpperCamelCase_   ,							encoder_hidden_states=UpperCamelCase_   ,							encoder_attention_mask=UpperCamelCase_   ,							)
        lowercase_       :List[str]     =       model(
            UpperCamelCase_   ,							attention_mask=UpperCamelCase_   ,							token_type_ids=UpperCamelCase_   ,							encoder_hidden_states=UpperCamelCase_   ,							)
        lowercase_       :Any     =       model(UpperCamelCase_   ,							attention_mask=UpperCamelCase_   ,							token_type_ids=UpperCamelCase_		)
        self.parent.assertEqual(result.last_hidden_state.shape   ,							(self.batch_size, self.seq_length, self.hidden_size)		)
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_		):
        lowercase_       :Tuple     =       MraForMaskedLM(config=UpperCamelCase_		)
        model.to(UpperCamelCase_		)
        model.eval()
        lowercase_       :Optional[int]     =       model(UpperCamelCase_   ,							attention_mask=UpperCamelCase_   ,							token_type_ids=UpperCamelCase_   ,							labels=UpperCamelCase_		)
        self.parent.assertEqual(result.logits.shape   ,							(self.batch_size, self.seq_length, self.vocab_size)		)
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_		):
        lowercase_       :List[str]     =       MraForQuestionAnswering(config=UpperCamelCase_		)
        model.to(UpperCamelCase_		)
        model.eval()
        lowercase_       :Union[str, Any]     =       model(
            UpperCamelCase_   ,							attention_mask=UpperCamelCase_   ,							token_type_ids=UpperCamelCase_   ,							start_positions=UpperCamelCase_   ,							end_positions=UpperCamelCase_   ,							)
        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   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_		):
        lowercase_       :List[Any]     =       self.num_labels
        lowercase_       :Any     =       MraForSequenceClassification(UpperCamelCase_		)
        model.to(UpperCamelCase_		)
        model.eval()
        lowercase_       :List[str]     =       model(UpperCamelCase_   ,							attention_mask=UpperCamelCase_   ,							token_type_ids=UpperCamelCase_   ,							labels=UpperCamelCase_		)
        self.parent.assertEqual(result.logits.shape   ,							(self.batch_size, self.num_labels)		)
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_		):
        lowercase_       :Tuple     =       self.num_labels
        lowercase_       :Any     =       MraForTokenClassification(config=UpperCamelCase_		)
        model.to(UpperCamelCase_		)
        model.eval()
        lowercase_       :Dict     =       model(UpperCamelCase_   ,							attention_mask=UpperCamelCase_   ,							token_type_ids=UpperCamelCase_   ,							labels=UpperCamelCase_		)
        self.parent.assertEqual(result.logits.shape   ,							(self.batch_size, self.seq_length, self.num_labels)		)
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_		):
        lowercase_       :List[str]     =       self.num_choices
        lowercase_       :List[str]     =       MraForMultipleChoice(config=UpperCamelCase_		)
        model.to(UpperCamelCase_		)
        model.eval()
        lowercase_       :int     =       input_ids.unsqueeze(1		).expand(-1   ,							self.num_choices   ,							-1		).contiguous()
        lowercase_       :Any     =       token_type_ids.unsqueeze(1		).expand(-1   ,							self.num_choices   ,							-1		).contiguous()
        lowercase_       :List[str]     =       input_mask.unsqueeze(1		).expand(-1   ,							self.num_choices   ,							-1		).contiguous()
        lowercase_       :Optional[int]     =       model(
            UpperCamelCase_   ,							attention_mask=UpperCamelCase_   ,							token_type_ids=UpperCamelCase_   ,							labels=UpperCamelCase_   ,							)
        self.parent.assertEqual(result.logits.shape   ,							(self.batch_size, self.num_choices)		)




   def 		UpperCamelCase  (     self		):
        lowercase_       :int     =       self.prepare_config_and_inputs()
        (
            (
            lowercase_
        )    ,      (
            lowercase_
        )    ,      (
            lowercase_
        )    ,      (
            lowercase_
        )    ,      (
            lowercase_
        )    ,      (
            lowercase_
        )    ,      (
            lowercase_
        )    ,      
        )       :str     =       config_and_inputs
        lowercase_       :Optional[Any]     =       {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
        return config, inputs_dict






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


   '''simple docstring'''







   lowercase					:	Dict									=(
       (
           MraModel,
           MraForMaskedLM,
           MraForMultipleChoice,
           MraForQuestionAnswering,
           MraForSequenceClassification,
           MraForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   lowercase					:	Optional[int]									=False
   lowercase					:	str									=False
   lowercase					:	List[Any]									=False
   lowercase					:	Union[str, Any]									=False

   lowercase					:	str									=()
   def 		UpperCamelCase  (     self		):
        lowercase_       :Optional[Any]     =       MraModelTester(self		)
        lowercase_       :Any     =       ConfigTester(self   ,							config_class=UpperCamelCase_   ,							hidden_size=37		)
   def 		UpperCamelCase  (     self		):
        self.config_tester.run_common_tests()
   def 		UpperCamelCase  (     self		):
        lowercase_       :Optional[int]     =       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*UpperCamelCase_		)
   def 		UpperCamelCase  (     self		):
        lowercase_       :List[Any]     =       self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
             lowercase_       :List[str]     =       type
             self.model_tester.create_and_check_model(*UpperCamelCase_		)
   def 		UpperCamelCase  (     self		):
        lowercase_       :Optional[Any]     =       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase_		)
   def 		UpperCamelCase  (     self		):
        lowercase_       :Tuple     =       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*UpperCamelCase_		)
   def 		UpperCamelCase  (     self		):
        lowercase_       :Optional[Any]     =       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*UpperCamelCase_		)
   def 		UpperCamelCase  (     self		):
        lowercase_       :List[Any]     =       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*UpperCamelCase_		)
   def 		UpperCamelCase  (     self		):
        lowercase_       :int     =       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*UpperCamelCase_		)
   @slow
   def 		UpperCamelCase  (     self		):
        for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             lowercase_       :Optional[int]     =       MraModel.from_pretrained(UpperCamelCase_		)
             self.assertIsNotNone(UpperCamelCase_		)




   @unittest.skip(reason='''MRA does not output attentions'''		)
   def 		UpperCamelCase  (     self		):
        return






@require_torch
class      UpperCamelCase					(		unittest.TestCase     ):


   '''simple docstring'''







   @slow
   def 		UpperCamelCase  (     self		):
        lowercase_       :Dict     =       MraModel.from_pretrained('''uw-madison/mra-base-512-4'''		)
        lowercase_       :int     =       torch.arange(256		).unsqueeze(0		)

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

        lowercase_       :Tuple     =       torch.Size((1, 256, 768)		)
        self.assertEqual(output.shape   ,							UpperCamelCase_		)

        lowercase_       :Tuple     =       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]   ,							UpperCamelCase_   ,							atol=1E-4		)		)
   @slow
   def 		UpperCamelCase  (     self		):
        lowercase_       :int     =       MraForMaskedLM.from_pretrained('''uw-madison/mra-base-512-4'''		)
        lowercase_       :Union[str, Any]     =       torch.arange(256		).unsqueeze(0		)

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

        lowercase_       :Any     =       5_0265

        lowercase_       :Tuple     =       torch.Size((1, 256, vocab_size)		)
        self.assertEqual(output.shape   ,							UpperCamelCase_		)

        lowercase_       :List[Any]     =       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]   ,							UpperCamelCase_   ,							atol=1E-4		)		)




   @slow
   def 		UpperCamelCase  (     self		):
        lowercase_       :Union[str, Any]     =       MraForMaskedLM.from_pretrained('''uw-madison/mra-base-4096-8-d3'''		)
        lowercase_       :int     =       torch.arange(4096		).unsqueeze(0		)

        with torch.no_grad():
             lowercase_       :List[Any]     =       model(UpperCamelCase_		)[0]

        lowercase_       :List[Any]     =       5_0265

        lowercase_       :str     =       torch.Size((1, 4096, vocab_size)		)
        self.assertEqual(output.shape   ,							UpperCamelCase_		)

        lowercase_       :List[Any]     =       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]   ,							UpperCamelCase_   ,							atol=1E-4		)		)





from torch import nn




def     UpperCamelCase    (      _a		)					->							str:
     '''simple docstring'''





     if act_fn in ["swish", "silu"]:
          return nn.SiLU()
     elif act_fn == "mish":
          return nn.Mish()
     elif act_fn == "gelu":
          return nn.GELU()
     else:
          raise ValueError(f"Unsupported activation function: {act_fn}"		)

