



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					:



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



						"""simple docstring"""






						_UpperCAmelCase	 =     parent
						_UpperCAmelCase	 =     batch_size
						_UpperCAmelCase	 =     encoder_seq_length
						_UpperCAmelCase	 =     decoder_seq_length
						# For common tests
						_UpperCAmelCase	 =     self.decoder_seq_length
						_UpperCAmelCase	 =     is_training
						_UpperCAmelCase	 =     use_attention_mask
						_UpperCAmelCase	 =     use_labels
						_UpperCAmelCase	 =     vocab_size
						_UpperCAmelCase	 =     hidden_size
						_UpperCAmelCase	 =     num_hidden_layers
						_UpperCAmelCase	 =     num_attention_heads
						_UpperCAmelCase	 =     d_ff
						_UpperCAmelCase	 =     relative_attention_num_buckets
						_UpperCAmelCase	 =     dropout_rate
						_UpperCAmelCase	 =     initializer_factor
						_UpperCAmelCase	 =     eos_token_id
						_UpperCAmelCase	 =     pad_token_id
						_UpperCAmelCase	 =     decoder_start_token_id
						_UpperCAmelCase	 =     None
						_UpperCAmelCase	 =     decoder_layers



			def 					_lowerCamelCase (		self     :						Optional[int])     ->					Any:



						"""simple docstring"""






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



			def 					_lowerCamelCase (		self     :						Tuple				,	A     :						int				,	A     :						Union[str, Any]				,	A     :						List[str]				,	A     :						Any=None				,	A     :						List[Any]=None				,	A     :						List[Any]=None				,	A     :						int=None				,	A     :						Optional[Any]=None				,	)     ->					int:



						"""simple docstring"""






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



						"""simple docstring"""






						_UpperCAmelCase	 =     ids_tensor([self.batch_size, self.encoder_seq_length]				,	self.vocab_size)
						_UpperCAmelCase	 =     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
						_UpperCAmelCase	 =     input_ids.clamp(self.pad_token_id + 1)
						_UpperCAmelCase	 =     decoder_input_ids.clamp(self.pad_token_id + 1)

						_UpperCAmelCase	 =     self.get_config()
						_UpperCAmelCase	 =     config.num_attention_heads
						_UpperCAmelCase	 =     self.prepare_inputs_dict(A				,	A				,	A)
						return config, input_dict



			def 					_lowerCamelCase (		self     :						Any)     ->					Dict:



						"""simple docstring"""






						_UpperCAmelCase ,  _UpperCAmelCase	 =     self.prepare_config_and_inputs()
						return config, inputs_dict



			def 					_lowerCamelCase (		self     :						Dict)     ->					Tuple:



						"""simple docstring"""






						return TaConfig(
						    vocab_size=1_66				,	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 					_lowerCamelCase (		self     :						Optional[Any])     ->					Optional[int]:



						"""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 					_lowerCamelCase (		self     :						str				,	A     :						List[Any]				,	A     :						Any				,	A     :						List[str]				,	A     :						str				,	A     :						Tuple				,	A     :						str				,	)     ->					Dict:



						"""simple docstring"""






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



						"""simple docstring"""






						_UpperCAmelCase	 =     UMTaModel(config=A).get_decoder().to(A).eval()
						# first forward pass
						_UpperCAmelCase	 =     model(A				,	use_cache=A)
						_UpperCAmelCase	 =     model(A)
						_UpperCAmelCase	 =     model(A				,	use_cache=A)

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

						_UpperCAmelCase ,  _UpperCAmelCase	 =     outputs.to_tuple()

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

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

						_UpperCAmelCase	 =     model(A)['last_hidden_state']
						_UpperCAmelCase	 =     model(A				,	past_key_values=A)['last_hidden_state']

						# select random slice
						_UpperCAmelCase	 =     ids_tensor((1,)				,	output_from_past.shape[-1]).item()
						_UpperCAmelCase	 =     output_from_no_past[:, -1, random_slice_idx].detach()
						_UpperCAmelCase	 =     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 					_lowerCamelCase (		self     :						Optional[Any]				,	A     :						str				,	A     :						Optional[Any]				,	)     ->					Tuple:



						"""simple docstring"""






						_UpperCAmelCase	 =     UMTaModel(config=A).to(A).half().eval()
						_UpperCAmelCase	 =     model(**A)['last_hidden_state']
						self.parent.assertFalse(torch.isnan(A).any().item())







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



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



						"""simple docstring"""






						_UpperCAmelCase	 =     UMTaModelTester(self)



			@unittest.skip('Test has a segmentation fault on torch 1.8.0')
			def 					_lowerCamelCase (		self     :						Union[str, Any])     ->					int:



						"""simple docstring"""






						_UpperCAmelCase	 =     self.model_tester.prepare_config_and_inputs()
						_UpperCAmelCase	 =     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 					_lowerCamelCase (		self     :						Tuple)     ->					Optional[Any]:



						"""simple docstring"""






						_UpperCAmelCase	 =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model_fpaa_forward(*A)



			def 					_lowerCamelCase (		self     :						Tuple)     ->					Tuple:



						"""simple docstring"""






						_UpperCAmelCase	 =     ['encoder_attentions', 'decoder_attentions', 'cross_attentions']
						_UpperCAmelCase	 =     self.model_tester.prepare_config_and_inputs()
						_UpperCAmelCase	 =     config_and_inputs[0]
						_UpperCAmelCase	 =     UMTaForConditionalGeneration(A).eval()
						model.to(A)

						_UpperCAmelCase	 =     {
						    '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()):
									_UpperCAmelCase	 =     {name: mask}
									# Explicitly pass decoder_head_mask as it is required from T5 model when head_mask specified
									if name == "head_mask":
												_UpperCAmelCase	 =     torch.ones(
												    config.num_decoder_layers				,	config.num_heads				,	device=A)

									_UpperCAmelCase	 =     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
									_UpperCAmelCase	 =     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 					_lowerCamelCase (		self     :						Union[str, Any])     ->					Union[str, Any]:



						"""simple docstring"""






						pass







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



			@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 					_lowerCamelCase (		self     :						str)     ->					Optional[int]:



						"""simple docstring"""






						_UpperCAmelCase	 =     UMTaForConditionalGeneration.from_pretrained('google/umt5-small'				,	return_dict=A).to(A)
						_UpperCAmelCase	 =     AutoTokenizer.from_pretrained('google/umt5-small'				,	use_fast=A				,	legacy=A)
						_UpperCAmelCase	 =     [
						    '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>.',
						]
						_UpperCAmelCase	 =     tokenizer(A				,	return_tensors='pt'				,	padding=A).input_ids
						# fmt: off
						_UpperCAmelCase	 =     torch.tensor(
						    [
						        [ 3_85_30, 21_07_03, 25_62_99, 14_10, 25_62_98, 2_74, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
						        [   8_26, 3_21, 6_71, 2_59_22, 25_62_99, 2_74, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
						        [  14_60, 3_39, 3_12, 1_90_14, 1_06_20, 7_58, 25_62_99, 23_55,2_74, 1, 0, 0, 0, 0, 0, 0,0, 0],
						        [   5_17, 25_62_99, 1_48_69, 2_81, 3_01, 25_62_98, 2_75, 11_99_83,1, 0, 0, 0, 0, 0, 0, 0,0, 0],
						        [   3_20, 25_62_99, 1_48_69, 2_81, 22_34, 2_89, 22_75, 3_33,6_13_91, 2_89, 25_62_98, 5_43, 25_62_97, 16_87_14, 3_29, 25_62_96,2_74, 1],
						    ])
						# fmt: on
						torch.testing.assert_allclose(A				,	A)

						_UpperCAmelCase	 =     model.generate(input_ids.to(A))
						_UpperCAmelCase	 =     [
						    '<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>',
						]
						_UpperCAmelCase	 =     tokenizer.batch_decode(A)
						self.assertEqual(A				,	A)





import os
from shutil import copyfile
from typing import List, Optional, Tuple

from ...tokenization_utils import AddedToken
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
							from .tokenization_xlnet import XLNetTokenizer
else:
							UpperCAmelCase__     = None


UpperCAmelCase__     = logging.get_logger(__name__)

UpperCAmelCase__     = {"vocab_file": "spiece.model", "tokenizer_file": "tokenizer.json"}

UpperCAmelCase__     = {
    "vocab_file": {
        "xlnet-base-cased": "https://huggingface.co/xlnet-base-cased/resolve/main/spiece.model",
        "xlnet-large-cased": "https://huggingface.co/xlnet-large-cased/resolve/main/spiece.model",
    },
    "tokenizer_file": {
        "xlnet-base-cased": "https://huggingface.co/xlnet-base-cased/resolve/main/tokenizer.json",
        "xlnet-large-cased": "https://huggingface.co/xlnet-large-cased/resolve/main/tokenizer.json",
    },
}

UpperCAmelCase__     = {
    "xlnet-base-cased": None,
    "xlnet-large-cased": None,
}

UpperCAmelCase__     = "▁"

# Segments (not really needed)
UpperCAmelCase__     = 0
UpperCAmelCase__     = 1
UpperCAmelCase__     = 2
UpperCAmelCase__     = 3
UpperCAmelCase__     = 4







class     __lowerCAmelCase					(    A		):
			UpperCamelCase           =  VOCAB_FILES_NAMES
			UpperCamelCase           =  PRETRAINED_VOCAB_FILES_MAP
			UpperCamelCase           =  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			UpperCamelCase           =  '''left'''
			UpperCamelCase           =  XLNetTokenizer



			def __init__(		self     :						Any				,	A     :						Union[str, Any]=None				,	A     :						str=None				,	A     :						Tuple=False				,	A     :						Tuple=True				,	A     :						Any=False				,	A     :						List[str]="<s>"				,	A     :						List[str]="</s>"				,	A     :						Optional[int]="<unk>"				,	A     :						Tuple="<sep>"				,	A     :						str="<pad>"				,	A     :						Dict="<cls>"				,	A     :						Dict="<mask>"				,	A     :						Optional[Any]=["<eop>", "<eod>"]				,	**A     :						Optional[Any]				,	)     ->					str:



						"""simple docstring"""






						_UpperCAmelCase	 =     AddedToken(A				,	lstrip=A				,	rstrip=A) if isinstance(A				,	A) else mask_token

						super().__init__(
						    vocab_file=A				,	tokenizer_file=A				,	do_lower_case=A				,	remove_space=A				,	keep_accents=A				,	bos_token=A				,	eos_token=A				,	unk_token=A				,	sep_token=A				,	pad_token=A				,	cls_token=A				,	mask_token=A				,	additional_special_tokens=A				,	**A				,	)

						_UpperCAmelCase	 =     3
						_UpperCAmelCase	 =     do_lower_case
						_UpperCAmelCase	 =     remove_space
						_UpperCAmelCase	 =     keep_accents
						_UpperCAmelCase	 =     vocab_file
						_UpperCAmelCase	 =     False if not self.vocab_file else True



			def 					_lowerCamelCase (		self     :						Tuple				,	A     :						List[int]				,	A     :						Optional[List[int]] = None)     ->					List[int]:



						"""simple docstring"""






						_UpperCAmelCase	 =     [self.sep_token_id]
						_UpperCAmelCase	 =     [self.cls_token_id]
						if token_ids_a is None:
									return token_ids_a + sep + cls
						return token_ids_a + sep + token_ids_a + sep + cls



			def 					_lowerCamelCase (		self     :						Tuple				,	A     :						List[int]				,	A     :						Optional[List[int]] = None)     ->					List[int]:



						"""simple docstring"""






						_UpperCAmelCase	 =     [self.sep_token_id]
						_UpperCAmelCase	 =     [2]

						if token_ids_a is None:
									return len(token_ids_a + sep) * [0] + cls_segment_id
						return len(token_ids_a + sep) * [0] + len(token_ids_a + sep) * [1] + cls_segment_id



			def 					_lowerCamelCase (		self     :						List[str]				,	A     :						str				,	A     :						Optional[str] = None)     ->					Tuple[str]:



						"""simple docstring"""






						if not self.can_save_slow_tokenizer:
									raise ValueError(
									    'Your fast tokenizer does not have the necessary information to save the vocabulary for a slow '
									    'tokenizer.')

						if not os.path.isdir(A):
									logger.error(F"Vocabulary path ({save_directory}) should be a directory")
									return
						_UpperCAmelCase	 =     os.path.join(
						    A				,	(filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file'])

						if os.path.abspath(self.vocab_file) != os.path.abspath(A):
									copyfile(self.vocab_file				,	A)

						return (out_vocab_file,)
