






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







     '''simple docstring'''




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


         """simple docstring"""



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




     def 		lowerCAmelCase_     (				self       :			Dict				):


         """simple docstring"""



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




     def 		lowerCAmelCase_     (				self       :			str ,	UpperCamelCase_       :			Dict ,	UpperCamelCase_       :			Any ,	UpperCamelCase_       :			List[str] ,	UpperCamelCase_       :			Optional[int]=None ,	UpperCamelCase_       :			Union[str, Any]=None ,	UpperCamelCase_       :			Tuple=None ,	UpperCamelCase_       :			int=None ,	UpperCamelCase_       :			Optional[int]=None ,	):


         """simple docstring"""



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


         """simple docstring"""



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

         __A         =	self.get_config()
         __A         =	config.num_attention_heads
         __A         =	self.prepare_inputs_dict(UpperCamelCase_ ,	UpperCamelCase_ ,	UpperCamelCase_				)
         return config, input_dict




     def 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



         __A         =	self.prepare_config_and_inputs()
         return config, inputs_dict




     def 		lowerCAmelCase_     (				self       :			List[str]				):


         """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 		lowerCAmelCase_     (				self       :			List[str]				):


         """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       :			Any ,	UpperCamelCase_       :			Optional[int] ,	UpperCamelCase_       :			str ,	UpperCamelCase_       :			str ,	UpperCamelCase_       :			Optional[Any] ,	UpperCamelCase_       :			Optional[Any] ,	UpperCamelCase_       :			Tuple ,	):


         """simple docstring"""



         __A         =	UMTaModel(config=UpperCamelCase_				)
         model.to(UpperCamelCase_				)
         model.eval()
         __A         =	model(
             input_ids=UpperCamelCase_ ,	decoder_input_ids=UpperCamelCase_ ,	attention_mask=UpperCamelCase_ ,	decoder_attention_mask=UpperCamelCase_ ,	)
         __A         =	model(input_ids=UpperCamelCase_ ,	decoder_input_ids=UpperCamelCase_				)
         __A         =	result.last_hidden_state
         __A         =	result.past_key_values
         __A         =	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(UpperCamelCase_				) ,	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       :			str ,	UpperCamelCase_       :			Any ,	UpperCamelCase_       :			Any ,	UpperCamelCase_       :			List[str] ,	UpperCamelCase_       :			List[str] ,	UpperCamelCase_       :			str ,	UpperCamelCase_       :			Any ,	):


         """simple docstring"""



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

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

         __A         =	outputs.to_tuple()

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

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

         __A         =	model(UpperCamelCase_				)['last_hidden_state']
         __A         =	model(UpperCamelCase_ ,	past_key_values=UpperCamelCase_				)['last_hidden_state']

         # select random slice
         __A         =	ids_tensor((1,) ,	output_from_past.shape[-1]				).item()
         __A         =	output_from_no_past[:, -1, random_slice_idx].detach()
         __A         =	output_from_past[:, 0, random_slice_idx].detach()

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





     def 		lowerCAmelCase_     (				self       :			Union[str, Any] ,	UpperCamelCase_       :			Any ,	UpperCamelCase_       :			Optional[int] ,	):


         """simple docstring"""



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


@require_torch
class       __lowercase      (	lowercase_			,      lowercase_			,      lowercase_			,      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 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



         __A         =	UMTaModelTester(self				)




     @unittest.skip("""Test has a segmentation fault on torch 1.8.0"""				)
     def 		lowerCAmelCase_     (				self       :			Union[str, Any]				):


         """simple docstring"""



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




     @unittest.skipIf(torch_device == """cpu""" ,	"""Cant do half precision"""				)
     def 		lowerCAmelCase_     (				self       :			str				):


         """simple docstring"""



         __A         =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model_fpaa_forward(*UpperCamelCase_				)




     def 		lowerCAmelCase_     (				self       :			Tuple				):


         """simple docstring"""



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

         __A         =	{
             'head_mask': torch.zeros(config.num_layers ,	config.num_heads ,	device=UpperCamelCase_				),
             'decoder_head_mask': torch.zeros(config.num_decoder_layers ,	config.num_heads ,	device=UpperCamelCase_				),
             'cross_attn_head_mask': torch.zeros(config.num_decoder_layers ,	config.num_heads ,	device=UpperCamelCase_				),
         }

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

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


         """simple docstring"""



         pass


@require_torch
@require_sentencepiece
@require_tokenizers
class       __lowercase      (	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 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



         __A         =	UMTaForConditionalGeneration.from_pretrained("""google/umt5-small""" ,	return_dict=UpperCamelCase_				).to(UpperCamelCase_				)
         __A         =	AutoTokenizer.from_pretrained("""google/umt5-small""" ,	use_fast=UpperCamelCase_ ,	legacy=UpperCamelCase_				)
         __A         =	[
             '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         =	tokenizer(UpperCamelCase_ ,	return_tensors="""pt""" ,	padding=UpperCamelCase_				).input_ids
         # fmt: off
         __A         =	torch.tensor(
             [
                 [ 38_530, 210_703, 256_299, 1_410, 256_298, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
                 [   826, 321, 671, 25_922, 256_299, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
                 [  1_460, 339, 312, 19_014, 10_620, 758, 256_299, 2_355,274, 1, 0, 0, 0, 0, 0, 0,0, 0],
                 [   517, 256_299, 14_869, 281, 301, 256_298, 275, 119_983,1, 0, 0, 0, 0, 0, 0, 0,0, 0],
                 [   320, 256_299, 14_869, 281, 2_234, 289, 2_275, 333,61_391, 289, 256_298, 543, 256_297, 168_714, 329, 256_296,274, 1],
             ]				)
         # fmt: on
         torch.testing.assert_allclose(UpperCamelCase_ ,	UpperCamelCase_				)

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












def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			str = "The quick brown fox jumps over the lazy dog"			,				)						->      bool:





    """simple docstring"""

    __A         =	set()

    # Replace all the whitespace in our sentence
    __A         =	input_str.replace(""" """			,				""""""    )
    for alpha in input_str:
        if "a" <= alpha.lower() <= "z":
            frequency.add(alpha.lower()    )
    return len(__lowercase    ) == 2_6



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			str = "The quick brown fox jumps over the lazy dog"			,				)						->      bool:





    """simple docstring"""

    __A         =	[False] * 2_6
    for char in input_str:
        if char.islower():
            __A         =	True
        elif char.isupper():
            __A         =	True
    return all(__lowercase    )



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			str = "The quick brown fox jumps over the lazy dog"			,				)						->      bool:





    """simple docstring"""

    return len({char for char in input_str.lower() if char.isalpha()}    ) == 2_6



def 			_SCREAMING_SNAKE_CASE						( )						->      None:





    """simple docstring"""

    from timeit import timeit

    __A         =	"""from __main__ import is_pangram, is_pangram_faster, is_pangram_fastest"""
    print(timeit("""is_pangram()"""			,				setup=__lowercase    )    )
    print(timeit("""is_pangram_faster()"""			,				setup=__lowercase    )    )
    print(timeit("""is_pangram_fastest()"""			,				setup=__lowercase    )    )
    # 5.348480500048026, 2.6477354579837993, 1.8470395830227062
    # 5.036091582966037, 2.644472333951853,  1.8869528750656173


if __name__ == "__main__":
  import doctest

  doctest.testmod()
  benchmark()




