



"""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 				lowerCamelCase  :






					'''simple docstring'''





					def __init__(       self:							Any						,     snake_case:							Union[str, Any]						,     snake_case:							str=99						,     snake_case:							Union[str, Any]=13						,     snake_case:							Any=7						,     snake_case:							int=9						,     snake_case:							str=True						,     snake_case:							str=True						,     snake_case:							Tuple=False						,     snake_case:							Optional[int]=32						,     snake_case:							str=5						,     snake_case:							Union[str, Any]=4						,     snake_case:							int=37						,     snake_case:							Any=8						,     snake_case:							List[str]=0.1						,     snake_case:							str=0.0_0_2						,     snake_case:							int=1						,     snake_case:							Any=0						,     snake_case:							str=0						,     snake_case:							Optional[Any]=None						,     snake_case:							List[Any]=None						,     )   ->	Any:
											snake_case_ :Tuple						=			parent
											snake_case_ :int						=			batch_size
											snake_case_ :Optional[Any]						=			encoder_seq_length
											snake_case_ :Union[str, Any]						=			decoder_seq_length
											# For common tests
											snake_case_ :Optional[Any]						=			self.decoder_seq_length
											snake_case_ :Dict						=			is_training
											snake_case_ :List[str]						=			use_attention_mask
											snake_case_ :List[str]						=			use_labels
											snake_case_ :str						=			vocab_size
											snake_case_ :Optional[int]						=			hidden_size
											snake_case_ :Any						=			num_hidden_layers
											snake_case_ :str						=			num_attention_heads
											snake_case_ :Optional[int]						=			d_ff
											snake_case_ :int						=			relative_attention_num_buckets
											snake_case_ :List[Any]						=			dropout_rate
											snake_case_ :Union[str, Any]						=			initializer_factor
											snake_case_ :str						=			eos_token_id
											snake_case_ :Any						=			pad_token_id
											snake_case_ :List[str]						=			decoder_start_token_id
											snake_case_ :Any						=			None
											snake_case_ :str						=			decoder_layers






					def    lowerCAmelCase_				(       self:							Optional[Any]						)   ->	Tuple:
											return TaConfig.from_pretrained("""google/umt5-base"""						)






					def    lowerCAmelCase_				(       self:							Optional[Any]						,     snake_case:							int						,     snake_case:							Dict						,     snake_case:							List[str]						,     snake_case:							Optional[int]=None						,     snake_case:							str=None						,     snake_case:							Any=None						,     snake_case:							Any=None						,     snake_case:							int=None						,     )   ->	Tuple:
											if attention_mask is None:
																	snake_case_ :Dict						=			input_ids.ne(config.pad_token_id						)
											if decoder_attention_mask is None:
																	snake_case_ :str						=			decoder_input_ids.ne(config.pad_token_id						)
											if head_mask is None:
																	snake_case_ :List[str]						=			torch.ones(config.num_hidden_layers						,     config.num_attention_heads						,     device=snake_case						)
											if decoder_head_mask is None:
																	snake_case_ :Union[str, Any]						=			torch.ones(config.num_decoder_layers						,     config.num_attention_heads						,     device=snake_case						)
											if cross_attn_head_mask is None:
																	snake_case_ :Optional[Any]						=			torch.ones(
																	    config.num_decoder_layers						,     config.num_attention_heads						,     device=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    lowerCAmelCase_				(       self:							List[Any]						)   ->	Union[str, Any]:
											snake_case_ :Dict						=			ids_tensor([self.batch_size, self.encoder_seq_length]						,     self.vocab_size						)
											snake_case_ :Any						=			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
											snake_case_ :Optional[int]						=			input_ids.clamp(self.pad_token_id + 1						)
											snake_case_ :int						=			decoder_input_ids.clamp(self.pad_token_id + 1						)

											snake_case_ :Union[str, Any]						=			self.get_config()
											snake_case_ :List[Any]						=			config.num_attention_heads
											snake_case_ :int						=			self.prepare_inputs_dict(snake_case						,     snake_case						,     snake_case						)
											return config, input_dict






					def    lowerCAmelCase_				(       self:							Any						)   ->	int:
											snake_case_ :str						=			self.prepare_config_and_inputs()
											return config, inputs_dict






					def    lowerCAmelCase_				(       self:							Tuple						)   ->	Union[str, Any]:
											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[Any]						)   ->	Any:
											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:							Union[str, Any]						,     snake_case:							Any						,     snake_case:							Dict						,     snake_case:							Any						,     snake_case:							str						,     snake_case:							str						,     snake_case:							Any						,     )   ->	List[str]:
											snake_case_ :str						=			UMTaModel(config=snake_case						)
											model.to(snake_case						)
											model.eval()
											snake_case_ :Dict						=			model(
											    input_ids=snake_case						,     decoder_input_ids=snake_case						,     attention_mask=snake_case						,     decoder_attention_mask=snake_case						,     )
											snake_case_ :Any						=			model(input_ids=snake_case						,     decoder_input_ids=snake_case						)
											snake_case_ :Tuple						=			result.last_hidden_state
											snake_case_ :Union[str, Any]						=			result.past_key_values
											snake_case_ :List[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(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    lowerCAmelCase_				(       self:							List[str]						,     snake_case:							str						,     snake_case:							Optional[int]						,     snake_case:							Tuple						,     snake_case:							Union[str, Any]						,     snake_case:							Optional[int]						,     snake_case:							Dict						,     )   ->	List[Any]:
											snake_case_ :List[Any]						=			UMTaModel(config=snake_case						).get_decoder().to(snake_case						).eval()
											# first forward pass
											snake_case_ :List[str]						=			model(snake_case						,     use_cache=snake_case						)
											snake_case_ :Optional[Any]						=			model(snake_case						)
											snake_case_ :List[str]						=			model(snake_case						,     use_cache=snake_case						)

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

											snake_case_ :List[str]						=			outputs.to_tuple()

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

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

											snake_case_ :str						=			model(snake_case						)["""last_hidden_state"""]
											snake_case_ :List[Any]						=			model(snake_case						,     past_key_values=snake_case						)["""last_hidden_state"""]

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

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






					def    lowerCAmelCase_				(       self:							Union[str, Any]						,     snake_case:							List[Any]						,     snake_case:							Optional[Any]						,     )   ->	Tuple:
											snake_case_ :Union[str, Any]						=			UMTaModel(config=snake_case						).to(snake_case						).half().eval()
											snake_case_ :Optional[Any]						=			model(**snake_case						)["""last_hidden_state"""]
											self.parent.assertFalse(torch.isnan(snake_case						).any().item()						)

@require_torch
class 				lowerCamelCase  (			_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						unittest.TestCase				):






					'''simple docstring'''





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






					def    lowerCAmelCase_				(       self:							Optional[int]						)   ->	Optional[Any]:
											snake_case_ :Dict						=			UMTaModelTester(self						)






					@unittest.skip("""Test has a segmentation fault on torch 1.8.0"""						)
					def    lowerCAmelCase_				(       self:							int						)   ->	str:
											snake_case_ :str						=			self.model_tester.prepare_config_and_inputs()
											snake_case_ :Any						=			UMTaModel(config_and_inputs[0]						).to(snake_case						)
											with tempfile.TemporaryDirectory() as tmpdirname:
																	torch.onnx.export(
																	    snake_case						,     (config_and_inputs[1], config_and_inputs[3], config_and_inputs[2])						,     f"""{tmpdirname}/t5_test.onnx"""						,     export_params=snake_case						,     opset_version=9						,     input_names=["""input_ids""", """decoder_input_ids"""]						,     )






					@unittest.skipIf(torch_device == """cpu"""						,     """Cant do half precision"""						)
					def    lowerCAmelCase_				(       self:							Optional[Any]						)   ->	Dict:
											snake_case_ :str						=			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model_fpaa_forward(*snake_case						)






					def    lowerCAmelCase_				(       self:							Optional[Any]						)   ->	Union[str, Any]:
											snake_case_ :List[str]						=			["""encoder_attentions""", """decoder_attentions""", """cross_attentions"""]
											snake_case_ :Optional[int]						=			self.model_tester.prepare_config_and_inputs()
											snake_case_ :Union[str, Any]						=			config_and_inputs[0]
											snake_case_ :Union[str, Any]						=			UMTaForConditionalGeneration(snake_case						).eval()
											model.to(snake_case						)

											snake_case_ :Tuple						=			{
											    """head_mask""": torch.zeros(config.num_layers						,     config.num_heads						,     device=snake_case						),
											    """decoder_head_mask""": torch.zeros(config.num_decoder_layers						,     config.num_heads						,     device=snake_case						),
											    """cross_attn_head_mask""": torch.zeros(config.num_decoder_layers						,     config.num_heads						,     device=snake_case						),
											}

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

																	snake_case_ :List[str]						=			model.generate(
																	    config_and_inputs[1]["""input_ids"""]						,     num_beams=1						,     max_length=3						,     output_attentions=snake_case						,     return_dict_in_generate=snake_case						,     **snake_case						,     )
																	# We check the state of decoder_attentions and cross_attentions just from the last step
																	snake_case_ :List[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    lowerCAmelCase_				(       self:							Union[str, Any]						)   ->	int:
											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    lowerCAmelCase_				(       self:							List[Any]						)   ->	Optional[Any]:
											snake_case_ :List[str]						=			UMTaForConditionalGeneration.from_pretrained("""google/umt5-small"""						,     return_dict=snake_case						).to(snake_case						)
											snake_case_ :Union[str, Any]						=			AutoTokenizer.from_pretrained("""google/umt5-small"""						,     use_fast=snake_case						,     legacy=snake_case						)
											snake_case_ :Optional[Any]						=			[
											    """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>.""",
											]
											snake_case_ :Optional[int]						=			tokenizer(snake_case						,     return_tensors="""pt"""						,     padding=snake_case						).input_ids
											# fmt: off
											snake_case_ :List[str]						=			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(snake_case						,     snake_case						)

											snake_case_ :List[str]						=			model.generate(input_ids.to(snake_case						)						)
											snake_case_ :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>""",
											]
											snake_case_ :Optional[Any]						=			tokenizer.batch_decode(snake_case						)
											self.assertEqual(snake_case						,     snake_case						)




'''simple docstring'''






def 				UpperCamelCase_			(    A__					:   int					):
  '''simple docstring'''





  lowerCAmelCase_			:      Union[str, Any]			=   (1 + 24 * n) ** 0.5
  return ((1 + root) / 6) % 1 == 0







def 				UpperCamelCase_			(    A__					:   int = 50_00					):
  '''simple docstring'''





  lowerCAmelCase_			:      Optional[int]			=   [(i * (3 * i - 1)) // 2 for i in range(1							,    A__					)]
  for i, pentagonal_i in enumerate(A__					):
    for j in range(A__							,    len(A__					)					):
      lowerCAmelCase_			:      int			=   pentagonal_nums[j]
      lowerCAmelCase_			:      Union[str, Any]			=   pentagonal_i + pentagonal_j
      lowerCAmelCase_			:      List[Any]			=   pentagonal_j - pentagonal_i
      if is_pentagonal(A__					) and is_pentagonal(A__					):
        return b

  return -1


if __name__ == "__main__":
  print(F'''{solution() = }''')



