


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

						"""simple docstring"""






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






						def 				_UpperCAmelCase       (  self							:      Optional[Any]	):
													A_			    =		ids_tensor([self.batch_size, self.seq_length] ,  self.vocab_size	)

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

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

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

													A_			    =		self.get_config()

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






						def 				_UpperCAmelCase       (  self							:      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=lowerCAmelCase ,  initializer_range=self.initializer_range ,  )






						def 				_UpperCAmelCase       (  self							:      Tuple	):
													A_			    =		self.get_config()
													A_			    =		3_00
													return config






						def 				_UpperCAmelCase       (  self							:      Dict	):
													(
													    (
													    A_
													)    ,					(
													    A_
													)    ,					(
													    A_
													)    ,					(
													    A_
													)    ,					(
													    A_
													)    ,					(
													    A_
													)    ,					(
													    A_
													)    ,					
													)			    =		self.prepare_config_and_inputs()

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






						def 				_UpperCAmelCase       (  self							:      Optional[int] ,  lowerCAmelCase							:      Optional[int] ,  lowerCAmelCase							:      str ,  lowerCAmelCase							:      List[Any] ,  lowerCAmelCase							:      Dict ,  lowerCAmelCase							:      Tuple ,  lowerCAmelCase							:      Optional[int] ,  lowerCAmelCase							:      int ,  lowerCAmelCase							:      int ,  lowerCAmelCase							:      str ,  ):
													A_			    =		True
													A_			    =		MraModel(lowerCAmelCase	)
													model.to(lowerCAmelCase	)
													model.eval()
													A_			    =		model(
													    lowerCAmelCase ,  attention_mask=lowerCAmelCase ,  token_type_ids=lowerCAmelCase ,  encoder_hidden_states=lowerCAmelCase ,  encoder_attention_mask=lowerCAmelCase ,  )
													A_			    =		model(
													    lowerCAmelCase ,  attention_mask=lowerCAmelCase ,  token_type_ids=lowerCAmelCase ,  encoder_hidden_states=lowerCAmelCase ,  )
													A_			    =		model(lowerCAmelCase ,  attention_mask=lowerCAmelCase ,  token_type_ids=lowerCAmelCase	)
													self.parent.assertEqual(result.last_hidden_state.shape ,  (self.batch_size, self.seq_length, self.hidden_size)	)






						def 				_UpperCAmelCase       (  self							:      Optional[int] ,  lowerCAmelCase							:      str ,  lowerCAmelCase							:      Any ,  lowerCAmelCase							:      Optional[int] ,  lowerCAmelCase							:      Optional[int] ,  lowerCAmelCase							:      List[str] ,  lowerCAmelCase							:      Union[str, Any] ,  lowerCAmelCase							:      Optional[int]	):
													A_			    =		MraForMaskedLM(config=lowerCAmelCase	)
													model.to(lowerCAmelCase	)
													model.eval()
													A_			    =		model(lowerCAmelCase ,  attention_mask=lowerCAmelCase ,  token_type_ids=lowerCAmelCase ,  labels=lowerCAmelCase	)
													self.parent.assertEqual(result.logits.shape ,  (self.batch_size, self.seq_length, self.vocab_size)	)






						def 				_UpperCAmelCase       (  self							:      Dict ,  lowerCAmelCase							:      Optional[int] ,  lowerCAmelCase							:      Union[str, Any] ,  lowerCAmelCase							:      List[Any] ,  lowerCAmelCase							:      Any ,  lowerCAmelCase							:      List[Any] ,  lowerCAmelCase							:      Tuple ,  lowerCAmelCase							:      List[str]	):
													A_			    =		MraForQuestionAnswering(config=lowerCAmelCase	)
													model.to(lowerCAmelCase	)
													model.eval()
													A_			    =		model(
													    lowerCAmelCase ,  attention_mask=lowerCAmelCase ,  token_type_ids=lowerCAmelCase ,  start_positions=lowerCAmelCase ,  end_positions=lowerCAmelCase ,  )
													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							:      List[str] ,  lowerCAmelCase							:      int ,  lowerCAmelCase							:      str ,  lowerCAmelCase							:      int ,  lowerCAmelCase							:      str ,  lowerCAmelCase							:      List[Any] ,  lowerCAmelCase							:      str ,  lowerCAmelCase							:      Dict	):
													A_			    =		self.num_labels
													A_			    =		MraForSequenceClassification(lowerCAmelCase	)
													model.to(lowerCAmelCase	)
													model.eval()
													A_			    =		model(lowerCAmelCase ,  attention_mask=lowerCAmelCase ,  token_type_ids=lowerCAmelCase ,  labels=lowerCAmelCase	)
													self.parent.assertEqual(result.logits.shape ,  (self.batch_size, self.num_labels)	)






						def 				_UpperCAmelCase       (  self							:      Union[str, Any] ,  lowerCAmelCase							:      str ,  lowerCAmelCase							:      str ,  lowerCAmelCase							:      Optional[Any] ,  lowerCAmelCase							:      str ,  lowerCAmelCase							:      str ,  lowerCAmelCase							:      int ,  lowerCAmelCase							:      Optional[Any]	):
													A_			    =		self.num_labels
													A_			    =		MraForTokenClassification(config=lowerCAmelCase	)
													model.to(lowerCAmelCase	)
													model.eval()
													A_			    =		model(lowerCAmelCase ,  attention_mask=lowerCAmelCase ,  token_type_ids=lowerCAmelCase ,  labels=lowerCAmelCase	)
													self.parent.assertEqual(result.logits.shape ,  (self.batch_size, self.seq_length, self.num_labels)	)






						def 				_UpperCAmelCase       (  self							:      Optional[int] ,  lowerCAmelCase							:      str ,  lowerCAmelCase							:      List[str] ,  lowerCAmelCase							:      List[Any] ,  lowerCAmelCase							:      int ,  lowerCAmelCase							:      Optional[Any] ,  lowerCAmelCase							:      str ,  lowerCAmelCase							:      List[Any]	):
													A_			    =		self.num_choices
													A_			    =		MraForMultipleChoice(config=lowerCAmelCase	)
													model.to(lowerCAmelCase	)
													model.eval()
													A_			    =		input_ids.unsqueeze(1	).expand(-1 ,  self.num_choices ,  -1	).contiguous()
													A_			    =		token_type_ids.unsqueeze(1	).expand(-1 ,  self.num_choices ,  -1	).contiguous()
													A_			    =		input_mask.unsqueeze(1	).expand(-1 ,  self.num_choices ,  -1	).contiguous()
													A_			    =		model(
													    lowerCAmelCase ,  attention_mask=lowerCAmelCase ,  token_type_ids=lowerCAmelCase ,  labels=lowerCAmelCase ,  )
													self.parent.assertEqual(result.logits.shape ,  (self.batch_size, self.num_choices)	)






						def 				_UpperCAmelCase       (  self							:      Tuple	):
													A_			    =		self.prepare_config_and_inputs()
													(
													    (
													    A_
													)    ,					(
													    A_
													)    ,					(
													    A_
													)    ,					(
													    A_
													)    ,					(
													    A_
													)    ,					(
													    A_
													)    ,					(
													    A_
													)    ,					
													)			    =		config_and_inputs
													A_			    =		{"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						):

						"""simple docstring"""






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

						_UpperCAmelCase		:       str   						=()






						def 				_UpperCAmelCase       (  self							:      Any	):
													A_			    =		MraModelTester(self	)
													A_			    =		ConfigTester(self ,  config_class=lowerCAmelCase ,  hidden_size=37	)






						def 				_UpperCAmelCase       (  self							:      List[Any]	):
													self.config_tester.run_common_tests()






						def 				_UpperCAmelCase       (  self							:      Optional[Any]	):
													A_			    =		self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*lowerCAmelCase	)






						def 				_UpperCAmelCase       (  self							:      List[Any]	):
													A_			    =		self.model_tester.prepare_config_and_inputs()
													for type in ["absolute", "relative_key", "relative_key_query"]:
																				A_			    =		type
																				self.model_tester.create_and_check_model(*lowerCAmelCase	)






						def 				_UpperCAmelCase       (  self							:      Dict	):
													A_			    =		self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_masked_lm(*lowerCAmelCase	)






						def 				_UpperCAmelCase       (  self							:      str	):
													A_			    =		self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_multiple_choice(*lowerCAmelCase	)






						def 				_UpperCAmelCase       (  self							:      Dict	):
													A_			    =		self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_question_answering(*lowerCAmelCase	)






						def 				_UpperCAmelCase       (  self							:      Optional[Any]	):
													A_			    =		self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_sequence_classification(*lowerCAmelCase	)






						def 				_UpperCAmelCase       (  self							:      List[str]	):
													A_			    =		self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase	)






						@slow
						def 				_UpperCAmelCase       (  self							:      List[Any]	):
													for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																				A_			    =		MraModel.from_pretrained(lowerCAmelCase	)
																				self.assertIsNotNone(lowerCAmelCase	)






						@unittest.skip(reason="MRA does not output attentions"	)
						def 				_UpperCAmelCase       (  self							:      List[str]	):
													return






@require_torch
class 					__lowerCAmelCase   (   unittest.TestCase						):

						"""simple docstring"""






						@slow
						def 				_UpperCAmelCase       (  self							:      Any	):
													A_			    =		MraModel.from_pretrained("uw-madison/mra-base-512-4"	)
													A_			    =		torch.arange(2_56	).unsqueeze(0	)

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

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

													A_			    =		torch.tensor(
													    [[[-0.0_1_4_0, 0.0_8_3_0, -0.0_3_8_1], [0.1_5_4_6, 0.1_4_0_2, 0.0_2_2_0], [0.1_1_6_2, 0.0_8_5_1, 0.0_1_6_5]]]	)

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






						@slow
						def 				_UpperCAmelCase       (  self							:      List[str]	):
													A_			    =		MraForMaskedLM.from_pretrained("uw-madison/mra-base-512-4"	)
													A_			    =		torch.arange(2_56	).unsqueeze(0	)

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

													A_			    =		5_02_65

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

													A_			    =		torch.tensor(
													    [[[9.2_5_9_5, -3.6_0_3_8, 1_1.8_8_1_9], [9.3_8_6_9, -3.2_6_9_3, 1_1.0_9_5_6], [1_1.8_5_2_4, -3.4_9_3_8, 1_3.1_2_1_0]]]	)

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






						@slow
						def 				_UpperCAmelCase       (  self							:      str	):
													A_			    =		MraForMaskedLM.from_pretrained("uw-madison/mra-base-4096-8-d3"	)
													A_			    =		torch.arange(40_96	).unsqueeze(0	)

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

													A_			    =		5_02_65

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

													A_			    =		torch.tensor(
													    [[[5.4_7_8_9, -2.3_5_6_4, 7.5_0_6_4], [7.9_0_6_7, -1.3_3_6_9, 9.9_6_6_8], [9.0_7_1_2, -1.8_1_0_6, 7.0_3_8_0]]]	)

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









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


__SCREAMING_SNAKE_CASE					:				Optional[int]									=		logging.get_logger(__name__)

__SCREAMING_SNAKE_CASE					:				Union[str, Any]									=		{
    '''tanreinama/GPTSAN-2.8B-spout_is_uniform''': (
        '''https://huggingface.co/tanreinama/GPTSAN-2.8B-spout_is_uniform/resolve/main/config.json'''
    ),
}






class 					__lowerCAmelCase   (   lowercase						):

						"""simple docstring"""






						_UpperCAmelCase		:       Optional[int]   						="gptsan-japanese"
						_UpperCAmelCase		:       Union[str, Any]   						=[
						    "past_key_values",
						]
						_UpperCAmelCase		:       int   						={
						    "hidden_size": "d_model",
						    "num_attention_heads": "num_heads",
						    "num_hidden_layers": "num_layers",
						}






						def __init__(  self							:      List[str] ,  lowerCAmelCase							:      Dict=3_60_00 ,  lowerCAmelCase							:      Union[str, Any]=12_80 ,  lowerCAmelCase							:      Tuple=10_24 ,  lowerCAmelCase							:      Any=81_92 ,  lowerCAmelCase							:      List[str]=40_96 ,  lowerCAmelCase							:      Dict=1_28 ,  lowerCAmelCase							:      int=10 ,  lowerCAmelCase							:      str=0 ,  lowerCAmelCase							:      int=16 ,  lowerCAmelCase							:      int=16 ,  lowerCAmelCase							:      str=1_28 ,  lowerCAmelCase							:      Dict=0.0 ,  lowerCAmelCase							:      List[str]=1e-5 ,  lowerCAmelCase							:      List[str]=False ,  lowerCAmelCase							:      int=0.0 ,  lowerCAmelCase							:      Tuple="float32" ,  lowerCAmelCase							:      Optional[int]=False ,  lowerCAmelCase							:      int=False ,  lowerCAmelCase							:      Tuple=False ,  lowerCAmelCase							:      Optional[Any]=0.0_0_2 ,  lowerCAmelCase							:      Tuple=False ,  lowerCAmelCase							:      str=True ,  lowerCAmelCase							:      Optional[int]=3_59_98 ,  lowerCAmelCase							:      Optional[Any]=3_59_95 ,  lowerCAmelCase							:      Tuple=3_59_99 ,  **lowerCAmelCase							:      Dict ,  ):
													A_			    =		vocab_size
													A_			    =		max_position_embeddings
													A_			    =		d_model
													A_			    =		d_ff
													A_			    =		d_ext
													A_			    =		d_spout
													A_			    =		num_switch_layers
													A_			    =		num_ext_layers
													A_			    =		num_switch_layers + num_ext_layers
													A_			    =		num_heads
													A_			    =		num_experts
													A_			    =		expert_capacity
													A_			    =		dropout_rate
													A_			    =		layer_norm_epsilon
													A_			    =		router_bias
													A_			    =		router_jitter_noise
													A_			    =		router_dtype
													A_			    =		router_ignore_padding_tokens
													A_			    =		output_hidden_states
													A_			    =		output_attentions
													A_			    =		initializer_factor
													A_			    =		output_router_logits
													A_			    =		use_cache

													super().__init__(
													    separator_token_id=lowerCAmelCase ,  pad_token_id=lowerCAmelCase ,  eos_token_id=lowerCAmelCase ,  **lowerCAmelCase ,  )





