





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      :					Optional[int]      ,				UpperCamelCase      :					Any      ,				UpperCamelCase      :					str=99      ,				UpperCamelCase      :					str=13      ,				UpperCamelCase      :					Optional[int]=7      ,				UpperCamelCase      :					Union[str, Any]=9      ,				UpperCamelCase      :					Tuple=True      ,				UpperCamelCase      :					Tuple=True      ,				UpperCamelCase      :					Optional[int]=False      ,				UpperCamelCase      :					Dict=32      ,				UpperCamelCase      :					Any=5      ,				UpperCamelCase      :					Union[str, Any]=4      ,				UpperCamelCase      :					str=37      ,				UpperCamelCase      :					Dict=8      ,				UpperCamelCase      :					Tuple=0.1      ,				UpperCamelCase      :					Any=0.0_02      ,				UpperCamelCase      :					str=1      ,				UpperCamelCase      :					Optional[int]=0      ,				UpperCamelCase      :					Optional[int]=0      ,				UpperCamelCase      :					Dict=None      ,				UpperCamelCase      :					Tuple=None      ,				):


    '''simple docstring'''





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







  def UpperCamelCase_ (      self      :					str							):


    '''simple docstring'''





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







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


    '''simple docstring'''





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


    '''simple docstring'''





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

    _snake_case   :				Tuple														=	self.get_config()
    _snake_case   :				int														=	config.num_attention_heads
    _snake_case   :				Optional[int]														=	self.prepare_inputs_dict(UpperCamelCase      ,				UpperCamelCase      ,				UpperCamelCase							)
    return config, input_dict







  def UpperCamelCase_ (      self      :					Any							):


    '''simple docstring'''





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







  def UpperCamelCase_ (      self      :					int							):


    '''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 UpperCamelCase_ (      self      :					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 UpperCamelCase_ (      self      :					Optional[int]      ,				UpperCamelCase      :					Any      ,				UpperCamelCase      :					Tuple      ,				UpperCamelCase      :					List[Any]      ,				UpperCamelCase      :					Optional[int]      ,				UpperCamelCase      :					int      ,				UpperCamelCase      :					Optional[int]      ,				):


    '''simple docstring'''





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


    '''simple docstring'''





    _snake_case   :				str														=	UMTaModel(config=UpperCamelCase							).get_decoder().to(UpperCamelCase							).eval()
    # first forward pass
    _snake_case   :				Optional[Any]														=	model(UpperCamelCase      ,				use_cache=UpperCamelCase							)
    _snake_case   :				Dict														=	model(UpperCamelCase							)
    _snake_case   :				Union[str, Any]														=	model(UpperCamelCase      ,				use_cache=UpperCamelCase							)

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

    _snake_case   :				List[Any]														=	outputs.to_tuple()

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

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

    _snake_case   :				Union[str, Any]														=	model(UpperCamelCase							)['last_hidden_state']
    _snake_case   :				Optional[Any]														=	model(UpperCamelCase      ,				past_key_values=UpperCamelCase							)['last_hidden_state']

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


    '''simple docstring'''





    _snake_case   :				Optional[int]														=	UMTaModel(config=UpperCamelCase							).to(UpperCamelCase							).half().eval()
    _snake_case   :				int														=	model(**UpperCamelCase							)['last_hidden_state']
    self.parent.assertFalse(torch.isnan(UpperCamelCase							).any().item()							)




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


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







  def UpperCamelCase_ (      self      :					Any							):


    '''simple docstring'''





    _snake_case   :				str														=	UMTaModelTester(self							)







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


    '''simple docstring'''





    _snake_case   :				Optional[int]														=	self.model_tester.prepare_config_and_inputs()
    _snake_case   :				Optional[int]														=	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 UpperCamelCase_ (      self      :					Optional[Any]							):


    '''simple docstring'''





    _snake_case   :				Optional[Any]														=	self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model_fpaa_forward(*UpperCamelCase							)







  def UpperCamelCase_ (      self      :					int							):


    '''simple docstring'''





    _snake_case   :				Any														=	['encoder_attentions', 'decoder_attentions', 'cross_attentions']
    _snake_case   :				Any														=	self.model_tester.prepare_config_and_inputs()
    _snake_case   :				List[str]														=	config_and_inputs[0]
    _snake_case   :				Tuple														=	UMTaForConditionalGeneration(UpperCamelCase							).eval()
    model.to(UpperCamelCase							)

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

      _snake_case   :				str														=	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
      _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 UpperCamelCase_ (      self      :					str							):


    '''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 UpperCamelCase_ (      self      :					List[Any]							):


    '''simple docstring'''





    _snake_case   :				List[Any]														=	UMTaForConditionalGeneration.from_pretrained('google/umt5-small'      ,				return_dict=UpperCamelCase							).to(UpperCamelCase							)
    _snake_case   :				int														=	AutoTokenizer.from_pretrained('google/umt5-small'      ,				use_fast=UpperCamelCase      ,				legacy=UpperCamelCase							)
    _snake_case   :				List[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(UpperCamelCase      ,				return_tensors='pt'      ,				padding=UpperCamelCase							).input_ids
    # fmt: off
    _snake_case   :				List[Any]														=	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(UpperCamelCase      ,				UpperCamelCase							)

    _snake_case   :				List[Any]														=	model.generate(input_ids.to(UpperCamelCase							)							)
    _snake_case   :				List[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   :				List[Any]														=	tokenizer.batch_decode(UpperCamelCase							)
    self.assertEqual(UpperCamelCase      ,				UpperCamelCase							)








from functools import reduce

lowerCAmelCase_						   =    (
    """73167176531330624919225119674426574742355349194934"""
    """96983520312774506326239578318016984801869478851843"""
    """85861560789112949495459501737958331952853208805511"""
    """12540698747158523863050715693290963295227443043557"""
    """66896648950445244523161731856403098711121722383113"""
    """62229893423380308135336276614282806444486645238749"""
    """30358907296290491560440772390713810515859307960866"""
    """70172427121883998797908792274921901699720888093776"""
    """65727333001053367881220235421809751254540594752243"""
    """52584907711670556013604839586446706324415722155397"""
    """53697817977846174064955149290862569321978468622482"""
    """83972241375657056057490261407972968652414535100474"""
    """82166370484403199890008895243450658541227588666881"""
    """16427171479924442928230863465674813919123162824586"""
    """17866458359124566529476545682848912883142607690042"""
    """24219022671055626321111109370544217506941658960408"""
    """07198403850962455444362981230987879927244284909188"""
    """84580156166097919133875499200524063689912560717606"""
    """05886116467109405077541002256983155200055935729725"""
    """71636269561882670428252483600823257530420752963450"""
)


def     lowerCamelCase_	(      lowerCAmelCase:	str = N				)->		int:
  return max(
      # mypy cannot properly interpret reduce
      int(reduce(lambda lowerCAmelCase    ,							lowerCAmelCase				: str(int(lowerCAmelCase				) * int(lowerCAmelCase				)				)    ,							n[i : i + 13]				)				)
      for i in range(len(lowerCAmelCase				) - 12				)				)


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




