
import unittest

from parameterized import parameterized

from transformers import OpenLlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers import OpenLlamaForCausalLM, OpenLlamaForSequenceClassification, OpenLlamaModel







class 				lowercase		:




							def __init__(				self   ,A__   ,A__=1_3   ,A__=7   ,A__=True   ,A__=True   ,A__=False   ,A__=True   ,A__=9_9   ,A__=3_2   ,A__=5   ,A__=4   ,A__=3_7   ,A__="gelu"   ,A__=0.1   ,A__=0.1   ,A__=5_1_2   ,A__=1_6   ,A__=2   ,A__=0.02   ,A__=3   ,A__=4   ,A__=None   ,):
												lowercase							=						parent
												lowercase							=						batch_size
												lowercase							=						seq_length
												lowercase							=						is_training
												lowercase							=						use_input_mask
												lowercase							=						use_token_type_ids
												lowercase							=						use_labels
												lowercase							=						vocab_size
												lowercase							=						hidden_size
												lowercase							=						num_hidden_layers
												lowercase							=						num_attention_heads
												lowercase							=						intermediate_size
												lowercase							=						hidden_act
												lowercase							=						hidden_dropout_prob
												lowercase							=						attention_probs_dropout_prob
												lowercase							=						max_position_embeddings
												lowercase							=						type_vocab_size
												lowercase							=						type_sequence_label_size
												lowercase							=						initializer_range
												lowercase							=						num_labels
												lowercase							=						num_choices
												lowercase							=						scope




							def  A__      (				self):
												lowercase							=						ids_tensor([self.batch_size, self.seq_length]   ,self.vocab_size)

												lowercase							=						None
												if self.use_input_mask:
																	lowercase							=						random_attention_mask([self.batch_size, self.seq_length])

												lowercase							=						None
												if self.use_token_type_ids:
																	lowercase							=						ids_tensor([self.batch_size, self.seq_length]   ,self.type_vocab_size)

												lowercase							=						None
												lowercase							=						None
												lowercase							=						None
												if self.use_labels:
																	lowercase							=						ids_tensor([self.batch_size]   ,self.type_sequence_label_size)
																	lowercase							=						ids_tensor([self.batch_size, self.seq_length]   ,self.num_labels)
																	lowercase							=						ids_tensor([self.batch_size]   ,self.num_choices)

												lowercase							=						self.get_config()

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




							def  A__      (				self):
												return OpenLlamaConfig(
												    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   ,use_stable_embedding=__snake_case   ,)




							def  A__      (				self   ,A__   ,A__   ,A__   ,A__   ,A__   ,A__   ,A__):
												lowercase							=						OpenLlamaModel(config=__snake_case)
												model.to(__snake_case)
												model.eval()
												lowercase							=						model(__snake_case   ,attention_mask=__snake_case)
												lowercase							=						model(__snake_case)
												self.parent.assertEqual(result.last_hidden_state.shape   ,(self.batch_size, self.seq_length, self.hidden_size))




							def  A__      (				self   ,A__   ,A__   ,A__   ,A__   ,A__   ,A__   ,A__   ,A__   ,A__   ,):
												lowercase							=						True
												lowercase							=						OpenLlamaModel(__snake_case)
												model.to(__snake_case)
												model.eval()
												lowercase							=						model(
												    __snake_case   ,attention_mask=__snake_case   ,encoder_hidden_states=__snake_case   ,encoder_attention_mask=__snake_case   ,)
												lowercase							=						model(
												    __snake_case   ,attention_mask=__snake_case   ,encoder_hidden_states=__snake_case   ,)
												lowercase							=						model(__snake_case   ,attention_mask=__snake_case)
												self.parent.assertEqual(result.last_hidden_state.shape   ,(self.batch_size, self.seq_length, self.hidden_size))




							def  A__      (				self   ,A__   ,A__   ,A__   ,A__   ,A__   ,A__   ,A__   ,A__   ,A__   ,):
												lowercase							=						OpenLlamaForCausalLM(config=__snake_case)
												model.to(__snake_case)
												model.eval()
												lowercase							=						model(__snake_case   ,attention_mask=__snake_case   ,labels=__snake_case)
												self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.seq_length, self.vocab_size))




							def  A__      (				self   ,A__   ,A__   ,A__   ,A__   ,A__   ,A__   ,A__   ,A__   ,A__   ,):
												lowercase							=						True
												lowercase							=						True
												lowercase							=						OpenLlamaForCausalLM(config=__snake_case)
												model.to(__snake_case)
												model.eval()

												# first forward pass
												lowercase							=						model(
												    __snake_case   ,attention_mask=__snake_case   ,encoder_hidden_states=__snake_case   ,encoder_attention_mask=__snake_case   ,use_cache=__snake_case   ,)
												lowercase							=						outputs.past_key_values

												# create hypothetical multiple next token and extent to next_input_ids
												lowercase							=						ids_tensor((self.batch_size, 3)   ,config.vocab_size)
												lowercase							=						ids_tensor((self.batch_size, 3)   ,vocab_size=2)

												# append to next input_ids and
												lowercase							=						torch.cat([input_ids, next_tokens]   ,dim=-1)
												lowercase							=						torch.cat([input_mask, next_mask]   ,dim=-1)

												lowercase							=						model(
												    __snake_case   ,attention_mask=__snake_case   ,encoder_hidden_states=__snake_case   ,encoder_attention_mask=__snake_case   ,output_hidden_states=__snake_case   ,)['''hidden_states'''][0]
												lowercase							=						model(
												    __snake_case   ,attention_mask=__snake_case   ,encoder_hidden_states=__snake_case   ,encoder_attention_mask=__snake_case   ,past_key_values=__snake_case   ,output_hidden_states=__snake_case   ,)['''hidden_states'''][0]

												# select random slice
												lowercase							=						ids_tensor((1,)   ,output_from_past.shape[-1]).item()
												lowercase							=						output_from_no_past[:, -3:, random_slice_idx].detach()
												lowercase							=						output_from_past[:, :, random_slice_idx].detach()

												self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1])

												# test that outputs are equal for slice
												self.parent.assertTrue(torch.allclose(__snake_case   ,__snake_case   ,atol=1E-3))





							def  A__      (				self):
												lowercase							=						self.prepare_config_and_inputs()
												(
												    lowercase		
												)							=						config_and_inputs
												lowercase							=						{'''input_ids''': input_ids, '''attention_mask''': input_mask}
												return config, inputs_dict







@require_torch
class 				lowercase		(					A__       ,  A__       ,  A__       ,  unittest.TestCase		):
							lowercase_       :       List[Any]       					=(
							    (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
							)
							lowercase_       :       str       					=(OpenLlamaForCausalLM,) if is_torch_available() else ()
							lowercase_       :       str       					=(
							    {
							        '''feature-extraction''': OpenLlamaModel,
							        '''text-classification''': OpenLlamaForSequenceClassification,
							        '''text-generation''': OpenLlamaForCausalLM,
							        '''zero-shot''': OpenLlamaForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							lowercase_       :       Tuple       					=False
							lowercase_       :       Dict       					=False




							def  A__      (				self):
												lowercase							=						OpenLlamaModelTester(self)
												lowercase							=						ConfigTester(self   ,config_class=__snake_case   ,hidden_size=3_7)




							def  A__      (				self):
												self.config_tester.run_common_tests()




							def  A__      (				self):
												lowercase							=						self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*__snake_case)




							def  A__      (				self):
												lowercase							=						self.model_tester.prepare_config_and_inputs()
												for type in ["absolute", "relative_key", "relative_key_query"]:
																	lowercase							=						type
																	self.model_tester.create_and_check_model(*__snake_case)




							def  A__      (				self):
												lowercase							=						self.model_tester.prepare_config_and_inputs_for_common()
												lowercase							=						3
												lowercase							=						input_dict['''input_ids''']
												lowercase							=						input_ids.ne(1).to(__snake_case)
												lowercase							=						ids_tensor([self.model_tester.batch_size]   ,self.model_tester.type_sequence_label_size)
												lowercase							=						OpenLlamaForSequenceClassification(__snake_case)
												model.to(__snake_case)
												model.eval()
												lowercase							=						model(__snake_case   ,attention_mask=__snake_case   ,labels=__snake_case)
												self.assertEqual(result.logits.shape   ,(self.model_tester.batch_size, self.model_tester.num_labels))




							def  A__      (				self):
												lowercase							=						self.model_tester.prepare_config_and_inputs_for_common()
												lowercase							=						3
												lowercase							=						'''single_label_classification'''
												lowercase							=						input_dict['''input_ids''']
												lowercase							=						input_ids.ne(1).to(__snake_case)
												lowercase							=						ids_tensor([self.model_tester.batch_size]   ,self.model_tester.type_sequence_label_size)
												lowercase							=						OpenLlamaForSequenceClassification(__snake_case)
												model.to(__snake_case)
												model.eval()
												lowercase							=						model(__snake_case   ,attention_mask=__snake_case   ,labels=__snake_case)
												self.assertEqual(result.logits.shape   ,(self.model_tester.batch_size, self.model_tester.num_labels))




							def  A__      (				self):
												lowercase							=						self.model_tester.prepare_config_and_inputs_for_common()
												lowercase							=						3
												lowercase							=						'''multi_label_classification'''
												lowercase							=						input_dict['''input_ids''']
												lowercase							=						input_ids.ne(1).to(__snake_case)
												lowercase							=						ids_tensor(
												    [self.model_tester.batch_size, config.num_labels]   ,self.model_tester.type_sequence_label_size).to(torch.float)
												lowercase							=						OpenLlamaForSequenceClassification(__snake_case)
												model.to(__snake_case)
												model.eval()
												lowercase							=						model(__snake_case   ,attention_mask=__snake_case   ,labels=__snake_case)
												self.assertEqual(result.logits.shape   ,(self.model_tester.batch_size, self.model_tester.num_labels))




							@unittest.skip('''Open-Llama buffers include complex numbers, which breaks this test''')
							def  A__      (				self):
												pass





							@parameterized.expand([('''linear''',), ('''dynamic''',)])
							def  A__      (				self   ,A__):
												lowercase							=						self.model_tester.prepare_config_and_inputs_for_common()
												lowercase							=						ids_tensor([1, 1_0]   ,config.vocab_size)
												lowercase							=						ids_tensor([1, int(config.max_position_embeddings * 1.5)]   ,config.vocab_size)

												set_seed(4_2)  # Fixed seed at init time so the two models get the same random weights
												lowercase							=						OpenLlamaModel(__snake_case)
												original_model.to(__snake_case)
												original_model.eval()
												lowercase							=						original_model(__snake_case).last_hidden_state
												lowercase							=						original_model(__snake_case).last_hidden_state

												set_seed(4_2)  # Fixed seed at init time so the two models get the same random weights
												lowercase							=						{'''type''': scaling_type, '''factor''': 10.0}
												lowercase							=						OpenLlamaModel(__snake_case)
												scaled_model.to(__snake_case)
												scaled_model.eval()
												lowercase							=						scaled_model(__snake_case).last_hidden_state
												lowercase							=						scaled_model(__snake_case).last_hidden_state

												# Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
												# maximum sequence length, so the outputs for the short input should match.
												if scaling_type == "dynamic":
																	self.assertTrue(torch.allclose(__snake_case   ,__snake_case   ,atol=1E-5))
												else:
																	self.assertFalse(torch.allclose(__snake_case   ,__snake_case   ,atol=1E-5))

												# The output should be different for long inputs
												self.assertFalse(torch.allclose(__snake_case   ,__snake_case   ,atol=1E-5))









'''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 				SCREAMING_SNAKE_CASE:







	"""simple docstring"""






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







	def      A	(						self			:     Tuple		)			->		Union[str, Any]:
								UpperCAmelCase :       Optional[int] 							=		ids_tensor([self.batch_size, self.seq_length]		,			self.vocab_size		)

								UpperCAmelCase :       Optional[int] 							=		None
								if self.use_input_mask:
															UpperCAmelCase :       Dict 							=		random_attention_mask([self.batch_size, self.seq_length]		)

								UpperCAmelCase :       Dict 							=		None
								if self.use_token_type_ids:
															UpperCAmelCase :       Union[str, Any] 							=		ids_tensor([self.batch_size, self.seq_length]		,			self.type_vocab_size		)

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

								UpperCAmelCase :       List[str] 							=		self.get_config()

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







	def      A	(						self			:     int		)			->		Tuple:
								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      A	(						self			:     Optional[Any]		)			->		Any:
								UpperCAmelCase :       Optional[Any] 							=		self.get_config()
								UpperCAmelCase :       int 							=		300
								return config







	def      A	(						self			:     Optional[Any]		)			->		Any:
								(
								    (
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		
								) :       Tuple 							=		self.prepare_config_and_inputs()

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







	def      A	(						self			:     Tuple		,			__snake_case			:     Optional[Any]		,			__snake_case			:     List[str]		,			__snake_case			:     List[str]		,			__snake_case			:     int		,			__snake_case			:     Union[str, Any]		,			__snake_case			:     Optional[Any]		,			__snake_case			:     Any		,			__snake_case			:     List[Any]		,			__snake_case			:     Optional[Any]		,			)			->		Tuple:
								UpperCAmelCase :       str 							=		True
								UpperCAmelCase :       Tuple 							=		MraModel(__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       Optional[int] 							=		model(
								    __snake_case		,			attention_mask=__snake_case		,			token_type_ids=__snake_case		,			encoder_hidden_states=__snake_case		,			encoder_attention_mask=__snake_case		,			)
								UpperCAmelCase :       Optional[Any] 							=		model(
								    __snake_case		,			attention_mask=__snake_case		,			token_type_ids=__snake_case		,			encoder_hidden_states=__snake_case		,			)
								UpperCAmelCase :       str 							=		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      A	(						self			:     Tuple		,			__snake_case			:     str		,			__snake_case			:     Optional[Any]		,			__snake_case			:     List[str]		,			__snake_case			:     Tuple		,			__snake_case			:     Optional[Any]		,			__snake_case			:     List[str]		,			__snake_case			:     int		)			->		Any:
								UpperCAmelCase :       Dict 							=		MraForMaskedLM(config=__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       Optional[int] 							=		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      A	(						self			:     Tuple		,			__snake_case			:     Tuple		,			__snake_case			:     Dict		,			__snake_case			:     Dict		,			__snake_case			:     Any		,			__snake_case			:     int		,			__snake_case			:     Optional[Any]		,			__snake_case			:     Tuple		)			->		Optional[int]:
								UpperCAmelCase :       List[str] 							=		MraForQuestionAnswering(config=__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       List[Any] 							=		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      A	(						self			:     str		,			__snake_case			:     Optional[int]		,			__snake_case			:     List[str]		,			__snake_case			:     str		,			__snake_case			:     int		,			__snake_case			:     Optional[Any]		,			__snake_case			:     Union[str, Any]		,			__snake_case			:     List[Any]		)			->		int:
								UpperCAmelCase :       int 							=		self.num_labels
								UpperCAmelCase :       Union[str, Any] 							=		MraForSequenceClassification(__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       List[str] 							=		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      A	(						self			:     str		,			__snake_case			:     Dict		,			__snake_case			:     Optional[Any]		,			__snake_case			:     Dict		,			__snake_case			:     Union[str, Any]		,			__snake_case			:     Union[str, Any]		,			__snake_case			:     Union[str, Any]		,			__snake_case			:     Dict		)			->		int:
								UpperCAmelCase :       Tuple 							=		self.num_labels
								UpperCAmelCase :       List[str] 							=		MraForTokenClassification(config=__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       str 							=		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      A	(						self			:     str		,			__snake_case			:     int		,			__snake_case			:     Any		,			__snake_case			:     Tuple		,			__snake_case			:     Optional[int]		,			__snake_case			:     List[str]		,			__snake_case			:     str		,			__snake_case			:     Union[str, Any]		)			->		Optional[Any]:
								UpperCAmelCase :       Tuple 							=		self.num_choices
								UpperCAmelCase :       int 							=		MraForMultipleChoice(config=__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       str 							=		input_ids.unsqueeze(1		).expand(-1		,			self.num_choices		,			-1		).contiguous()
								UpperCAmelCase :       List[Any] 							=		token_type_ids.unsqueeze(1		).expand(-1		,			self.num_choices		,			-1		).contiguous()
								UpperCAmelCase :       List[str] 							=		input_mask.unsqueeze(1		).expand(-1		,			self.num_choices		,			-1		).contiguous()
								UpperCAmelCase :       List[str] 							=		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      A	(						self			:     str		)			->		Dict:
								UpperCAmelCase :       Any 							=		self.prepare_config_and_inputs()
								(
								    (
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		
								) :       List[str] 							=		config_and_inputs
								UpperCAmelCase :       Any 							=		{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
								return config, inputs_dict

@require_torch
class 				SCREAMING_SNAKE_CASE(							A__							,		unittest.TestCase			):







	"""simple docstring"""






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

	lowerCamelCase__			 =       ()







	def      A	(						self			:     int		)			->		Union[str, Any]:
								UpperCAmelCase :       List[str] 							=		MraModelTester(self		)
								UpperCAmelCase :       Optional[int] 							=		ConfigTester(self		,			config_class=__snake_case		,			hidden_size=37		)







	def      A	(						self			:     Optional[Any]		)			->		str:
								self.config_tester.run_common_tests()







	def      A	(						self			:     Tuple		)			->		Optional[Any]:
								UpperCAmelCase :       Optional[int] 							=		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*__snake_case		)







	def      A	(						self			:     List[Any]		)			->		Optional[Any]:
								UpperCAmelCase :       Tuple 							=		self.model_tester.prepare_config_and_inputs()
								for type in ["absolute", "relative_key", "relative_key_query"]:
															UpperCAmelCase :       List[Any] 							=		type
															self.model_tester.create_and_check_model(*__snake_case		)







	def      A	(						self			:     Tuple		)			->		Dict:
								UpperCAmelCase :       List[Any] 							=		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*__snake_case		)







	def      A	(						self			:     Tuple		)			->		List[str]:
								UpperCAmelCase :       Tuple 							=		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_multiple_choice(*__snake_case		)







	def      A	(						self			:     int		)			->		Dict:
								UpperCAmelCase :       Optional[int] 							=		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_question_answering(*__snake_case		)







	def      A	(						self			:     Dict		)			->		Optional[int]:
								UpperCAmelCase :       Any 							=		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_sequence_classification(*__snake_case		)







	def      A	(						self			:     Any		)			->		Optional[int]:
								UpperCAmelCase :       List[str] 							=		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*__snake_case		)







	@slow
	def      A	(						self			:     Dict		)			->		Any:
								for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															UpperCAmelCase :       str 							=		MraModel.from_pretrained(__snake_case		)
															self.assertIsNotNone(__snake_case		)







	@unittest.skip(reason='''MRA does not output attentions'''		)
	def      A	(						self			:     str		)			->		Optional[Any]:
								return

@require_torch
class 				SCREAMING_SNAKE_CASE(							unittest.TestCase			):







	"""simple docstring"""






	@slow
	def      A	(						self			:     Tuple		)			->		List[Any]:
								UpperCAmelCase :       int 							=		MraModel.from_pretrained('''uw-madison/mra-base-512-4'''		)
								UpperCAmelCase :       Optional[Any] 							=		torch.arange(256		).unsqueeze(0		)

								with torch.no_grad():
															UpperCAmelCase :       List[Any] 							=		model(__snake_case		)[0]

								UpperCAmelCase :       Optional[Any] 							=		torch.Size((1, 256, 768)		)
								self.assertEqual(output.shape		,			__snake_case		)

								UpperCAmelCase :       Any 							=		torch.tensor(
								    [[[-0.01_40, 0.08_30, -0.03_81], [0.15_46, 0.14_02, 0.02_20], [0.11_62, 0.08_51, 0.01_65]]]		)

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







	@slow
	def      A	(						self			:     Optional[Any]		)			->		Any:
								UpperCAmelCase :       Optional[int] 							=		MraForMaskedLM.from_pretrained('''uw-madison/mra-base-512-4'''		)
								UpperCAmelCase :       Dict 							=		torch.arange(256		).unsqueeze(0		)

								with torch.no_grad():
															UpperCAmelCase :       List[Any] 							=		model(__snake_case		)[0]

								UpperCAmelCase :       int 							=		50265

								UpperCAmelCase :       int 							=		torch.Size((1, 256, vocab_size)		)
								self.assertEqual(output.shape		,			__snake_case		)

								UpperCAmelCase :       Union[str, Any] 							=		torch.tensor(
								    [[[9.25_95, -3.60_38, 11.88_19], [9.38_69, -3.26_93, 11.09_56], [11.85_24, -3.49_38, 13.12_10]]]		)

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







	@slow
	def      A	(						self			:     str		)			->		List[Any]:
								UpperCAmelCase :       List[Any] 							=		MraForMaskedLM.from_pretrained('''uw-madison/mra-base-4096-8-d3'''		)
								UpperCAmelCase :       List[Any] 							=		torch.arange(4096		).unsqueeze(0		)

								with torch.no_grad():
															UpperCAmelCase :       Tuple 							=		model(__snake_case		)[0]

								UpperCAmelCase :       Optional[int] 							=		50265

								UpperCAmelCase :       Tuple 							=		torch.Size((1, 4096, vocab_size)		)
								self.assertEqual(output.shape		,			__snake_case		)

								UpperCAmelCase :       Optional[int] 							=		torch.tensor(
								    [[[5.47_89, -2.35_64, 7.50_64], [7.90_67, -1.33_69, 9.96_68], [9.07_12, -1.81_06, 7.03_80]]]		)

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





