

"""simple docstring"""




import unittest

from transformers import YosoConfig, 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
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import (
						    YosoForMaskedLM,
						    YosoForMultipleChoice,
						    YosoForQuestionAnswering,
						    YosoForSequenceClassification,
						    YosoForTokenClassification,
						    YosoModel,
						)
						from transformers.models.yoso.modeling_yoso import YOSO_PRETRAINED_MODEL_ARCHIVE_LIST





class 				lowerCamelCase  :

			def __init__(   self			:  Optional[Any]      ,			__UpperCAmelCase			:  Any      ,			__UpperCAmelCase			:  str=1_3      ,			__UpperCAmelCase			:  Optional[Any]=7      ,			__UpperCAmelCase			:  str=True      ,			__UpperCAmelCase			:  List[str]=True      ,			__UpperCAmelCase			:  str=True      ,			__UpperCAmelCase			:  Union[str, Any]=True      ,			__UpperCAmelCase			:  int=9_9      ,			__UpperCAmelCase			:  int=3_2      ,			__UpperCAmelCase			:  int=5      ,			__UpperCAmelCase			:  Tuple=4      ,			__UpperCAmelCase			:  Union[str, Any]=3_7      ,			__UpperCAmelCase			:  Any="gelu"      ,			__UpperCAmelCase			:  List[Any]=0.1      ,			__UpperCAmelCase			:  str=0.1      ,			__UpperCAmelCase			:  List[str]=5_1_2      ,			__UpperCAmelCase			:  List[Any]=1_6      ,			__UpperCAmelCase			:  List[str]=2      ,			__UpperCAmelCase			:  Any=0.02      ,			__UpperCAmelCase			:  Tuple=3      ,			__UpperCAmelCase			:  str=4      ,			__UpperCAmelCase			:  Tuple=None      ,			)				->      Union[str, Any]:
						SCREAMING_SNAKE_CASE__									= parent
						SCREAMING_SNAKE_CASE__									= batch_size
						SCREAMING_SNAKE_CASE__									= seq_length
						SCREAMING_SNAKE_CASE__									= is_training
						SCREAMING_SNAKE_CASE__									= use_input_mask
						SCREAMING_SNAKE_CASE__									= use_token_type_ids
						SCREAMING_SNAKE_CASE__									= use_labels
						SCREAMING_SNAKE_CASE__									= vocab_size
						SCREAMING_SNAKE_CASE__									= hidden_size
						SCREAMING_SNAKE_CASE__									= num_hidden_layers
						SCREAMING_SNAKE_CASE__									= num_attention_heads
						SCREAMING_SNAKE_CASE__									= intermediate_size
						SCREAMING_SNAKE_CASE__									= hidden_act
						SCREAMING_SNAKE_CASE__									= hidden_dropout_prob
						SCREAMING_SNAKE_CASE__									= attention_probs_dropout_prob
						SCREAMING_SNAKE_CASE__									= max_position_embeddings
						SCREAMING_SNAKE_CASE__									= type_vocab_size
						SCREAMING_SNAKE_CASE__									= type_sequence_label_size
						SCREAMING_SNAKE_CASE__									= initializer_range
						SCREAMING_SNAKE_CASE__									= num_labels
						SCREAMING_SNAKE_CASE__									= num_choices
						SCREAMING_SNAKE_CASE__									= scope

			def    SCREAMING_SNAKE_CASE			(   self			:  List[str]    )				->      List[Any]:
						SCREAMING_SNAKE_CASE__									= ids_tensor([self.batch_size, self.seq_length]      ,			self.vocab_size    )

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

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

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

						SCREAMING_SNAKE_CASE__									= self.get_config()

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

			def    SCREAMING_SNAKE_CASE			(   self			:  List[str]    )				->      str:
						return YosoConfig(
						    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=__UpperCAmelCase      ,			initializer_range=self.initializer_range      ,			)

			def    SCREAMING_SNAKE_CASE			(   self			:  Tuple    )				->      Any:
						SCREAMING_SNAKE_CASE__									= self.get_config()
						SCREAMING_SNAKE_CASE__									= 3_0_0
						return config

			def    SCREAMING_SNAKE_CASE			(   self			:  int    )				->      Any:
						(
						    (
						    SCREAMING_SNAKE_CASE__						
						)							,						(
						    SCREAMING_SNAKE_CASE__						
						)							,						(
						    SCREAMING_SNAKE_CASE__						
						)							,						(
						    SCREAMING_SNAKE_CASE__						
						)							,						(
						    SCREAMING_SNAKE_CASE__						
						)							,						(
						    SCREAMING_SNAKE_CASE__						
						)							,						(
						    SCREAMING_SNAKE_CASE__						
						)							,												
						)									= self.prepare_config_and_inputs()

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

			def    SCREAMING_SNAKE_CASE			(   self			:  str      ,			__UpperCAmelCase			:  Tuple      ,			__UpperCAmelCase			:  Optional[int]      ,			__UpperCAmelCase			:  Optional[Any]      ,			__UpperCAmelCase			:  Union[str, Any]      ,			__UpperCAmelCase			:  Optional[Any]      ,			__UpperCAmelCase			:  Optional[int]      ,			__UpperCAmelCase			:  Union[str, Any]      ,			__UpperCAmelCase			:  Optional[int]      ,			__UpperCAmelCase			:  str      ,			)				->      Tuple:
						SCREAMING_SNAKE_CASE__									= True
						SCREAMING_SNAKE_CASE__									= YosoModel(__UpperCAmelCase    )
						model.to(__UpperCAmelCase    )
						model.eval()
						SCREAMING_SNAKE_CASE__									= model(
						    __UpperCAmelCase      ,			attention_mask=__UpperCAmelCase      ,			token_type_ids=__UpperCAmelCase      ,			encoder_hidden_states=__UpperCAmelCase      ,			encoder_attention_mask=__UpperCAmelCase      ,			)
						SCREAMING_SNAKE_CASE__									= model(
						    __UpperCAmelCase      ,			attention_mask=__UpperCAmelCase      ,			token_type_ids=__UpperCAmelCase      ,			encoder_hidden_states=__UpperCAmelCase      ,			)
						SCREAMING_SNAKE_CASE__									= model(__UpperCAmelCase      ,			attention_mask=__UpperCAmelCase      ,			token_type_ids=__UpperCAmelCase    )
						self.parent.assertEqual(result.last_hidden_state.shape      ,			(self.batch_size, self.seq_length, self.hidden_size)    )

			def    SCREAMING_SNAKE_CASE			(   self			:  str      ,			__UpperCAmelCase			:  List[Any]      ,			__UpperCAmelCase			:  Optional[Any]      ,			__UpperCAmelCase			:  List[str]      ,			__UpperCAmelCase			:  str      ,			__UpperCAmelCase			:  Union[str, Any]      ,			__UpperCAmelCase			:  Dict      ,			__UpperCAmelCase			:  Optional[Any]    )				->      Optional[int]:
						SCREAMING_SNAKE_CASE__									= YosoForMaskedLM(config=__UpperCAmelCase    )
						model.to(__UpperCAmelCase    )
						model.eval()
						SCREAMING_SNAKE_CASE__									= model(__UpperCAmelCase      ,			attention_mask=__UpperCAmelCase      ,			token_type_ids=__UpperCAmelCase      ,			labels=__UpperCAmelCase    )
						self.parent.assertEqual(result.logits.shape      ,			(self.batch_size, self.seq_length, self.vocab_size)    )

			def    SCREAMING_SNAKE_CASE			(   self			:  Dict      ,			__UpperCAmelCase			:  List[Any]      ,			__UpperCAmelCase			:  List[str]      ,			__UpperCAmelCase			:  int      ,			__UpperCAmelCase			:  Any      ,			__UpperCAmelCase			:  str      ,			__UpperCAmelCase			:  Optional[int]      ,			__UpperCAmelCase			:  Optional[int]    )				->      List[Any]:
						SCREAMING_SNAKE_CASE__									= YosoForQuestionAnswering(config=__UpperCAmelCase    )
						model.to(__UpperCAmelCase    )
						model.eval()
						SCREAMING_SNAKE_CASE__									= model(
						    __UpperCAmelCase      ,			attention_mask=__UpperCAmelCase      ,			token_type_ids=__UpperCAmelCase      ,			start_positions=__UpperCAmelCase      ,			end_positions=__UpperCAmelCase      ,			)
						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    SCREAMING_SNAKE_CASE			(   self			:  List[Any]      ,			__UpperCAmelCase			:  int      ,			__UpperCAmelCase			:  List[Any]      ,			__UpperCAmelCase			:  Optional[int]      ,			__UpperCAmelCase			:  int      ,			__UpperCAmelCase			:  Optional[Any]      ,			__UpperCAmelCase			:  Dict      ,			__UpperCAmelCase			:  Union[str, Any]    )				->      Union[str, Any]:
						SCREAMING_SNAKE_CASE__									= self.num_labels
						SCREAMING_SNAKE_CASE__									= YosoForSequenceClassification(__UpperCAmelCase    )
						model.to(__UpperCAmelCase    )
						model.eval()
						SCREAMING_SNAKE_CASE__									= model(__UpperCAmelCase      ,			attention_mask=__UpperCAmelCase      ,			token_type_ids=__UpperCAmelCase      ,			labels=__UpperCAmelCase    )
						self.parent.assertEqual(result.logits.shape      ,			(self.batch_size, self.num_labels)    )

			def    SCREAMING_SNAKE_CASE			(   self			:  Optional[int]      ,			__UpperCAmelCase			:  Optional[int]      ,			__UpperCAmelCase			:  List[Any]      ,			__UpperCAmelCase			:  Union[str, Any]      ,			__UpperCAmelCase			:  Optional[int]      ,			__UpperCAmelCase			:  Any      ,			__UpperCAmelCase			:  Any      ,			__UpperCAmelCase			:  List[Any]    )				->      Tuple:
						SCREAMING_SNAKE_CASE__									= self.num_labels
						SCREAMING_SNAKE_CASE__									= YosoForTokenClassification(config=__UpperCAmelCase    )
						model.to(__UpperCAmelCase    )
						model.eval()
						SCREAMING_SNAKE_CASE__									= model(__UpperCAmelCase      ,			attention_mask=__UpperCAmelCase      ,			token_type_ids=__UpperCAmelCase      ,			labels=__UpperCAmelCase    )
						self.parent.assertEqual(result.logits.shape      ,			(self.batch_size, self.seq_length, self.num_labels)    )

			def    SCREAMING_SNAKE_CASE			(   self			:  Optional[int]      ,			__UpperCAmelCase			:  Optional[Any]      ,			__UpperCAmelCase			:  Union[str, Any]      ,			__UpperCAmelCase			:  Dict      ,			__UpperCAmelCase			:  Optional[Any]      ,			__UpperCAmelCase			:  Optional[Any]      ,			__UpperCAmelCase			:  Dict      ,			__UpperCAmelCase			:  Union[str, Any]    )				->      Optional[int]:
						SCREAMING_SNAKE_CASE__									= self.num_choices
						SCREAMING_SNAKE_CASE__									= YosoForMultipleChoice(config=__UpperCAmelCase    )
						model.to(__UpperCAmelCase    )
						model.eval()
						SCREAMING_SNAKE_CASE__									= input_ids.unsqueeze(1    ).expand(-1      ,			self.num_choices      ,			-1    ).contiguous()
						SCREAMING_SNAKE_CASE__									= token_type_ids.unsqueeze(1    ).expand(-1      ,			self.num_choices      ,			-1    ).contiguous()
						SCREAMING_SNAKE_CASE__									= input_mask.unsqueeze(1    ).expand(-1      ,			self.num_choices      ,			-1    ).contiguous()
						SCREAMING_SNAKE_CASE__									= model(
						    __UpperCAmelCase      ,			attention_mask=__UpperCAmelCase      ,			token_type_ids=__UpperCAmelCase      ,			labels=__UpperCAmelCase      ,			)
						self.parent.assertEqual(result.logits.shape      ,			(self.batch_size, self.num_choices)    )





			def    SCREAMING_SNAKE_CASE			(   self			:  List[Any]    )				->      Union[str, Any]:
						SCREAMING_SNAKE_CASE__									= self.prepare_config_and_inputs()
						(
						    (
						    SCREAMING_SNAKE_CASE__						
						)							,						(
						    SCREAMING_SNAKE_CASE__						
						)							,						(
						    SCREAMING_SNAKE_CASE__						
						)							,						(
						    SCREAMING_SNAKE_CASE__						
						)							,						(
						    SCREAMING_SNAKE_CASE__						
						)							,						(
						    SCREAMING_SNAKE_CASE__						
						)							,						(
						    SCREAMING_SNAKE_CASE__						
						)							,												
						)									= config_and_inputs
						SCREAMING_SNAKE_CASE__									= {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
						return config, inputs_dict





@require_torch
class 				lowerCamelCase  (A__     ,A__     ,unittest.TestCase     ):
			lowerCamelCase__				:						List[str]   						=							(
			    (
			        YosoModel,
			        YosoForMaskedLM,
			        YosoForMultipleChoice,
			        YosoForQuestionAnswering,
			        YosoForSequenceClassification,
			        YosoForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			lowerCamelCase__				:						Optional[int]   						=							False
			lowerCamelCase__				:						Optional[int]   						=							False
			lowerCamelCase__				:						Optional[int]   						=							False

			lowerCamelCase__				:						Union[str, Any]   						=							()
			lowerCamelCase__				:						Union[str, Any]   						=							(
			    {
			        'feature-extraction': YosoModel,
			        'fill-mask': YosoForMaskedLM,
			        'question-answering': YosoForQuestionAnswering,
			        'text-classification': YosoForSequenceClassification,
			        'token-classification': YosoForTokenClassification,
			        'zero-shot': YosoForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)

			def    SCREAMING_SNAKE_CASE			(   self			:  Any    )				->      Tuple:
						SCREAMING_SNAKE_CASE__									= YosoModelTester(self    )
						SCREAMING_SNAKE_CASE__									= ConfigTester(self      ,			config_class=__UpperCAmelCase      ,			hidden_size=3_7    )

			def    SCREAMING_SNAKE_CASE			(   self			:  Any    )				->      Optional[int]:
						self.config_tester.run_common_tests()

			def    SCREAMING_SNAKE_CASE			(   self			:  Any    )				->      Optional[int]:
						SCREAMING_SNAKE_CASE__									= self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*__UpperCAmelCase    )

			def    SCREAMING_SNAKE_CASE			(   self			:  Union[str, Any]    )				->      Any:
						SCREAMING_SNAKE_CASE__									= self.model_tester.prepare_config_and_inputs()
						for type in ["absolute", "relative_key", "relative_key_query"]:
									SCREAMING_SNAKE_CASE__									= type
									self.model_tester.create_and_check_model(*__UpperCAmelCase    )

			def    SCREAMING_SNAKE_CASE			(   self			:  Optional[Any]    )				->      Optional[int]:
						SCREAMING_SNAKE_CASE__									= self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*__UpperCAmelCase    )

			def    SCREAMING_SNAKE_CASE			(   self			:  Dict    )				->      str:
						SCREAMING_SNAKE_CASE__									= self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_multiple_choice(*__UpperCAmelCase    )

			def    SCREAMING_SNAKE_CASE			(   self			:  List[Any]    )				->      Dict:
						SCREAMING_SNAKE_CASE__									= self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_question_answering(*__UpperCAmelCase    )

			def    SCREAMING_SNAKE_CASE			(   self			:  int    )				->      Tuple:
						SCREAMING_SNAKE_CASE__									= self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_sequence_classification(*__UpperCAmelCase    )

			def    SCREAMING_SNAKE_CASE			(   self			:  List[Any]    )				->      Dict:
						SCREAMING_SNAKE_CASE__									= self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*__UpperCAmelCase    )

			@slow
			def    SCREAMING_SNAKE_CASE			(   self			:  Any    )				->      Tuple:
						for model_name in YOSO_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									SCREAMING_SNAKE_CASE__									= YosoModel.from_pretrained(__UpperCAmelCase    )
									self.assertIsNotNone(__UpperCAmelCase    )





			def    SCREAMING_SNAKE_CASE			(   self			:  Optional[int]    )				->      Union[str, Any]:
						return





@require_torch
class 				lowerCamelCase  (unittest.TestCase     ):

			@slow
			def    SCREAMING_SNAKE_CASE			(   self			:  str    )				->      Optional[Any]:
						SCREAMING_SNAKE_CASE__									= YosoModel.from_pretrained("""uw-madison/yoso-4096"""    )
						SCREAMING_SNAKE_CASE__									= torch.tensor([[0, 1, 2, 3, 4, 5]]    )

						with torch.no_grad():
									SCREAMING_SNAKE_CASE__									= model(__UpperCAmelCase    )[0]

						SCREAMING_SNAKE_CASE__									= torch.Size((1, 6, 7_6_8)    )
						self.assertEqual(output.shape      ,			__UpperCAmelCase    )

						SCREAMING_SNAKE_CASE__									= torch.tensor(
						    [[[-0.0_611, 0.1_242, 0.0_840], [0.0_280, -0.0_048, 0.1_125], [0.0_106, 0.0_226, 0.0_751]]]    )

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

			@slow
			def    SCREAMING_SNAKE_CASE			(   self			:  List[str]    )				->      Dict:
						SCREAMING_SNAKE_CASE__									= YosoForMaskedLM.from_pretrained("""uw-madison/yoso-4096"""    )
						SCREAMING_SNAKE_CASE__									= torch.tensor([[0, 1, 2, 3, 4, 5]]    )

						with torch.no_grad():
									SCREAMING_SNAKE_CASE__									= model(__UpperCAmelCase    )[0]

						SCREAMING_SNAKE_CASE__									= 5_0_2_6_5

						SCREAMING_SNAKE_CASE__									= torch.Size((1, 6, vocab_size)    )
						self.assertEqual(output.shape      ,			__UpperCAmelCase    )

						SCREAMING_SNAKE_CASE__									= torch.tensor(
						    [[[-2.1_313, -3.7_285, -2.2_407], [-2.7_047, -3.3_314, -2.6_408], [0.0_629, -2.5_166, -0.3_356]]]    )

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





			@slow
			def    SCREAMING_SNAKE_CASE			(   self			:  Dict    )				->      Union[str, Any]:
						SCREAMING_SNAKE_CASE__									= YosoForMaskedLM.from_pretrained("""uw-madison/yoso-4096"""    )
						SCREAMING_SNAKE_CASE__									= torch.arange(4_0_9_6    ).unsqueeze(0    )

						with torch.no_grad():
									SCREAMING_SNAKE_CASE__									= model(__UpperCAmelCase    )[0]

						SCREAMING_SNAKE_CASE__									= 5_0_2_6_5

						SCREAMING_SNAKE_CASE__									= torch.Size((1, 4_0_9_6, vocab_size)    )
						self.assertEqual(output.shape      ,			__UpperCAmelCase    )

						SCREAMING_SNAKE_CASE__									= torch.tensor(
						    [[[-2.3_914, -4.3_742, -5.0_956], [-4.0_988, -4.2_384, -7.0_406], [-3.1_427, -3.7_192, -6.6_800]]]    )

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





"""simple docstring"""




def 		A    (       snake_case__						,			snake_case__      ):




			'''simple docstring'''


			return a + b


if __name__ == "__main__":
						A_							:     List[str]     =	5
						A_							:     Optional[int]     =	6
						print(F'The sum of {a} + {b} is {add(a, b)}')


