





"""simple docstring"""
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							:



 def __init__(     self				,       _snake_case				,       _snake_case=2				,       _snake_case=8				,       _snake_case=True				,       _snake_case=True				,       _snake_case=True				,       _snake_case=True				,       _snake_case=99				,       _snake_case=16				,       _snake_case=5				,       _snake_case=2				,       _snake_case=36				,       _snake_case="gelu"				,       _snake_case=0.0				,       _snake_case=0.0				,       _snake_case=5_12				,       _snake_case=16				,       _snake_case=2				,       _snake_case=0.02				,       _snake_case=3				,       _snake_case=4				,       _snake_case=None				,       ):







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







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







  """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=_snake_case				,       initializer_range=self.initializer_range				,       )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.get_config()
  lowerCAmelCase    				= 3_00
  return config



 def    UpperCamelCase__  (     self     ):







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







  """simple docstring"""






  lowerCAmelCase    				= MraModel(config=_snake_case     )
  model.to(_snake_case     )
  model.eval()
  lowerCAmelCase    				= model(_snake_case				,       attention_mask=_snake_case				,       token_type_ids=_snake_case     )
  lowerCAmelCase    				= model(_snake_case				,       token_type_ids=_snake_case     )
  lowerCAmelCase    				= model(_snake_case     )
  self.parent.assertEqual(result.last_hidden_state.shape				,       (self.batch_size, self.seq_length, self.hidden_size)     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       ):







  """simple docstring"""






  lowerCAmelCase    				= True
  lowerCAmelCase    				= MraModel(_snake_case     )
  model.to(_snake_case     )
  model.eval()
  lowerCAmelCase    				= model(
      _snake_case				,       attention_mask=_snake_case				,       token_type_ids=_snake_case				,       encoder_hidden_states=_snake_case				,       encoder_attention_mask=_snake_case				,       )
  lowerCAmelCase    				= model(
      _snake_case				,       attention_mask=_snake_case				,       token_type_ids=_snake_case				,       encoder_hidden_states=_snake_case				,       )
  lowerCAmelCase    				= model(_snake_case				,       attention_mask=_snake_case				,       token_type_ids=_snake_case     )
  self.parent.assertEqual(result.last_hidden_state.shape				,       (self.batch_size, self.seq_length, self.hidden_size)     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= MraForMaskedLM(config=_snake_case     )
  model.to(_snake_case     )
  model.eval()
  lowerCAmelCase    				= model(_snake_case				,       attention_mask=_snake_case				,       token_type_ids=_snake_case				,       labels=_snake_case     )
  self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.seq_length, self.vocab_size)     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= MraForQuestionAnswering(config=_snake_case     )
  model.to(_snake_case     )
  model.eval()
  lowerCAmelCase    				= model(
      _snake_case				,       attention_mask=_snake_case				,       token_type_ids=_snake_case				,       start_positions=_snake_case				,       end_positions=_snake_case				,       )
  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				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= self.num_labels
  lowerCAmelCase    				= MraForSequenceClassification(_snake_case     )
  model.to(_snake_case     )
  model.eval()
  lowerCAmelCase    				= model(_snake_case				,       attention_mask=_snake_case				,       token_type_ids=_snake_case				,       labels=_snake_case     )
  self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.num_labels)     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= self.num_labels
  lowerCAmelCase    				= MraForTokenClassification(config=_snake_case     )
  model.to(_snake_case     )
  model.eval()
  lowerCAmelCase    				= model(_snake_case				,       attention_mask=_snake_case				,       token_type_ids=_snake_case				,       labels=_snake_case     )
  self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.seq_length, self.num_labels)     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= self.num_choices
  lowerCAmelCase    				= MraForMultipleChoice(config=_snake_case     )
  model.to(_snake_case     )
  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(
      _snake_case				,       attention_mask=_snake_case				,       token_type_ids=_snake_case				,       labels=_snake_case				,       )
  self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.num_choices)     )



 def    UpperCamelCase__  (     self     ):







  """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							(   a__   , unittest.TestCase  ):
 snake_case__             =       (
     (
         MraModel,
         MraForMaskedLM,
         MraForMultipleChoice,
         MraForQuestionAnswering,
         MraForSequenceClassification,
         MraForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 snake_case__             =       False
 snake_case__             =       False
 snake_case__             =       False
 snake_case__             =       False

 snake_case__             =       ()



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= MraModelTester(self     )
  lowerCAmelCase    				= ConfigTester(self				,       config_class=_snake_case				,       hidden_size=37     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  self.config_tester.run_common_tests()



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_model(*_snake_case     )



 def    UpperCamelCase__  (     self     ):







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



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_for_masked_lm(*_snake_case     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_for_multiple_choice(*_snake_case     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_for_question_answering(*_snake_case     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_for_sequence_classification(*_snake_case     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_for_token_classification(*_snake_case     )



 @slow
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






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



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







  """simple docstring"""






  return







@require_torch
class       a							(   unittest.TestCase  ):



 @slow
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= MraModel.from_pretrained('uw-madison/mra-base-512-4'     )
  lowerCAmelCase    				= torch.arange(2_56     ).unsqueeze(0     )

  with torch.no_grad():
   lowerCAmelCase    				= model(_snake_case     )[0]

  lowerCAmelCase    				= torch.Size((1, 2_56, 7_68)     )
  self.assertEqual(output.shape				,       _snake_case     )

  lowerCAmelCase    				= torch.tensor(
      [[[-0.0_140, 0.0_830, -0.0_381], [0.1_546, 0.1_402, 0.0_220], [0.1_162, 0.0_851, 0.0_165]]]     )

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



 @slow
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= MraForMaskedLM.from_pretrained('uw-madison/mra-base-512-4'     )
  lowerCAmelCase    				= torch.arange(2_56     ).unsqueeze(0     )

  with torch.no_grad():
   lowerCAmelCase    				= model(_snake_case     )[0]

  lowerCAmelCase    				= 5_02_65

  lowerCAmelCase    				= torch.Size((1, 2_56, vocab_size)     )
  self.assertEqual(output.shape				,       _snake_case     )

  lowerCAmelCase    				= torch.tensor(
      [[[9.2_595, -3.6_038, 11.8_819], [9.3_869, -3.2_693, 11.0_956], [11.8_524, -3.4_938, 13.1_210]]]     )

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



 @slow
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






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

  with torch.no_grad():
   lowerCAmelCase    				= model(_snake_case     )[0]

  lowerCAmelCase    				= 5_02_65

  lowerCAmelCase    				= torch.Size((1, 40_96, vocab_size)     )
  self.assertEqual(output.shape				,       _snake_case     )

  lowerCAmelCase    				= torch.tensor(
      [[[5.4_789, -2.3_564, 7.5_064], [7.9_067, -1.3_369, 9.9_668], [9.0_712, -1.8_106, 7.0_380]]]     )

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












"""simple docstring"""
from ...configuration_utils import PretrainedConfig
from ...utils import logging


__UpperCamelCase     :						Optional[Any]						       =  logging.get_logger(__name__)

__UpperCamelCase     :						List[Any]						       =  {
    # See all MEGATRON_BERT models at https://huggingface.co/models?filter=bert
}







class       a							(   a__  ):
 snake_case__             =       '''megatron-bert'''



 def __init__(     self				,       _snake_case=2_90_56				,       _snake_case=10_24				,       _snake_case=24				,       _snake_case=16				,       _snake_case=40_96				,       _snake_case="gelu"				,       _snake_case=0.1				,       _snake_case=0.1				,       _snake_case=5_12				,       _snake_case=2				,       _snake_case=0.02				,       _snake_case=1E-12				,       _snake_case=0				,       _snake_case="absolute"				,       _snake_case=True				,       **_snake_case				,       ):







  """simple docstring"""






  super().__init__(pad_token_id=_snake_case				,       **_snake_case     )

  lowerCAmelCase    				= vocab_size
  lowerCAmelCase    				= hidden_size
  lowerCAmelCase    				= num_hidden_layers
  lowerCAmelCase    				= num_attention_heads
  lowerCAmelCase    				= hidden_act
  lowerCAmelCase    				= intermediate_size
  lowerCAmelCase    				= hidden_dropout_prob
  lowerCAmelCase    				= attention_probs_dropout_prob
  lowerCAmelCase    				= max_position_embeddings
  lowerCAmelCase    				= type_vocab_size
  lowerCAmelCase    				= initializer_range
  lowerCAmelCase    				= layer_norm_eps
  lowerCAmelCase    				= position_embedding_type
  lowerCAmelCase    				= use_cache





