




import inspect
from typing import Callable, List, Optional, Union

import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextModel,
    CLIPTokenizer,
    WhisperForConditionalGeneration,
    WhisperProcessor,
)

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.utils import logging


UpperCAmelCase__     :   Union[str, Any]				     = logging.get_logger(__name__)  # pylint: disable=invalid-name







class a (     UpperCamelCase_  ):
						"""simple docstring"""

						def __init__( self :     Optional[int]				,							lowerCamelCase__ :     WhisperForConditionalGeneration				,							lowerCamelCase__ :     WhisperProcessor				,							lowerCamelCase__ :     AutoencoderKL				,							lowerCamelCase__ :     CLIPTextModel				,							lowerCamelCase__ :     CLIPTokenizer				,							lowerCamelCase__ :     UNetaDConditionModel				,							lowerCamelCase__ :     Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]				,							lowerCamelCase__ :     StableDiffusionSafetyChecker				,							lowerCamelCase__ :     CLIPImageProcessor				,							)    ->						Union[str, Any]:
													"""simple docstring"""


													super().__init__()

													if safety_checker is None:
																				logger.warning(
																				    f'You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure'
																				    ''' that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered'''
																				    ''' results in services or applications open to the public. Both the diffusers team and Hugging Face'''
																				    ''' strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling'''
																				    ''' it only for use-cases that involve analyzing network behavior or auditing its results. For more'''
																				    ''' information, please have a look at https://github.com/huggingface/diffusers/pull/254 .'''				)

													self.register_modules(
													    speech_model=UpperCamelCase__				,							speech_processor=UpperCamelCase__				,							vae=UpperCamelCase__				,							text_encoder=UpperCamelCase__				,							tokenizer=UpperCamelCase__				,							unet=UpperCamelCase__				,							scheduler=UpperCamelCase__				,							feature_extractor=UpperCamelCase__				,							)






						def     UpperCAmelCase_      ( self :     Any				,							lowerCamelCase__ :     Optional[Union[str, int]] = "auto"				)    ->						Dict:
													"""simple docstring"""


													if slice_size == "auto":
																				__lowercase      =    self.unet.config.attention_head_dim // 2
													self.unet.set_attention_slice(UpperCamelCase__				)






						def     UpperCAmelCase_      ( self :     Dict				)    ->						int:
													"""simple docstring"""


													self.enable_attention_slicing(UpperCamelCase__				)






						@torch.no_grad()
						def __call__( self :     List[str]				,							lowerCamelCase__ :     Union[str, Any]				,							lowerCamelCase__ :     Any=16_000				,							lowerCamelCase__ :     int = 512				,							lowerCamelCase__ :     int = 512				,							lowerCamelCase__ :     int = 50				,							lowerCamelCase__ :     float = 7.5				,							lowerCamelCase__ :     Optional[Union[str, List[str]]] = None				,							lowerCamelCase__ :     Optional[int] = 1				,							lowerCamelCase__ :     float = 0.0				,							lowerCamelCase__ :     Optional[torch.Generator] = None				,							lowerCamelCase__ :     Optional[torch.FloatTensor] = None				,							lowerCamelCase__ :     Optional[str] = "pil"				,							lowerCamelCase__ :     bool = True				,							lowerCamelCase__ :     Optional[Callable[[int, int, torch.FloatTensor], None]] = None				,							lowerCamelCase__ :     int = 1				,							**lowerCamelCase__ :     int				,							)    ->						Optional[Any]:
													"""simple docstring"""


													__lowercase      =    self.speech_processor.feature_extractor(
													    UpperCamelCase__				,							return_tensors='''pt'''				,							sampling_rate=UpperCamelCase__				).input_features.to(self.device				)
													__lowercase      =    self.speech_model.generate(UpperCamelCase__				,							max_length=480_000				)

													__lowercase      =    self.speech_processor.tokenizer.batch_decode(UpperCamelCase__				,							skip_special_tokens=UpperCamelCase__				,							normalize=UpperCamelCase__				)[
													    0
													]

													if isinstance(UpperCamelCase__				,							UpperCamelCase__				):
																				__lowercase      =    1
													elif isinstance(UpperCamelCase__				,							UpperCamelCase__				):
																				__lowercase      =    len(UpperCamelCase__				)
													else:
																				raise ValueError(f'`prompt` has to be of type `str` or `list` but is {type(UpperCamelCase__				)}'				)

													if height % 8 != 0 or width % 8 != 0:
																				raise ValueError(f'`height` and `width` have to be divisible by 8 but are {height} and {width}.'				)

													if (callback_steps is None) or (
													    callback_steps is not None and (not isinstance(UpperCamelCase__				,							UpperCamelCase__				) or callback_steps <= 0)
													):
																				raise ValueError(
																				    f'`callback_steps` has to be a positive integer but is {callback_steps} of type'
																				    f' {type(UpperCamelCase__				)}.'				)

													# get prompt text embeddings
													__lowercase      =    self.tokenizer(
													    UpperCamelCase__				,							padding='''max_length'''				,							max_length=self.tokenizer.model_max_length				,							return_tensors='''pt'''				,							)
													__lowercase      =    text_inputs.input_ids

													if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
																				__lowercase      =    self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :]				)
																				logger.warning(
																				    '''The following part of your input was truncated because CLIP can only handle sequences up to'''
																				    f' {self.tokenizer.model_max_length} tokens: {removed_text}'				)
																				__lowercase      =    text_input_ids[:, : self.tokenizer.model_max_length]
													__lowercase      =    self.text_encoder(text_input_ids.to(self.device				)				)[0]

													# duplicate text embeddings for each generation per prompt, using mps friendly method
													__lowercase      =    text_embeddings.shape
													__lowercase      =    text_embeddings.repeat(1				,							UpperCamelCase__				,							1				)
													__lowercase      =    text_embeddings.view(bs_embed * num_images_per_prompt				,							UpperCamelCase__				,							-1				)

													# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
													# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
													# corresponds to doing no classifier free guidance.
													__lowercase      =    guidance_scale > 1.0
													# get unconditional embeddings for classifier free guidance
													if do_classifier_free_guidance:
																				__lowercase      =    42
																				if negative_prompt is None:
																											__lowercase      =    [''''''] * batch_size
																				elif type(UpperCamelCase__				) is not type(UpperCamelCase__				):
																											raise TypeError(
																											    f'`negative_prompt` should be the same type to `prompt`, but got {type(UpperCamelCase__				)} !='
																											    f' {type(UpperCamelCase__				)}.'				)
																				elif isinstance(UpperCamelCase__				,							UpperCamelCase__				):
																											__lowercase      =    [negative_prompt]
																				elif batch_size != len(UpperCamelCase__				):
																											raise ValueError(
																											    f'`negative_prompt`: {negative_prompt} has batch size {len(UpperCamelCase__				)}, but `prompt`:'
																											    f' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches'
																											    ''' the batch size of `prompt`.'''				)
																				else:
																											__lowercase      =    negative_prompt

																				__lowercase      =    text_input_ids.shape[-1]
																				__lowercase      =    self.tokenizer(
																				    UpperCamelCase__				,							padding='''max_length'''				,							max_length=UpperCamelCase__				,							truncation=UpperCamelCase__				,							return_tensors='''pt'''				,							)
																				__lowercase      =    self.text_encoder(uncond_input.input_ids.to(self.device				)				)[0]

																				# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
																				__lowercase      =    uncond_embeddings.shape[1]
																				__lowercase      =    uncond_embeddings.repeat(1				,							UpperCamelCase__				,							1				)
																				__lowercase      =    uncond_embeddings.view(batch_size * num_images_per_prompt				,							UpperCamelCase__				,							-1				)

																				# For classifier free guidance, we need to do two forward passes.
																				# Here we concatenate the unconditional and text embeddings into a single batch
																				# to avoid doing two forward passes
																				__lowercase      =    torch.cat([uncond_embeddings, text_embeddings]				)

													# get the initial random noise unless the user supplied it

													# Unlike in other pipelines, latents need to be generated in the target device
													# for 1-to-1 results reproducibility with the CompVis implementation.
													# However this currently doesn't work in `mps`.
													__lowercase      =    (batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
													__lowercase      =    text_embeddings.dtype
													if latents is None:
																				if self.device.type == "mps":
																											# randn does not exist on mps
																											__lowercase      =    torch.randn(UpperCamelCase__				,							generator=UpperCamelCase__				,							device='''cpu'''				,							dtype=UpperCamelCase__				).to(
																											    self.device				)
																				else:
																											__lowercase      =    torch.randn(UpperCamelCase__				,							generator=UpperCamelCase__				,							device=self.device				,							dtype=UpperCamelCase__				)
													else:
																				if latents.shape != latents_shape:
																											raise ValueError(f'Unexpected latents shape, got {latents.shape}, expected {latents_shape}'				)
																				__lowercase      =    latents.to(self.device				)

													# set timesteps
													self.scheduler.set_timesteps(UpperCamelCase__				)

													# Some schedulers like PNDM have timesteps as arrays
													# It's more optimized to move all timesteps to correct device beforehand
													__lowercase      =    self.scheduler.timesteps.to(self.device				)

													# scale the initial noise by the standard deviation required by the scheduler
													__lowercase      =    latents * self.scheduler.init_noise_sigma

													# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
													# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
													# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
													# and should be between [0, 1]
													__lowercase      =    '''eta''' in set(inspect.signature(self.scheduler.step				).parameters.keys()				)
													__lowercase      =    {}
													if accepts_eta:
																				__lowercase      =    eta

													for i, t in enumerate(self.progress_bar(UpperCamelCase__				)				):
																				# expand the latents if we are doing classifier free guidance
																				__lowercase      =    torch.cat([latents] * 2				) if do_classifier_free_guidance else latents
																				__lowercase      =    self.scheduler.scale_model_input(UpperCamelCase__				,							UpperCamelCase__				)

																				# predict the noise residual
																				__lowercase      =    self.unet(UpperCamelCase__				,							UpperCamelCase__				,							encoder_hidden_states=UpperCamelCase__				).sample

																				# perform guidance
																				if do_classifier_free_guidance:
																											__lowercase      =    noise_pred.chunk(2				)
																											__lowercase      =    noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

																				# compute the previous noisy sample x_t -> x_t-1
																				__lowercase      =    self.scheduler.step(UpperCamelCase__				,							UpperCamelCase__				,							UpperCamelCase__				,							**UpperCamelCase__				).prev_sample

																				# call the callback, if provided
																				if callback is not None and i % callback_steps == 0:
																											callback(UpperCamelCase__				,							UpperCamelCase__				,							UpperCamelCase__				)

													__lowercase      =    1 / 0.1_8_2_1_5 * latents
													__lowercase      =    self.vae.decode(UpperCamelCase__				).sample

													__lowercase      =    (image / 2 + 0.5).clamp(0				,							1				)

													# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
													__lowercase      =    image.cpu().permute(0				,							2				,							3				,							1				).float().numpy()

													if output_type == "pil":
																				__lowercase      =    self.numpy_to_pil(UpperCamelCase__				)

													if not return_dict:
																				return image

													return StableDiffusionPipelineOutput(images=UpperCamelCase__				,							nsfw_content_detected=UpperCamelCase__				)












import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotConfig, is_flax_available
from transformers.testing_utils import jax_device, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
			import os

			# The slow tests are often failing with OOM error on GPU
			# This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
			# but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
			UpperCAmelCase__				     = "platform"

			import jax
			import jax.numpy as jnp

			from transformers import BlenderbotTokenizer
			from transformers.models.blenderbot.modeling_flax_blenderbot import (
			    FlaxBlenderbotForConditionalGeneration,
			    FlaxBlenderbotModel,
			    shift_tokens_right,
			)







def 		_A(       UpperCamelCase__							:			Tuple				,       UpperCamelCase__							:			Optional[int]				,       UpperCamelCase__							:			Dict=None				,       UpperCamelCase__							:			Optional[int]=None				,       UpperCamelCase__							:			int=None				,       UpperCamelCase__							:			Union[str, Any]=None				,       UpperCamelCase__							:			Optional[Any]=None				,       UpperCamelCase__							:			List[str]=None				,       )    -> int:
							'''simple docstring'''


							if attention_mask is None:
														__lowercase      =    np.where(input_ids != config.pad_token_id				,       1				,       0				)
							if decoder_attention_mask is None:
														__lowercase      =    np.where(decoder_input_ids != config.pad_token_id				,       1				,       0				)
							if head_mask is None:
														__lowercase      =    np.ones((config.encoder_layers, config.encoder_attention_heads)				)
							if decoder_head_mask is None:
														__lowercase      =    np.ones((config.decoder_layers, config.decoder_attention_heads)				)
							if cross_attn_head_mask is None:
														__lowercase      =    np.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": attention_mask,
							}







class a :
						"""simple docstring"""

						def __init__( self :     Dict				,							lowerCamelCase__ :     Dict				,							lowerCamelCase__ :     Union[str, Any]=13				,							lowerCamelCase__ :     Union[str, Any]=7				,							lowerCamelCase__ :     Dict=True				,							lowerCamelCase__ :     int=False				,							lowerCamelCase__ :     Any=99				,							lowerCamelCase__ :     str=16				,							lowerCamelCase__ :     Union[str, Any]=2				,							lowerCamelCase__ :     Union[str, Any]=4				,							lowerCamelCase__ :     Tuple=4				,							lowerCamelCase__ :     Optional[Any]="gelu"				,							lowerCamelCase__ :     int=0.1				,							lowerCamelCase__ :     Tuple=0.1				,							lowerCamelCase__ :     Any=32				,							lowerCamelCase__ :     List[str]=2				,							lowerCamelCase__ :     Optional[Any]=1				,							lowerCamelCase__ :     List[Any]=0				,							lowerCamelCase__ :     int=0.0_2				,							)    ->						Optional[int]:
													"""simple docstring"""


													__lowercase      =    parent
													__lowercase      =    batch_size
													__lowercase      =    seq_length
													__lowercase      =    is_training
													__lowercase      =    use_labels
													__lowercase      =    vocab_size
													__lowercase      =    hidden_size
													__lowercase      =    num_hidden_layers
													__lowercase      =    num_attention_heads
													__lowercase      =    intermediate_size
													__lowercase      =    hidden_act
													__lowercase      =    hidden_dropout_prob
													__lowercase      =    attention_probs_dropout_prob
													__lowercase      =    max_position_embeddings
													__lowercase      =    eos_token_id
													__lowercase      =    pad_token_id
													__lowercase      =    bos_token_id
													__lowercase      =    initializer_range






						def     UpperCAmelCase_      ( self :     int				)    ->						Any:
													"""simple docstring"""


													__lowercase      =    np.clip(ids_tensor([self.batch_size, self.seq_length - 1]				,							self.vocab_size				)				,							3				,							self.vocab_size				)
													__lowercase      =    np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)				,							dtype=np.intaa				))				,							-1				)

													__lowercase      =    shift_tokens_right(lowerCamelCase__				,							1				,							2				)

													__lowercase      =    BlenderbotConfig(
													    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_id=self.eos_token_id				,							bos_token_id=self.bos_token_id				,							pad_token_id=self.pad_token_id				,							initializer_range=self.initializer_range				,							use_cache=lowerCamelCase__				,							)
													__lowercase      =    prepare_blenderbot_inputs_dict(lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				)
													return config, inputs_dict






						def     UpperCAmelCase_      ( self :     Dict				)    ->						Optional[int]:
													"""simple docstring"""


													__lowercase     ,     __lowercase      =    self.prepare_config_and_inputs()
													return config, inputs_dict






						def     UpperCAmelCase_      ( self :     Optional[Any]				,							lowerCamelCase__ :     str				,							lowerCamelCase__ :     List[str]				,							lowerCamelCase__ :     Tuple				)    ->						Any:
													"""simple docstring"""


													__lowercase      =    20
													__lowercase      =    model_class_name(lowerCamelCase__				)

													__lowercase      =    model.encode(inputs_dict['''input_ids''']				)

													__lowercase     ,     __lowercase      =    (
													    inputs_dict['''decoder_input_ids'''],
													    inputs_dict['''decoder_attention_mask'''],
													)

													__lowercase      =    model.init_cache(decoder_input_ids.shape[0]				,							lowerCamelCase__				,							lowerCamelCase__				)
													__lowercase      =    jnp.ones((decoder_input_ids.shape[0], max_decoder_length)				,							dtype='''i4'''				)

													__lowercase      =    jnp.broadcast_to(
													    jnp.arange(decoder_input_ids.shape[-1] - 1				)[None, :]				,							(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)				,							)
													__lowercase      =    model.decode(
													    decoder_input_ids[:, :-1]				,							lowerCamelCase__				,							decoder_attention_mask=lowerCamelCase__				,							past_key_values=lowerCamelCase__				,							decoder_position_ids=lowerCamelCase__				,							)

													__lowercase      =    jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]				,							dtype='''i4'''				)
													__lowercase      =    model.decode(
													    decoder_input_ids[:, -1:]				,							lowerCamelCase__				,							decoder_attention_mask=lowerCamelCase__				,							past_key_values=outputs_cache.past_key_values				,							decoder_position_ids=lowerCamelCase__				,							)

													__lowercase      =    model.decode(lowerCamelCase__				,							lowerCamelCase__				)

													__lowercase      =    np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])				)				)
													self.parent.assertTrue(diff < 1e-3				,							msg=f'Max diff is {diff}'				)






						def     UpperCAmelCase_      ( self :     int				,							lowerCamelCase__ :     Optional[Any]				,							lowerCamelCase__ :     str				,							lowerCamelCase__ :     List[Any]				)    ->						Optional[Any]:
													"""simple docstring"""


													__lowercase      =    20
													__lowercase      =    model_class_name(lowerCamelCase__				)

													__lowercase      =    model.encode(inputs_dict['''input_ids''']				)

													__lowercase     ,     __lowercase      =    (
													    inputs_dict['''decoder_input_ids'''],
													    inputs_dict['''decoder_attention_mask'''],
													)

													__lowercase      =    jnp.concatenate(
													    [
													        decoder_attention_mask,
													        jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])				),
													    ]				,							axis=-1				,							)

													__lowercase      =    model.init_cache(decoder_input_ids.shape[0]				,							lowerCamelCase__				,							lowerCamelCase__				)
													__lowercase      =    jnp.broadcast_to(
													    jnp.arange(decoder_input_ids.shape[-1] - 1				)[None, :]				,							(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)				,							)

													__lowercase      =    model.decode(
													    decoder_input_ids[:, :-1]				,							lowerCamelCase__				,							decoder_attention_mask=lowerCamelCase__				,							past_key_values=lowerCamelCase__				,							decoder_position_ids=lowerCamelCase__				,							)
													__lowercase      =    jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]				,							dtype='''i4'''				)
													__lowercase      =    model.decode(
													    decoder_input_ids[:, -1:]				,							lowerCamelCase__				,							past_key_values=outputs_cache.past_key_values				,							decoder_attention_mask=lowerCamelCase__				,							decoder_position_ids=lowerCamelCase__				,							)

													__lowercase      =    model.decode(lowerCamelCase__				,							lowerCamelCase__				,							decoder_attention_mask=lowerCamelCase__				)

													__lowercase      =    np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])				)				)
													self.parent.assertTrue(diff < 1e-3				,							msg=f'Max diff is {diff}'				)







@require_flax
class a (     unittest.TestCase  ):
						"""simple docstring"""

						UpperCamelCase_	:       List[Any]       	=		99






						def     UpperCAmelCase_      ( self :     List[Any]				)    ->						Dict:
													"""simple docstring"""


													__lowercase      =    np.array(
													    [
													        [71, 82, 18, 33, 46, 91, 2],
													        [68, 34, 26, 58, 30, 82, 2],
													        [5, 97, 17, 39, 94, 40, 2],
													        [76, 83, 94, 25, 70, 78, 2],
													        [87, 59, 41, 35, 48, 66, 2],
													        [55, 13, 16, 58, 5, 2, 1],  # note padding
													        [64, 27, 31, 51, 12, 75, 2],
													        [52, 64, 86, 17, 83, 39, 2],
													        [48, 61, 9, 24, 71, 82, 2],
													        [26, 1, 60, 48, 22, 13, 2],
													        [21, 5, 62, 28, 14, 76, 2],
													        [45, 98, 37, 86, 59, 48, 2],
													        [70, 70, 50, 9, 28, 0, 2],
													    ]				,							dtype=np.intaa				,							)

													__lowercase      =    input_ids.shape[0]
													__lowercase      =    BlenderbotConfig(
													    vocab_size=self.vocab_size				,							d_model=24				,							encoder_layers=2				,							decoder_layers=2				,							encoder_attention_heads=2				,							decoder_attention_heads=2				,							encoder_ffn_dim=32				,							decoder_ffn_dim=32				,							max_position_embeddings=48				,							eos_token_id=2				,							pad_token_id=1				,							bos_token_id=0				,							)
													return config, input_ids, batch_size






						def     UpperCAmelCase_      ( self :     Tuple				)    ->						Optional[int]:
													"""simple docstring"""


													__lowercase     ,     __lowercase     ,     __lowercase      =    self._get_config_and_data()
													__lowercase      =    FlaxBlenderbotForConditionalGeneration(lowerCamelCase__				)
													__lowercase      =    lm_model(input_ids=lowerCamelCase__				)
													__lowercase      =    (batch_size, input_ids.shape[1], config.vocab_size)
													self.assertEqual(outputs['''logits'''].shape				,							lowerCamelCase__				)






						def     UpperCAmelCase_      ( self :     Union[str, Any]				)    ->						List[Any]:
													"""simple docstring"""


													__lowercase      =    BlenderbotConfig(
													    vocab_size=self.vocab_size				,							d_model=14				,							encoder_layers=2				,							decoder_layers=2				,							encoder_attention_heads=2				,							decoder_attention_heads=2				,							encoder_ffn_dim=8				,							decoder_ffn_dim=8				,							max_position_embeddings=48				,							)
													__lowercase      =    FlaxBlenderbotForConditionalGeneration(lowerCamelCase__				)
													__lowercase      =    np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]				,							dtype=np.intaa				)
													__lowercase      =    np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]				,							dtype=np.intaa				)
													__lowercase      =    lm_model(input_ids=lowerCamelCase__				,							decoder_input_ids=lowerCamelCase__				)
													__lowercase      =    (*summary.shape, config.vocab_size)
													self.assertEqual(outputs['''logits'''].shape				,							lowerCamelCase__				)






						def     UpperCAmelCase_      ( self :     List[str]				)    ->						List[Any]:
													"""simple docstring"""


													__lowercase      =    np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]				,							dtype=np.intaa				)
													__lowercase      =    shift_tokens_right(lowerCamelCase__				,							1				,							2				)
													__lowercase      =    np.equal(lowerCamelCase__				,							1				).astype(np.floataa				).sum()
													__lowercase      =    np.equal(lowerCamelCase__				,							1				).astype(np.floataa				).sum()
													self.assertEqual(shifted.shape				,							input_ids.shape				)
													self.assertEqual(lowerCamelCase__				,							n_pad_before - 1				)
													self.assertTrue(np.equal(shifted[:, 0]				,							2				).all()				)







@require_flax
class a (     __SCREAMING_SNAKE_CASE       ,						unittest.TestCase       ,						__SCREAMING_SNAKE_CASE  ):
						"""simple docstring"""

						UpperCamelCase_	:       Optional[Any]       	=		True
						UpperCamelCase_	:       List[str]       	=		(
						    (
						        FlaxBlenderbotModel,
						        FlaxBlenderbotForConditionalGeneration,
						    )
						    if is_flax_available()
						    else ()
						)
						UpperCamelCase_	:       Union[str, Any]       	=		(FlaxBlenderbotForConditionalGeneration,) if is_flax_available() else ()






						def     UpperCAmelCase_      ( self :     Optional[int]				)    ->						List[str]:
													"""simple docstring"""


													__lowercase      =    FlaxBlenderbotModelTester(self				)






						def     UpperCAmelCase_      ( self :     Any				)    ->						List[str]:
													"""simple docstring"""


													__lowercase     ,     __lowercase      =    self.model_tester.prepare_config_and_inputs()
													for model_class in self.all_model_classes:
																				self.model_tester.check_use_cache_forward(lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				)






						def     UpperCAmelCase_      ( self :     Optional[Any]				)    ->						Optional[Any]:
													"""simple docstring"""


													__lowercase     ,     __lowercase      =    self.model_tester.prepare_config_and_inputs()
													for model_class in self.all_model_classes:
																				self.model_tester.check_use_cache_forward_with_attn_mask(lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				)






						def     UpperCAmelCase_      ( self :     List[str]				)    ->						int:
													"""simple docstring"""


													__lowercase     ,     __lowercase      =    self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																				with self.subTest(model_class.__name__				):
																											__lowercase      =    self._prepare_for_class(lowerCamelCase__				,							lowerCamelCase__				)
																											__lowercase      =    model_class(lowerCamelCase__				)

																											@jax.jit
																											def encode_jitted(lowerCamelCase__ :     List[str]				,							lowerCamelCase__ :     Any=None				,							**lowerCamelCase__ :     Dict				):
																																		return model.encode(input_ids=lowerCamelCase__				,							attention_mask=lowerCamelCase__				)

																											with self.subTest('''JIT Enabled'''				):
																																		__lowercase      =    encode_jitted(**lowerCamelCase__				).to_tuple()

																											with self.subTest('''JIT Disabled'''				):
																																		with jax.disable_jit():
																																									__lowercase      =    encode_jitted(**lowerCamelCase__				).to_tuple()

																											self.assertEqual(len(lowerCamelCase__				)				,							len(lowerCamelCase__				)				)
																											for jitted_output, output in zip(lowerCamelCase__				,							lowerCamelCase__				):
																																		self.assertEqual(jitted_output.shape				,							output.shape				)






						def     UpperCAmelCase_      ( self :     Dict				)    ->						Optional[Any]:
													"""simple docstring"""


													__lowercase     ,     __lowercase      =    self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																				with self.subTest(model_class.__name__				):
																											__lowercase      =    model_class(lowerCamelCase__				)
																											__lowercase      =    model.encode(inputs_dict['''input_ids''']				,							inputs_dict['''attention_mask''']				)

																											__lowercase      =    {
																											    '''decoder_input_ids''': inputs_dict['''decoder_input_ids'''],
																											    '''decoder_attention_mask''': inputs_dict['''decoder_attention_mask'''],
																											    '''encoder_outputs''': encoder_outputs,
																											}

																											@jax.jit
																											def decode_jitted(lowerCamelCase__ :     List[str]				,							lowerCamelCase__ :     List[str]				,							lowerCamelCase__ :     Tuple				):
																																		return model.decode(
																																		    decoder_input_ids=lowerCamelCase__				,							decoder_attention_mask=lowerCamelCase__				,							encoder_outputs=lowerCamelCase__				,							)

																											with self.subTest('''JIT Enabled'''				):
																																		__lowercase      =    decode_jitted(**lowerCamelCase__				).to_tuple()

																											with self.subTest('''JIT Disabled'''				):
																																		with jax.disable_jit():
																																									__lowercase      =    decode_jitted(**lowerCamelCase__				).to_tuple()

																											self.assertEqual(len(lowerCamelCase__				)				,							len(lowerCamelCase__				)				)
																											for jitted_output, output in zip(lowerCamelCase__				,							lowerCamelCase__				):
																																		self.assertEqual(jitted_output.shape				,							output.shape				)






						@slow
						def     UpperCAmelCase_      ( self :     int				)    ->						Union[str, Any]:
													"""simple docstring"""


													for model_class_name in self.all_model_classes:
																				__lowercase      =    model_class_name.from_pretrained('''facebook/blenderbot-400M-distill'''				)
																				# FlaxBlenderbotForSequenceClassification expects eos token in input_ids
																				__lowercase      =    np.ones((1, 1)				) * model.config.eos_token_id
																				__lowercase      =    model(lowerCamelCase__				)
																				self.assertIsNotNone(lowerCamelCase__				)






						@unittest.skipUnless(jax_device != '''cpu'''				,							'''3B test too slow on CPU.'''				)
						@slow
						def     UpperCAmelCase_      ( self :     Tuple				)    ->						List[str]:
													"""simple docstring"""


													__lowercase      =    {'''num_beams''': 1, '''early_stopping''': True, '''min_length''': 15, '''max_length''': 25}
													__lowercase      =    {'''skip_special_tokens''': True, '''clean_up_tokenization_spaces''': True}

													__lowercase      =    FlaxBlenderbotForConditionalGeneration.from_pretrained('''facebook/blenderbot-3B'''				,							from_pt=lowerCamelCase__				)
													__lowercase      =    BlenderbotTokenizer.from_pretrained('''facebook/blenderbot-3B'''				)

													__lowercase      =    ['''Sam''']
													__lowercase      =    tokenizer(lowerCamelCase__				,							return_tensors='''jax'''				)

													__lowercase      =    model.generate(**lowerCamelCase__				,							**lowerCamelCase__				)
													__lowercase      =    '''Sam is a great name. It means "sun" in Gaelic.'''

													__lowercase      =    tokenizer.batch_decode(lowerCamelCase__				,							**lowerCamelCase__				)
													assert generated_txt[0].strip() == tgt_text






