



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




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



    '''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 				lowerCamelCase					(  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 				lowerCamelCase					(  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=_a		, initializer_range=self.initializer_range		, )




 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    __lowerCamelCase  							=			self.get_config()
    __lowerCamelCase  							=			300
    return config




 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    (
        __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 				lowerCamelCase					(  self		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase ):



    '''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 				lowerCamelCase					(  self		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, ):



    '''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 				lowerCamelCase					(  self		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase ):



    '''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 				lowerCamelCase					(  self		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase ):



    '''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 				lowerCamelCase					(  self		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase ):



    '''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 				lowerCamelCase					(  self		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase ):



    '''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 				lowerCamelCase					(  self		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase ):



    '''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 				lowerCamelCase					(  self ):



    '''simple docstring'''
    __lowerCamelCase  							=			self.prepare_config_and_inputs()
    (
        __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        __lowerCAmelCase   (		__lowercase     ,		unittest.TestCase      ):
 lowerCAmelCase__				    =     (
     (
         MraModel,
         MraForMaskedLM,
         MraForMultipleChoice,
         MraForQuestionAnswering,
         MraForSequenceClassification,
         MraForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 lowerCAmelCase__				    =     False
 lowerCAmelCase__				    =     False
 lowerCAmelCase__				    =     False
 lowerCAmelCase__				    =     False

 lowerCAmelCase__				    =     ()




 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    __lowerCamelCase  							=			MraModelTester(self )
    __lowerCamelCase  							=			ConfigTester(self		, config_class=_a		, hidden_size=37 )




 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    self.config_tester.run_common_tests()




 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    __lowerCamelCase  							=			self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*_a )




 def 				lowerCamelCase					(  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(*_a )




 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    __lowerCamelCase  							=			self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_masked_lm(*_a )




 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    __lowerCamelCase  							=			self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_multiple_choice(*_a )




 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    __lowerCamelCase  							=			self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_question_answering(*_a )




 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    __lowerCamelCase  							=			self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_sequence_classification(*_a )




 def 				lowerCamelCase					(  self ):



    '''simple docstring'''
    __lowerCamelCase  							=			self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_token_classification(*_a )




 @slow
 def 				lowerCamelCase					(  self ):



    '''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 				lowerCamelCase					(  self ):



    '''simple docstring'''
    return



@require_torch
class        __lowerCAmelCase   (		unittest.TestCase      ):




 @slow
 def 				lowerCamelCase					(  self ):



    '''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.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]		, _a		, atol=1E-4 ) )




 @slow
 def 				lowerCamelCase					(  self ):



    '''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.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]		, _a		, atol=1E-4 ) )




 @slow
 def 				lowerCamelCase					(  self ):



    '''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.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]		, _a		, atol=1E-4 ) )









import operator as op


a__				=		'''scaler.pt'''
a__				=		'''pytorch_model'''
a__				=		'''random_states'''
a__				=		'''optimizer'''
a__				=		'''scheduler'''
a__				=		'''pytorch_model.bin'''
a__				=		'''pytorch_model.bin.index.json'''
a__				=		'''model.safetensors'''
a__				=		'''model.safetensors.index.json'''
a__				=		'''1.10.2'''
a__				=		'''py38'''
a__				=		'''4.17.0'''
a__				=		['''ml.p3.16xlarge''', '''ml.p3dn.24xlarge''', '''ml.p4dn.24xlarge''']
a__				=		['''FULL_SHARD''', '''SHARD_GRAD_OP''', '''NO_SHARD''', '''HYBRID_SHARD''', '''HYBRID_SHARD_ZERO2''']
a__				=		['''TRANSFORMER_BASED_WRAP''', '''SIZE_BASED_WRAP''', '''NO_WRAP''']
a__				=		['''BACKWARD_PRE''', '''BACKWARD_POST''', '''NO_PREFETCH''']
a__				=		['''FULL_STATE_DICT''', '''LOCAL_STATE_DICT''', '''SHARDED_STATE_DICT''']
a__				=		'''2.0.1'''
a__				=		['''pdsh''', '''standard''', '''openmpi''', '''mvapich''']
a__				=		['''default''', '''reduce-overhead''', '''max-autotune''']

a__				=		{'''>''': op.gt, '''>=''': op.ge, '''==''': op.eq, '''!=''': op.ne, '''<=''': op.le, '''<''': op.lt}

# These are the args for `torch.distributed.launch` for pytorch < 1.9
a__				=		[
    '''nnodes''',
    '''nproc_per_node''',
    '''rdzv_backend''',
    '''rdzv_endpoint''',
    '''rdzv_id''',
    '''rdzv_conf''',
    '''standalone''',
    '''max_restarts''',
    '''monitor_interval''',
    '''start_method''',
    '''role''',
    '''module''',
    '''m''',
    '''no_python''',
    '''run_path''',
    '''log_dir''',
    '''r''',
    '''redirects''',
    '''t''',
    '''tee''',
    '''node_rank''',
    '''master_addr''',
    '''master_port''',
]

a__				=		['''DEEPSPEED''', '''MULTI_GPU''', '''FSDP''', '''MEGATRON_LM''']
a__				=		['''DEEPSPEED''', '''MULTI_XPU''', '''FSDP''']



