



"""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    _lowerCAmelCase    :





		"""simple docstring"""






		_lowerCamelCase			  =       BlenderbotConfig
		_lowerCamelCase			  =       {}
		_lowerCamelCase			  =       '''gelu'''





		def __init__(					self							,    _lowercase							,    _lowercase=1_3							,    _lowercase=7							,    _lowercase=True							,    _lowercase=False							,    _lowercase=9_9							,    _lowercase=3_2							,    _lowercase=2							,    _lowercase=4							,    _lowercase=3_7							,    _lowercase=0.1							,    _lowercase=0.1							,    _lowercase=2_0							,    _lowercase=2							,    _lowercase=1							,    _lowercase=0							,    ) ->						Any:







								'''simple docstring'''

								snake_case_   : Optional[int]						=			parent
								snake_case_   : Any						=			batch_size
								snake_case_   : Tuple						=			seq_length
								snake_case_   : Optional[int]						=			is_training
								snake_case_   : str						=			use_labels
								snake_case_   : Tuple						=			vocab_size
								snake_case_   : Union[str, Any]						=			hidden_size
								snake_case_   : Tuple						=			num_hidden_layers
								snake_case_   : List[Any]						=			num_attention_heads
								snake_case_   : List[str]						=			intermediate_size

								snake_case_   : List[str]						=			hidden_dropout_prob
								snake_case_   : List[str]						=			attention_probs_dropout_prob
								snake_case_   : str						=			max_position_embeddings
								snake_case_   : List[str]						=			eos_token_id
								snake_case_   : Dict						=			pad_token_id
								snake_case_   : int						=			bos_token_id





		def    UpperCAmelCase__							(					self   ) ->						int:







								'''simple docstring'''

								snake_case_   : int						=			ids_tensor([self.batch_size, self.seq_length - 1]							,    self.vocab_size   )
								snake_case_   : int						=			tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size   )							,    1   )
								snake_case_   : Union[str, Any]						=			tf.concat([input_ids, eos_tensor]							,    axis=1   )

								snake_case_   : List[Any]						=			ids_tensor([self.batch_size, self.seq_length]							,    self.vocab_size   )

								snake_case_   : Tuple						=			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							,    )
								snake_case_   : List[str]						=			prepare_blenderbot_inputs_dict(_lowercase							,    _lowercase							,    _lowercase   )
								return config, inputs_dict





		def    UpperCAmelCase__							(					self							,    _lowercase							,    _lowercase   ) ->						List[Any]:







								'''simple docstring'''

								snake_case_   : Optional[Any]						=			TFBlenderbotModel(config=_lowercase   ).get_decoder()
								snake_case_   : Optional[int]						=			inputs_dict["""input_ids"""]

								snake_case_   : Optional[Any]						=			input_ids[:1, :]
								snake_case_   : List[Any]						=			inputs_dict["""attention_mask"""][:1, :]
								snake_case_   : Dict						=			inputs_dict["""head_mask"""]
								snake_case_   : Tuple						=			1

								# first forward pass
								snake_case_   : Union[str, Any]						=			model(_lowercase							,    attention_mask=_lowercase							,    head_mask=_lowercase							,    use_cache=_lowercase   )

								snake_case_			,		snake_case_   : List[str]						=			outputs.to_tuple()

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

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

								snake_case_   : Dict						=			model(_lowercase							,    attention_mask=_lowercase   )[0]
								snake_case_   : str						=			model(_lowercase							,    attention_mask=_lowercase							,    past_key_values=_lowercase   )[0]

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

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

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




def 	__lowerCAmelCase						(		__UpperCamelCase      :			List[Any]					, __UpperCamelCase      :			Optional[int]					, __UpperCamelCase      :			List[Any]					, __UpperCamelCase      :			Optional[int]=None					, __UpperCamelCase      :			Union[str, Any]=None					, __UpperCamelCase      :			List[Any]=None					, __UpperCamelCase      :			Tuple=None					, __UpperCamelCase      :			Union[str, Any]=None					, ):


						'''simple docstring'''






						if attention_mask is None:
												snake_case_   : Tuple						=			tf.cast(tf.math.not_equal(__UpperCamelCase					, config.pad_token_id    )					, tf.inta    )
						if decoder_attention_mask is None:
												snake_case_   : Tuple						=			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:
												snake_case_   : Optional[int]						=			tf.ones((config.encoder_layers, config.encoder_attention_heads)    )
						if decoder_head_mask is None:
												snake_case_   : Tuple						=			tf.ones((config.decoder_layers, config.decoder_attention_heads)    )
						if cross_attn_head_mask is None:
												snake_case_   : Dict						=			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    _lowerCAmelCase    (				SCREAMING_SNAKE_CASE__  ,				SCREAMING_SNAKE_CASE__  ,				unittest.TestCase		):





		"""simple docstring"""






		_lowerCamelCase			  =       (TFBlenderbotForConditionalGeneration, TFBlenderbotModel) if is_tf_available() else ()
		_lowerCamelCase			  =       (TFBlenderbotForConditionalGeneration,) if is_tf_available() else ()
		_lowerCamelCase			  =       (
		    {
		        '''conversational''': TFBlenderbotForConditionalGeneration,
		        '''feature-extraction''': TFBlenderbotModel,
		        '''summarization''': TFBlenderbotForConditionalGeneration,
		        '''text2text-generation''': TFBlenderbotForConditionalGeneration,
		        '''translation''': TFBlenderbotForConditionalGeneration,
		    }
		    if is_tf_available()
		    else {}
		)
		_lowerCamelCase			  =       True
		_lowerCamelCase			  =       False
		_lowerCamelCase			  =       False





		def    UpperCAmelCase__							(					self   ) ->						Union[str, Any]:







								'''simple docstring'''

								snake_case_   : Optional[Any]						=			TFBlenderbotModelTester(self   )
								snake_case_   : Tuple						=			ConfigTester(self							,    config_class=_lowercase   )





		def    UpperCAmelCase__							(					self   ) ->						Union[str, Any]:







								'''simple docstring'''

								self.config_tester.run_common_tests()





		def    UpperCAmelCase__							(					self   ) ->						Tuple:







								'''simple docstring'''

								snake_case_   : int						=			self.model_tester.prepare_config_and_inputs_for_common()
								self.model_tester.check_decoder_model_past_large_inputs(*_lowercase   )




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





		"""simple docstring"""






		_lowerCamelCase			  =       ['''My friends are cool but they eat too many carbs.''']
		_lowerCamelCase			  =       '''facebook/blenderbot-400M-distill'''





		@cached_property
		def    UpperCAmelCase__							(					self   ) ->						List[str]:







								'''simple docstring'''

								return BlenderbotTokenizer.from_pretrained(self.model_name   )





		@cached_property
		def    UpperCAmelCase__							(					self   ) ->						Tuple:







								'''simple docstring'''

								snake_case_   : Optional[int]						=			TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name   )
								return model





		@slow
		def    UpperCAmelCase__							(					self   ) ->						Optional[Any]:







								'''simple docstring'''

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






"""simple docstring"""





import unittest

from transformers import is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, require_torch, slow


if is_flax_available():
					import optax
					from flax.training.common_utils import onehot

					from transformers import AutoTokenizer, FlaxMTaForConditionalGeneration
					from transformers.models.ta.modeling_flax_ta import shift_tokens_right



@require_torch
@require_sentencepiece
@require_tokenizers
@require_flax
class    _lowerCAmelCase    (				unittest.TestCase		):





		"""simple docstring"""






		@slow
		def    UpperCAmelCase__							(					self   ) ->						Optional[int]:







								'''simple docstring'''

								snake_case_   : int						=			FlaxMTaForConditionalGeneration.from_pretrained("""google/mt5-small"""   )
								snake_case_   : List[Any]						=			AutoTokenizer.from_pretrained("""google/mt5-small"""   )

								snake_case_   : Dict						=			tokenizer("""Hello there"""							,    return_tensors="""np"""   ).input_ids
								snake_case_   : List[str]						=			tokenizer("""Hi I am"""							,    return_tensors="""np"""   ).input_ids

								snake_case_   : Optional[Any]						=			shift_tokens_right(_lowercase							,    model.config.pad_token_id							,    model.config.decoder_start_token_id   )

								snake_case_   : Tuple						=			model(_lowercase							,    decoder_input_ids=_lowercase   ).logits
								snake_case_   : Tuple						=			optax.softmax_cross_entropy(_lowercase							,    onehot(_lowercase							,    logits.shape[-1]   )   ).mean()

								snake_case_   : List[str]						=			-(labels.shape[-1] * loss.item())

								snake_case_   : Optional[int]						=			-84.9127
								self.assertTrue(abs(mtf_score - EXPECTED_SCORE   ) < 1E-4   )

