
"""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_						:


				"""simple docstring"""



				def __init__(self       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__=13       ,	SCREAMING_SNAKE_CASE__=7       ,	SCREAMING_SNAKE_CASE__=True       ,	SCREAMING_SNAKE_CASE__=True       ,	SCREAMING_SNAKE_CASE__=True       ,	SCREAMING_SNAKE_CASE__=True       ,	SCREAMING_SNAKE_CASE__=99       ,	SCREAMING_SNAKE_CASE__=32       ,	SCREAMING_SNAKE_CASE__=5       ,	SCREAMING_SNAKE_CASE__=4       ,	SCREAMING_SNAKE_CASE__=37       ,	SCREAMING_SNAKE_CASE__="gelu"       ,	SCREAMING_SNAKE_CASE__=0.1       ,	SCREAMING_SNAKE_CASE__=0.1       ,	SCREAMING_SNAKE_CASE__=5_12       ,	SCREAMING_SNAKE_CASE__=16       ,	SCREAMING_SNAKE_CASE__=2       ,	SCREAMING_SNAKE_CASE__=0.02       ,	SCREAMING_SNAKE_CASE__=3       ,	SCREAMING_SNAKE_CASE__=4       ,	SCREAMING_SNAKE_CASE__=None       ,	) ->		Optional[int]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	List[str]									=       parent
											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       batch_size
											SCREAMING_SNAKE_CASE__						:	Any									=       seq_length
											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       is_training
											SCREAMING_SNAKE_CASE__						:	str									=       use_input_mask
											SCREAMING_SNAKE_CASE__						:	Tuple									=       use_token_type_ids
											SCREAMING_SNAKE_CASE__						:	List[Any]									=       use_labels
											SCREAMING_SNAKE_CASE__						:	List[Any]									=       vocab_size
											SCREAMING_SNAKE_CASE__						:	Tuple									=       hidden_size
											SCREAMING_SNAKE_CASE__						:	Dict									=       num_hidden_layers
											SCREAMING_SNAKE_CASE__						:	List[Any]									=       num_attention_heads
											SCREAMING_SNAKE_CASE__						:	Tuple									=       intermediate_size
											SCREAMING_SNAKE_CASE__						:	Tuple									=       hidden_act
											SCREAMING_SNAKE_CASE__						:	Any									=       hidden_dropout_prob
											SCREAMING_SNAKE_CASE__						:	Optional[int]									=       attention_probs_dropout_prob
											SCREAMING_SNAKE_CASE__						:	Tuple									=       max_position_embeddings
											SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       type_vocab_size
											SCREAMING_SNAKE_CASE__						:	Any									=       type_sequence_label_size
											SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       initializer_range
											SCREAMING_SNAKE_CASE__						:	List[Any]									=       num_labels
											SCREAMING_SNAKE_CASE__						:	Any									=       num_choices
											SCREAMING_SNAKE_CASE__						:	Any									=       scope



				def      __magic_name__   (self   ) ->		str:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	int									=       ids_tensor([self.batch_size, self.seq_length]       ,	self.vocab_size   )

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

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

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

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

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



				def      __magic_name__   (self   ) ->		Tuple:



											"""simple docstring"""




											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=SCREAMING_SNAKE_CASE__       ,	initializer_range=self.initializer_range       ,	)



				def      __magic_name__   (self   ) ->		Optional[int]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	str									=       self.get_config()
											SCREAMING_SNAKE_CASE__						:	Any									=       3_00
											return config



				def      __magic_name__   (self   ) ->		List[str]:



											"""simple docstring"""




											(
											    (
											    SCREAMING_SNAKE_CASE__
											)				,				(
											    SCREAMING_SNAKE_CASE__
											)				,				(
											    SCREAMING_SNAKE_CASE__
											)				,				(
											    SCREAMING_SNAKE_CASE__
											)				,				(
											    SCREAMING_SNAKE_CASE__
											)				,				(
											    SCREAMING_SNAKE_CASE__
											)				,				(
											    SCREAMING_SNAKE_CASE__
											)				,				
											)						:	List[Any]									=       self.prepare_config_and_inputs()

											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       True
											SCREAMING_SNAKE_CASE__						:	int									=       floats_tensor([self.batch_size, self.seq_length, self.hidden_size]   )
											SCREAMING_SNAKE_CASE__						:	List[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      __magic_name__   (self       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   ) ->		Optional[Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	str									=       YosoModel(config=SCREAMING_SNAKE_CASE__   )
											model.to(SCREAMING_SNAKE_CASE__   )
											model.eval()
											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       model(SCREAMING_SNAKE_CASE__       ,	attention_mask=SCREAMING_SNAKE_CASE__       ,	token_type_ids=SCREAMING_SNAKE_CASE__   )
											SCREAMING_SNAKE_CASE__						:	Dict									=       model(SCREAMING_SNAKE_CASE__       ,	token_type_ids=SCREAMING_SNAKE_CASE__   )
											SCREAMING_SNAKE_CASE__						:	Optional[int]									=       model(SCREAMING_SNAKE_CASE__   )
											self.parent.assertEqual(result.last_hidden_state.shape       ,	(self.batch_size, self.seq_length, self.hidden_size)   )



				def      __magic_name__   (self       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	) ->		Union[str, Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	int									=       True
											SCREAMING_SNAKE_CASE__						:	Dict									=       YosoModel(SCREAMING_SNAKE_CASE__   )
											model.to(SCREAMING_SNAKE_CASE__   )
											model.eval()
											SCREAMING_SNAKE_CASE__						:	Any									=       model(
											    SCREAMING_SNAKE_CASE__       ,	attention_mask=SCREAMING_SNAKE_CASE__       ,	token_type_ids=SCREAMING_SNAKE_CASE__       ,	encoder_hidden_states=SCREAMING_SNAKE_CASE__       ,	encoder_attention_mask=SCREAMING_SNAKE_CASE__       ,	)
											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       model(
											    SCREAMING_SNAKE_CASE__       ,	attention_mask=SCREAMING_SNAKE_CASE__       ,	token_type_ids=SCREAMING_SNAKE_CASE__       ,	encoder_hidden_states=SCREAMING_SNAKE_CASE__       ,	)
											SCREAMING_SNAKE_CASE__						:	List[str]									=       model(SCREAMING_SNAKE_CASE__       ,	attention_mask=SCREAMING_SNAKE_CASE__       ,	token_type_ids=SCREAMING_SNAKE_CASE__   )
											self.parent.assertEqual(result.last_hidden_state.shape       ,	(self.batch_size, self.seq_length, self.hidden_size)   )



				def      __magic_name__   (self       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   ) ->		List[Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       YosoForMaskedLM(config=SCREAMING_SNAKE_CASE__   )
											model.to(SCREAMING_SNAKE_CASE__   )
											model.eval()
											SCREAMING_SNAKE_CASE__						:	Any									=       model(SCREAMING_SNAKE_CASE__       ,	attention_mask=SCREAMING_SNAKE_CASE__       ,	token_type_ids=SCREAMING_SNAKE_CASE__       ,	labels=SCREAMING_SNAKE_CASE__   )
											self.parent.assertEqual(result.logits.shape       ,	(self.batch_size, self.seq_length, self.vocab_size)   )



				def      __magic_name__   (self       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   ) ->		Optional[int]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       YosoForQuestionAnswering(config=SCREAMING_SNAKE_CASE__   )
											model.to(SCREAMING_SNAKE_CASE__   )
											model.eval()
											SCREAMING_SNAKE_CASE__						:	Optional[int]									=       model(
											    SCREAMING_SNAKE_CASE__       ,	attention_mask=SCREAMING_SNAKE_CASE__       ,	token_type_ids=SCREAMING_SNAKE_CASE__       ,	start_positions=SCREAMING_SNAKE_CASE__       ,	end_positions=SCREAMING_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      __magic_name__   (self       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   ) ->		Optional[Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       self.num_labels
											SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       YosoForSequenceClassification(SCREAMING_SNAKE_CASE__   )
											model.to(SCREAMING_SNAKE_CASE__   )
											model.eval()
											SCREAMING_SNAKE_CASE__						:	List[Any]									=       model(SCREAMING_SNAKE_CASE__       ,	attention_mask=SCREAMING_SNAKE_CASE__       ,	token_type_ids=SCREAMING_SNAKE_CASE__       ,	labels=SCREAMING_SNAKE_CASE__   )
											self.parent.assertEqual(result.logits.shape       ,	(self.batch_size, self.num_labels)   )



				def      __magic_name__   (self       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   ) ->		Any:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	int									=       self.num_labels
											SCREAMING_SNAKE_CASE__						:	Optional[int]									=       YosoForTokenClassification(config=SCREAMING_SNAKE_CASE__   )
											model.to(SCREAMING_SNAKE_CASE__   )
											model.eval()
											SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       model(SCREAMING_SNAKE_CASE__       ,	attention_mask=SCREAMING_SNAKE_CASE__       ,	token_type_ids=SCREAMING_SNAKE_CASE__       ,	labels=SCREAMING_SNAKE_CASE__   )
											self.parent.assertEqual(result.logits.shape       ,	(self.batch_size, self.seq_length, self.num_labels)   )



				def      __magic_name__   (self       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   ) ->		Tuple:



											"""simple docstring"""




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



				def      __magic_name__   (self   ) ->		Tuple:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	int									=       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__
											)				,				
											)						:	Optional[int]									=       config_and_inputs
											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       {"""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  ):


				"""simple docstring"""



				__UpperCamelCase				:						Union[str, Any] 	=       (
				    (
				        YosoModel,
				        YosoForMaskedLM,
				        YosoForMultipleChoice,
				        YosoForQuestionAnswering,
				        YosoForSequenceClassification,
				        YosoForTokenClassification,
				    )
				    if is_torch_available()
				    else ()
				)
				__UpperCamelCase				:						Tuple 	=       False
				__UpperCamelCase				:						List[str] 	=       False
				__UpperCamelCase				:						Tuple 	=       False

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



				def      __magic_name__   (self   ) ->		Tuple:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	int									=       YosoModelTester(self   )
											SCREAMING_SNAKE_CASE__						:	Optional[int]									=       ConfigTester(self       ,	config_class=SCREAMING_SNAKE_CASE__       ,	hidden_size=37   )



				def      __magic_name__   (self   ) ->		int:



											"""simple docstring"""




											self.config_tester.run_common_tests()



				def      __magic_name__   (self   ) ->		str:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	str									=       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE__   )



				def      __magic_name__   (self   ) ->		Tuple:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Tuple									=       self.model_tester.prepare_config_and_inputs()
											for type in ["absolute", "relative_key", "relative_key_query"]:
																		SCREAMING_SNAKE_CASE__						:	int									=       type
																		self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE__   )



				def      __magic_name__   (self   ) ->		int:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	List[Any]									=       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE__   )



				def      __magic_name__   (self   ) ->		Dict:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_multiple_choice(*SCREAMING_SNAKE_CASE__   )



				def      __magic_name__   (self   ) ->		Any:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Optional[int]									=       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE__   )



				def      __magic_name__   (self   ) ->		Union[str, Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Optional[int]									=       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_sequence_classification(*SCREAMING_SNAKE_CASE__   )



				def      __magic_name__   (self   ) ->		List[Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	List[Any]									=       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE__   )



				@slow
				def      __magic_name__   (self   ) ->		Optional[Any]:



											"""simple docstring"""




											for model_name in YOSO_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		SCREAMING_SNAKE_CASE__						:	Tuple									=       YosoModel.from_pretrained(SCREAMING_SNAKE_CASE__   )
																		self.assertIsNotNone(SCREAMING_SNAKE_CASE__   )



				def      __magic_name__   (self   ) ->		List[Any]:



											"""simple docstring"""




											return






@require_torch
class      lowerCAmelCase_						(unittest.TestCase  ):


				"""simple docstring"""



				@slow
				def      __magic_name__   (self   ) ->		Optional[Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       YosoModel.from_pretrained("""uw-madison/yoso-4096"""   )
											SCREAMING_SNAKE_CASE__						:	Any									=       torch.tensor([[0, 1, 2, 3, 4, 5]]   )

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

											SCREAMING_SNAKE_CASE__						:	str									=       torch.Size((1, 6, 7_68)   )
											self.assertEqual(output.shape       ,	SCREAMING_SNAKE_CASE__   )

											SCREAMING_SNAKE_CASE__						:	Optional[int]									=       torch.tensor(
											    [[[-0.0611, 0.1242, 0.0840], [0.0280, -0.0048, 0.1125], [0.0106, 0.0226, 0.0751]]]   )

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



				@slow
				def      __magic_name__   (self   ) ->		Optional[int]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Dict									=       YosoForMaskedLM.from_pretrained("""uw-madison/yoso-4096"""   )
											SCREAMING_SNAKE_CASE__						:	int									=       torch.tensor([[0, 1, 2, 3, 4, 5]]   )

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

											SCREAMING_SNAKE_CASE__						:	int									=       5_02_65

											SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       torch.Size((1, 6, vocab_size)   )
											self.assertEqual(output.shape       ,	SCREAMING_SNAKE_CASE__   )

											SCREAMING_SNAKE_CASE__						:	Any									=       torch.tensor(
											    [[[-2.1313, -3.7285, -2.2407], [-2.7047, -3.3314, -2.6408], [0.0629, -2.5166, -0.3356]]]   )

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



				@slow
				def      __magic_name__   (self   ) ->		Tuple:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       YosoForMaskedLM.from_pretrained("""uw-madison/yoso-4096"""   )
											SCREAMING_SNAKE_CASE__						:	Tuple									=       torch.arange(40_96   ).unsqueeze(0   )

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

											SCREAMING_SNAKE_CASE__						:	str									=       5_02_65

											SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       torch.Size((1, 40_96, vocab_size)   )
											self.assertEqual(output.shape       ,	SCREAMING_SNAKE_CASE__   )

											SCREAMING_SNAKE_CASE__						:	Dict									=       torch.tensor(
											    [[[-2.3914, -4.3742, -5.0956], [-4.0988, -4.2384, -7.0406], [-3.1427, -3.7192, -6.6800]]]   )

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






"""simple docstring"""




from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


UpperCAmelCase__				:   Optional[Any]   							=			{'configuration_umt5': ['UMT5Config', 'UMT5OnnxConfig']}


try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					UpperCAmelCase__				:   Tuple   							=			[
					    'UMT5EncoderModel',
					    'UMT5ForConditionalGeneration',
					    'UMT5ForQuestionAnswering',
					    'UMT5Model',
					    'UMT5PreTrainedModel',
					]

if TYPE_CHECKING:
					from .configuration_umta import UMTaConfig, UMTaOnnxConfig

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_umta import (
										    UMTaEncoderModel,
										    UMTaForConditionalGeneration,
										    UMTaForQuestionAnswering,
										    UMTaModel,
										    UMTaPreTrainedModel,
										)
else:
					import sys

					UpperCAmelCase__				:   Optional[Any]   							=			_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




