
'''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    UpperCAmelCase_				:





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





 def        snake_case__		(    self):
     return TaConfig.from_pretrained("google/umt5-base")





 def        snake_case__		(    self ,       lowercase_ ,       lowercase_ ,       lowercase_ ,       lowercase_=None ,       lowercase_=None ,       lowercase_=None ,       lowercase_=None ,       lowercase_=None ,       ):





     if attention_mask is None:
         snake_case_ :    Any							=						input_ids.ne(config.pad_token_id)
     if decoder_attention_mask is None:
         snake_case_ :    List[str]							=						decoder_input_ids.ne(config.pad_token_id)
     if head_mask is None:
         snake_case_ :    Dict							=						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_ :    List[str]							=						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        snake_case__		(    self):
     snake_case_ :    List[str]							=						ids_tensor([self.batch_size, self.encoder_seq_length] ,       self.vocab_size)
     snake_case_ :    Dict							=						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_ :    Tuple							=						decoder_input_ids.clamp(self.pad_token_id + 1)

     snake_case_ :    Optional[Any]							=						self.get_config()
     snake_case_ :    List[Any]							=						config.num_attention_heads
     snake_case_ :    int							=						self.prepare_inputs_dict(UpperCamelCase_ ,       UpperCamelCase_ ,       UpperCamelCase_)
     return config, input_dict





 def        snake_case__		(    self):
     snake_case_ :    Optional[int]							=						self.prepare_config_and_inputs()
     return config, inputs_dict





 def        snake_case__		(    self):
     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        snake_case__		(    self):
     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 ,       lowercase_ ,       lowercase_ ,       lowercase_ ,       lowercase_ ,       lowercase_ ,       lowercase_ ,       ):
     snake_case_ :    Dict							=						UMTaModel(config=UpperCamelCase_)
     model.to(UpperCamelCase_)
     model.eval()
     snake_case_ :    int							=						model(
         input_ids=UpperCamelCase_ ,       decoder_input_ids=UpperCamelCase_ ,       attention_mask=UpperCamelCase_ ,       decoder_attention_mask=UpperCamelCase_ ,       )
     snake_case_ :    Any							=						model(input_ids=UpperCamelCase_ ,       decoder_input_ids=UpperCamelCase_)
     snake_case_ :    Tuple							=						result.last_hidden_state
     snake_case_ :    Optional[int]							=						result.past_key_values
     snake_case_ :    str							=						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        snake_case__		(    self ,       lowercase_ ,       lowercase_ ,       lowercase_ ,       lowercase_ ,       lowercase_ ,       lowercase_ ,       ):
     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_ :    Dict							=						model(UpperCamelCase_ ,       use_cache=UpperCamelCase_)

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

     snake_case_ :    Any							=						outputs.to_tuple()

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

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

     snake_case_ :    Optional[int]							=						model(UpperCamelCase_)['last_hidden_state']
     snake_case_ :    str							=						model(UpperCamelCase_ ,       past_key_values=UpperCamelCase_)['last_hidden_state']

     # select random slice
     snake_case_ :    List[Any]							=						ids_tensor((1,) ,       output_from_past.shape[-1]).item()
     snake_case_ :    List[str]							=						output_from_no_past[:, -1, random_slice_idx].detach()
     snake_case_ :    Optional[int]							=						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        snake_case__		(    self ,       lowercase_ ,       lowercase_ ,       ):
     snake_case_ :    Optional[int]							=						UMTaModel(config=UpperCamelCase_).to(UpperCamelCase_).half().eval()
     snake_case_ :    Tuple							=						model(**UpperCamelCase_)['last_hidden_state']
     self.parent.assertFalse(torch.isnan(UpperCamelCase_).any().item())






@require_torch
class    UpperCAmelCase_				(	snake_case__    ,						snake_case__    ,						snake_case__    ,						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        snake_case__		(    self):
     snake_case_ :    List[str]							=						UMTaModelTester(self)





 @unittest.skip("Test has a segmentation fault on torch 1.8.0")
 def        snake_case__		(    self):
     snake_case_ :    Any							=						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        snake_case__		(    self):
     snake_case_ :    str							=						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model_fpaa_forward(*UpperCamelCase_)





 def        snake_case__		(    self):
     snake_case_ :    List[str]							=						['encoder_attentions', 'decoder_attentions', 'cross_attentions']
     snake_case_ :    List[Any]							=						self.model_tester.prepare_config_and_inputs()
     snake_case_ :    Dict							=						config_and_inputs[0]
     snake_case_ :    Union[str, Any]							=						UMTaForConditionalGeneration(UpperCamelCase_).eval()
     model.to(UpperCamelCase_)

     snake_case_ :    List[Any]							=						{
         '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_ :    Dict							=						torch.ones(
                 config.num_decoder_layers ,       config.num_heads ,       device=UpperCamelCase_)

         snake_case_ :    Optional[Any]							=						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_ :    Dict							=						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):
     pass






@require_torch
@require_sentencepiece
@require_tokenizers
class    UpperCAmelCase_				(	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        snake_case__		(    self):
     snake_case_ :    int							=						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[str]							=						[
         '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_ :    Dict							=						tokenizer(UpperCamelCase_ ,       return_tensors="pt" ,       padding=UpperCamelCase_).input_ids
     # fmt: off
     snake_case_ :    Optional[int]							=						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_ :    Optional[int]							=						model.generate(input_ids.to(UpperCamelCase_))
     snake_case_ :    Tuple							=						[
         '<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[int]							=						tokenizer.batch_decode(UpperCamelCase_)
     self.assertEqual(UpperCamelCase_ ,       UpperCamelCase_)






'''simple docstring'''
from __future__ import annotations

from collections.abc import Callable
from typing import Generic, TypeVar

a_											=						TypeVar("T")
a_											=						TypeVar("U")






class    UpperCAmelCase_				(	Generic[T, U]					):





 def __init__(    self ,       lowercase_ ,       lowercase_):
     snake_case_ :    Any     			=						key
     snake_case_ :    List[Any]     			=						val
     snake_case_ :    DoubleLinkedListNode[T, U] | None     			=						None
     snake_case_ :    DoubleLinkedListNode[T, U] | None     			=						None





 def __repr__(    self):
     return (
         F'Node: key: {self.key}, val: {self.val}, '
         F'has next: {bool(self.next)}, has prev: {bool(self.prev)}'
     )






class    UpperCAmelCase_				(	Generic[T, U]					):





 def __init__(    self):
     snake_case_ :    DoubleLinkedListNode[T, U]     			=						DoubleLinkedListNode(lowercase_ ,       lowercase_)
     snake_case_ :    DoubleLinkedListNode[T, U]     			=						DoubleLinkedListNode(lowercase_ ,       lowercase_)
     snake_case_			, snake_case_ :    Union[str, Any]     			=						self.rear, self.head





 def __repr__(    self):
     snake_case_ :    Dict     			=						["DoubleLinkedList"]
     snake_case_ :    Dict     			=						self.head
     while node.next is not None:
         rep.append(str(lowercase_))
         snake_case_ :    List[str]     			=						node.next
     rep.append(str(self.rear))
     return ",\n    ".join(lowercase_)





 def        snake_case__		(    self ,       lowercase_):
     snake_case_ :    List[str]     			=						self.rear.prev

     # All nodes other than self.head are guaranteed to have non-None previous
     assert previous is not None

     snake_case_ :    Tuple     			=						node
     snake_case_ :    str     			=						previous
     snake_case_ :    Optional[Any]     			=						node
     snake_case_ :    Any     			=						self.rear





 def        snake_case__		(    self ,       lowercase_):

     if node.prev is None or node.next is None:
         return None

     snake_case_ :    Union[str, Any]     			=						node.next
     snake_case_ :    Optional[int]     			=						node.prev
     snake_case_ :    str     			=						None
     snake_case_ :    int     			=						None
     return node






class    UpperCAmelCase_				(	Generic[T, U]					):
 UpperCAmelCase_		    =						{}





 def __init__(    self ,       lowercase_):
     snake_case_ :    DoubleLinkedList[T, U]     			=						DoubleLinkedList()
     snake_case_ :    List[str]     			=						capacity
     snake_case_ :    Any     			=						0
     snake_case_ :    Dict     			=						0
     snake_case_ :    Union[str, Any]     			=						0
     snake_case_ :    dict[T, DoubleLinkedListNode[T, U]]     			=						{}





 def __repr__(    self):
     return (
         F'CacheInfo(hits={self.hits}, misses={self.miss}, '
         F'capacity={self.capacity}, current size={self.num_keys})'
     )





 def __contains__(    self ,       lowercase_):
     return key in self.cache





 def        snake_case__		(    self ,       lowercase_):
     # Note: pythonic interface would throw KeyError rather than return None

     if key in self.cache:
         self.hits += 1
         snake_case_ :    DoubleLinkedListNode[T, U]     			=						self.cache[key]
         snake_case_ :    Optional[Any]     			=						self.list.remove(self.cache[key])
         assert node == value_node

         # node is guaranteed not None because it is in self.cache
         assert node is not None
         self.list.add(lowercase_)
         return node.val
     self.miss += 1
     return None





 def        snake_case__		(    self ,       lowercase_ ,       lowercase_):

     if key not in self.cache:
         if self.num_keys >= self.capacity:
             # delete first node (oldest) when over capacity
             snake_case_ :    Optional[Any]     			=						self.list.head.next

             # guaranteed to have a non-None first node when num_keys > 0
             # explain to type checker via assertions
             assert first_node is not None
             assert first_node.key is not None
             assert (
                 self.list.remove(lowercase_) is not None
             )  # node guaranteed to be in list assert node.key is not None

             del self.cache[first_node.key]
             self.num_keys -= 1
         snake_case_ :    Dict     			=						DoubleLinkedListNode(lowercase_ ,       lowercase_)
         self.list.add(self.cache[key])
         self.num_keys += 1

     else:
         # bump node to the end of the list, update value
         snake_case_ :    List[Any]     			=						self.list.remove(self.cache[key])
         assert node is not None  # node guaranteed to be in list
         snake_case_ :    Optional[Any]     			=						value
         self.list.add(lowercase_)





 @classmethod
 def        snake_case__		(    cls ,       lowercase_ = 1_28):

     def cache_decorator_inner(lowercase_) -> Callable[..., U]:
         def cache_decorator_wrapper(*lowercase_) -> U:
             if func not in cls.decorator_function_to_instance_map:
                 snake_case_ :    List[str]     			=						LRUCache(lowercase_)

             snake_case_ :    Optional[Any]     			=						cls.decorator_function_to_instance_map[func].get(args[0])
             if result is None:
                 snake_case_ :    Any     			=						func(*lowercase_)
                 cls.decorator_function_to_instance_map[func].put(args[0] ,       lowercase_)
             return result

         def cache_info() -> LRUCache[T, U]:
             return cls.decorator_function_to_instance_map[func]

         setattr(lowercase_ ,       "cache_info" ,       lowercase_)  # noqa: B010

         return cache_decorator_wrapper

     return cache_decorator_inner


if __name__ == "__main__":
  import doctest

  doctest.testmod()




