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				:





      def __init__(		self						:					Dict					,       __lowerCamelCase						:					int					,       __lowerCamelCase						:					str=9_9					,       __lowerCamelCase						:					Any=1_3					,       __lowerCamelCase						:					Any=7					,       __lowerCamelCase						:					Union[str, Any]=9					,       __lowerCamelCase						:					Any=True					,       __lowerCamelCase						:					Optional[Any]=True					,       __lowerCamelCase						:					str=False					,       __lowerCamelCase						:					Tuple=3_2					,       __lowerCamelCase						:					Dict=5					,       __lowerCamelCase						:					Optional[Any]=4					,       __lowerCamelCase						:					Union[str, Any]=3_7					,       __lowerCamelCase						:					int=8					,       __lowerCamelCase						:					Optional[int]=0.1					,       __lowerCamelCase						:					Tuple=0.002					,       __lowerCamelCase						:					str=1					,       __lowerCamelCase						:					Optional[int]=0					,       __lowerCamelCase						:					Tuple=0					,       __lowerCamelCase						:					int=None					,       __lowerCamelCase						:					List[str]=None					,       )    ->   Any:






         """simple docstring"""







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





      def        _lowercase				(		self						:					Optional[Any]       )    ->   Tuple:






         """simple docstring"""







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





      def        _lowercase				(		self						:					Optional[int]					,       __lowerCamelCase						:					Tuple					,       __lowerCamelCase						:					Dict					,       __lowerCamelCase						:					Any					,       __lowerCamelCase						:					Any=None					,       __lowerCamelCase						:					List[Any]=None					,       __lowerCamelCase						:					str=None					,       __lowerCamelCase						:					Union[str, Any]=None					,       __lowerCamelCase						:					Optional[int]=None					,       )    ->   Any:






         """simple docstring"""







         if attention_mask is None:
            UpperCAmelCase							=     input_ids.ne(config.pad_token_id       )
         if decoder_attention_mask is None:
            UpperCAmelCase							=     decoder_input_ids.ne(config.pad_token_id       )
         if head_mask is None:
            UpperCAmelCase							=     torch.ones(config.num_hidden_layers					,       config.num_attention_heads					,       device=_a       )
         if decoder_head_mask is None:
            UpperCAmelCase							=     torch.ones(config.num_decoder_layers					,       config.num_attention_heads					,       device=_a       )
         if cross_attn_head_mask is None:
            UpperCAmelCase							=     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        _lowercase				(		self						:					List[Any]       )    ->   Tuple:






         """simple docstring"""







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

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





      def        _lowercase				(		self						:					Dict       )    ->   Any:






         """simple docstring"""







         UpperCAmelCase							=     self.prepare_config_and_inputs()
         return config, inputs_dict





      def        _lowercase				(		self						:					str       )    ->   List[str]:






         """simple docstring"""







         return TaConfig(
             vocab_size=1_6_6					,       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        _lowercase				(		self						:					Optional[int]       )    ->   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        _lowercase				(		self						:					str					,       __lowerCamelCase						:					Tuple					,       __lowerCamelCase						:					Optional[Any]					,       __lowerCamelCase						:					Any					,       __lowerCamelCase						:					Optional[int]					,       __lowerCamelCase						:					Tuple					,       __lowerCamelCase						:					Optional[Any]					,       )    ->   Optional[Any]:






         """simple docstring"""







         UpperCAmelCase							=     UMTaModel(config=_a       )
         model.to(_a       )
         model.eval()
         UpperCAmelCase							=     model(
             input_ids=_a					,       decoder_input_ids=_a					,       attention_mask=_a					,       decoder_attention_mask=_a					,       )
         UpperCAmelCase							=     model(input_ids=_a					,       decoder_input_ids=_a       )
         UpperCAmelCase							=     result.last_hidden_state
         UpperCAmelCase							=     result.past_key_values
         UpperCAmelCase							=     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        _lowercase				(		self						:					Dict					,       __lowerCamelCase						:					int					,       __lowerCamelCase						:					Any					,       __lowerCamelCase						:					Dict					,       __lowerCamelCase						:					Any					,       __lowerCamelCase						:					Optional[Any]					,       __lowerCamelCase						:					Union[str, Any]					,       )    ->   Union[str, Any]:






         """simple docstring"""







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

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

         UpperCAmelCase							=     outputs.to_tuple()

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

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

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

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






         """simple docstring"""







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






@require_torch
class   __lowercase				(  UpperCamelCase_ ,					UpperCamelCase_ ,					UpperCamelCase_ ,					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        _lowercase				(		self						:					Dict       )    ->   Optional[int]:






         """simple docstring"""







         UpperCAmelCase							=     UMTaModelTester(self       )





      @unittest.skip("""Test has a segmentation fault on torch 1.8.0"""       )
      def        _lowercase				(		self						:					int       )    ->   Any:






         """simple docstring"""







         UpperCAmelCase							=     self.model_tester.prepare_config_and_inputs()
         UpperCAmelCase							=     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        _lowercase				(		self						:					str       )    ->   str:






         """simple docstring"""







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





      def        _lowercase				(		self						:					Dict       )    ->   List[str]:






         """simple docstring"""







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

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

            UpperCAmelCase							=     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
            UpperCAmelCase							=     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        _lowercase				(		self						:					List[Any]       )    ->   Dict:






         """simple docstring"""







         pass






@require_torch
@require_sentencepiece
@require_tokenizers
class   __lowercase				(  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        _lowercase				(		self						:					Tuple       )    ->   Union[str, Any]:






         """simple docstring"""







         UpperCAmelCase							=     UMTaForConditionalGeneration.from_pretrained("""google/umt5-small"""					,       return_dict=_a       ).to(_a       )
         UpperCAmelCase							=     AutoTokenizer.from_pretrained("""google/umt5-small"""					,       use_fast=_a					,       legacy=_a       )
         UpperCAmelCase							=     [
             """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>.""",
         ]
         UpperCAmelCase							=     tokenizer(_a					,       return_tensors="""pt"""					,       padding=_a       ).input_ids
         # fmt: off
         UpperCAmelCase							=     torch.tensor(
             [
                 [ 3_8_5_3_0, 2_1_0_7_0_3, 2_5_6_2_9_9, 1_4_1_0, 2_5_6_2_9_8, 2_7_4, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
                 [   8_2_6, 3_2_1, 6_7_1, 2_5_9_2_2, 2_5_6_2_9_9, 2_7_4, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
                 [  1_4_6_0, 3_3_9, 3_1_2, 1_9_0_1_4, 1_0_6_2_0, 7_5_8, 2_5_6_2_9_9, 2_3_5_5,2_7_4, 1, 0, 0, 0, 0, 0, 0,0, 0],
                 [   5_1_7, 2_5_6_2_9_9, 1_4_8_6_9, 2_8_1, 3_0_1, 2_5_6_2_9_8, 2_7_5, 1_1_9_9_8_3,1, 0, 0, 0, 0, 0, 0, 0,0, 0],
                 [   3_2_0, 2_5_6_2_9_9, 1_4_8_6_9, 2_8_1, 2_2_3_4, 2_8_9, 2_2_7_5, 3_3_3,6_1_3_9_1, 2_8_9, 2_5_6_2_9_8, 5_4_3, 2_5_6_2_9_7, 1_6_8_7_1_4, 3_2_9, 2_5_6_2_9_6,2_7_4, 1],
             ]       )
         # fmt: on
         torch.testing.assert_allclose(_a					,       _a       )

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







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available


__a     =			{"""tokenization_herbert""": ["""HerbertTokenizer"""]}

try:
 if not is_tokenizers_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __a     =			["""HerbertTokenizerFast"""]


if TYPE_CHECKING:
 from .tokenization_herbert import HerbertTokenizer

 try:
  if not is_tokenizers_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_herbert_fast import HerbertTokenizerFast

else:
 import sys

 __a     =			_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



