

import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    PNDMScheduler,
    StableDiffusionLDMaDPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS


enable_full_determinism()




class   UpperCamelCase			(  unittest.TestCase ):
    __UpperCamelCase    =StableDiffusionLDMaDPipeline
    __UpperCamelCase    =TEXT_TO_IMAGE_PARAMS
    __UpperCamelCase    =TEXT_TO_IMAGE_BATCH_PARAMS
    __UpperCamelCase    =TEXT_TO_IMAGE_IMAGE_PARAMS

    def UpperCamelCase      (      self			:   int     ):

     """simple docstring"""






     torch.manual_seed(0     )
     SCREAMING_SNAKE_CASE		    =      UNetaDConditionModel(
         block_out_channels=(3_2, 6_4)				,     layers_per_block=2				,     sample_size=3_2				,     in_channels=4				,     out_channels=4				,     down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')				,     up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')				,     cross_attention_dim=3_2				,     )
     SCREAMING_SNAKE_CASE		    =      DDIMScheduler(
         beta_start=0.00_085				,     beta_end=0.012				,     beta_schedule='scaled_linear'				,     clip_sample=snake_case__				,     set_alpha_to_one=snake_case__				,     )
     torch.manual_seed(0     )
     SCREAMING_SNAKE_CASE		    =      AutoencoderKL(
         block_out_channels=[3_2, 6_4]				,     in_channels=6				,     out_channels=6				,     down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']				,     up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']				,     latent_channels=4				,     )
     torch.manual_seed(0     )
     SCREAMING_SNAKE_CASE		    =      CLIPTextConfig(
         bos_token_id=0				,     eos_token_id=2				,     hidden_size=3_2				,     intermediate_size=3_7				,     layer_norm_eps=1E-05				,     num_attention_heads=4				,     num_hidden_layers=5				,     pad_token_id=1				,     vocab_size=1_0_0_0				,     )
     SCREAMING_SNAKE_CASE		    =      CLIPTextModel(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'     )

     SCREAMING_SNAKE_CASE		    =      {
         'unet': unet,
         'scheduler': scheduler,
         'vae': vae,
         'text_encoder': text_encoder,
         'tokenizer': tokenizer,
         'safety_checker': None,
         'feature_extractor': None,
     }
     return components

    def UpperCamelCase      (      self			:   str				,     snake_case__			:   Any				,     snake_case__			:   Tuple=0     ):

     """simple docstring"""






     if str(snake_case__     ).startswith('mps'     ):
      SCREAMING_SNAKE_CASE		    =      torch.manual_seed(snake_case__     )
     else:
      SCREAMING_SNAKE_CASE		    =      torch.Generator(device=snake_case__     ).manual_seed(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      {
         'prompt': 'A painting of a squirrel eating a burger',
         'generator': generator,
         'num_inference_steps': 2,
         'guidance_scale': 6.0,
         'output_type': 'numpy',
     }
     return inputs

    def UpperCamelCase      (      self			:   int     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      'cpu'  # ensure determinism for the device-dependent torch.Generator

     SCREAMING_SNAKE_CASE		    =      self.get_dummy_components()
     SCREAMING_SNAKE_CASE		    =      StableDiffusionLDMaDPipeline(**snake_case__     )
     SCREAMING_SNAKE_CASE		    =      ldmad_pipe.to(snake_case__     )
     ldmad_pipe.set_progress_bar_config(disable=snake_case__     )

     SCREAMING_SNAKE_CASE		    =      self.get_dummy_inputs(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      ldmad_pipe(**snake_case__     )
     SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE		    =      output.rgb, output.depth

     SCREAMING_SNAKE_CASE		    =      rgb[0, -3:, -3:, -1]
     SCREAMING_SNAKE_CASE		    =      depth[0, -3:, -1]

     assert rgb.shape == (1, 6_4, 6_4, 3)
     assert depth.shape == (1, 6_4, 6_4)

     SCREAMING_SNAKE_CASE		    =      np.array(
         [0.37_338_176, 0.70_247, 0.74_203_193, 0.51_643_604, 0.58_256_793, 0.60_932_136, 0.4_181_095, 0.48_355_877, 0.46_535_262]     )
     SCREAMING_SNAKE_CASE		    =      np.array([103.46_727, 85.812_004, 87.849_236]     )

     assert np.abs(image_slice_rgb.flatten() - expected_slice_rgb     ).max() < 1E-2
     assert np.abs(image_slice_depth.flatten() - expected_slice_depth     ).max() < 1E-2

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

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      self.get_dummy_components()
     SCREAMING_SNAKE_CASE		    =      StableDiffusionLDMaDPipeline(**snake_case__     )
     SCREAMING_SNAKE_CASE		    =      ldmad_pipe.to(snake_case__     )
     ldmad_pipe.set_progress_bar_config(disable=snake_case__     )

     SCREAMING_SNAKE_CASE		    =      self.get_dummy_inputs(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      3 * [inputs['prompt']]

     # forward
     SCREAMING_SNAKE_CASE		    =      ldmad_pipe(**snake_case__     )
     SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE		    =      output.rgb, output.depth
     SCREAMING_SNAKE_CASE		    =      rgb_slice_a[0, -3:, -3:, -1]
     SCREAMING_SNAKE_CASE		    =      depth_slice_a[0, -3:, -1]

     SCREAMING_SNAKE_CASE		    =      self.get_dummy_inputs(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      3 * [inputs.pop('prompt'     )]

     SCREAMING_SNAKE_CASE		    =      ldmad_pipe.tokenizer(
         snake_case__				,     padding='max_length'				,     max_length=ldmad_pipe.tokenizer.model_max_length				,     truncation=snake_case__				,     return_tensors='pt'				,     )
     SCREAMING_SNAKE_CASE		    =      text_inputs['input_ids'].to(snake_case__     )

     SCREAMING_SNAKE_CASE		    =      ldmad_pipe.text_encoder(snake_case__     )[0]

     SCREAMING_SNAKE_CASE		    =      prompt_embeds

     # forward
     SCREAMING_SNAKE_CASE		    =      ldmad_pipe(**snake_case__     )
     SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE		    =      output.rgb, output.depth
     SCREAMING_SNAKE_CASE		    =      rgb_slice_a[0, -3:, -3:, -1]
     SCREAMING_SNAKE_CASE		    =      depth_slice_a[0, -3:, -1]

     assert np.abs(rgb_slice_a.flatten() - rgb_slice_a.flatten()     ).max() < 1E-4
     assert np.abs(depth_slice_a.flatten() - depth_slice_a.flatten()     ).max() < 1E-4




    def UpperCamelCase      (      self			:   List[str]     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      'cpu'  # ensure determinism for the device-dependent torch.Generator
     SCREAMING_SNAKE_CASE		    =      self.get_dummy_components()
     SCREAMING_SNAKE_CASE		    =      PNDMScheduler(skip_prk_steps=snake_case__     )
     SCREAMING_SNAKE_CASE		    =      StableDiffusionLDMaDPipeline(**snake_case__     )
     SCREAMING_SNAKE_CASE		    =      ldmad_pipe.to(snake_case__     )
     ldmad_pipe.set_progress_bar_config(disable=snake_case__     )

     SCREAMING_SNAKE_CASE		    =      self.get_dummy_inputs(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      'french fries'
     SCREAMING_SNAKE_CASE		    =      ldmad_pipe(**snake_case__				,     negative_prompt=snake_case__     )

     SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE		    =      output.rgb, output.depth
     SCREAMING_SNAKE_CASE		    =      rgb[0, -3:, -3:, -1]
     SCREAMING_SNAKE_CASE		    =      depth[0, -3:, -1]

     assert rgb.shape == (1, 6_4, 6_4, 3)
     assert depth.shape == (1, 6_4, 6_4)

     SCREAMING_SNAKE_CASE		    =      np.array(
         [0.37_044, 0.71_811_503, 0.7_223_251, 0.48_603_675, 0.5_638_391, 0.6_364_948, 0.42_833_704, 0.4_901_315, 0.47_926_217]     )
     SCREAMING_SNAKE_CASE		    =      np.array([107.84_738, 84.62_802, 89.962_135]     )
     assert np.abs(rgb_slice.flatten() - expected_slice_rgb     ).max() < 1E-2
     assert np.abs(depth_slice.flatten() - expected_slice_depth     ).max() < 1E-2




@slow
@require_torch_gpu
class   UpperCamelCase			(  unittest.TestCase ):

    def UpperCamelCase      (      self			:   Optional[int]     ):

     """simple docstring"""






     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()

    def UpperCamelCase      (      self			:   Optional[int]				,     snake_case__			:   List[str]				,     snake_case__			:   Dict="cpu"				,     snake_case__			:   Tuple=torch.floataa				,     snake_case__			:   Tuple=0     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      torch.Generator(device=snake_case__     ).manual_seed(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      np.random.RandomState(snake_case__     ).standard_normal((1, 4, 6_4, 6_4)     )
     SCREAMING_SNAKE_CASE		    =      torch.from_numpy(snake_case__     ).to(device=snake_case__				,     dtype=snake_case__     )
     SCREAMING_SNAKE_CASE		    =      {
         'prompt': 'a photograph of an astronaut riding a horse',
         'latents': latents,
         'generator': generator,
         'num_inference_steps': 3,
         'guidance_scale': 7.5,
         'output_type': 'numpy',
     }
     return inputs




    def UpperCamelCase      (      self			:   Any     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'     )
     SCREAMING_SNAKE_CASE		    =      ldmad_pipe.to(snake_case__     )
     ldmad_pipe.set_progress_bar_config(disable=snake_case__     )

     SCREAMING_SNAKE_CASE		    =      self.get_inputs(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      ldmad_pipe(**snake_case__     )
     SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE		    =      output.rgb, output.depth
     SCREAMING_SNAKE_CASE		    =      rgb[0, -3:, -3:, -1].flatten()
     SCREAMING_SNAKE_CASE		    =      rgb[0, -3:, -1].flatten()

     assert rgb.shape == (1, 5_1_2, 5_1_2, 3)
     assert depth.shape == (1, 5_1_2, 5_1_2)

     SCREAMING_SNAKE_CASE		    =      np.array(
         [0.53_805_465, 0.56_707_305, 0.5_486_515, 0.57_012_236, 0.5_814_511, 0.56_253_487, 0.54_843_014, 0.55_092_263, 0.6_459_706]     )
     SCREAMING_SNAKE_CASE		    =      np.array(
         [0.9_263_781, 0.6_678_672, 0.5_486_515, 0.92_202_145, 0.67_831_135, 0.56_253_487, 0.9_241_694, 0.7_551_478, 0.6_459_706]     )
     assert np.abs(rgb_slice - expected_slice_rgb     ).max() < 3E-3
     assert np.abs(depth_slice - expected_slice_depth     ).max() < 3E-3




@nightly
@require_torch_gpu
class   UpperCamelCase			(  unittest.TestCase ):

    def UpperCamelCase      (      self			:   Tuple     ):

     """simple docstring"""






     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()

    def UpperCamelCase      (      self			:   Tuple				,     snake_case__			:   Optional[int]				,     snake_case__			:   Any="cpu"				,     snake_case__			:   Optional[int]=torch.floataa				,     snake_case__			:   int=0     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      torch.Generator(device=snake_case__     ).manual_seed(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      np.random.RandomState(snake_case__     ).standard_normal((1, 4, 6_4, 6_4)     )
     SCREAMING_SNAKE_CASE		    =      torch.from_numpy(snake_case__     ).to(device=snake_case__				,     dtype=snake_case__     )
     SCREAMING_SNAKE_CASE		    =      {
         'prompt': 'a photograph of an astronaut riding a horse',
         'latents': latents,
         'generator': generator,
         'num_inference_steps': 5_0,
         'guidance_scale': 7.5,
         'output_type': 'numpy',
     }
     return inputs

    def UpperCamelCase      (      self			:   List[str]     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'     ).to(snake_case__     )
     ldmad_pipe.set_progress_bar_config(disable=snake_case__     )

     SCREAMING_SNAKE_CASE		    =      self.get_inputs(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      ldmad_pipe(**snake_case__     )
     SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE		    =      output.rgb, output.depth

     SCREAMING_SNAKE_CASE		    =      0.495_586
     SCREAMING_SNAKE_CASE		    =      0.33_795_515
     SCREAMING_SNAKE_CASE		    =      112.48_518
     SCREAMING_SNAKE_CASE		    =      98.489_746
     assert np.abs(expected_rgb_mean - rgb.mean()     ) < 1E-3
     assert np.abs(expected_rgb_std - rgb.std()     ) < 1E-3
     assert np.abs(expected_depth_mean - depth.mean()     ) < 1E-3
     assert np.abs(expected_depth_std - depth.std()     ) < 1E-3




    def UpperCamelCase      (      self			:   List[str]     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d-4c'     ).to(snake_case__     )
     ldmad_pipe.set_progress_bar_config(disable=snake_case__     )

     SCREAMING_SNAKE_CASE		    =      self.get_inputs(snake_case__     )
     SCREAMING_SNAKE_CASE		    =      ldmad_pipe(**snake_case__     )
     SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE		    =      output.rgb, output.depth

     SCREAMING_SNAKE_CASE		    =      0.4_194_127
     SCREAMING_SNAKE_CASE		    =      0.35_375_586
     SCREAMING_SNAKE_CASE		    =      0.5_638_502
     SCREAMING_SNAKE_CASE		    =      0.34_686_103

     assert rgb.shape == (1, 5_1_2, 5_1_2, 3)
     assert depth.shape == (1, 5_1_2, 5_1_2, 1)
     assert np.abs(expected_rgb_mean - rgb.mean()     ) < 1E-3
     assert np.abs(expected_rgb_std - rgb.std()     ) < 1E-3
     assert np.abs(expected_depth_mean - depth.mean()     ) < 1E-3
     assert np.abs(expected_depth_std - depth.std()     ) < 1E-3




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			:
    __UpperCamelCase    =BlenderbotSmallConfig
    __UpperCamelCase    ={}
    __UpperCamelCase    ="gelu"

    def __init__(      self			:   Optional[int]				,     snake_case__			:   Dict				,     snake_case__			:   Optional[int]=1_3				,     snake_case__			:   Optional[int]=7				,     snake_case__			:   Optional[Any]=True				,     snake_case__			:   List[str]=False				,     snake_case__			:   int=9_9				,     snake_case__			:   Any=3_2				,     snake_case__			:   Tuple=2				,     snake_case__			:   Optional[int]=4				,     snake_case__			:   int=3_7				,     snake_case__			:   Optional[Any]=0.1				,     snake_case__			:   int=0.1				,     snake_case__			:   Any=2_0				,     snake_case__			:   List[Any]=2				,     snake_case__			:   int=1				,     snake_case__			:   int=0				,     ):

     """simple docstring"""






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

     SCREAMING_SNAKE_CASE		    =      hidden_dropout_prob
     SCREAMING_SNAKE_CASE		    =      attention_probs_dropout_prob
     SCREAMING_SNAKE_CASE		    =      max_position_embeddings
     SCREAMING_SNAKE_CASE		    =      eos_token_id
     SCREAMING_SNAKE_CASE		    =      pad_token_id
     SCREAMING_SNAKE_CASE		    =      bos_token_id

    def UpperCamelCase      (      self			:   str     ):

     """simple docstring"""






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

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

     SCREAMING_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				,     )
     SCREAMING_SNAKE_CASE		    =      prepare_blenderbot_small_inputs_dict(snake_case__				,     snake_case__				,     snake_case__     )
     return config, inputs_dict




    def UpperCamelCase      (      self			:   Dict				,     snake_case__			:   Tuple				,     snake_case__			:   Any     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      TFBlenderbotSmallModel(config=snake_case__     ).get_decoder()
     SCREAMING_SNAKE_CASE		    =      inputs_dict['input_ids']

     SCREAMING_SNAKE_CASE		    =      input_ids[:1, :]
     SCREAMING_SNAKE_CASE		    =      inputs_dict['attention_mask'][:1, :]
     SCREAMING_SNAKE_CASE		    =      inputs_dict['head_mask']
     SCREAMING_SNAKE_CASE		    =      1

     # first forward pass
     SCREAMING_SNAKE_CASE		    =      model(snake_case__				,     attention_mask=snake_case__				,     head_mask=snake_case__				,     use_cache=snake_case__     )

     SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE		    =      outputs.to_tuple()

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

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

     SCREAMING_SNAKE_CASE		    =      model(snake_case__				,     attention_mask=snake_case__     )[0]
     SCREAMING_SNAKE_CASE		    =      model(snake_case__				,     attention_mask=snake_case__				,     past_key_values=snake_case__     )[0]

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

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

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


def 							__lowerCAmelCase     (							_UpperCamelCase		:	Dict		,				_UpperCamelCase		:	Optional[int]		,				_UpperCamelCase		:	str		,				_UpperCamelCase		:	Dict=None		,				_UpperCamelCase		:	Optional[int]=None		,				_UpperCamelCase		:	Any=None		,				_UpperCamelCase		:	List[str]=None		,				_UpperCamelCase		:	Optional[Any]=None		,				)		->				str:
 '''simple docstring'''


 if attention_mask is None:
  SCREAMING_SNAKE_CASE		    =      tf.cast(tf.math.not_equal(_UpperCamelCase		,				config.pad_token_id						)		,				tf.inta						)
 if decoder_attention_mask is None:
  SCREAMING_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:
  SCREAMING_SNAKE_CASE		    =      tf.ones((config.encoder_layers, config.encoder_attention_heads)						)
 if decoder_head_mask is None:
  SCREAMING_SNAKE_CASE		    =      tf.ones((config.decoder_layers, config.decoder_attention_heads)						)
 if cross_attn_head_mask is None:
  SCREAMING_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			(  SCREAMING_SNAKE_CASE						,      SCREAMING_SNAKE_CASE						,      unittest.TestCase ):
    __UpperCamelCase    =(
        (TFBlenderbotSmallForConditionalGeneration, TFBlenderbotSmallModel) if is_tf_available() else ()
    )
    __UpperCamelCase    =(TFBlenderbotSmallForConditionalGeneration,) if is_tf_available() else ()
    __UpperCamelCase    =(
        {
            "conversational": TFBlenderbotSmallForConditionalGeneration,
            "feature-extraction": TFBlenderbotSmallModel,
            "summarization": TFBlenderbotSmallForConditionalGeneration,
            "text2text-generation": TFBlenderbotSmallForConditionalGeneration,
            "translation": TFBlenderbotSmallForConditionalGeneration,
        }
        if is_tf_available()
        else {}
    )
    __UpperCamelCase    =True
    __UpperCamelCase    =False
    __UpperCamelCase    =False

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

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      TFBlenderbotSmallModelTester(self     )
     SCREAMING_SNAKE_CASE		    =      ConfigTester(self				,     config_class=snake_case__     )

    def UpperCamelCase      (      self			:   Any     ):

     """simple docstring"""






     self.config_tester.run_common_tests()




    def UpperCamelCase      (      self			:   List[Any]     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      self.model_tester.prepare_config_and_inputs_for_common()
     self.model_tester.check_decoder_model_past_large_inputs(*snake_case__     )




@require_tokenizers
@require_tf
class   UpperCamelCase			(  unittest.TestCase ):
    __UpperCamelCase    =[
        "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?"
    ]
    __UpperCamelCase    ="facebook/blenderbot_small-90M"

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

     """simple docstring"""






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

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

     """simple docstring"""






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




    @slow
    def UpperCamelCase      (      self			:   Tuple     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      self.tokenizer(self.src_text				,     return_tensors='tf'     )
     SCREAMING_SNAKE_CASE		    =      self.model.generate(
         model_inputs.input_ids				,     attention_mask=model_inputs.attention_mask				,     num_beams=2				,     use_cache=snake_case__				,     )
     SCREAMING_SNAKE_CASE		    =      self.tokenizer.batch_decode(generated_ids.numpy()				,     skip_special_tokens=snake_case__     )[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.",
     )

