'''simple docstring'''




from __future__ import annotations

import unittest

from transformers import BlenderbotConfig, BlenderbotTokenizer, is_tf_available
from transformers.testing_utils import require_tf, require_tokenizers, slow
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
						import tensorflow as tf

						from transformers import TFAutoModelForSeqaSeqLM, TFBlenderbotForConditionalGeneration, TFBlenderbotModel






@require_tf
class 				a__:


			'''simple docstring'''






			UpperCAmelCase_			:     Tuple      =       BlenderbotConfig
			UpperCAmelCase_			:     str      =       {}
			UpperCAmelCase_			:     Optional[Any]      =       '''gelu'''



			def __init__(   self       ,    __lowerCAmelCase       ,    __lowerCAmelCase=13       ,    __lowerCAmelCase=7       ,    __lowerCAmelCase=True       ,    __lowerCAmelCase=False       ,    __lowerCAmelCase=99       ,    __lowerCAmelCase=32       ,    __lowerCAmelCase=2       ,    __lowerCAmelCase=4       ,    __lowerCAmelCase=37       ,    __lowerCAmelCase=0.1       ,    __lowerCAmelCase=0.1       ,    __lowerCAmelCase=20       ,    __lowerCAmelCase=2       ,    __lowerCAmelCase=1       ,    __lowerCAmelCase=0       ,    ):




					"""simple docstring"""


					lowerCAmelCase   					=  parent
					lowerCAmelCase   					=  batch_size
					lowerCAmelCase   					=  seq_length
					lowerCAmelCase   					=  is_training
					lowerCAmelCase   					=  use_labels
					lowerCAmelCase   					=  vocab_size
					lowerCAmelCase   					=  hidden_size
					lowerCAmelCase   					=  num_hidden_layers
					lowerCAmelCase   					=  num_attention_heads
					lowerCAmelCase   					=  intermediate_size

					lowerCAmelCase   					=  hidden_dropout_prob
					lowerCAmelCase   					=  attention_probs_dropout_prob
					lowerCAmelCase   					=  max_position_embeddings
					lowerCAmelCase   					=  eos_token_id
					lowerCAmelCase   					=  pad_token_id
					lowerCAmelCase   					=  bos_token_id



			def 							a_       (   self):




					"""simple docstring"""


					lowerCAmelCase   					=  ids_tensor([self.batch_size, self.seq_length - 1]       ,    self.vocab_size)
					lowerCAmelCase   					=  tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size)       ,    1)
					lowerCAmelCase   					=  tf.concat([input_ids, eos_tensor]       ,    axis=1)

					lowerCAmelCase   					=  ids_tensor([self.batch_size, self.seq_length]       ,    self.vocab_size)

					lowerCAmelCase   					=  self.config_cls(
					    vocab_size=self.vocab_size       ,    d_model=self.hidden_size       ,    encoder_layers=self.num_hidden_layers       ,    decoder_layers=self.num_hidden_layers       ,    encoder_attention_heads=self.num_attention_heads       ,    decoder_attention_heads=self.num_attention_heads       ,    encoder_ffn_dim=self.intermediate_size       ,    decoder_ffn_dim=self.intermediate_size       ,    dropout=self.hidden_dropout_prob       ,    attention_dropout=self.attention_probs_dropout_prob       ,    max_position_embeddings=self.max_position_embeddings       ,    eos_token_ids=[2]       ,    bos_token_id=self.bos_token_id       ,    pad_token_id=self.pad_token_id       ,    decoder_start_token_id=self.pad_token_id       ,    **self.config_updates       ,    )
					lowerCAmelCase   					=  prepare_blenderbot_inputs_dict(__lowerCAmelCase       ,    __lowerCAmelCase       ,    __lowerCAmelCase)
					return config, inputs_dict







			def 							a_       (   self       ,    __lowerCAmelCase       ,    __lowerCAmelCase):




					"""simple docstring"""


					lowerCAmelCase   					=  TFBlenderbotModel(config=__lowerCAmelCase).get_decoder()
					lowerCAmelCase   					=  inputs_dict["""input_ids"""]

					lowerCAmelCase   					=  input_ids[:1, :]
					lowerCAmelCase   					=  inputs_dict["""attention_mask"""][:1, :]
					lowerCAmelCase   					=  inputs_dict["""head_mask"""]
					lowerCAmelCase   					=  1

					# first forward pass
					lowerCAmelCase   					=  model(__lowerCAmelCase       ,    attention_mask=__lowerCAmelCase       ,    head_mask=__lowerCAmelCase       ,    use_cache=__lowerCAmelCase)

					lowerCAmelCase					,     lowerCAmelCase   					=  outputs.to_tuple()

					# create hypothetical next token and extent to next_input_ids
					lowerCAmelCase   					=  ids_tensor((self.batch_size, 3)       ,    config.vocab_size)
					lowerCAmelCase   					=  tf.cast(ids_tensor((self.batch_size, 3)       ,    2)       ,    tf.inta)

					# append to next input_ids and
					lowerCAmelCase   					=  tf.concat([input_ids, next_tokens]       ,    axis=-1)
					lowerCAmelCase   					=  tf.concat([attention_mask, next_attn_mask]       ,    axis=-1)

					lowerCAmelCase   					=  model(__lowerCAmelCase       ,    attention_mask=__lowerCAmelCase)[0]
					lowerCAmelCase   					=  model(__lowerCAmelCase       ,    attention_mask=__lowerCAmelCase       ,    past_key_values=__lowerCAmelCase)[0]

					self.parent.assertEqual(next_tokens.shape[1]       ,    output_from_past.shape[1])

					# select random slice
					lowerCAmelCase   					=  int(ids_tensor((1,)       ,    output_from_past.shape[-1]))
					lowerCAmelCase   					=  output_from_no_past[:, -3:, random_slice_idx]
					lowerCAmelCase   					=  output_from_past[:, :, random_slice_idx]

					# test that outputs are equal for slice
					tf.debugging.assert_near(__lowerCAmelCase       ,    __lowerCAmelCase       ,    rtol=1E-3)







def 				snake_case__							(			_A:					Any      ,   _A:					Optional[int]      ,   _A:					Optional[int]      ,   _A:					List[str]=None      ,   _A:					str=None      ,   _A:					Union[str, Any]=None      ,   _A:					Tuple=None      ,   _A:					Optional[Any]=None      ,   )    ->	Optional[Any]:

		'''simple docstring'''


		if attention_mask is None:
				lowerCAmelCase   					=  tf.cast(tf.math.not_equal(_A      ,   config.pad_token_id )      ,   tf.inta )
		if decoder_attention_mask is None:
				lowerCAmelCase   					=  tf.concat(
				    [
				        tf.ones(decoder_input_ids[:, :1].shape      ,   dtype=tf.inta ),
				        tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]      ,   config.pad_token_id )      ,   tf.inta ),
				    ]      ,   axis=-1      ,   )
		if head_mask is None:
				lowerCAmelCase   					=  tf.ones((config.encoder_layers, config.encoder_attention_heads) )
		if decoder_head_mask is None:
				lowerCAmelCase   					=  tf.ones((config.decoder_layers, config.decoder_attention_heads) )
		if cross_attn_head_mask is None:
				lowerCAmelCase   					=  tf.ones((config.decoder_layers, config.decoder_attention_heads) )
		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,
		}






@require_tf
class 				a__(					lowerCAmelCase__ ,						lowerCAmelCase__ ,						unittest.TestCase	):


			'''simple docstring'''






			UpperCAmelCase_			:     str      =       (TFBlenderbotForConditionalGeneration, TFBlenderbotModel) if is_tf_available() else ()
			UpperCAmelCase_			:     Any      =       (TFBlenderbotForConditionalGeneration,) if is_tf_available() else ()
			UpperCAmelCase_			:     Optional[int]      =       (
			    {
			        '''conversational''': TFBlenderbotForConditionalGeneration,
			        '''feature-extraction''': TFBlenderbotModel,
			        '''summarization''': TFBlenderbotForConditionalGeneration,
			        '''text2text-generation''': TFBlenderbotForConditionalGeneration,
			        '''translation''': TFBlenderbotForConditionalGeneration,
			    }
			    if is_tf_available()
			    else {}
			)
			UpperCAmelCase_			:     Any      =       True
			UpperCAmelCase_			:     Union[str, Any]      =       False
			UpperCAmelCase_			:     List[str]      =       False



			def 							a_       (   self):




					"""simple docstring"""


					lowerCAmelCase   					=  TFBlenderbotModelTester(self)
					lowerCAmelCase   					=  ConfigTester(self       ,    config_class=__lowerCAmelCase)



			def 							a_       (   self):




					"""simple docstring"""


					self.config_tester.run_common_tests()







			def 							a_       (   self):




					"""simple docstring"""


					lowerCAmelCase   					=  self.model_tester.prepare_config_and_inputs_for_common()
					self.model_tester.check_decoder_model_past_large_inputs(*__lowerCAmelCase)







@require_tokenizers
@require_tf
class 				a__(					unittest.TestCase	):


			'''simple docstring'''






			UpperCAmelCase_			:     List[str]      =       ['''My friends are cool but they eat too many carbs.''']
			UpperCAmelCase_			:     Union[str, Any]      =       '''facebook/blenderbot-400M-distill'''



			@cached_property
			def 							a_       (   self):




					"""simple docstring"""


					return BlenderbotTokenizer.from_pretrained(self.model_name)



			@cached_property
			def 							a_       (   self):




					"""simple docstring"""


					lowerCAmelCase   					=  TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name)
					return model







			@slow
			def 							a_       (   self):




					"""simple docstring"""


					lowerCAmelCase   					=  self.tokenizer(self.src_text       ,    return_tensors="""tf""")
					lowerCAmelCase   					=  self.model.generate(
					    model_inputs.input_ids       ,    )
					lowerCAmelCase   					=  self.tokenizer.batch_decode(generated_ids.numpy()       ,    skip_special_tokens=__lowerCAmelCase)[0]
					assert (
					    generated_words
					    == " That's unfortunate. Are they trying to lose weight or are they just trying to be healthier?"
					)





'''simple docstring'''




from ...utils import logging
from ..ta.modeling_tf_ta import TFTaEncoderModel, TFTaForConditionalGeneration, TFTaModel
from .configuration_mta import MTaConfig


__lowercase          =  logging.get_logger(__name__)

__lowercase          =  '''T5Config'''






class 				a__(					lowerCAmelCase__	):


			'''simple docstring'''






			UpperCAmelCase_			:     List[str]      =       '''mt5'''
			UpperCAmelCase_			:     Tuple      =       MTaConfig







class 				a__(					lowerCAmelCase__	):


			'''simple docstring'''






			UpperCAmelCase_			:     List[Any]      =       '''mt5'''
			UpperCAmelCase_			:     int      =       MTaConfig







class 				a__(					lowerCAmelCase__	):


			'''simple docstring'''






			UpperCAmelCase_			:     Tuple      =       '''mt5'''
			UpperCAmelCase_			:     Union[str, Any]      =       MTaConfig




