




import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotSmallConfig, is_flax_available
from transformers.testing_utils import 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_		:				Optional[int]									=     "platform"

							import jax
							import jax.numpy as jnp

							from transformers.models.blenderbot_small.modeling_flax_blenderbot_small import (
							    FlaxBlenderbotSmallForConditionalGeneration,
							    FlaxBlenderbotSmallModel,
							    shift_tokens_right,
							)




def 		lowerCAmelCase_   (	lowerCamelCase	,     lowerCamelCase	,     lowerCamelCase=None	,     lowerCamelCase=None	,     lowerCamelCase=None	,     lowerCamelCase=None	,     lowerCamelCase=None	,     lowerCamelCase=None	,     ):
				if attention_mask is None:
								__magic_name__	:	Tuple								=np.where(input_ids != config.pad_token_id	,     1	,     0   )
				if decoder_attention_mask is None:
								__magic_name__	:	Optional[Any]								=np.where(decoder_input_ids != config.pad_token_id	,     1	,     0   )
				if head_mask is None:
								__magic_name__	:	str								=np.ones((config.encoder_layers, config.encoder_attention_heads)   )
				if decoder_head_mask is None:
								__magic_name__	:	List[str]								=np.ones((config.decoder_layers, config.decoder_attention_heads)   )
				if cross_attn_head_mask is None:
								__magic_name__	:	Tuple								=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      :

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

										__magic_name__	:	Optional[Any]								=parent
										__magic_name__	:	Dict								=batch_size
										__magic_name__	:	Optional[Any]								=seq_length
										__magic_name__	:	Union[str, Any]								=is_training
										__magic_name__	:	List[str]								=use_labels
										__magic_name__	:	Any								=vocab_size
										__magic_name__	:	str								=hidden_size
										__magic_name__	:	Union[str, Any]								=num_hidden_layers
										__magic_name__	:	Optional[Any]								=num_attention_heads
										__magic_name__	:	Any								=intermediate_size
										__magic_name__	:	Optional[Any]								=hidden_act
										__magic_name__	:	Any								=hidden_dropout_prob
										__magic_name__	:	Tuple								=attention_probs_dropout_prob
										__magic_name__	:	List[Any]								=max_position_embeddings
										__magic_name__	:	Optional[int]								=eos_token_id
										__magic_name__	:	List[str]								=pad_token_id
										__magic_name__	:	Union[str, Any]								=bos_token_id
										__magic_name__	:	Optional[int]								=initializer_range

						def     A__						(		self   :Optional[Any] ):
										'''simple docstring'''

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

										__magic_name__	:	Any								=shift_tokens_right(__snake_case    , 1    , 2 )

										__magic_name__	:	Optional[int]								=BlenderbotSmallConfig(
										    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=__snake_case    , )
										__magic_name__	:	Dict								=prepare_blenderbot_inputs_dict(__snake_case    , __snake_case    , __snake_case )
										return config, inputs_dict

						def     A__						(		self   :List[str] ):
										'''simple docstring'''

										__magic_name__				, __magic_name__	:	str								=self.prepare_config_and_inputs()
										return config, inputs_dict

						def     A__						(		self   :List[str]    , __snake_case   :Optional[int]    , __snake_case   :Tuple    , __snake_case   :str ):
										'''simple docstring'''

										__magic_name__	:	Optional[int]								=20
										__magic_name__	:	List[Any]								=model_class_name(__snake_case )

										__magic_name__	:	int								=model.encode(inputs_dict["""input_ids"""] )

										__magic_name__				, __magic_name__	:	str								=(
										    inputs_dict["""decoder_input_ids"""],
										    inputs_dict["""decoder_attention_mask"""],
										)

										__magic_name__	:	Optional[int]								=model.init_cache(decoder_input_ids.shape[0]    , __snake_case    , __snake_case )
										__magic_name__	:	List[Any]								=jnp.ones((decoder_input_ids.shape[0], max_decoder_length)    , dtype="""i4""" )

										__magic_name__	:	Optional[int]								=jnp.broadcast_to(
										    jnp.arange(decoder_input_ids.shape[-1] - 1 )[None, :]    , (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)    , )
										__magic_name__	:	List[Any]								=model.decode(
										    decoder_input_ids[:, :-1]    , __snake_case    , decoder_attention_mask=__snake_case    , past_key_values=__snake_case    , decoder_position_ids=__snake_case    , )

										__magic_name__	:	List[str]								=jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]    , dtype="""i4""" )
										__magic_name__	:	List[Any]								=model.decode(
										    decoder_input_ids[:, -1:]    , __snake_case    , decoder_attention_mask=__snake_case    , past_key_values=outputs_cache.past_key_values    , decoder_position_ids=__snake_case    , )

										__magic_name__	:	List[Any]								=model.decode(__snake_case    , __snake_case )

										__magic_name__	:	Tuple								=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     A__						(		self   :List[str]    , __snake_case   :Union[str, Any]    , __snake_case   :List[str]    , __snake_case   :Optional[Any] ):
										'''simple docstring'''

										__magic_name__	:	Union[str, Any]								=20
										__magic_name__	:	Optional[Any]								=model_class_name(__snake_case )

										__magic_name__	:	Dict								=model.encode(inputs_dict["""input_ids"""] )

										__magic_name__				, __magic_name__	:	List[Any]								=(
										    inputs_dict["""decoder_input_ids"""],
										    inputs_dict["""decoder_attention_mask"""],
										)

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

										__magic_name__	:	List[Any]								=model.init_cache(decoder_input_ids.shape[0]    , __snake_case    , __snake_case )
										__magic_name__	:	Union[str, Any]								=jnp.broadcast_to(
										    jnp.arange(decoder_input_ids.shape[-1] - 1 )[None, :]    , (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)    , )

										__magic_name__	:	str								=model.decode(
										    decoder_input_ids[:, :-1]    , __snake_case    , decoder_attention_mask=__snake_case    , past_key_values=__snake_case    , decoder_position_ids=__snake_case    , )
										__magic_name__	:	Optional[Any]								=jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]    , dtype="""i4""" )
										__magic_name__	:	Dict								=model.decode(
										    decoder_input_ids[:, -1:]    , __snake_case    , past_key_values=outputs_cache.past_key_values    , decoder_attention_mask=__snake_case    , decoder_position_ids=__snake_case    , )

										__magic_name__	:	List[Any]								=model.decode(__snake_case    , __snake_case    , decoder_attention_mask=__snake_case )

										__magic_name__	:	List[Any]								=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			):
						UpperCamelCase     						=					99

						def     A__						(		self   :Tuple ):
										'''simple docstring'''

										__magic_name__	:	List[str]								=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    , )

										__magic_name__	:	Any								=input_ids.shape[0]
										__magic_name__	:	Tuple								=BlenderbotSmallConfig(
										    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     A__						(		self   :Dict ):
										'''simple docstring'''

										__magic_name__				, __magic_name__				, __magic_name__	:	int								=self._get_config_and_data()
										__magic_name__	:	Dict								=FlaxBlenderbotSmallForConditionalGeneration(__snake_case )
										__magic_name__	:	Optional[Any]								=lm_model(input_ids=__snake_case )
										__magic_name__	:	int								=(batch_size, input_ids.shape[1], config.vocab_size)
										self.assertEqual(outputs["""logits"""].shape    , __snake_case )

						def     A__						(		self   :Optional[Any] ):
										'''simple docstring'''

										__magic_name__	:	int								=BlenderbotSmallConfig(
										    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    , )
										__magic_name__	:	Optional[int]								=FlaxBlenderbotSmallForConditionalGeneration(__snake_case )
										__magic_name__	:	str								=np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]    , dtype=np.intaa )
										__magic_name__	:	str								=np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]    , dtype=np.intaa )
										__magic_name__	:	Optional[Any]								=lm_model(input_ids=__snake_case    , decoder_input_ids=__snake_case )
										__magic_name__	:	Union[str, Any]								=(*summary.shape, config.vocab_size)
										self.assertEqual(outputs["""logits"""].shape    , __snake_case )


						def     A__						(		self   :int ):
										'''simple docstring'''

										__magic_name__	:	Optional[Any]								=np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]    , dtype=np.intaa )
										__magic_name__	:	List[str]								=shift_tokens_right(__snake_case    , 1    , 2 )
										__magic_name__	:	Union[str, Any]								=np.equal(__snake_case    , 1 ).astype(np.floataa ).sum()
										__magic_name__	:	int								=np.equal(__snake_case    , 1 ).astype(np.floataa ).sum()
										self.assertEqual(shifted.shape    , input_ids.shape )
										self.assertEqual(__snake_case    , n_pad_before - 1 )
										self.assertTrue(np.equal(shifted[:, 0]    , 2 ).all() )




@require_flax
class 		__A      (			UpperCamelCase__					,						unittest.TestCase					,						UpperCamelCase__			):
						UpperCamelCase     						=					True
						UpperCamelCase     						=					(
						    (
						        FlaxBlenderbotSmallModel,
						        FlaxBlenderbotSmallForConditionalGeneration,
						    )
						    if is_flax_available()
						    else ()
						)
						UpperCamelCase     						=					(FlaxBlenderbotSmallForConditionalGeneration,) if is_flax_available() else ()

						def     A__						(		self   :Optional[int] ):
										'''simple docstring'''

										__magic_name__	:	int								=FlaxBlenderbotSmallModelTester(self )

						def     A__						(		self   :List[str] ):
										'''simple docstring'''

										__magic_name__				, __magic_name__	:	Optional[int]								=self.model_tester.prepare_config_and_inputs()
										for model_class in self.all_model_classes:
														self.model_tester.check_use_cache_forward(__snake_case    , __snake_case    , __snake_case )

						def     A__						(		self   :str ):
										'''simple docstring'''

										__magic_name__				, __magic_name__	:	Dict								=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(__snake_case    , __snake_case    , __snake_case )

						def     A__						(		self   :Union[str, Any] ):
										'''simple docstring'''

										__magic_name__				, __magic_name__	:	Optional[Any]								=self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
														with self.subTest(model_class.__name__ ):
																		__magic_name__	:	List[Any]								=self._prepare_for_class(__snake_case    , __snake_case )
																		__magic_name__	:	Tuple								=model_class(__snake_case )

																		@jax.jit
																		def encode_jitted(__snake_case   :Union[str, Any]    , __snake_case   :Tuple=None    , **__snake_case   :Any ):
																						return model.encode(input_ids=__snake_case    , attention_mask=__snake_case )

																		with self.subTest("""JIT Enabled""" ):
																						__magic_name__	:	List[Any]								=encode_jitted(**__snake_case ).to_tuple()

																		with self.subTest("""JIT Disabled""" ):
																						with jax.disable_jit():
																										__magic_name__	:	Tuple								=encode_jitted(**__snake_case ).to_tuple()

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

						def     A__						(		self   :Any ):
										'''simple docstring'''

										__magic_name__				, __magic_name__	:	Optional[int]								=self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
														with self.subTest(model_class.__name__ ):
																		__magic_name__	:	Optional[int]								=model_class(__snake_case )
																		__magic_name__	:	List[str]								=model.encode(inputs_dict["""input_ids"""]    , inputs_dict["""attention_mask"""] )

																		__magic_name__	:	List[str]								={
																		    """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(__snake_case   :Optional[int]    , __snake_case   :List[Any]    , __snake_case   :Optional[Any] ):
																						return model.decode(
																						    decoder_input_ids=__snake_case    , decoder_attention_mask=__snake_case    , encoder_outputs=__snake_case    , )

																		with self.subTest("""JIT Enabled""" ):
																						__magic_name__	:	Dict								=decode_jitted(**__snake_case ).to_tuple()

																		with self.subTest("""JIT Disabled""" ):
																						with jax.disable_jit():
																										__magic_name__	:	int								=decode_jitted(**__snake_case ).to_tuple()

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


						@slow
						def     A__						(		self   :List[str] ):
										'''simple docstring'''

										for model_class_name in self.all_model_classes:
														__magic_name__	:	Dict								=model_class_name.from_pretrained("""facebook/blenderbot_small-90M""" )
														# FlaxBlenderbotForSequenceClassification expects eos token in input_ids
														__magic_name__	:	Optional[int]								=np.ones((1, 1) ) * model.config.eos_token_id
														__magic_name__	:	Any								=model(__snake_case )
														self.assertIsNotNone(__snake_case )










from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
							from PIL import Image

							from ..image_utils import load_image

if is_tf_available():
							import tensorflow as tf

							from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
							from ..tf_utils import stable_softmax

if is_torch_available():
							from ..models.auto.modeling_auto import MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING

UpperCAmelCase_		:				str									=     logging.get_logger(__name__)




@add_end_docstrings(UpperCamelCase__			)
class 		__A      (			UpperCamelCase__			):

						def __init__(		self   :Any    , *__snake_case   :List[Any]    , **__snake_case   :Union[str, Any] ):
										'''simple docstring'''

										super().__init__(*__snake_case    , **__snake_case )
										requires_backends(self    , """vision""" )
										self.check_model_type(
										    TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
										    if self.framework == """tf"""
										    else MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING )

						def     A__						(		self   :int    , __snake_case   :Optional[Any]=None ):
										'''simple docstring'''

										__magic_name__	:	Any								={}
										if top_k is not None:
														__magic_name__	:	str								=top_k
										return {}, {}, postprocess_params

						def __call__(		self   :Tuple    , __snake_case   :Union[str, List[str], "Image.Image", List["Image.Image"]]    , **__snake_case   :Union[str, Any] ):
										'''simple docstring'''

										return super().__call__(__snake_case    , **__snake_case )

						def     A__						(		self   :str    , __snake_case   :Any ):
										'''simple docstring'''

										__magic_name__	:	Optional[Any]								=load_image(__snake_case )
										__magic_name__	:	Any								=self.image_processor(images=__snake_case    , return_tensors=self.framework )
										return model_inputs

						def     A__						(		self   :int    , __snake_case   :Optional[int] ):
										'''simple docstring'''

										__magic_name__	:	str								=self.model(**__snake_case )
										return model_outputs


						def     A__						(		self   :List[str]    , __snake_case   :int    , __snake_case   :Union[str, Any]=5 ):
										'''simple docstring'''

										if top_k > self.model.config.num_labels:
														__magic_name__	:	Dict								=self.model.config.num_labels

										if self.framework == "pt":
														__magic_name__	:	int								=model_outputs.logits.softmax(-1 )[0]
														__magic_name__				, __magic_name__	:	str								=probs.topk(__snake_case )
										elif self.framework == "tf":
														__magic_name__	:	Union[str, Any]								=stable_softmax(model_outputs.logits    , axis=-1 )[0]
														__magic_name__	:	List[str]								=tf.math.top_k(__snake_case    , k=__snake_case )
														__magic_name__				, __magic_name__	:	str								=topk.values.numpy(), topk.indices.numpy()
										else:
														raise ValueError(f"Unsupported framework: {self.framework}" )

										__magic_name__	:	str								=scores.tolist()
										__magic_name__	:	Union[str, Any]								=ids.tolist()
										return [{"score": score, "label": self.model.config.idalabel[_id]} for score, _id in zip(__snake_case    , __snake_case )]




