




from __future__ import annotations

import unittest

from transformers import BlenderbotSmallConfig, BlenderbotSmallTokenizer, 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, TFBlenderbotSmallForConditionalGeneration, TFBlenderbotSmallModel
@require_tf
class 						UpperCAmelCase  :







						'''simple docstring'''





						lowerCAmelCase_							=  BlenderbotSmallConfig
						lowerCAmelCase_							=  {}
						lowerCAmelCase_							=  '''gelu'''






						def __init__(		self				:	Union[str, Any]			,      __lowercase				:	Dict			,      __lowercase				:	int=13			,      __lowercase				:	str=7			,      __lowercase				:	Tuple=True			,      __lowercase				:	int=False			,      __lowercase				:	Optional[int]=99			,      __lowercase				:	List[str]=32			,      __lowercase				:	List[Any]=2			,      __lowercase				:	Dict=4			,      __lowercase				:	str=37			,      __lowercase				:	Optional[int]=0.1			,      __lowercase				:	str=0.1			,      __lowercase				:	List[Any]=20			,      __lowercase				:	List[Any]=2			,      __lowercase				:	Tuple=1			,      __lowercase				:	Tuple=0			,      ):






							"""simple docstring"""






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

							snake_case_						       =	hidden_dropout_prob
							snake_case_						       =	attention_probs_dropout_prob
							snake_case_						       =	max_position_embeddings
							snake_case_						       =	eos_token_id
							snake_case_						       =	pad_token_id
							snake_case_						       =	bos_token_id






						def 			snake_case__		(		self				:	List[str]		):






							"""simple docstring"""






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

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

							snake_case_						       =	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_						       =	prepare_blenderbot_small_inputs_dict(__lowercase			,      __lowercase			,      __lowercase		)
							return config, inputs_dict






						def 			snake_case__		(		self				:	Optional[Any]			,      __lowercase				:	List[str]			,      __lowercase				:	List[str]		):






							"""simple docstring"""






							snake_case_						       =	TFBlenderbotSmallModel(config=__lowercase		).get_decoder()
							snake_case_						       =	inputs_dict["input_ids"]

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

							# first forward pass
							snake_case_						       =	model(__lowercase			,      attention_mask=__lowercase			,      head_mask=__lowercase			,      use_cache=__lowercase		)

							snake_case_       , snake_case_						       =	outputs.to_tuple()

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

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

							snake_case_						       =	model(__lowercase			,      attention_mask=__lowercase		)[0]
							snake_case_						       =	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_						       =	int(ids_tensor((1,)			,      output_from_past.shape[-1]		)		)
							snake_case_						       =	output_from_no_past[:, -3:, random_slice_idx]
							snake_case_						       =	output_from_past[:, :, random_slice_idx]

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







def 		lowerCamelCase__       (						_A	,      _A	,      _A	,      _A=None	,      _A=None	,      _A=None	,      _A=None	,      _A=None	,      ):







	'''simple docstring'''





	if attention_mask is None:
		snake_case_						       =	tf.cast(tf.math.not_equal(_A	,      config.pad_token_id				)	,      tf.inta				)
	if decoder_attention_mask is None:
		snake_case_						       =	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_						       =	tf.ones((config.encoder_layers, config.encoder_attention_heads)				)
	if decoder_head_mask is None:
		snake_case_						       =	tf.ones((config.decoder_layers, config.decoder_attention_heads)				)
	if cross_attn_head_mask is None:
		snake_case_						       =	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 						UpperCAmelCase  (							UpperCAmelCase__			,  UpperCAmelCase__			,  unittest.TestCase     ):







						'''simple docstring'''





						lowerCAmelCase_							=  (
						    (TFBlenderbotSmallForConditionalGeneration, TFBlenderbotSmallModel) if is_tf_available() else ()
						)
						lowerCAmelCase_							=  (TFBlenderbotSmallForConditionalGeneration,) if is_tf_available() else ()
						lowerCAmelCase_							=  (
						    {
						        '''conversational''': TFBlenderbotSmallForConditionalGeneration,
						        '''feature-extraction''': TFBlenderbotSmallModel,
						        '''summarization''': TFBlenderbotSmallForConditionalGeneration,
						        '''text2text-generation''': TFBlenderbotSmallForConditionalGeneration,
						        '''translation''': TFBlenderbotSmallForConditionalGeneration,
						    }
						    if is_tf_available()
						    else {}
						)
						lowerCAmelCase_							=  True
						lowerCAmelCase_							=  False
						lowerCAmelCase_							=  False






						def 			snake_case__		(		self				:	Optional[int]		):






							"""simple docstring"""






							snake_case_						       =	TFBlenderbotSmallModelTester(self		)
							snake_case_						       =	ConfigTester(self			,      config_class=__lowercase		)






						def 			snake_case__		(		self				:	Tuple		):






							"""simple docstring"""






							self.config_tester.run_common_tests()






						def 			snake_case__		(		self				:	List[str]		):






							"""simple docstring"""






							snake_case_						       =	self.model_tester.prepare_config_and_inputs_for_common()
							self.model_tester.check_decoder_model_past_large_inputs(*__lowercase		)







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







						'''simple docstring'''





						lowerCAmelCase_							=  [
						    '''Social anxiety\nWow, I am never shy. Do you have anxiety?\nYes. I end up sweating and blushing and feel like '''
						    '''  i\'m going to throw up.\nand why is that?'''
						]
						lowerCAmelCase_							=  '''facebook/blenderbot_small-90M'''






						@cached_property
						def 			snake_case__		(		self				:	Union[str, Any]		):






							"""simple docstring"""






							return BlenderbotSmallTokenizer.from_pretrained("facebook/blenderbot-90M"		)






						@cached_property
						def 			snake_case__		(		self				:	Optional[int]		):






							"""simple docstring"""






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






						@slow
						def 			snake_case__		(		self				:	Tuple		):






							"""simple docstring"""






							snake_case_						       =	self.tokenizer(self.src_text			,      return_tensors="tf"		)
							snake_case_						       =	self.model.generate(
							    model_inputs.input_ids			,      attention_mask=model_inputs.attention_mask			,      num_beams=2			,      use_cache=__lowercase			,      )
							snake_case_						       =	self.tokenizer.batch_decode(generated_ids.numpy()			,      skip_special_tokens=__lowercase		)[0]
							assert generated_words in (
							    "i don't know. i just feel like i'm going to throw up. it's not fun.",
							    "i'm not sure. i just feel like i've been feeling like i have to be in a certain place",
							    "i'm not sure. i just feel like i've been in a bad situation.",
							)






import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np
class 						UpperCAmelCase  :







						'''simple docstring'''





						def 			snake_case__		(		self				:	List[Any]		):






							"""simple docstring"""






							torch.manual_seed(0		)
							snake_case_						       =	TaEncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5"		)

							torch.manual_seed(0		)
							snake_case_						       =	AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"		)

							torch.manual_seed(0		)
							snake_case_						       =	UNetaDConditionModel(
							    sample_size=32			,      layers_per_block=1			,      block_out_channels=[32, 64]			,      down_block_types=[
							        "ResnetDownsampleBlock2D",
							        "SimpleCrossAttnDownBlock2D",
							    ]			,      mid_block_type="UNetMidBlock2DSimpleCrossAttn"			,      up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"]			,      in_channels=3			,      out_channels=6			,      cross_attention_dim=32			,      encoder_hid_dim=32			,      attention_head_dim=8			,      addition_embed_type="text"			,      addition_embed_type_num_heads=2			,      cross_attention_norm="group_norm"			,      resnet_time_scale_shift="scale_shift"			,      act_fn="gelu"			,      )
							unet.set_attn_processor(AttnAddedKVProcessor()		)  # For reproducibility tests

							torch.manual_seed(0		)
							snake_case_						       =	DDPMScheduler(
							    num_train_timesteps=10_00			,      beta_schedule="squaredcos_cap_v2"			,      beta_start=0.0001			,      beta_end=0.02			,      thresholding=__lowercase			,      dynamic_thresholding_ratio=0.95			,      sample_max_value=1.0			,      prediction_type="epsilon"			,      variance_type="learned_range"			,      )

							torch.manual_seed(0		)
							snake_case_						       =	IFWatermarker()

							return {
							    "text_encoder": text_encoder,
							    "tokenizer": tokenizer,
							    "unet": unet,
							    "scheduler": scheduler,
							    "watermarker": watermarker,
							    "safety_checker": None,
							    "feature_extractor": None,
							}






						def 			snake_case__		(		self				:	List[str]		):






							"""simple docstring"""






							torch.manual_seed(0		)
							snake_case_						       =	TaEncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5"		)

							torch.manual_seed(0		)
							snake_case_						       =	AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"		)

							torch.manual_seed(0		)
							snake_case_						       =	UNetaDConditionModel(
							    sample_size=32			,      layers_per_block=[1, 2]			,      block_out_channels=[32, 64]			,      down_block_types=[
							        "ResnetDownsampleBlock2D",
							        "SimpleCrossAttnDownBlock2D",
							    ]			,      mid_block_type="UNetMidBlock2DSimpleCrossAttn"			,      up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"]			,      in_channels=6			,      out_channels=6			,      cross_attention_dim=32			,      encoder_hid_dim=32			,      attention_head_dim=8			,      addition_embed_type="text"			,      addition_embed_type_num_heads=2			,      cross_attention_norm="group_norm"			,      resnet_time_scale_shift="scale_shift"			,      act_fn="gelu"			,      class_embed_type="timestep"			,      mid_block_scale_factor=1.414			,      time_embedding_act_fn="gelu"			,      time_embedding_dim=32			,      )
							unet.set_attn_processor(AttnAddedKVProcessor()		)  # For reproducibility tests

							torch.manual_seed(0		)
							snake_case_						       =	DDPMScheduler(
							    num_train_timesteps=10_00			,      beta_schedule="squaredcos_cap_v2"			,      beta_start=0.0001			,      beta_end=0.02			,      thresholding=__lowercase			,      dynamic_thresholding_ratio=0.95			,      sample_max_value=1.0			,      prediction_type="epsilon"			,      variance_type="learned_range"			,      )

							torch.manual_seed(0		)
							snake_case_						       =	DDPMScheduler(
							    num_train_timesteps=10_00			,      beta_schedule="squaredcos_cap_v2"			,      beta_start=0.0001			,      beta_end=0.02			,      )

							torch.manual_seed(0		)
							snake_case_						       =	IFWatermarker()

							return {
							    "text_encoder": text_encoder,
							    "tokenizer": tokenizer,
							    "unet": unet,
							    "scheduler": scheduler,
							    "image_noising_scheduler": image_noising_scheduler,
							    "watermarker": watermarker,
							    "safety_checker": None,
							    "feature_extractor": None,
							}






						def 			snake_case__		(		self				:	List[str]		):






							"""simple docstring"""






							snake_case_						       =	self.get_dummy_components()
							snake_case_						       =	self.pipeline_class(**__lowercase		)
							pipe.to(__lowercase		)
							pipe.set_progress_bar_config(disable=__lowercase		)

							snake_case_						       =	self.get_dummy_inputs(__lowercase		)

							snake_case_						       =	inputs["prompt"]
							snake_case_						       =	inputs["generator"]
							snake_case_						       =	inputs["num_inference_steps"]
							snake_case_						       =	inputs["output_type"]

							if "image" in inputs:
								snake_case_						       =	inputs["image"]
							else:
								snake_case_						       =	None

							if "mask_image" in inputs:
								snake_case_						       =	inputs["mask_image"]
							else:
								snake_case_						       =	None

							if "original_image" in inputs:
								snake_case_						       =	inputs["original_image"]
							else:
								snake_case_						       =	None

							snake_case_       , snake_case_						       =	pipe.encode_prompt(__lowercase		)

							# inputs with prompt converted to embeddings
							snake_case_						       =	{
							    "prompt_embeds": prompt_embeds,
							    "negative_prompt_embeds": negative_prompt_embeds,
							    "generator": generator,
							    "num_inference_steps": num_inference_steps,
							    "output_type": output_type,
							}

							if image is not None:
								snake_case_						       =	image

							if mask_image is not None:
								snake_case_						       =	mask_image

							if original_image is not None:
								snake_case_						       =	original_image

							# set all optional components to None
							for optional_component in pipe._optional_components:
								setattr(__lowercase			,      __lowercase			,      __lowercase		)

							snake_case_						       =	pipe(**__lowercase		)[0]

							with tempfile.TemporaryDirectory() as tmpdir:
								pipe.save_pretrained(__lowercase		)
								snake_case_						       =	self.pipeline_class.from_pretrained(__lowercase		)
								pipe_loaded.to(__lowercase		)
								pipe_loaded.set_progress_bar_config(disable=__lowercase		)

							pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()		)  # For reproducibility tests

							for optional_component in pipe._optional_components:
								self.assertTrue(
								    getattr(__lowercase			,      __lowercase		) is None			,      f"`{optional_component}` did not stay set to None after loading."			,      )

							snake_case_						       =	self.get_dummy_inputs(__lowercase		)

							snake_case_						       =	inputs["generator"]
							snake_case_						       =	inputs["num_inference_steps"]
							snake_case_						       =	inputs["output_type"]

							# inputs with prompt converted to embeddings
							snake_case_						       =	{
							    "prompt_embeds": prompt_embeds,
							    "negative_prompt_embeds": negative_prompt_embeds,
							    "generator": generator,
							    "num_inference_steps": num_inference_steps,
							    "output_type": output_type,
							}

							if image is not None:
								snake_case_						       =	image

							if mask_image is not None:
								snake_case_						       =	mask_image

							if original_image is not None:
								snake_case_						       =	original_image

							snake_case_						       =	pipe_loaded(**__lowercase		)[0]

							snake_case_						       =	np.abs(to_np(__lowercase		) - to_np(__lowercase		)		).max()
							self.assertLess(__lowercase			,      1E-4		)






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






							"""simple docstring"""






							snake_case_						       =	self.get_dummy_components()
							snake_case_						       =	self.pipeline_class(**__lowercase		)
							pipe.to(__lowercase		)
							pipe.set_progress_bar_config(disable=__lowercase		)

							snake_case_						       =	self.get_dummy_inputs(__lowercase		)
							snake_case_						       =	pipe(**__lowercase		)[0]

							with tempfile.TemporaryDirectory() as tmpdir:
								pipe.save_pretrained(__lowercase		)
								snake_case_						       =	self.pipeline_class.from_pretrained(__lowercase		)
								pipe_loaded.to(__lowercase		)
								pipe_loaded.set_progress_bar_config(disable=__lowercase		)

							pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()		)  # For reproducibility tests

							snake_case_						       =	self.get_dummy_inputs(__lowercase		)
							snake_case_						       =	pipe_loaded(**__lowercase		)[0]

							snake_case_						       =	np.abs(to_np(__lowercase		) - to_np(__lowercase		)		).max()
							self.assertLess(__lowercase			,      1E-4		)
