



"""simple docstring"""






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      A__  :




	'''simple docstring'''





	def __init__(	self:						int			,     _SCREAMING_SNAKE_CASE:						Union[str, Any]			,     _SCREAMING_SNAKE_CASE:						List[str]=99			,     _SCREAMING_SNAKE_CASE:						Dict=13			,     _SCREAMING_SNAKE_CASE:						Tuple=7			,     _SCREAMING_SNAKE_CASE:						Union[str, Any]=9			,     _SCREAMING_SNAKE_CASE:						str=True			,     _SCREAMING_SNAKE_CASE:						int=True			,     _SCREAMING_SNAKE_CASE:						Optional[Any]=False			,     _SCREAMING_SNAKE_CASE:						Union[str, Any]=32			,     _SCREAMING_SNAKE_CASE:						Optional[int]=5			,     _SCREAMING_SNAKE_CASE:						Optional[Any]=4			,     _SCREAMING_SNAKE_CASE:						Union[str, Any]=37			,     _SCREAMING_SNAKE_CASE:						List[Any]=8			,     _SCREAMING_SNAKE_CASE:						Optional[int]=0.1			,     _SCREAMING_SNAKE_CASE:						str=0.002			,     _SCREAMING_SNAKE_CASE:						List[Any]=1			,     _SCREAMING_SNAKE_CASE:						List[str]=0			,     _SCREAMING_SNAKE_CASE:						int=0			,     _SCREAMING_SNAKE_CASE:						Optional[int]=None			,     _SCREAMING_SNAKE_CASE:						str=None			,     )							->			Tuple:

				"""simple docstring"""

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

	def 				_SCREAMING_SNAKE_CASE  (	self:						Dict)							->			Tuple:

				"""simple docstring"""

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

	def 				_SCREAMING_SNAKE_CASE  (	self:						int			,     _SCREAMING_SNAKE_CASE:						Optional[Any]			,     _SCREAMING_SNAKE_CASE:						Optional[int]			,     _SCREAMING_SNAKE_CASE:						int			,     _SCREAMING_SNAKE_CASE:						Any=None			,     _SCREAMING_SNAKE_CASE:						Optional[Any]=None			,     _SCREAMING_SNAKE_CASE:						Optional[Any]=None			,     _SCREAMING_SNAKE_CASE:						Optional[Any]=None			,     _SCREAMING_SNAKE_CASE:						Any=None			,     )							->			Optional[Any]:

				"""simple docstring"""

				if attention_mask is None:
							__lowerCAmelCase :    List[str]				=			input_ids.ne(config.pad_token_id)
				if decoder_attention_mask is None:
							__lowerCAmelCase :    Dict				=			decoder_input_ids.ne(config.pad_token_id)
				if head_mask is None:
							__lowerCAmelCase :    List[str]				=			torch.ones(config.num_hidden_layers			,     config.num_attention_heads			,     device=A_)
				if decoder_head_mask is None:
							__lowerCAmelCase :    str				=			torch.ones(config.num_decoder_layers			,     config.num_attention_heads			,     device=A_)
				if cross_attn_head_mask is None:
							__lowerCAmelCase :    Any				=			torch.ones(
							    config.num_decoder_layers			,     config.num_attention_heads			,     device=A_)
				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 				_SCREAMING_SNAKE_CASE  (	self:						List[Any])							->			Dict:

				"""simple docstring"""

				__lowerCAmelCase :    Tuple				=			ids_tensor([self.batch_size, self.encoder_seq_length]			,     self.vocab_size)
				__lowerCAmelCase :    List[str]				=			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
				__lowerCAmelCase :    str				=			input_ids.clamp(self.pad_token_id + 1)
				__lowerCAmelCase :    Dict				=			decoder_input_ids.clamp(self.pad_token_id + 1)

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

	def 				_SCREAMING_SNAKE_CASE  (	self:						str)							->			Optional[int]:

				"""simple docstring"""

				__lowerCAmelCase		,      __lowerCAmelCase :    Tuple				=			self.prepare_config_and_inputs()
				return config, inputs_dict

	def 				_SCREAMING_SNAKE_CASE  (	self:						Union[str, Any])							->			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 				_SCREAMING_SNAKE_CASE  (	self:						List[str])							->			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 				_SCREAMING_SNAKE_CASE  (	self:						List[Any]			,     _SCREAMING_SNAKE_CASE:						List[Any]			,     _SCREAMING_SNAKE_CASE:						Tuple			,     _SCREAMING_SNAKE_CASE:						Any			,     _SCREAMING_SNAKE_CASE:						Tuple			,     _SCREAMING_SNAKE_CASE:						int			,     _SCREAMING_SNAKE_CASE:						Union[str, Any]			,     )							->			Dict:

				"""simple docstring"""

				__lowerCAmelCase :    Optional[Any]				=			UMTaModel(config=A_)
				model.to(A_)
				model.eval()
				__lowerCAmelCase :    List[str]				=			model(
				    input_ids=A_			,     decoder_input_ids=A_			,     attention_mask=A_			,     decoder_attention_mask=A_			,     )
				__lowerCAmelCase :    Tuple				=			model(input_ids=A_			,     decoder_input_ids=A_)
				__lowerCAmelCase :    Tuple				=			result.last_hidden_state
				__lowerCAmelCase :    List[str]				=			result.past_key_values
				__lowerCAmelCase :    Optional[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(A_)			,     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 				_SCREAMING_SNAKE_CASE  (	self:						Dict			,     _SCREAMING_SNAKE_CASE:						Optional[Any]			,     _SCREAMING_SNAKE_CASE:						int			,     _SCREAMING_SNAKE_CASE:						Dict			,     _SCREAMING_SNAKE_CASE:						Union[str, Any]			,     _SCREAMING_SNAKE_CASE:						str			,     _SCREAMING_SNAKE_CASE:						List[str]			,     )							->			Union[str, Any]:

				"""simple docstring"""

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

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

				__lowerCAmelCase		,      __lowerCAmelCase :    Optional[int]				=			outputs.to_tuple()

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

				# append to next input_ids and
				__lowerCAmelCase :    Union[str, Any]				=			torch.cat([input_ids, next_tokens]			,     dim=-1)

				__lowerCAmelCase :    Dict				=			model(A_)["last_hidden_state"]
				__lowerCAmelCase :    Tuple				=			model(A_			,     past_key_values=A_)["last_hidden_state"]

				# select random slice
				__lowerCAmelCase :    Optional[int]				=			ids_tensor((1,)			,     output_from_past.shape[-1]).item()
				__lowerCAmelCase :    List[str]				=			output_from_no_past[:, -1, random_slice_idx].detach()
				__lowerCAmelCase :    Tuple				=			output_from_past[:, 0, random_slice_idx].detach()

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







	def 				_SCREAMING_SNAKE_CASE  (	self:						str			,     _SCREAMING_SNAKE_CASE:						Dict			,     _SCREAMING_SNAKE_CASE:						Any			,     )							->			List[str]:

				"""simple docstring"""

				__lowerCAmelCase :    Any				=			UMTaModel(config=A_).to(A_).half().eval()
				__lowerCAmelCase :    Optional[int]				=			model(**A_)["last_hidden_state"]
				self.parent.assertFalse(torch.isnan(A_).any().item())






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




	'''simple docstring'''





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

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

				"""simple docstring"""

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

	@unittest.skip("Test has a segmentation fault on torch 1.8.0")
	def 				_SCREAMING_SNAKE_CASE  (	self:						Optional[Any])							->			List[str]:

				"""simple docstring"""

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

	@unittest.skipIf(torch_device == "cpu"			,     "Cant do half precision")
	def 				_SCREAMING_SNAKE_CASE  (	self:						Any)							->			Optional[int]:

				"""simple docstring"""

				__lowerCAmelCase :    List[Any]				=			self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model_fpaa_forward(*A_)

	def 				_SCREAMING_SNAKE_CASE  (	self:						Any)							->			int:

				"""simple docstring"""

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

				__lowerCAmelCase :    Optional[Any]				=			{
				    "head_mask": torch.zeros(config.num_layers			,     config.num_heads			,     device=A_),
				    "decoder_head_mask": torch.zeros(config.num_decoder_layers			,     config.num_heads			,     device=A_),
				    "cross_attn_head_mask": torch.zeros(config.num_decoder_layers			,     config.num_heads			,     device=A_),
				}

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

							__lowerCAmelCase :    Any				=			model.generate(
							    config_and_inputs[1]["input_ids"]			,     num_beams=1			,     max_length=3			,     output_attentions=A_			,     return_dict_in_generate=A_			,     **A_			,     )
							# We check the state of decoder_attentions and cross_attentions just from the last step
							__lowerCAmelCase :    Optional[Any]				=			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 				_SCREAMING_SNAKE_CASE  (	self:						Tuple)							->			Any:

				"""simple docstring"""

				pass






@require_torch
@require_sentencepiece
@require_tokenizers
class      A__  (    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 				_SCREAMING_SNAKE_CASE  (	self:						List[str])							->			Optional[Any]:

				"""simple docstring"""

				__lowerCAmelCase :    str				=			UMTaForConditionalGeneration.from_pretrained("google/umt5-small"			,     return_dict=A_).to(A_)
				__lowerCAmelCase :    List[Any]				=			AutoTokenizer.from_pretrained("google/umt5-small"			,     use_fast=A_			,     legacy=A_)
				__lowerCAmelCase :    List[str]				=			[
				    "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>.",
				]
				__lowerCAmelCase :    Dict				=			tokenizer(A_			,     return_tensors="pt"			,     padding=A_).input_ids
				# fmt: off
				__lowerCAmelCase :    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(A_			,     A_)

				__lowerCAmelCase :    Optional[int]				=			model.generate(input_ids.to(A_))
				__lowerCAmelCase :    Dict				=			[
				    "<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>",
				]
				__lowerCAmelCase :    Optional[Any]				=			tokenizer.batch_decode(A_)
				self.assertEqual(A_			,     A_)



import pytest

import datasets


# Import fixture modules as plugins
__snake_case   							=						['''tests.fixtures.files''', '''tests.fixtures.hub''', '''tests.fixtures.fsspec''']






def 		_A		(		_lowercase					,  _lowercase )					-> Tuple:


     """simple docstring"""


     for item in items:
          if any(marker in item.keywords for marker in ['integration', 'unit'] ):
               continue
          item.add_marker(pytest.mark.unit )






def 		_A		(		_lowercase )					-> str:


     """simple docstring"""


     config.addinivalue_line('markers'					,  'torchaudio_latest: mark test to run with torchaudio>=0.12' )






@pytest.fixture(autouse=_lowercase )
def 		_A		(		_lowercase					,  _lowercase )					-> Any:


     """simple docstring"""


     __UpperCamelCase =			tmp_path_factory.getbasetemp() / 'cache'
     __UpperCamelCase =			test_hf_cache_home / 'datasets'
     __UpperCamelCase =			test_hf_cache_home / 'metrics'
     __UpperCamelCase =			test_hf_cache_home / 'modules'
     monkeypatch.setattr('datasets.config.HF_DATASETS_CACHE'					,  str(_lowercase ) )
     monkeypatch.setattr('datasets.config.HF_METRICS_CACHE'					,  str(_lowercase ) )
     monkeypatch.setattr('datasets.config.HF_MODULES_CACHE'					,  str(_lowercase ) )
     __UpperCamelCase =			test_hf_datasets_cache / 'downloads'
     monkeypatch.setattr('datasets.config.DOWNLOADED_DATASETS_PATH'					,  str(_lowercase ) )
     __UpperCamelCase =			test_hf_datasets_cache / 'downloads' / 'extracted'
     monkeypatch.setattr('datasets.config.EXTRACTED_DATASETS_PATH'					,  str(_lowercase ) )






@pytest.fixture(autouse=_lowercase					,  scope='session' )
def 		_A		(		)					-> Dict:


     """simple docstring"""


     datasets.disable_progress_bar()






@pytest.fixture(autouse=_lowercase )
def 		_A		(		_lowercase )					-> Tuple:


     """simple docstring"""


     monkeypatch.setattr('datasets.config.HF_UPDATE_DOWNLOAD_COUNTS'					,  _lowercase )






@pytest.fixture
def 		_A		(		_lowercase )					-> Any:


     """simple docstring"""


     monkeypatch.setattr('sqlalchemy.util.deprecations.SILENCE_UBER_WARNING'					,  _lowercase )


