



"""simple docstring"""







import collections
import inspect
import unittest

from transformers import SwinvaConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch
							from torch import nn

							from transformers import SwinvaForImageClassification, SwinvaForMaskedImageModeling, SwinvaModel
							from transformers.models.swinva.modeling_swinva import SWINV2_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
							from PIL import Image

							from transformers import AutoImageProcessor
class        __A       :







		'''simple docstring'''





		def __init__(	self						:  Optional[int]			,_snake_case						:  Optional[Any]			,_snake_case						:  Union[str, Any]=13			,_snake_case						:  Any=32			,_snake_case						:  int=2			,_snake_case						:  str=3			,_snake_case						:  Optional[Any]=16			,_snake_case						:  List[Any]=[1, 2, 1]			,_snake_case						:  Dict=[2, 2, 4]			,_snake_case						:  List[Any]=2			,_snake_case						:  Any=2.0			,_snake_case						:  Optional[int]=True			,_snake_case						:  Optional[int]=0.0			,_snake_case						:  Union[str, Any]=0.0			,_snake_case						:  str=0.1			,_snake_case						:  List[Any]="gelu"			,_snake_case						:  Tuple=False			,_snake_case						:  Optional[int]=True			,_snake_case						:  str=0.02			,_snake_case						:  List[str]=1e-5			,_snake_case						:  int=True			,_snake_case						:  Dict=None			,_snake_case						:  str=True			,_snake_case						:  List[Any]=10			,_snake_case						:  Any=8			,)					->		Union[str, Any]:






				"""simple docstring"""



				lowercase__   :    Dict				=			parent
				lowercase__   :    Any				=			batch_size
				lowercase__   :    Union[str, Any]				=			image_size
				lowercase__   :    Dict				=			patch_size
				lowercase__   :    int				=			num_channels
				lowercase__   :    Any				=			embed_dim
				lowercase__   :    int				=			depths
				lowercase__   :    Dict				=			num_heads
				lowercase__   :    List[Any]				=			window_size
				lowercase__   :    int				=			mlp_ratio
				lowercase__   :    Optional[int]				=			qkv_bias
				lowercase__   :    str				=			hidden_dropout_prob
				lowercase__   :    List[Any]				=			attention_probs_dropout_prob
				lowercase__   :    Dict				=			drop_path_rate
				lowercase__   :    int				=			hidden_act
				lowercase__   :    Tuple				=			use_absolute_embeddings
				lowercase__   :    Tuple				=			patch_norm
				lowercase__   :    Tuple				=			layer_norm_eps
				lowercase__   :    Optional[Any]				=			initializer_range
				lowercase__   :    int				=			is_training
				lowercase__   :    Optional[int]				=			scope
				lowercase__   :    str				=			use_labels
				lowercase__   :    Dict				=			type_sequence_label_size
				lowercase__   :    Union[str, Any]				=			encoder_stride







		def 					UpperCAmelCase    (	self						:  Dict						)					->		Any:






				"""simple docstring"""



				lowercase__   :    str				=			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

				lowercase__   :    Optional[Any]				=			None
				if self.use_labels:
						lowercase__   :    Optional[int]				=			ids_tensor([self.batch_size]			,self.type_sequence_label_size						)

				lowercase__   :    Union[str, Any]				=			self.get_config()

				return config, pixel_values, labels







		def 					UpperCAmelCase    (	self						:  Optional[Any]						)					->		List[str]:






				"""simple docstring"""



				return SwinvaConfig(
				    image_size=self.image_size			,patch_size=self.patch_size			,num_channels=self.num_channels			,embed_dim=self.embed_dim			,depths=self.depths			,num_heads=self.num_heads			,window_size=self.window_size			,mlp_ratio=self.mlp_ratio			,qkv_bias=self.qkv_bias			,hidden_dropout_prob=self.hidden_dropout_prob			,attention_probs_dropout_prob=self.attention_probs_dropout_prob			,drop_path_rate=self.drop_path_rate			,hidden_act=self.hidden_act			,use_absolute_embeddings=self.use_absolute_embeddings			,path_norm=self.patch_norm			,layer_norm_eps=self.layer_norm_eps			,initializer_range=self.initializer_range			,encoder_stride=self.encoder_stride			,)







		def 					UpperCAmelCase    (	self						:  str			,_snake_case						:  Dict			,_snake_case						:  List[str]			,_snake_case						:  Optional[int]						)					->		Optional[int]:






				"""simple docstring"""



				lowercase__   :    Any				=			SwinvaModel(config=_snake_case						)
				model.to(_snake_case						)
				model.eval()
				lowercase__   :    str				=			model(_snake_case						)

				lowercase__   :    List[Any]				=			((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths						) - 1))
				lowercase__   :    Tuple				=			int(config.embed_dim * 2 ** (len(config.depths						) - 1)						)

				self.parent.assertEqual(result.last_hidden_state.shape			,(self.batch_size, expected_seq_len, expected_dim)						)







		def 					UpperCAmelCase    (	self						:  Union[str, Any]			,_snake_case						:  List[str]			,_snake_case						:  Optional[Any]			,_snake_case						:  int						)					->		Any:






				"""simple docstring"""



				lowercase__   :    Union[str, Any]				=			SwinvaForMaskedImageModeling(config=_snake_case						)
				model.to(_snake_case						)
				model.eval()
				lowercase__   :    Tuple				=			model(_snake_case						)
				self.parent.assertEqual(
				    result.logits.shape			,(self.batch_size, self.num_channels, self.image_size, self.image_size)						)

				# test greyscale images
				lowercase__   :    Optional[int]				=			1
				lowercase__   :    List[Any]				=			SwinvaForMaskedImageModeling(_snake_case						)
				model.to(_snake_case						)
				model.eval()

				lowercase__   :    Union[str, Any]				=			floats_tensor([self.batch_size, 1, self.image_size, self.image_size]						)
				lowercase__   :    str				=			model(_snake_case						)
				self.parent.assertEqual(result.logits.shape			,(self.batch_size, 1, self.image_size, self.image_size)						)







		def 					UpperCAmelCase    (	self						:  str			,_snake_case						:  str			,_snake_case						:  str			,_snake_case						:  Tuple						)					->		Any:






				"""simple docstring"""



				lowercase__   :    Tuple				=			self.type_sequence_label_size
				lowercase__   :    Dict				=			SwinvaForImageClassification(_snake_case						)
				model.to(_snake_case						)
				model.eval()
				lowercase__   :    str				=			model(_snake_case			,labels=_snake_case						)
				self.parent.assertEqual(result.logits.shape			,(self.batch_size, self.type_sequence_label_size)						)







		def 					UpperCAmelCase    (	self						:  Dict						)					->		Dict:






				"""simple docstring"""



				lowercase__   :    Optional[int]				=			self.prepare_config_and_inputs()
				lowercase__		,			lowercase__		,			lowercase__   :    Union[str, Any]				=			config_and_inputs
				lowercase__   :    List[str]				=			{'''pixel_values''': pixel_values}
				return config, inputs_dict





@require_torch
class        __A       (	A_     ,A_     ,unittest.TestCase    ):







		'''simple docstring'''





		lowerCAmelCase			:					Union[str, Any]   	=				(
		    (SwinvaModel, SwinvaForImageClassification, SwinvaForMaskedImageModeling) if is_torch_available() else ()
		)
		lowerCAmelCase			:					Optional[int]   	=				(
		    {"feature-extraction": SwinvaModel, "image-classification": SwinvaForImageClassification}
		    if is_torch_available()
		    else {}
		)

		lowerCAmelCase			:					List[Any]   	=				False
		lowerCAmelCase			:					Dict   	=				False
		lowerCAmelCase			:					List[Any]   	=				False
		lowerCAmelCase			:					Any   	=				False







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






				"""simple docstring"""



				lowercase__   :    Optional[Any]				=			SwinvaModelTester(self						)
				lowercase__   :    List[str]				=			ConfigTester(self			,config_class=_snake_case			,embed_dim=37						)







		def 					UpperCAmelCase    (	self						:  int						)					->		Any:






				"""simple docstring"""



				self.config_tester.create_and_test_config_to_json_string()
				self.config_tester.create_and_test_config_to_json_file()
				self.config_tester.create_and_test_config_from_and_save_pretrained()
				self.config_tester.create_and_test_config_with_num_labels()
				self.config_tester.check_config_can_be_init_without_params()
				self.config_tester.check_config_arguments_init()







		def 					UpperCAmelCase    (	self						:  str						)					->		List[Any]:






				"""simple docstring"""



				lowercase__   :    int				=			self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*_snake_case						)







		@unittest.skip(reason='''Got `CUDA error: misaligned address` with PyTorch 2.0.0.'''						)
		def 					UpperCAmelCase    (	self						:  Optional[Any]						)					->		Optional[int]:






				"""simple docstring"""



				pass







		@unittest.skip(reason='''Swinv2 does not use inputs_embeds'''						)
		def 					UpperCAmelCase    (	self						:  List[str]						)					->		str:






				"""simple docstring"""



				pass







		def 					UpperCAmelCase    (	self						:  Optional[int]						)					->		Tuple:






				"""simple docstring"""



				lowercase__		,			lowercase__   :    Optional[int]				=			self.model_tester.prepare_config_and_inputs_for_common()

				for model_class in self.all_model_classes:
						lowercase__   :    List[Any]				=			model_class(_snake_case						)
						self.assertIsInstance(model.get_input_embeddings()			,(nn.Module)						)
						lowercase__   :    str				=			model.get_output_embeddings()
						self.assertTrue(x is None or isinstance(_snake_case			,nn.Linear						)						)







		def 					UpperCAmelCase    (	self						:  int						)					->		List[Any]:






				"""simple docstring"""



				lowercase__		,			lowercase__   :    Union[str, Any]				=			self.model_tester.prepare_config_and_inputs_for_common()

				for model_class in self.all_model_classes:
						lowercase__   :    str				=			model_class(_snake_case						)
						lowercase__   :    List[Any]				=			inspect.signature(model.forward						)
						# signature.parameters is an OrderedDict => so arg_names order is deterministic
						lowercase__   :    Optional[Any]				=			[*signature.parameters.keys()]

						lowercase__   :    Tuple				=			['''pixel_values''']
						self.assertListEqual(arg_names[:1]			,_snake_case						)







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






				"""simple docstring"""



				lowercase__		,			lowercase__   :    List[str]				=			self.model_tester.prepare_config_and_inputs_for_common()
				lowercase__   :    Tuple				=			True

				for model_class in self.all_model_classes:
						lowercase__   :    Optional[int]				=			True
						lowercase__   :    str				=			False
						lowercase__   :    Union[str, Any]				=			True
						lowercase__   :    Optional[Any]				=			model_class(_snake_case						)
						model.to(_snake_case						)
						model.eval()
						with torch.no_grad():
								lowercase__   :    str				=			model(**self._prepare_for_class(_snake_case			,_snake_case						)						)
						lowercase__   :    Dict				=			outputs.attentions
						lowercase__   :    Any				=			len(self.model_tester.depths						)
						self.assertEqual(len(_snake_case						)			,_snake_case						)

						# check that output_attentions also work using config
						del inputs_dict["output_attentions"]
						lowercase__   :    List[Any]				=			True
						lowercase__   :    Optional[Any]				=			config.window_size**2
						lowercase__   :    Any				=			model_class(_snake_case						)
						model.to(_snake_case						)
						model.eval()
						with torch.no_grad():
								lowercase__   :    List[str]				=			model(**self._prepare_for_class(_snake_case			,_snake_case						)						)
						lowercase__   :    Optional[Any]				=			outputs.attentions
						self.assertEqual(len(_snake_case						)			,_snake_case						)

						self.assertListEqual(
						    list(attentions[0].shape[-3:]						)			,[self.model_tester.num_heads[0], window_size_squared, window_size_squared]			,)
						lowercase__   :    Optional[Any]				=			len(_snake_case						)

						# Check attention is always last and order is fine
						lowercase__   :    Optional[int]				=			True
						lowercase__   :    Tuple				=			True
						lowercase__   :    Optional[Any]				=			model_class(_snake_case						)
						model.to(_snake_case						)
						model.eval()
						with torch.no_grad():
								lowercase__   :    Optional[Any]				=			model(**self._prepare_for_class(_snake_case			,_snake_case						)						)

						if hasattr(self.model_tester			,'''num_hidden_states_types'''						):
								lowercase__   :    int				=			self.model_tester.num_hidden_states_types
						else:
								# also another +1 for reshaped_hidden_states
								lowercase__   :    List[str]				=			2
						self.assertEqual(out_len + added_hidden_states			,len(_snake_case						)						)

						lowercase__   :    Optional[int]				=			outputs.attentions

						self.assertEqual(len(_snake_case						)			,_snake_case						)

						self.assertListEqual(
						    list(self_attentions[0].shape[-3:]						)			,[self.model_tester.num_heads[0], window_size_squared, window_size_squared]			,)







		def 					UpperCAmelCase    (	self						:  List[str]			,_snake_case						:  int			,_snake_case						:  List[str]			,_snake_case						:  Optional[int]			,_snake_case						:  List[Any]						)					->		Union[str, Any]:






				"""simple docstring"""



				lowercase__   :    List[Any]				=			model_class(_snake_case						)
				model.to(_snake_case						)
				model.eval()

				with torch.no_grad():
						lowercase__   :    int				=			model(**self._prepare_for_class(_snake_case			,_snake_case						)						)

				lowercase__   :    Optional[int]				=			outputs.hidden_states

				lowercase__   :    List[Any]				=			getattr(
				    self.model_tester			,'''expected_num_hidden_layers'''			,len(self.model_tester.depths						) + 1						)
				self.assertEqual(len(_snake_case						)			,_snake_case						)

				# Swinv2 has a different seq_length
				lowercase__   :    Dict				=			(
				    config.patch_size
				    if isinstance(config.patch_size			,collections.abc.Iterable						)
				    else (config.patch_size, config.patch_size)
				)

				lowercase__   :    int				=			(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

				self.assertListEqual(
				    list(hidden_states[0].shape[-2:]						)			,[num_patches, self.model_tester.embed_dim]			,)

				lowercase__   :    Tuple				=			outputs.reshaped_hidden_states
				self.assertEqual(len(_snake_case						)			,_snake_case						)

				lowercase__		,			lowercase__		,			lowercase__		,			lowercase__   :    List[str]				=			reshaped_hidden_states[0].shape
				lowercase__   :    int				=			(
				    reshaped_hidden_states[0].view(_snake_case			,_snake_case			,height * width						).permute(0			,2			,1						)
				)
				self.assertListEqual(
				    list(reshaped_hidden_states.shape[-2:]						)			,[num_patches, self.model_tester.embed_dim]			,)







		def 					UpperCAmelCase    (	self						:  Tuple						)					->		int:






				"""simple docstring"""



				lowercase__		,			lowercase__   :    Optional[int]				=			self.model_tester.prepare_config_and_inputs_for_common()

				lowercase__   :    str				=			(
				    self.model_tester.image_size
				    if isinstance(self.model_tester.image_size			,collections.abc.Iterable						)
				    else (self.model_tester.image_size, self.model_tester.image_size)
				)

				for model_class in self.all_model_classes:
						lowercase__   :    List[str]				=			True
						self.check_hidden_states_output(_snake_case			,_snake_case			,_snake_case			,_snake_case						)

						# check that output_hidden_states also work using config
						del inputs_dict["output_hidden_states"]
						lowercase__   :    str				=			True

						self.check_hidden_states_output(_snake_case			,_snake_case			,_snake_case			,_snake_case						)







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






				"""simple docstring"""



				lowercase__		,			lowercase__   :    List[str]				=			self.model_tester.prepare_config_and_inputs_for_common()
				lowercase__   :    List[Any]				=			3

				lowercase__   :    Tuple				=			(
				    self.model_tester.image_size
				    if isinstance(self.model_tester.image_size			,collections.abc.Iterable						)
				    else (self.model_tester.image_size, self.model_tester.image_size)
				)
				lowercase__   :    Optional[int]				=			(
				    config.patch_size
				    if isinstance(config.patch_size			,collections.abc.Iterable						)
				    else (config.patch_size, config.patch_size)
				)

				lowercase__   :    Dict				=			image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
				lowercase__   :    Dict				=			image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

				for model_class in self.all_model_classes:
						lowercase__   :    str				=			True
						self.check_hidden_states_output(_snake_case			,_snake_case			,_snake_case			,(padded_height, padded_width)						)

						# check that output_hidden_states also work using config
						del inputs_dict["output_hidden_states"]
						lowercase__   :    Dict				=			True
						self.check_hidden_states_output(_snake_case			,_snake_case			,_snake_case			,(padded_height, padded_width)						)







		def 					UpperCAmelCase    (	self						:  Tuple						)					->		List[Any]:






				"""simple docstring"""



				lowercase__   :    Union[str, Any]				=			self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_masked_image_modeling(*_snake_case						)







		def 					UpperCAmelCase    (	self						:  Dict						)					->		Any:






				"""simple docstring"""



				lowercase__   :    Tuple				=			self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_image_classification(*_snake_case						)







		@slow
		def 					UpperCAmelCase    (	self						:  List[str]						)					->		Optional[Any]:






				"""simple docstring"""



				for model_name in SWINV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
						lowercase__   :    Union[str, Any]				=			SwinvaModel.from_pretrained(_snake_case						)
						self.assertIsNotNone(_snake_case						)







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






				"""simple docstring"""



				lowercase__		,			lowercase__   :    str				=			self.model_tester.prepare_config_and_inputs_for_common()

				lowercase__   :    Tuple				=			_config_zero_init(_snake_case						)
				for model_class in self.all_model_classes:
						lowercase__   :    Optional[int]				=			model_class(config=_snake_case						)
						for name, param in model.named_parameters():
								if "embeddings" not in name and "logit_scale" not in name and param.requires_grad:
										self.assertIn(
										    ((param.data.mean() * 1e9).round() / 1e9).item()			,[0.0, 1.0]			,msg=f"""Parameter {name} of model {model_class} seems not properly initialized"""			,)





@require_vision
@require_torch
class        __A       (	unittest.TestCase    ):







		'''simple docstring'''





		@cached_property
		def 					UpperCAmelCase    (	self						:  Optional[Any]						)					->		Tuple:






				"""simple docstring"""



				return (
				    AutoImageProcessor.from_pretrained('''microsoft/swinv2-tiny-patch4-window8-256'''						)
				    if is_vision_available()
				    else None
				)







		@slow
		def 					UpperCAmelCase    (	self						:  Any						)					->		List[str]:






				"""simple docstring"""



				lowercase__   :    str				=			SwinvaForImageClassification.from_pretrained('''microsoft/swinv2-tiny-patch4-window8-256'''						).to(
				    _snake_case						)
				lowercase__   :    Union[str, Any]				=			self.default_image_processor

				lowercase__   :    List[str]				=			Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''						)
				lowercase__   :    Dict				=			image_processor(images=_snake_case			,return_tensors='''pt'''						).to(_snake_case						)

				# forward pass
				with torch.no_grad():
						lowercase__   :    Optional[Any]				=			model(**_snake_case						)

				# verify the logits
				lowercase__   :    str				=			torch.Size((1, 1_000)						)
				self.assertEqual(outputs.logits.shape			,_snake_case						)
				lowercase__   :    Dict				=			torch.tensor([-0.3947, -0.4306, 0.0026]						).to(_snake_case						)
				self.assertTrue(torch.allclose(outputs.logits[0, :3]			,_snake_case			,atol=1e-4						)						)




import tempfile
import unittest

from transformers import TaConfig, is_torch_available
from transformers.testing_utils import (
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
    torch_device,
)

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


if is_torch_available():
     import torch

     from transformers import AutoTokenizer, UMTaForConditionalGeneration, UMTaForQuestionAnswering, UMTaModel






class   _lowerCamelCase :




       """simple docstring"""



       def __init__(      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE=99				,       _SCREAMING_SNAKE_CASE=13				,       _SCREAMING_SNAKE_CASE=7				,       _SCREAMING_SNAKE_CASE=9				,       _SCREAMING_SNAKE_CASE=True				,       _SCREAMING_SNAKE_CASE=True				,       _SCREAMING_SNAKE_CASE=False				,       _SCREAMING_SNAKE_CASE=32				,       _SCREAMING_SNAKE_CASE=5				,       _SCREAMING_SNAKE_CASE=4				,       _SCREAMING_SNAKE_CASE=37				,       _SCREAMING_SNAKE_CASE=8				,       _SCREAMING_SNAKE_CASE=0.1				,       _SCREAMING_SNAKE_CASE=0.0_0_2				,       _SCREAMING_SNAKE_CASE=1				,       _SCREAMING_SNAKE_CASE=0				,       _SCREAMING_SNAKE_CASE=0				,       _SCREAMING_SNAKE_CASE=None				,       _SCREAMING_SNAKE_CASE=None				,       )->str:




         '''simple docstring'''





         A_      :				Union[str, Any]		     =					parent
         A_      :				Optional[Any]		     =					batch_size
         A_      :				int		     =					encoder_seq_length
         A_      :				List[Any]		     =					decoder_seq_length
         # For common tests
         A_      :				str		     =					self.decoder_seq_length
         A_      :				int		     =					is_training
         A_      :				Tuple		     =					use_attention_mask
         A_      :				Any		     =					use_labels
         A_      :				List[str]		     =					vocab_size
         A_      :				Dict		     =					hidden_size
         A_      :				Optional[Any]		     =					num_hidden_layers
         A_      :				List[Any]		     =					num_attention_heads
         A_      :				Tuple		     =					d_ff
         A_      :				List[str]		     =					relative_attention_num_buckets
         A_      :				Tuple		     =					dropout_rate
         A_      :				Optional[Any]		     =					initializer_factor
         A_      :				List[str]		     =					eos_token_id
         A_      :				Union[str, Any]		     =					pad_token_id
         A_      :				Any		     =					decoder_start_token_id
         A_      :				int		     =					None
         A_      :				Tuple		     =					decoder_layers



       def        _snake_case       (      self		)->str:




         '''simple docstring'''





         return TaConfig.from_pretrained('''google/umt5-base'''		)



       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE=None				,       _SCREAMING_SNAKE_CASE=None				,       _SCREAMING_SNAKE_CASE=None				,       _SCREAMING_SNAKE_CASE=None				,       _SCREAMING_SNAKE_CASE=None				,       )->List[str]:




         '''simple docstring'''





         if attention_mask is None:
           A_      :				Union[str, Any]		     =					input_ids.ne(config.pad_token_id		)
         if decoder_attention_mask is None:
           A_      :				Any		     =					decoder_input_ids.ne(config.pad_token_id		)
         if head_mask is None:
           A_      :				str		     =					torch.ones(config.num_hidden_layers				,       config.num_attention_heads				,       device=_SCREAMING_SNAKE_CASE		)
         if decoder_head_mask is None:
           A_      :				List[str]		     =					torch.ones(config.num_decoder_layers				,       config.num_attention_heads				,       device=_SCREAMING_SNAKE_CASE		)
         if cross_attn_head_mask is None:
           A_      :				Union[str, Any]		     =					torch.ones(
               config.num_decoder_layers				,       config.num_attention_heads				,       device=_SCREAMING_SNAKE_CASE		)
         return {
             "input_ids": input_ids,
             "decoder_input_ids": decoder_input_ids,
             "attention_mask": attention_mask,
             "decoder_attention_mask": decoder_attention_mask,
             "head_mask": head_mask,
             "decoder_head_mask": decoder_head_mask,
             "cross_attn_head_mask": cross_attn_head_mask,
         }



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




         '''simple docstring'''





         A_      :				List[Any]		     =					ids_tensor([self.batch_size, self.encoder_seq_length]				,       self.vocab_size		)
         A_      :				int		     =					ids_tensor([self.batch_size, self.decoder_seq_length]				,       self.vocab_size		)

         # we need to clamp the input ids here to avoid having pad token in between
         # this is because for NllbMoe the position_ids are prepared such that
         # all pad tokens have pos id = 2 and rest are between 2..seq_length
         # and the seq_length here is seq_length - num_pad_tokens
         # but when using past, there is no way of knowing if the past input ids had
         # pad tokens in them, which results in incorrect seq_lenth and which in turn results in
         # position_ids being off by num_pad_tokens in past input
         A_      :				Tuple		     =					input_ids.clamp(self.pad_token_id + 1		)
         A_      :				int		     =					decoder_input_ids.clamp(self.pad_token_id + 1		)

         A_      :				Optional[Any]		     =					self.get_config()
         A_      :				int		     =					config.num_attention_heads
         A_      :				Dict		     =					self.prepare_inputs_dict(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)
         return config, input_dict



       def        _snake_case       (      self		)->Optional[Any]:




         '''simple docstring'''





         A_     ,       A_      :				str		     =					self.prepare_config_and_inputs()
         return config, inputs_dict



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




         '''simple docstring'''





         return TaConfig(
             vocab_size=166				,       d_model=self.hidden_size				,       d_ff=self.d_ff				,       d_kv=self.hidden_size // self.num_attention_heads				,       num_layers=self.num_hidden_layers				,       num_decoder_layers=self.decoder_layers				,       num_heads=self.num_attention_heads				,       relative_attention_num_buckets=self.relative_attention_num_buckets				,       dropout_rate=self.dropout_rate				,       initializer_factor=self.initializer_factor				,       eos_token_id=self.eos_token_id				,       bos_token_id=self.pad_token_id				,       pad_token_id=self.pad_token_id				,       decoder_start_token_id=self.decoder_start_token_id				,       )



       def        _snake_case       (      self		)->Dict:




         '''simple docstring'''





         return TaConfig(
             vocab_size=self.vocab_size				,       d_model=self.hidden_size				,       d_ff=self.d_ff				,       d_kv=self.hidden_size // self.num_attention_heads				,       num_layers=self.num_hidden_layers				,       num_decoder_layers=self.decoder_layers				,       num_heads=self.num_attention_heads				,       relative_attention_num_buckets=self.relative_attention_num_buckets				,       dropout_rate=self.dropout_rate				,       initializer_factor=self.initializer_factor				,       eos_token_id=self.eos_token_id				,       bos_token_id=self.pad_token_id				,       pad_token_id=self.pad_token_id				,       decoder_start_token_id=self.decoder_start_token_id				,       )



       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       )->Dict:




         '''simple docstring'''





         A_      :				Dict		     =					UMTaModel(config=_SCREAMING_SNAKE_CASE		)
         model.to(_SCREAMING_SNAKE_CASE		)
         model.eval()
         A_      :				Tuple		     =					model(
             input_ids=_SCREAMING_SNAKE_CASE				,       decoder_input_ids=_SCREAMING_SNAKE_CASE				,       attention_mask=_SCREAMING_SNAKE_CASE				,       decoder_attention_mask=_SCREAMING_SNAKE_CASE				,       )
         A_      :				List[Any]		     =					model(input_ids=_SCREAMING_SNAKE_CASE				,       decoder_input_ids=_SCREAMING_SNAKE_CASE		)
         A_      :				str		     =					result.last_hidden_state
         A_      :				List[Any]		     =					result.past_key_values
         A_      :				Any		     =					result.encoder_last_hidden_state

         self.parent.assertEqual(encoder_output.size()				,       (self.batch_size, self.encoder_seq_length, self.hidden_size)		)
         self.parent.assertEqual(decoder_output.size()				,       (self.batch_size, self.decoder_seq_length, self.hidden_size)		)
         # There should be `num_layers` key value embeddings stored in decoder_past
         self.parent.assertEqual(len(_SCREAMING_SNAKE_CASE		)				,       config.num_layers		)
         # There should be a self attn key, a self attn value, a cross attn key and a cross attn value stored in each decoder_past tuple
         self.parent.assertEqual(len(decoder_past[0]		)				,       4		)



       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       )->Optional[Any]:




         '''simple docstring'''





         A_      :				List[Any]		     =					UMTaModel(config=_SCREAMING_SNAKE_CASE		).get_decoder().to(_SCREAMING_SNAKE_CASE		).eval()
         # first forward pass
         A_      :				Union[str, Any]		     =					model(_SCREAMING_SNAKE_CASE				,       use_cache=_SCREAMING_SNAKE_CASE		)
         A_      :				List[str]		     =					model(_SCREAMING_SNAKE_CASE		)
         A_      :				int		     =					model(_SCREAMING_SNAKE_CASE				,       use_cache=_SCREAMING_SNAKE_CASE		)

         self.parent.assertTrue(len(_SCREAMING_SNAKE_CASE		) == len(_SCREAMING_SNAKE_CASE		)		)
         self.parent.assertTrue(len(_SCREAMING_SNAKE_CASE		) == len(_SCREAMING_SNAKE_CASE		) + 1		)

         A_     ,       A_      :				int		     =					outputs.to_tuple()

         # create hypothetical next token and extent to next_input_ids
         A_      :				str		     =					ids_tensor((self.batch_size, 1)				,       config.vocab_size		)

         # append to next input_ids and
         A_      :				Tuple		     =					torch.cat([input_ids, next_tokens]				,       dim=-1		)

         A_      :				Tuple		     =					model(_SCREAMING_SNAKE_CASE		)['''last_hidden_state''']
         A_      :				Union[str, Any]		     =					model(_SCREAMING_SNAKE_CASE				,       past_key_values=_SCREAMING_SNAKE_CASE		)['''last_hidden_state''']

         # select random slice
         A_      :				List[Any]		     =					ids_tensor((1,)				,       output_from_past.shape[-1]		).item()
         A_      :				Union[str, Any]		     =					output_from_no_past[:, -1, random_slice_idx].detach()
         A_      :				List[str]		     =					output_from_past[:, 0, random_slice_idx].detach()

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







       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       )->Optional[Any]:




         '''simple docstring'''





         A_      :				List[str]		     =					UMTaModel(config=_SCREAMING_SNAKE_CASE		).to(_SCREAMING_SNAKE_CASE		).half().eval()
         A_      :				Tuple		     =					model(**_SCREAMING_SNAKE_CASE		)['''last_hidden_state''']
         self.parent.assertFalse(torch.isnan(_SCREAMING_SNAKE_CASE		).any().item()		)






@require_torch
class   _lowerCamelCase (					UpperCamelCase					,							UpperCamelCase					,							UpperCamelCase					,							unittest.TestCase			):




       """simple docstring"""
       snake_case      =    (
           (UMTaModel, UMTaForConditionalGeneration, UMTaForQuestionAnswering) if is_torch_available() else ()
       )
       snake_case      =    (UMTaForConditionalGeneration,) if is_torch_available() else ()
       snake_case      =    (
           {
               "conversational": UMTaForConditionalGeneration,
               "feature-extraction": UMTaModel,
               "summarization": UMTaForConditionalGeneration,
               "text2text-generation": UMTaForConditionalGeneration,
               "translation": UMTaForConditionalGeneration,
               "question-answering": UMTaForQuestionAnswering,
           }
           if is_torch_available()
           else {}
       )
       snake_case      =    True
       snake_case      =    False
       snake_case      =    False
       snake_case      =    True
       snake_case      =    True
       # The small UMT5 model needs higher percentages for CPU/MP tests
       snake_case      =    [0.8, 0.9]



       def        _snake_case       (      self		)->str:




         '''simple docstring'''





         A_      :				Union[str, Any]		     =					UMTaModelTester(self		)



       @unittest.skip('''Test has a segmentation fault on torch 1.8.0'''		)
       def        _snake_case       (      self		)->List[str]:




         '''simple docstring'''





         A_      :				Any		     =					self.model_tester.prepare_config_and_inputs()
         A_      :				str		     =					UMTaModel(config_and_inputs[0]		).to(_SCREAMING_SNAKE_CASE		)
         with tempfile.TemporaryDirectory() as tmpdirname:
           torch.onnx.export(
               _SCREAMING_SNAKE_CASE				,       (config_and_inputs[1], config_and_inputs[3], config_and_inputs[2])				,       F'''{tmpdirname}/t5_test.onnx'''				,       export_params=_SCREAMING_SNAKE_CASE				,       opset_version=9				,       input_names=['''input_ids''', '''decoder_input_ids''']				,       )



       @unittest.skipIf(torch_device == '''cpu'''				,       '''Cant do half precision'''		)
       def        _snake_case       (      self		)->str:




         '''simple docstring'''





         A_      :				Dict		     =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model_fpaa_forward(*_SCREAMING_SNAKE_CASE		)



       def        _snake_case       (      self		)->int:




         '''simple docstring'''





         A_      :				Optional[int]		     =					['''encoder_attentions''', '''decoder_attentions''', '''cross_attentions''']
         A_      :				List[Any]		     =					self.model_tester.prepare_config_and_inputs()
         A_      :				Union[str, Any]		     =					config_and_inputs[0]
         A_      :				int		     =					UMTaForConditionalGeneration(_SCREAMING_SNAKE_CASE		).eval()
         model.to(_SCREAMING_SNAKE_CASE		)

         A_      :				Any		     =					{
             '''head_mask''': torch.zeros(config.num_layers				,       config.num_heads				,       device=_SCREAMING_SNAKE_CASE		),
             '''decoder_head_mask''': torch.zeros(config.num_decoder_layers				,       config.num_heads				,       device=_SCREAMING_SNAKE_CASE		),
             '''cross_attn_head_mask''': torch.zeros(config.num_decoder_layers				,       config.num_heads				,       device=_SCREAMING_SNAKE_CASE		),
         }

         for attn_name, (name, mask) in zip(_SCREAMING_SNAKE_CASE				,       head_masking.items()		):
           A_      :				Optional[int]		     =					{name: mask}
           # Explicitly pass decoder_head_mask as it is required from T5 model when head_mask specified
           if name == "head_mask":
             A_      :				Dict		     =					torch.ones(
                 config.num_decoder_layers				,       config.num_heads				,       device=_SCREAMING_SNAKE_CASE		)

           A_      :				Union[str, Any]		     =					model.generate(
               config_and_inputs[1]['''input_ids''']				,       num_beams=1				,       max_length=3				,       output_attentions=_SCREAMING_SNAKE_CASE				,       return_dict_in_generate=_SCREAMING_SNAKE_CASE				,       **_SCREAMING_SNAKE_CASE				,       )
           # We check the state of decoder_attentions and cross_attentions just from the last step
           A_      :				Tuple		     =					out[attn_name] if attn_name == attention_names[0] else out[attn_name][-1]
           self.assertEqual(sum([w.sum().item() for w in attn_weights]		)				,       0.0		)







       @unittest.skip('''Does not work on the tiny model as we keep hitting edge cases.'''		)
       def        _snake_case       (      self		)->Any:




         '''simple docstring'''





         pass






@require_torch
@require_sentencepiece
@require_tokenizers
class   _lowerCamelCase (					unittest.TestCase			):




       """simple docstring"""



       @slow
       @unittest.skip(
           '''Unless we stop stripping left and right by default for all special tokens, the expected ids obtained here will not match the original ones. Wait for https://github.com/huggingface/transformers/pull/23909 to be merged'''		)
       def        _snake_case       (      self		)->int:




         '''simple docstring'''





         A_      :				Optional[Any]		     =					UMTaForConditionalGeneration.from_pretrained('''google/umt5-small'''				,       return_dict=_SCREAMING_SNAKE_CASE		).to(_SCREAMING_SNAKE_CASE		)
         A_      :				Union[str, Any]		     =					AutoTokenizer.from_pretrained('''google/umt5-small'''				,       use_fast=_SCREAMING_SNAKE_CASE				,       legacy=_SCREAMING_SNAKE_CASE		)
         A_      :				Tuple		     =					[
             '''Bonjour monsieur <extra_id_0> bien <extra_id_1>.''',
             '''No se como puedo <extra_id_0>.''',
             '''This is the reason why we <extra_id_0> them.''',
             '''The <extra_id_0> walks in <extra_id_1>, seats''',
             '''A <extra_id_0> walks into a bar and orders a <extra_id_1> with <extra_id_2> pinch of <extra_id_3>.''',
         ]
         A_      :				Union[str, Any]		     =					tokenizer(_SCREAMING_SNAKE_CASE				,       return_tensors='''pt'''				,       padding=_SCREAMING_SNAKE_CASE		).input_ids
         # fmt: off
         A_      :				Dict		     =					torch.tensor(
             [
                 [ 3_8530, 21_0703, 25_6299, 1410, 25_6298, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
                 [   826, 321, 671, 2_5922, 25_6299, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
                 [  1460, 339, 312, 1_9014, 1_0620, 758, 25_6299, 2355,274, 1, 0, 0, 0, 0, 0, 0,0, 0],
                 [   517, 25_6299, 1_4869, 281, 301, 25_6298, 275, 11_9983,1, 0, 0, 0, 0, 0, 0, 0,0, 0],
                 [   320, 25_6299, 1_4869, 281, 2234, 289, 2275, 333,6_1391, 289, 25_6298, 543, 25_6297, 16_8714, 329, 25_6296,274, 1],
             ]		)
         # fmt: on
         torch.testing.assert_allclose(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)

         A_      :				List[Any]		     =					model.generate(input_ids.to(_SCREAMING_SNAKE_CASE		)		)
         A_      :				str		     =					[
             '''<pad><extra_id_0> et<extra_id_1> [eod] <extra_id_2><extra_id_55>.. [eod] 💐 💐 💐 💐 💐 💐 💐 💐 💐 💐 💐 <extra_id_56>ajšietosto<extra_id_56>lleux<extra_id_19><extra_id_6>ajšie</s>''',
             '''<pad><extra_id_0>.<extra_id_1>.,<0x0A>...spech <0x0A><extra_id_20> <extra_id_21></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>''',
             '''<pad><extra_id_0> are not going to be a part of the world. We are not going to be a part of<extra_id_1> and<extra_id_2><0x0A><extra_id_48>.<extra_id_48></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>''',
             '''<pad><extra_id_0> door<extra_id_1>, the door<extra_id_2> 피해[/</s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>''',
             '''<pad><extra_id_0>nyone who<extra_id_1> drink<extra_id_2> a<extra_id_3> alcohol<extra_id_4> A<extra_id_5> A. This<extra_id_6> I<extra_id_7><extra_id_52><extra_id_53></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>''',
         ]
         A_      :				str		     =					tokenizer.batch_decode(_SCREAMING_SNAKE_CASE		)
         self.assertEqual(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)




