



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     :List[Any]   ,__lowercase     :Union[str, Any]   ,__lowercase     :Optional[Any]=2   ,__lowercase     :List[Any]=8   ,__lowercase     :int=True   ,__lowercase     :Optional[Any]=True   ,__lowercase     :List[Any]=True   ,__lowercase     :List[Any]=True   ,__lowercase     :List[Any]=9_9   ,__lowercase     :Dict=1_6   ,__lowercase     :List[str]=5   ,__lowercase     :Tuple=2   ,__lowercase     :Dict=3_6   ,__lowercase     :List[str]="gelu"   ,__lowercase     :Optional[Any]=0.0   ,__lowercase     :Any=0.0   ,__lowercase     :Dict=5_1_2   ,__lowercase     :Optional[int]=1_6   ,__lowercase     :List[Any]=2   ,__lowercase     :int=0.02   ,__lowercase     :Optional[int]=3   ,__lowercase     :List[str]=4   ,__lowercase     :Any=None   ,):
						snake_case__	:      str     		=	parent
						snake_case__	:      str     		=	batch_size
						snake_case__	:      List[str]     		=	seq_length
						snake_case__	:      str     		=	is_training
						snake_case__	:      int     		=	use_input_mask
						snake_case__	:      List[Any]     		=	use_token_type_ids
						snake_case__	:      int     		=	use_labels
						snake_case__	:      Any     		=	vocab_size
						snake_case__	:      Any     		=	hidden_size
						snake_case__	:      List[Any]     		=	num_hidden_layers
						snake_case__	:      int     		=	num_attention_heads
						snake_case__	:      Tuple     		=	intermediate_size
						snake_case__	:      int     		=	hidden_act
						snake_case__	:      Any     		=	hidden_dropout_prob
						snake_case__	:      List[Any]     		=	attention_probs_dropout_prob
						snake_case__	:      Tuple     		=	max_position_embeddings
						snake_case__	:      str     		=	type_vocab_size
						snake_case__	:      Optional[Any]     		=	type_sequence_label_size
						snake_case__	:      Dict     		=	initializer_range
						snake_case__	:      List[str]     		=	num_labels
						snake_case__	:      Any     		=	num_choices
						snake_case__	:      Optional[int]     		=	scope



		def 	__lowerCamelCase			(     self     :List[Any]  ):
						snake_case__	:      List[Any]     		=	ids_tensor([self.batch_size, self.seq_length]   ,self.vocab_size  )

						snake_case__	:      Optional[Any]     		=	None
						if self.use_input_mask:
										snake_case__	:      List[Any]     		=	random_attention_mask([self.batch_size, self.seq_length]  )

						snake_case__	:      Tuple     		=	None
						if self.use_token_type_ids:
										snake_case__	:      Optional[int]     		=	ids_tensor([self.batch_size, self.seq_length]   ,self.type_vocab_size  )

						snake_case__	:      List[str]     		=	None
						snake_case__	:      int     		=	None
						snake_case__	:      Any     		=	None
						if self.use_labels:
										snake_case__	:      Dict     		=	ids_tensor([self.batch_size]   ,self.type_sequence_label_size  )
										snake_case__	:      List[Any]     		=	ids_tensor([self.batch_size, self.seq_length]   ,self.num_labels  )
										snake_case__	:      Any     		=	ids_tensor([self.batch_size]   ,self.num_choices  )

						snake_case__	:      Dict     		=	self.get_config()

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



		def 	__lowerCamelCase			(     self     :Optional[Any]  ):
						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=__lowercase   ,initializer_range=self.initializer_range   ,)



		def 	__lowerCamelCase			(     self     :Any  ):
						snake_case__	:      List[Any]     		=	self.get_config()
						snake_case__	:      Dict     		=	3_0_0
						return config



		def 	__lowerCamelCase			(     self     :Any  ):
						(
						    (
						    snake_case__
						)		, (
						    snake_case__
						)		, (
						    snake_case__
						)		, (
						    snake_case__
						)		, (
						    snake_case__
						)		, (
						    snake_case__
						)		, (
						    snake_case__
						)		, 
						)	:      Any     		=	self.prepare_config_and_inputs()

						snake_case__	:      int     		=	True
						snake_case__	:      Union[str, Any]     		=	floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
						snake_case__	:      Tuple     		=	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     :List[str]   ,__lowercase     :str   ,__lowercase     :List[str]   ,__lowercase     :List[str]   ,__lowercase     :Optional[int]   ,__lowercase     :Tuple   ,__lowercase     :List[Any]   ,__lowercase     :List[str]  ):
						snake_case__	:      Dict     		=	MraModel(config=__lowercase  )
						model.to(__lowercase  )
						model.eval()
						snake_case__	:      Optional[int]     		=	model(__lowercase   ,attention_mask=__lowercase   ,token_type_ids=__lowercase  )
						snake_case__	:      List[str]     		=	model(__lowercase   ,token_type_ids=__lowercase  )
						snake_case__	:      List[Any]     		=	model(__lowercase  )
						self.parent.assertEqual(result.last_hidden_state.shape   ,(self.batch_size, self.seq_length, self.hidden_size)  )



		def 	__lowerCamelCase			(     self     :List[Any]   ,__lowercase     :Dict   ,__lowercase     :str   ,__lowercase     :Optional[Any]   ,__lowercase     :List[str]   ,__lowercase     :List[str]   ,__lowercase     :Any   ,__lowercase     :Optional[Any]   ,__lowercase     :Tuple   ,__lowercase     :int   ,):
						snake_case__	:      Any     		=	True
						snake_case__	:      Any     		=	MraModel(__lowercase  )
						model.to(__lowercase  )
						model.eval()
						snake_case__	:      Tuple     		=	model(
						    __lowercase   ,attention_mask=__lowercase   ,token_type_ids=__lowercase   ,encoder_hidden_states=__lowercase   ,encoder_attention_mask=__lowercase   ,)
						snake_case__	:      Optional[int]     		=	model(
						    __lowercase   ,attention_mask=__lowercase   ,token_type_ids=__lowercase   ,encoder_hidden_states=__lowercase   ,)
						snake_case__	:      Dict     		=	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 	__lowerCamelCase			(     self     :Any   ,__lowercase     :int   ,__lowercase     :int   ,__lowercase     :List[str]   ,__lowercase     :Optional[int]   ,__lowercase     :Dict   ,__lowercase     :Union[str, Any]   ,__lowercase     :Tuple  ):
						snake_case__	:      Tuple     		=	MraForMaskedLM(config=__lowercase  )
						model.to(__lowercase  )
						model.eval()
						snake_case__	:      List[Any]     		=	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 	__lowerCamelCase			(     self     :Tuple   ,__lowercase     :Any   ,__lowercase     :Dict   ,__lowercase     :List[str]   ,__lowercase     :Tuple   ,__lowercase     :int   ,__lowercase     :Tuple   ,__lowercase     :List[Any]  ):
						snake_case__	:      Dict     		=	MraForQuestionAnswering(config=__lowercase  )
						model.to(__lowercase  )
						model.eval()
						snake_case__	:      Union[str, Any]     		=	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 	__lowerCamelCase			(     self     :Optional[int]   ,__lowercase     :Union[str, Any]   ,__lowercase     :int   ,__lowercase     :Union[str, Any]   ,__lowercase     :Dict   ,__lowercase     :List[Any]   ,__lowercase     :Any   ,__lowercase     :str  ):
						snake_case__	:      Any     		=	self.num_labels
						snake_case__	:      Optional[Any]     		=	MraForSequenceClassification(__lowercase  )
						model.to(__lowercase  )
						model.eval()
						snake_case__	:      int     		=	model(__lowercase   ,attention_mask=__lowercase   ,token_type_ids=__lowercase   ,labels=__lowercase  )
						self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.num_labels)  )



		def 	__lowerCamelCase			(     self     :Optional[int]   ,__lowercase     :Any   ,__lowercase     :Optional[int]   ,__lowercase     :str   ,__lowercase     :Any   ,__lowercase     :Union[str, Any]   ,__lowercase     :Union[str, Any]   ,__lowercase     :Tuple  ):
						snake_case__	:      int     		=	self.num_labels
						snake_case__	:      str     		=	MraForTokenClassification(config=__lowercase  )
						model.to(__lowercase  )
						model.eval()
						snake_case__	:      int     		=	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 	__lowerCamelCase			(     self     :Any   ,__lowercase     :str   ,__lowercase     :str   ,__lowercase     :Tuple   ,__lowercase     :Tuple   ,__lowercase     :List[str]   ,__lowercase     :str   ,__lowercase     :str  ):
						snake_case__	:      str     		=	self.num_choices
						snake_case__	:      Dict     		=	MraForMultipleChoice(config=__lowercase  )
						model.to(__lowercase  )
						model.eval()
						snake_case__	:      Dict     		=	input_ids.unsqueeze(1  ).expand(-1   ,self.num_choices   ,-1  ).contiguous()
						snake_case__	:      List[str]     		=	token_type_ids.unsqueeze(1  ).expand(-1   ,self.num_choices   ,-1  ).contiguous()
						snake_case__	:      Optional[Any]     		=	input_mask.unsqueeze(1  ).expand(-1   ,self.num_choices   ,-1  ).contiguous()
						snake_case__	:      Any     		=	model(
						    __lowercase   ,attention_mask=__lowercase   ,token_type_ids=__lowercase   ,labels=__lowercase   ,)
						self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.num_choices)  )






		def 	__lowerCamelCase			(     self     :Dict  ):
						snake_case__	:      Union[str, Any]     		=	self.prepare_config_and_inputs()
						(
						    (
						    snake_case__
						)		, (
						    snake_case__
						)		, (
						    snake_case__
						)		, (
						    snake_case__
						)		, (
						    snake_case__
						)		, (
						    snake_case__
						)		, (
						    snake_case__
						)		, 
						)	:      List[str]     		=	config_and_inputs
						snake_case__	:      int     		=	{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
						return config, inputs_dict


@require_torch
class       a							(			__lowerCamelCase  ,							unittest.TestCase				):
		__lowerCAmelCase	:    int											=					(
		    (
		        MraModel,
		        MraForMaskedLM,
		        MraForMultipleChoice,
		        MraForQuestionAnswering,
		        MraForSequenceClassification,
		        MraForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		__lowerCAmelCase	:    Union[str, Any]											=					False
		__lowerCAmelCase	:    Tuple											=					False
		__lowerCAmelCase	:    Optional[Any]											=					False
		__lowerCAmelCase	:    Dict											=					False

		__lowerCAmelCase	:    str											=					()



		def 	__lowerCamelCase			(     self     :Optional[int]  ):
						snake_case__	:      Optional[int]     		=	MraModelTester(self  )
						snake_case__	:      Optional[int]     		=	ConfigTester(self   ,config_class=__lowercase   ,hidden_size=3_7  )



		def 	__lowerCamelCase			(     self     :Tuple  ):
						self.config_tester.run_common_tests()



		def 	__lowerCamelCase			(     self     :Union[str, Any]  ):
						snake_case__	:      Optional[int]     		=	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*__lowercase  )



		def 	__lowerCamelCase			(     self     :Optional[Any]  ):
						snake_case__	:      int     		=	self.model_tester.prepare_config_and_inputs()
						for type in ["absolute", "relative_key", "relative_key_query"]:
										snake_case__	:      Tuple     		=	type
										self.model_tester.create_and_check_model(*__lowercase  )



		def 	__lowerCamelCase			(     self     :Optional[int]  ):
						snake_case__	:      List[str]     		=	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*__lowercase  )



		def 	__lowerCamelCase			(     self     :Union[str, Any]  ):
						snake_case__	:      Any     		=	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_multiple_choice(*__lowercase  )



		def 	__lowerCamelCase			(     self     :Union[str, Any]  ):
						snake_case__	:      List[str]     		=	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_question_answering(*__lowercase  )



		def 	__lowerCamelCase			(     self     :Any  ):
						snake_case__	:      int     		=	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_sequence_classification(*__lowercase  )



		def 	__lowerCamelCase			(     self     :List[str]  ):
						snake_case__	:      Tuple     		=	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*__lowercase  )



		@slow
		def 	__lowerCamelCase			(     self     :int  ):
						for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										snake_case__	:      int     		=	MraModel.from_pretrained(__lowercase  )
										self.assertIsNotNone(__lowercase  )






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


@require_torch
class       a							(			unittest.TestCase				):



		@slow
		def 	__lowerCamelCase			(     self     :Dict  ):
						snake_case__	:      List[Any]     		=	MraModel.from_pretrained('''uw-madison/mra-base-512-4'''  )
						snake_case__	:      Dict     		=	torch.arange(2_5_6  ).unsqueeze(0  )

						with torch.no_grad():
										snake_case__	:      Optional[Any]     		=	model(__lowercase  )[0]

						snake_case__	:      Tuple     		=	torch.Size((1, 2_5_6, 7_6_8)  )
						self.assertEqual(output.shape   ,__lowercase  )

						snake_case__	:      Any     		=	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]   ,__lowercase   ,atol=1e-4  )  )



		@slow
		def 	__lowerCamelCase			(     self     :List[str]  ):
						snake_case__	:      List[str]     		=	MraForMaskedLM.from_pretrained('''uw-madison/mra-base-512-4'''  )
						snake_case__	:      Optional[Any]     		=	torch.arange(2_5_6  ).unsqueeze(0  )

						with torch.no_grad():
										snake_case__	:      Union[str, Any]     		=	model(__lowercase  )[0]

						snake_case__	:      Union[str, Any]     		=	5_0_2_6_5

						snake_case__	:      Union[str, Any]     		=	torch.Size((1, 2_5_6, vocab_size)  )
						self.assertEqual(output.shape   ,__lowercase  )

						snake_case__	:      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]   ,__lowercase   ,atol=1e-4  )  )






		@slow
		def 	__lowerCamelCase			(     self     :List[Any]  ):
						snake_case__	:      Dict     		=	MraForMaskedLM.from_pretrained('''uw-madison/mra-base-4096-8-d3'''  )
						snake_case__	:      int     		=	torch.arange(4_0_9_6  ).unsqueeze(0  )

						with torch.no_grad():
										snake_case__	:      Optional[int]     		=	model(__lowercase  )[0]

						snake_case__	:      str     		=	5_0_2_6_5

						snake_case__	:      Optional[Any]     		=	torch.Size((1, 4_0_9_6, vocab_size)  )
						self.assertEqual(output.shape   ,__lowercase  )

						snake_case__	:      List[str]     		=	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]   ,__lowercase   ,atol=1e-4  )  )






from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


A__								=						logging.get_logger(__name__)

A__								=						{
    '''distilbert-base-uncased''': '''https://huggingface.co/distilbert-base-uncased/resolve/main/config.json''',
    '''distilbert-base-uncased-distilled-squad''': (
        '''https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/config.json'''
    ),
    '''distilbert-base-cased''': '''https://huggingface.co/distilbert-base-cased/resolve/main/config.json''',
    '''distilbert-base-cased-distilled-squad''': (
        '''https://huggingface.co/distilbert-base-cased-distilled-squad/resolve/main/config.json'''
    ),
    '''distilbert-base-german-cased''': '''https://huggingface.co/distilbert-base-german-cased/resolve/main/config.json''',
    '''distilbert-base-multilingual-cased''': (
        '''https://huggingface.co/distilbert-base-multilingual-cased/resolve/main/config.json'''
    ),
    '''distilbert-base-uncased-finetuned-sst-2-english''': (
        '''https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english/resolve/main/config.json'''
    ),
}


class       a							(			__lowerCamelCase				):
		__lowerCAmelCase	:    str											=					"""distilbert"""
		__lowerCAmelCase	:    str											=					{
		    """hidden_size""": """dim""",
		    """num_attention_heads""": """n_heads""",
		    """num_hidden_layers""": """n_layers""",
		}



		def __init__(     self     :Dict   ,__lowercase     :Tuple=3_0_5_2_2   ,__lowercase     :Optional[Any]=5_1_2   ,__lowercase     :List[str]=False   ,__lowercase     :List[str]=6   ,__lowercase     :Optional[Any]=1_2   ,__lowercase     :Tuple=7_6_8   ,__lowercase     :int=4 * 7_6_8   ,__lowercase     :List[Any]=0.1   ,__lowercase     :List[str]=0.1   ,__lowercase     :Union[str, Any]="gelu"   ,__lowercase     :List[str]=0.02   ,__lowercase     :Optional[int]=0.1   ,__lowercase     :Dict=0.2   ,__lowercase     :Union[str, Any]=0   ,**__lowercase     :Optional[Any]   ,):
						snake_case__	:      List[str]     		=	vocab_size
						snake_case__	:      Tuple     		=	max_position_embeddings
						snake_case__	:      Optional[int]     		=	sinusoidal_pos_embds
						snake_case__	:      str     		=	n_layers
						snake_case__	:      List[Any]     		=	n_heads
						snake_case__	:      Tuple     		=	dim
						snake_case__	:      str     		=	hidden_dim
						snake_case__	:      int     		=	dropout
						snake_case__	:      Dict     		=	attention_dropout
						snake_case__	:      Tuple     		=	activation
						snake_case__	:      int     		=	initializer_range
						snake_case__	:      Optional[Any]     		=	qa_dropout
						snake_case__	:      Union[str, Any]     		=	seq_classif_dropout
						super().__init__(**__lowercase   ,pad_token_id=__lowercase  )


class       a							(			__lowerCamelCase				):



		@property
		def 	__lowerCamelCase			(     self     :Union[str, Any]  ):
						if self.task == "multiple-choice":
										snake_case__	:      str     		=	{0: '''batch''', 1: '''choice''', 2: '''sequence'''}
						else:
										snake_case__	:      Optional[int]     		=	{0: '''batch''', 1: '''sequence'''}
						return OrderedDict(
						    [
						        ('''input_ids''', dynamic_axis),
						        ('''attention_mask''', dynamic_axis),
						    ]  )

