




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
					__A        =    '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 			lowerCamelCase_		(       UpperCamelCase__				:      Optional[int] ,     UpperCamelCase__				:      Optional[Any] ,     UpperCamelCase__				:      Optional[int]=None ,     UpperCamelCase__				:      Dict=None ,     UpperCamelCase__				:      int=None ,     UpperCamelCase__				:      Optional[int]=None ,     UpperCamelCase__				:      Any=None ,     UpperCamelCase__				:      List[Any]=None ,     )     ->       Union[str, Any]:





				"""simple docstring"""


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

			"""simple docstring"""







			def __init__(      self			,						lowerCamelCase__			,						lowerCamelCase__=13			,						lowerCamelCase__=7			,						lowerCamelCase__=True			,						lowerCamelCase__=False			,						lowerCamelCase__=99			,						lowerCamelCase__=16			,						lowerCamelCase__=2			,						lowerCamelCase__=4			,						lowerCamelCase__=4			,						lowerCamelCase__="gelu"			,						lowerCamelCase__=0.1			,						lowerCamelCase__=0.1			,						lowerCamelCase__=32			,						lowerCamelCase__=2			,						lowerCamelCase__=1			,						lowerCamelCase__=0			,						lowerCamelCase__=0.02			,						)    ->      Any:



							'''simple docstring'''
							__lowerCamelCase		       =    parent
							__lowerCamelCase		       =    batch_size
							__lowerCamelCase		       =    seq_length
							__lowerCamelCase		       =    is_training
							__lowerCamelCase		       =    use_labels
							__lowerCamelCase		       =    vocab_size
							__lowerCamelCase		       =    hidden_size
							__lowerCamelCase		       =    num_hidden_layers
							__lowerCamelCase		       =    num_attention_heads
							__lowerCamelCase		       =    intermediate_size
							__lowerCamelCase		       =    hidden_act
							__lowerCamelCase		       =    hidden_dropout_prob
							__lowerCamelCase		       =    attention_probs_dropout_prob
							__lowerCamelCase		       =    max_position_embeddings
							__lowerCamelCase		       =    eos_token_id
							__lowerCamelCase		       =    pad_token_id
							__lowerCamelCase		       =    bos_token_id
							__lowerCamelCase		       =    initializer_range


			def       lowercase_    (      self							)    ->      Dict:



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

							__lowerCamelCase		       =    shift_tokens_right(lowerCamelCase__			,						1			,						2							)

							__lowerCamelCase		       =    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__			,						)
							__lowerCamelCase		       =    prepare_blenderbot_inputs_dict(lowerCamelCase__			,						lowerCamelCase__			,						lowerCamelCase__							)
							return config, inputs_dict


			def       lowercase_    (      self							)    ->      Optional[Any]:



							'''simple docstring'''
							__lowerCamelCase       ,						__lowerCamelCase		       =    self.prepare_config_and_inputs()
							return config, inputs_dict


			def       lowercase_    (      self			,						lowerCamelCase__			,						lowerCamelCase__			,						lowerCamelCase__							)    ->      str:



							'''simple docstring'''
							__lowerCamelCase		       =    20
							__lowerCamelCase		       =    model_class_name(lowerCamelCase__							)

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

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

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

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

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

							__lowerCamelCase		       =    model.decode(lowerCamelCase__			,						lowerCamelCase__							)

							__lowerCamelCase		       =    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       lowercase_    (      self			,						lowerCamelCase__			,						lowerCamelCase__			,						lowerCamelCase__							)    ->      Optional[int]:



							'''simple docstring'''
							__lowerCamelCase		       =    20
							__lowerCamelCase		       =    model_class_name(lowerCamelCase__							)

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

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

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

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

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

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

							__lowerCamelCase		       =    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     __lowerCAmelCase (       unittest.TestCase ):

			"""simple docstring"""







			snake_case_   		=       99


			def       lowercase_    (      self							)    ->      Dict:



							'''simple docstring'''
							__lowerCamelCase		       =    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			,						)

							__lowerCamelCase		       =    input_ids.shape[0]
							__lowerCamelCase		       =    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       lowercase_    (      self							)    ->      int:



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


			def       lowercase_    (      self							)    ->      List[str]:



							'''simple docstring'''
							__lowerCamelCase		       =    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			,						)
							__lowerCamelCase		       =    FlaxBlenderbotForConditionalGeneration(lowerCamelCase__							)
							__lowerCamelCase		       =    np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]			,						dtype=np.intaa							)
							__lowerCamelCase		       =    np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]			,						dtype=np.intaa							)
							__lowerCamelCase		       =    lm_model(input_ids=lowerCamelCase__			,						decoder_input_ids=lowerCamelCase__							)
							__lowerCamelCase		       =    (*summary.shape, config.vocab_size)
							self.assertEqual(outputs['logits'].shape			,						lowerCamelCase__							)






			def       lowercase_    (      self							)    ->      int:



							'''simple docstring'''
							__lowerCamelCase		       =    np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]			,						dtype=np.intaa							)
							__lowerCamelCase		       =    shift_tokens_right(lowerCamelCase__			,						1			,						2							)
							__lowerCamelCase		       =    np.equal(lowerCamelCase__			,						1							).astype(np.floataa							).sum()
							__lowerCamelCase		       =    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     __lowerCAmelCase (       lowercase_			,				unittest.TestCase			,				lowercase_ ):

			"""simple docstring"""







			snake_case_   		=       True
			snake_case_   		=       (
			    (
			        FlaxBlenderbotModel,
			        FlaxBlenderbotForConditionalGeneration,
			    )
			    if is_flax_available()
			    else ()
			)
			snake_case_   		=       (FlaxBlenderbotForConditionalGeneration,) if is_flax_available() else ()


			def       lowercase_    (      self							)    ->      int:



							'''simple docstring'''
							__lowerCamelCase		       =    FlaxBlenderbotModelTester(self							)


			def       lowercase_    (      self							)    ->      str:



							'''simple docstring'''
							__lowerCamelCase       ,						__lowerCamelCase		       =    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       lowercase_    (      self							)    ->      Union[str, Any]:



							'''simple docstring'''
							__lowerCamelCase       ,						__lowerCamelCase		       =    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       lowercase_    (      self							)    ->      Dict:



							'''simple docstring'''
							__lowerCamelCase       ,						__lowerCamelCase		       =    self.model_tester.prepare_config_and_inputs_for_common()

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

															@jax.jit
															def encode_jitted(lowerCamelCase__			,						lowerCamelCase__=None			,						**lowerCamelCase__							):
																			return model.encode(input_ids=lowerCamelCase__			,						attention_mask=lowerCamelCase__							)

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

															with self.subTest('JIT Disabled'							):
																			with jax.disable_jit():
																							__lowerCamelCase		       =    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       lowercase_    (      self							)    ->      List[Any]:



							'''simple docstring'''
							__lowerCamelCase       ,						__lowerCamelCase		       =    self.model_tester.prepare_config_and_inputs_for_common()

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

															__lowerCamelCase		       =    {
															    '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__			,						lowerCamelCase__			,						lowerCamelCase__							):
																			return model.decode(
																			    decoder_input_ids=lowerCamelCase__			,						decoder_attention_mask=lowerCamelCase__			,						encoder_outputs=lowerCamelCase__			,						)

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

															with self.subTest('JIT Disabled'							):
																			with jax.disable_jit():
																							__lowerCamelCase		       =    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       lowercase_    (      self							)    ->      List[str]:



							'''simple docstring'''







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






			@unittest.skipUnless(jax_device != 'cpu'			,						'3B test too slow on CPU.'							)
			@slow
			def       lowercase_    (      self							)    ->      Dict:



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

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

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

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

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






'''simple docstring'''

import warnings

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase							:   str  						=	logging.get_logger(__name__)

lowerCamelCase							:   int  						=	{
    'RUCAIBox/mvp': 'https://huggingface.co/RUCAIBox/mvp/resolve/main/config.json',
}
class   __lowerCAmelCase						(lowercase_     ):





      '''simple docstring'''






      lowerCAmelCase__					:     Union[str, Any]  					=							"""mvp"""
      lowerCAmelCase__					:     Optional[Any]  					=							["""past_key_values"""]
      lowerCAmelCase__					:     List[str]  					=							{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}



      def __init__(self    :    Any   ,						UpperCamelCase    :    Optional[int]=50267   ,						UpperCamelCase    :    Tuple=1024   ,						UpperCamelCase    :    int=12   ,						UpperCamelCase    :    Tuple=4096   ,						UpperCamelCase    :    Dict=16   ,						UpperCamelCase    :    int=12   ,						UpperCamelCase    :    Optional[int]=4096   ,						UpperCamelCase    :    Optional[int]=16   ,						UpperCamelCase    :    Tuple=0.0   ,						UpperCamelCase    :    Tuple=0.0   ,						UpperCamelCase    :    List[Any]="gelu"   ,						UpperCamelCase    :    Union[str, Any]=1024   ,						UpperCamelCase    :    Optional[Any]=0.1   ,						UpperCamelCase    :    str=0.0   ,						UpperCamelCase    :    str=0.0   ,						UpperCamelCase    :    Optional[Any]=0.02   ,						UpperCamelCase    :    List[str]=0.0   ,						UpperCamelCase    :    List[str]=False   ,						UpperCamelCase    :    Optional[int]=True   ,						UpperCamelCase    :    Any=1   ,						UpperCamelCase    :    int=0   ,						UpperCamelCase    :    int=2   ,						UpperCamelCase    :    Any=True   ,						UpperCamelCase    :    Optional[Any]=2   ,						UpperCamelCase    :    Optional[Any]=2   ,						UpperCamelCase    :    Tuple=False   ,						UpperCamelCase    :    int=100   ,						UpperCamelCase    :    Optional[Any]=800   ,						**UpperCamelCase    :    str   ,						):

             '''simple docstring'''







             lowercase__      			= vocab_size
             lowercase__      			= max_position_embeddings
             lowercase__      			= d_model
             lowercase__      			= encoder_ffn_dim
             lowercase__      			= encoder_layers
             lowercase__      			= encoder_attention_heads
             lowercase__      			= decoder_ffn_dim
             lowercase__      			= decoder_layers
             lowercase__      			= decoder_attention_heads
             lowercase__      			= dropout
             lowercase__      			= attention_dropout
             lowercase__      			= activation_dropout
             lowercase__      			= activation_function
             lowercase__      			= init_std
             lowercase__      			= encoder_layerdrop
             lowercase__      			= decoder_layerdrop
             lowercase__      			= classifier_dropout
             lowercase__      			= use_cache
             lowercase__      			= encoder_layers
             lowercase__      			= scale_embedding  # scale factor will be sqrt(d_model) if True
             lowercase__      			= use_prompt
             lowercase__      			= prompt_length
             lowercase__      			= prompt_mid_dim

             super().__init__(
                 pad_token_id=UpperCamelCase   ,						bos_token_id=UpperCamelCase   ,						eos_token_id=UpperCamelCase   ,						is_encoder_decoder=UpperCamelCase   ,						decoder_start_token_id=UpperCamelCase   ,						forced_eos_token_id=UpperCamelCase   ,						**UpperCamelCase   ,						)

             if self.forced_bos_token_id is None and kwargs.get('''force_bos_token_to_be_generated'''   ,						UpperCamelCase ):
                    lowercase__      			= self.bos_token_id
                    warnings.warn(
                        f"Please make sure the config includes `forced_bos_token_id={self.bos_token_id}` in future versions. "
                        '''The config can simply be saved and uploaded again to be fixed.''' )




