



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	:       str			,     lowercase	:       List[str]			,     lowercase	:       Dict=99			,     lowercase	:       List[Any]=13			,     lowercase	:       str=7			,     lowercase	:       Optional[Any]=9			,     lowercase	:       Optional[int]=True			,     lowercase	:       str=True			,     lowercase	:       Optional[Any]=False			,     lowercase	:       Any=32			,     lowercase	:       Optional[int]=5			,     lowercase	:       List[str]=4			,     lowercase	:       Optional[int]=37			,     lowercase	:       Tuple=8			,     lowercase	:       Optional[Any]=0.1			,     lowercase	:       Tuple=0.002			,     lowercase	:       Tuple=1			,     lowercase	:       Any=0			,     lowercase	:       Any=0			,     lowercase	:       Dict=None			,     lowercase	:       Any=None			,     )						->					Any:
         """simple docstring"""







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





   def 		snake_case__			(				self	:       str   )						->					Tuple:
         """simple docstring"""







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





   def 		snake_case__			(				self	:       Any			,     lowercase	:       Any			,     lowercase	:       int			,     lowercase	:       str			,     lowercase	:       int=None			,     lowercase	:       List[Any]=None			,     lowercase	:       Optional[int]=None			,     lowercase	:       Tuple=None			,     lowercase	:       Union[str, Any]=None			,     )						->					Any:
         """simple docstring"""







         if attention_mask is None:
               __lowercase					=				input_ids.ne(config.pad_token_id   )
         if decoder_attention_mask is None:
               __lowercase					=				decoder_input_ids.ne(config.pad_token_id   )
         if head_mask is None:
               __lowercase					=				torch.ones(config.num_hidden_layers			,     config.num_attention_heads			,     device=_a   )
         if decoder_head_mask is None:
               __lowercase					=				torch.ones(config.num_decoder_layers			,     config.num_attention_heads			,     device=_a   )
         if cross_attn_head_mask is None:
               __lowercase					=				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 		snake_case__			(				self	:       Union[str, Any]   )						->					Tuple:
         """simple docstring"""







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

         __lowercase					=				self.get_config()
         __lowercase					=				config.num_attention_heads
         __lowercase					=				self.prepare_inputs_dict(_a			,     _a			,     _a   )
         return config, input_dict





   def 		snake_case__			(				self	:       Tuple   )						->					Any:
         """simple docstring"""







         __lowercase					=				self.prepare_config_and_inputs()
         return config, inputs_dict





   def 		snake_case__			(				self	:       int   )						->					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 		snake_case__			(				self	:       Tuple   )						->					List[Any]:
         """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	:       Dict			,     lowercase	:       List[Any]			,     lowercase	:       List[str]			,     lowercase	:       Tuple			,     lowercase	:       Any			,     lowercase	:       Tuple			,     lowercase	:       Union[str, Any]			,     )						->					Optional[Any]:
         """simple docstring"""







         __lowercase					=				UMTaModel(config=_a   )
         model.to(_a   )
         model.eval()
         __lowercase					=				model(
             input_ids=_a			,     decoder_input_ids=_a			,     attention_mask=_a			,     decoder_attention_mask=_a			,     )
         __lowercase					=				model(input_ids=_a			,     decoder_input_ids=_a   )
         __lowercase					=				result.last_hidden_state
         __lowercase					=				result.past_key_values
         __lowercase					=				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 		snake_case__			(				self	:       Dict			,     lowercase	:       Tuple			,     lowercase	:       List[Any]			,     lowercase	:       int			,     lowercase	:       Any			,     lowercase	:       List[Any]			,     lowercase	:       Any			,     )						->					Union[str, Any]:
         """simple docstring"""







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

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

         __lowercase					=				outputs.to_tuple()

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

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

         __lowercase					=				model(_a   )["""last_hidden_state"""]
         __lowercase					=				model(_a			,     past_key_values=_a   )["""last_hidden_state"""]

         # select random slice
         __lowercase					=				ids_tensor((1,)			,     output_from_past.shape[-1]   ).item()
         __lowercase					=				output_from_no_past[:, -1, random_slice_idx].detach()
         __lowercase					=				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 		snake_case__			(				self	:       Optional[Any]			,     lowercase	:       int			,     lowercase	:       Union[str, Any]			,     )						->					List[Any]:
         """simple docstring"""







         __lowercase					=				UMTaModel(config=_a   ).to(_a   ).half().eval()
         __lowercase					=				model(**_a   )["""last_hidden_state"""]
         self.parent.assertFalse(torch.isnan(_a   ).any().item()   )





@require_torch
class 			_lowerCAmelCase   ( UpperCamelCase_    ,					UpperCamelCase_    ,					UpperCamelCase_    ,					unittest.TestCase			):






   """simple docstring"""





   lowercase__      :  Union[str, Any]      =			(
       (UMTaModel, UMTaForConditionalGeneration, UMTaForQuestionAnswering) if is_torch_available() else ()
   )
   lowercase__      :  Optional[int]      =			(UMTaForConditionalGeneration,) if is_torch_available() else ()
   lowercase__      :  List[str]      =			(
       {
           """conversational""": UMTaForConditionalGeneration,
           """feature-extraction""": UMTaModel,
           """summarization""": UMTaForConditionalGeneration,
           """text2text-generation""": UMTaForConditionalGeneration,
           """translation""": UMTaForConditionalGeneration,
           """question-answering""": UMTaForQuestionAnswering,
       }
       if is_torch_available()
       else {}
   )
   lowercase__      :  Union[str, Any]      =			True
   lowercase__      :  Tuple      =			False
   lowercase__      :  Optional[Any]      =			False
   lowercase__      :  List[Any]      =			True
   lowercase__      :  List[str]      =			True
   # The small UMT5 model needs higher percentages for CPU/MP tests
   lowercase__      :  Union[str, Any]      =			[0.8, 0.9]





   def 		snake_case__			(				self	:       int   )						->					Optional[int]:
         """simple docstring"""







         __lowercase					=				UMTaModelTester(self   )





   @unittest.skip("""Test has a segmentation fault on torch 1.8.0"""   )
   def 		snake_case__			(				self	:       str   )						->					Any:
         """simple docstring"""







         __lowercase					=				self.model_tester.prepare_config_and_inputs()
         __lowercase					=				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 		snake_case__			(				self	:       Optional[Any]   )						->					str:
         """simple docstring"""







         __lowercase					=				self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model_fpaa_forward(*_a   )





   def 		snake_case__			(				self	:       Any   )						->					List[str]:
         """simple docstring"""







         __lowercase					=				["""encoder_attentions""", """decoder_attentions""", """cross_attentions"""]
         __lowercase					=				self.model_tester.prepare_config_and_inputs()
         __lowercase					=				config_and_inputs[0]
         __lowercase					=				UMTaForConditionalGeneration(_a   ).eval()
         model.to(_a   )

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

               __lowercase					=				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
               __lowercase					=				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	:       List[Any]   )						->					Dict:
         """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	:       str   )						->					Union[str, Any]:
         """simple docstring"""







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

         __lowercase					=				model.generate(input_ids.to(_a   )   )
         __lowercase					=				[
             """<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>""",
         ]
         __lowercase					=				tokenizer.batch_decode(_a   )
         self.assertEqual(_a			,     _a   )










def  UpperCAmelCase__    (    lowercase__						,       lowercase__						,       lowercase__   ) ->      float:
						__lowercase					=				(num_of_terms / 2) * (2 * first_term + (num_of_terms - 1) * common_diff)
						# formula for sum of series
						return total






def  UpperCAmelCase__    (    ) ->      List[str]:
						print(sum_of_series(1						,       1						,       10   )   )


if __name__ == "__main__":
			import doctest

			doctest.testmod()





