


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			:







       def __init__(			self: int,A_: Union[str, Any],A_: List[str]=99,A_: Dict=13,A_: Tuple=7,A_: Union[str, Any]=9,A_: str=True,A_: int=True,A_: Optional[Any]=False,A_: Union[str, Any]=32,A_: Optional[int]=5,A_: Optional[Any]=4,A_: Union[str, Any]=37,A_: List[Any]=8,A_: Optional[int]=0.1,A_: str=0.0_0_2,A_: List[Any]=1,A_: List[str]=0,A_: int=0,A_: Optional[int]=None,A_: str=None,):






            '''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 		snake_case_	(			self: Dict					):






            '''simple docstring'''




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







       def 		snake_case_	(			self: int,A_: Optional[Any],A_: Optional[int],A_: int,A_: Any=None,A_: Optional[Any]=None,A_: Optional[Any]=None,A_: Optional[Any]=None,A_: Any=None,):






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






            '''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 		snake_case_	(			self: str					):






            '''simple docstring'''




            __UpperCamelCase,    __UpperCamelCase =			self.prepare_config_and_inputs()
            return config, inputs_dict







       def 		snake_case_	(			self: Union[str, Any]					):






            '''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: List[str]					):






            '''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: List[Any],A_: List[Any],A_: Tuple,A_: Any,A_: Tuple,A_: int,A_: Union[str, 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 		snake_case_	(			self: Dict,A_: Optional[Any],A_: int,A_: Dict,A_: Union[str, Any],A_: str,A_: List[str],):






            '''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,    __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 		snake_case_	(			self: str,A_: Dict,A_: 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 	__lowerCamelCase			(_a				,      _a				,      _a				,      unittest.TestCase						):
       _lowercase											=     (
           (UMTaModel, UMTaForConditionalGeneration, UMTaForQuestionAnswering) if is_torch_available() else ()
       )
       _lowercase											=     (UMTaForConditionalGeneration,) if is_torch_available() else ()
       _lowercase											=     (
           {
               """conversational""": UMTaForConditionalGeneration,
               """feature-extraction""": UMTaModel,
               """summarization""": UMTaForConditionalGeneration,
               """text2text-generation""": UMTaForConditionalGeneration,
               """translation""": UMTaForConditionalGeneration,
               """question-answering""": UMTaForQuestionAnswering,
           }
           if is_torch_available()
           else {}
       )
       _lowercase											=     True
       _lowercase											=     False
       _lowercase											=     False
       _lowercase											=     True
       _lowercase											=     True
       # The small UMT5 model needs higher percentages for CPU/MP tests
       _lowercase											=     [0.8, 0.9]







       def 		snake_case_	(			self: Any					):






            '''simple docstring'''




            __UpperCamelCase =			UMTaModelTester(self					)







       @unittest.skip('Test has a segmentation fault on torch 1.8.0'					)
       def 		snake_case_	(			self: Optional[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 		snake_case_	(			self: Any					):






            '''simple docstring'''




            __UpperCamelCase =			self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model_fpaa_forward(*A_					)







       def 		snake_case_	(			self: Any					):






            '''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 		snake_case_	(			self: Tuple					):






            '''simple docstring'''




            pass



@require_torch
@require_sentencepiece
@require_tokenizers
class 	__lowerCamelCase			(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: List[str]					):






            '''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_8530, 21_0703, 25_6299, 1410, 25_6298, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
                    [   826, 321, 671, 2_5922, 25_6299, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
                    [  1460, 339, 312, 1_9014, 1_0620, 758, 25_6299, 2355,274, 1, 0, 0, 0, 0, 0, 0,0, 0],
                    [   517, 25_6299, 1_4869, 281, 301, 25_6298, 275, 11_9983,1, 0, 0, 0, 0, 0, 0, 0,0, 0],
                    [   320, 25_6299, 1_4869, 281, 2234, 289, 2275, 333,6_1391, 289, 25_6298, 543, 25_6297, 16_8714, 329, 25_6296,274, 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 collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__snake_case   							=						logging.get_logger(__name__)

__snake_case   							=						{
    '''bert-base-uncased''': '''https://huggingface.co/bert-base-uncased/resolve/main/config.json''',
    '''bert-large-uncased''': '''https://huggingface.co/bert-large-uncased/resolve/main/config.json''',
    '''bert-base-cased''': '''https://huggingface.co/bert-base-cased/resolve/main/config.json''',
    '''bert-large-cased''': '''https://huggingface.co/bert-large-cased/resolve/main/config.json''',
    '''bert-base-multilingual-uncased''': '''https://huggingface.co/bert-base-multilingual-uncased/resolve/main/config.json''',
    '''bert-base-multilingual-cased''': '''https://huggingface.co/bert-base-multilingual-cased/resolve/main/config.json''',
    '''bert-base-chinese''': '''https://huggingface.co/bert-base-chinese/resolve/main/config.json''',
    '''bert-base-german-cased''': '''https://huggingface.co/bert-base-german-cased/resolve/main/config.json''',
    '''bert-large-uncased-whole-word-masking''': (
        '''https://huggingface.co/bert-large-uncased-whole-word-masking/resolve/main/config.json'''
    ),
    '''bert-large-cased-whole-word-masking''': (
        '''https://huggingface.co/bert-large-cased-whole-word-masking/resolve/main/config.json'''
    ),
    '''bert-large-uncased-whole-word-masking-finetuned-squad''': (
        '''https://huggingface.co/bert-large-uncased-whole-word-masking-finetuned-squad/resolve/main/config.json'''
    ),
    '''bert-large-cased-whole-word-masking-finetuned-squad''': (
        '''https://huggingface.co/bert-large-cased-whole-word-masking-finetuned-squad/resolve/main/config.json'''
    ),
    '''bert-base-cased-finetuned-mrpc''': '''https://huggingface.co/bert-base-cased-finetuned-mrpc/resolve/main/config.json''',
    '''bert-base-german-dbmdz-cased''': '''https://huggingface.co/bert-base-german-dbmdz-cased/resolve/main/config.json''',
    '''bert-base-german-dbmdz-uncased''': '''https://huggingface.co/bert-base-german-dbmdz-uncased/resolve/main/config.json''',
    '''cl-tohoku/bert-base-japanese''': '''https://huggingface.co/cl-tohoku/bert-base-japanese/resolve/main/config.json''',
    '''cl-tohoku/bert-base-japanese-whole-word-masking''': (
        '''https://huggingface.co/cl-tohoku/bert-base-japanese-whole-word-masking/resolve/main/config.json'''
    ),
    '''cl-tohoku/bert-base-japanese-char''': (
        '''https://huggingface.co/cl-tohoku/bert-base-japanese-char/resolve/main/config.json'''
    ),
    '''cl-tohoku/bert-base-japanese-char-whole-word-masking''': (
        '''https://huggingface.co/cl-tohoku/bert-base-japanese-char-whole-word-masking/resolve/main/config.json'''
    ),
    '''TurkuNLP/bert-base-finnish-cased-v1''': (
        '''https://huggingface.co/TurkuNLP/bert-base-finnish-cased-v1/resolve/main/config.json'''
    ),
    '''TurkuNLP/bert-base-finnish-uncased-v1''': (
        '''https://huggingface.co/TurkuNLP/bert-base-finnish-uncased-v1/resolve/main/config.json'''
    ),
    '''wietsedv/bert-base-dutch-cased''': '''https://huggingface.co/wietsedv/bert-base-dutch-cased/resolve/main/config.json''',
    # See all BERT models at https://huggingface.co/models?filter=bert
}



class 	__lowerCamelCase			(_a						):
       _lowercase											=     """bert"""







       def __init__(			self: Any,A_: Dict=3_0522,A_: Optional[Any]=768,A_: Union[str, Any]=12,A_: List[Any]=12,A_: Optional[int]=3072,A_: Union[str, Any]="gelu",A_: List[str]=0.1,A_: Dict=0.1,A_: Optional[int]=512,A_: Optional[Any]=2,A_: Union[str, Any]=0.0_2,A_: List[Any]=1E-12,A_: Optional[int]=0,A_: List[Any]="absolute",A_: str=True,A_: Union[str, Any]=None,**A_: int,):






            '''simple docstring'''




            super().__init__(pad_token_id=A_,**A_					)

            __UpperCamelCase =			vocab_size
            __UpperCamelCase =			hidden_size
            __UpperCamelCase =			num_hidden_layers
            __UpperCamelCase =			num_attention_heads
            __UpperCamelCase =			hidden_act
            __UpperCamelCase =			intermediate_size
            __UpperCamelCase =			hidden_dropout_prob
            __UpperCamelCase =			attention_probs_dropout_prob
            __UpperCamelCase =			max_position_embeddings
            __UpperCamelCase =			type_vocab_size
            __UpperCamelCase =			initializer_range
            __UpperCamelCase =			layer_norm_eps
            __UpperCamelCase =			position_embedding_type
            __UpperCamelCase =			use_cache
            __UpperCamelCase =			classifier_dropout



class 	__lowerCamelCase			(_a						):







       @property
       def 		snake_case_	(			self: Optional[int]					):






            '''simple docstring'''




            if self.task == "multiple-choice":
                 __UpperCamelCase =			{0: 'batch', 1: 'choice', 2: 'sequence'}
            else:
                 __UpperCamelCase =			{0: 'batch', 1: 'sequence'}
            return OrderedDict(
                [
                    ('input_ids', dynamic_axis),
                    ('attention_mask', dynamic_axis),
                    ('token_type_ids', dynamic_axis),
                ]					)


