
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 						lowercase_				:
 A__  : int      					=			BlenderbotConfig
 A__  : int      					=			{}
 A__  : Optional[int]      					=			"""gelu"""

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




        """simple docstring"""


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

        UpperCamelCase_											=				hidden_dropout_prob
        UpperCamelCase_											=				attention_probs_dropout_prob
        UpperCamelCase_											=				max_position_embeddings
        UpperCamelCase_											=				eos_token_id
        UpperCamelCase_											=				pad_token_id
        UpperCamelCase_											=				bos_token_id

 def lowerCamelCase_ (					self   ):




        """simple docstring"""


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

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

        UpperCamelCase_											=				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   ,		)
        UpperCamelCase_											=				prepare_blenderbot_inputs_dict(__a   ,		__a   ,		__a   )
        return config, inputs_dict

 def lowerCamelCase_ (					self   ,		__UpperCamelCase   ,		__UpperCamelCase   ):




        """simple docstring"""


        UpperCamelCase_											=				TFBlenderbotModel(config=__a   ).get_decoder()
        UpperCamelCase_											=				inputs_dict["""input_ids"""]

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

        # first forward pass
        UpperCamelCase_											=				model(__a   ,		attention_mask=__a   ,		head_mask=__a   ,		use_cache=__a   )

        UpperCamelCase_				, UpperCamelCase_											=				outputs.to_tuple()

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

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

        UpperCamelCase_											=				model(__a   ,		attention_mask=__a   )[0]
        UpperCamelCase_											=				model(__a   ,		attention_mask=__a   ,		past_key_values=__a   )[0]

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

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

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




def  lowerCamelCase__ (  a__     :     Union[str, Any]		,					a__     :     Union[str, Any]		,					a__     :     Dict		,					a__     :     Union[str, Any]=None		,					a__     :     Optional[int]=None		,					a__     :     Tuple=None		,					a__     :     Any=None		,					a__     :     Union[str, Any]=None		,					)						->		List[Any]:
       if attention_mask is None:
              UpperCamelCase_											=				tf.cast(tf.math.not_equal(a__		,					config.pad_token_id   )		,					tf.inta   )
       if decoder_attention_mask is None:
              UpperCamelCase_											=				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:
              UpperCamelCase_											=				tf.ones((config.encoder_layers, config.encoder_attention_heads)   )
       if decoder_head_mask is None:
              UpperCamelCase_											=				tf.ones((config.decoder_layers, config.decoder_attention_heads)   )
       if cross_attn_head_mask is None:
              UpperCamelCase_											=				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 						lowercase_				(				lowerCamelCase_							,		lowerCamelCase_							,		unittest.TestCase ):
 A__  : Any      					=			(TFBlenderbotForConditionalGeneration, TFBlenderbotModel) if is_tf_available() else ()
 A__  : Any      					=			(TFBlenderbotForConditionalGeneration,) if is_tf_available() else ()
 A__  : Optional[Any]      					=			(
     {
         """conversational""": TFBlenderbotForConditionalGeneration,
         """feature-extraction""": TFBlenderbotModel,
         """summarization""": TFBlenderbotForConditionalGeneration,
         """text2text-generation""": TFBlenderbotForConditionalGeneration,
         """translation""": TFBlenderbotForConditionalGeneration,
     }
     if is_tf_available()
     else {}
 )
 A__  : Tuple      					=			True
 A__  : Any      					=			False
 A__  : Tuple      					=			False

 def lowerCamelCase_ (					self   ):




        """simple docstring"""


        UpperCamelCase_											=				TFBlenderbotModelTester(self   )
        UpperCamelCase_											=				ConfigTester(self   ,		config_class=__a   )

 def lowerCamelCase_ (					self   ):




        """simple docstring"""


        self.config_tester.run_common_tests()

 def lowerCamelCase_ (					self   ):




        """simple docstring"""


        UpperCamelCase_											=				self.model_tester.prepare_config_and_inputs_for_common()
        self.model_tester.check_decoder_model_past_large_inputs(*__a   )



@require_tokenizers
@require_tf
class 						lowercase_				(				unittest.TestCase ):
 A__  : Tuple      					=			["""My friends are cool but they eat too many carbs."""]
 A__  : List[str]      					=			"""facebook/blenderbot-400M-distill"""

 @cached_property
 def lowerCamelCase_ (					self   ):




        """simple docstring"""


        return BlenderbotTokenizer.from_pretrained(self.model_name   )

 @cached_property
 def lowerCamelCase_ (					self   ):




        """simple docstring"""


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

 @slow
 def lowerCamelCase_ (					self   ):




        """simple docstring"""


        UpperCamelCase_											=				self.tokenizer(self.src_text   ,		return_tensors="""tf"""   )
        UpperCamelCase_											=				self.model.generate(
            model_inputs.input_ids   ,		)
        UpperCamelCase_											=				self.tokenizer.batch_decode(generated_ids.numpy()   ,		skip_special_tokens=__a   )[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 diffusers.pipelines.pipeline_utils import is_safetensors_compatible
class 				__SCREAMING_SNAKE_CASE  (unittest.TestCase      ):






				"""simple docstring"""







				def   UpperCamelCase__   (  self     :					str					):
								_a								= [
								    "safety_checker/pytorch_model.bin",
								    "safety_checker/model.safetensors",
								    "vae/diffusion_pytorch_model.bin",
								    "vae/diffusion_pytorch_model.safetensors",
								    "text_encoder/pytorch_model.bin",
								    "text_encoder/model.safetensors",
								    "unet/diffusion_pytorch_model.bin",
								    "unet/diffusion_pytorch_model.safetensors",
								]
								self.assertTrue(is_safetensors_compatible(__a					)					)




				def   UpperCamelCase__   (  self     :					List[str]					):
								_a								= [
								    "unet/diffusion_pytorch_model.bin",
								    "unet/diffusion_pytorch_model.safetensors",
								]
								self.assertTrue(is_safetensors_compatible(__a					)					)




				def   UpperCamelCase__   (  self     :					List[str]					):
								_a								= [
								    "safety_checker/pytorch_model.bin",
								    "safety_checker/model.safetensors",
								    "vae/diffusion_pytorch_model.bin",
								    "vae/diffusion_pytorch_model.safetensors",
								    "text_encoder/pytorch_model.bin",
								    "text_encoder/model.safetensors",
								    "unet/diffusion_pytorch_model.bin",
								    # Removed: 'unet/diffusion_pytorch_model.safetensors',
								]
								self.assertFalse(is_safetensors_compatible(__a					)					)




				def   UpperCamelCase__   (  self     :					List[str]					):
								_a								= [
								    "text_encoder/pytorch_model.bin",
								    "text_encoder/model.safetensors",
								]
								self.assertTrue(is_safetensors_compatible(__a					)					)




				def   UpperCamelCase__   (  self     :					Optional[Any]					):
								_a								= [
								    "safety_checker/pytorch_model.bin",
								    "safety_checker/model.safetensors",
								    "vae/diffusion_pytorch_model.bin",
								    "vae/diffusion_pytorch_model.safetensors",
								    "text_encoder/pytorch_model.bin",
								    # Removed: 'text_encoder/model.safetensors',
								    "unet/diffusion_pytorch_model.bin",
								    "unet/diffusion_pytorch_model.safetensors",
								]
								self.assertFalse(is_safetensors_compatible(__a					)					)




				def   UpperCamelCase__   (  self     :					str					):
								_a								= [
								    "safety_checker/pytorch_model.fp16.bin",
								    "safety_checker/model.fp16.safetensors",
								    "vae/diffusion_pytorch_model.fp16.bin",
								    "vae/diffusion_pytorch_model.fp16.safetensors",
								    "text_encoder/pytorch_model.fp16.bin",
								    "text_encoder/model.fp16.safetensors",
								    "unet/diffusion_pytorch_model.fp16.bin",
								    "unet/diffusion_pytorch_model.fp16.safetensors",
								]
								_a								= "fp16"
								self.assertTrue(is_safetensors_compatible(__a					,				variant=__a					)					)




				def   UpperCamelCase__   (  self     :					Any					):
								_a								= [
								    "unet/diffusion_pytorch_model.fp16.bin",
								    "unet/diffusion_pytorch_model.fp16.safetensors",
								]
								_a								= "fp16"
								self.assertTrue(is_safetensors_compatible(__a					,				variant=__a					)					)




				def   UpperCamelCase__   (  self     :					Any					):
								# pass variant but use the non-variant filenames
								_a								= [
								    "unet/diffusion_pytorch_model.bin",
								    "unet/diffusion_pytorch_model.safetensors",
								]
								_a								= "fp16"
								self.assertTrue(is_safetensors_compatible(__a					,				variant=__a					)					)




				def   UpperCamelCase__   (  self     :					Optional[Any]					):
								_a								= [
								    "safety_checker/pytorch_model.fp16.bin",
								    "safety_checker/model.fp16.safetensors",
								    "vae/diffusion_pytorch_model.fp16.bin",
								    "vae/diffusion_pytorch_model.fp16.safetensors",
								    "text_encoder/pytorch_model.fp16.bin",
								    "text_encoder/model.fp16.safetensors",
								    "unet/diffusion_pytorch_model.fp16.bin",
								    # Removed: 'unet/diffusion_pytorch_model.fp16.safetensors',
								]
								_a								= "fp16"
								self.assertFalse(is_safetensors_compatible(__a					,				variant=__a					)					)




				def   UpperCamelCase__   (  self     :					Dict					):
								_a								= [
								    "text_encoder/pytorch_model.fp16.bin",
								    "text_encoder/model.fp16.safetensors",
								]
								_a								= "fp16"
								self.assertTrue(is_safetensors_compatible(__a					,				variant=__a					)					)




				def   UpperCamelCase__   (  self     :					List[str]					):
								# pass variant but use the non-variant filenames
								_a								= [
								    "text_encoder/pytorch_model.bin",
								    "text_encoder/model.safetensors",
								]
								_a								= "fp16"
								self.assertTrue(is_safetensors_compatible(__a					,				variant=__a					)					)




				def   UpperCamelCase__   (  self     :					Optional[int]					):
								_a								= [
								    "safety_checker/pytorch_model.fp16.bin",
								    "safety_checker/model.fp16.safetensors",
								    "vae/diffusion_pytorch_model.fp16.bin",
								    "vae/diffusion_pytorch_model.fp16.safetensors",
								    "text_encoder/pytorch_model.fp16.bin",
								    # 'text_encoder/model.fp16.safetensors',
								    "unet/diffusion_pytorch_model.fp16.bin",
								    "unet/diffusion_pytorch_model.fp16.safetensors",
								]
								_a								= "fp16"
								self.assertFalse(is_safetensors_compatible(__a					,				variant=__a					)					)
