
'''simple docstring'''







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
							_lowercase :  List[str]    							=       "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 					snake_case_		(				__SCREAMING_SNAKE_CASE					: Optional[Any]      ,      __SCREAMING_SNAKE_CASE					: int      ,      __SCREAMING_SNAKE_CASE					: str=None      ,      __SCREAMING_SNAKE_CASE					: Union[str, Any]=None      ,      __SCREAMING_SNAKE_CASE					: List[Any]=None      ,      __SCREAMING_SNAKE_CASE					: Any=None      ,      __SCREAMING_SNAKE_CASE					: Any=None      ,      __SCREAMING_SNAKE_CASE					: Union[str, Any]=None      ,      ):



						"""simple docstring"""





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






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


											"""simple docstring"""





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






					def 	_snake_case    (   self		):


											"""simple docstring"""





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

											lowercase_	:     Optional[int]													=       shift_tokens_right(__SCREAMING_SNAKE_CASE				,							1				,							2		)

											lowercase_	:     Union[str, Any]													=       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=__SCREAMING_SNAKE_CASE				,							)
											lowercase_	:     Optional[int]													=       prepare_blenderbot_inputs_dict(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		)
											return config, inputs_dict






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_						, lowercase_	:     Any													=       self.prepare_config_and_inputs()
											return config, inputs_dict






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											lowercase_	:     List[str]													=       20
											lowercase_	:     Optional[Any]													=       model_class_name(__SCREAMING_SNAKE_CASE		)

											lowercase_	:     List[str]													=       model.encode(inputs_dict['''input_ids''']		)

											lowercase_						, lowercase_	:     List[Any]													=       (
											    inputs_dict['''decoder_input_ids'''],
											    inputs_dict['''decoder_attention_mask'''],
											)

											lowercase_	:     int													=       model.init_cache(decoder_input_ids.shape[0]				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		)
											lowercase_	:     str													=       jnp.ones((decoder_input_ids.shape[0], max_decoder_length)				,							dtype='''i4'''		)

											lowercase_	:     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)				,							)
											lowercase_	:     Optional[int]													=       model.decode(
											    decoder_input_ids[:, :-1]				,							__SCREAMING_SNAKE_CASE				,							decoder_attention_mask=__SCREAMING_SNAKE_CASE				,							past_key_values=__SCREAMING_SNAKE_CASE				,							decoder_position_ids=__SCREAMING_SNAKE_CASE				,							)

											lowercase_	:     int													=       jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]				,							dtype='''i4'''		)
											lowercase_	:     List[str]													=       model.decode(
											    decoder_input_ids[:, -1:]				,							__SCREAMING_SNAKE_CASE				,							decoder_attention_mask=__SCREAMING_SNAKE_CASE				,							past_key_values=outputs_cache.past_key_values				,							decoder_position_ids=__SCREAMING_SNAKE_CASE				,							)

											lowercase_	:     Tuple													=       model.decode(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		)

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


											"""simple docstring"""





											lowercase_	:     int													=       20
											lowercase_	:     Any													=       model_class_name(__SCREAMING_SNAKE_CASE		)

											lowercase_	:     int													=       model.encode(inputs_dict['''input_ids''']		)

											lowercase_						, lowercase_	:     Tuple													=       (
											    inputs_dict['''decoder_input_ids'''],
											    inputs_dict['''decoder_attention_mask'''],
											)

											lowercase_	:     List[str]													=       jnp.concatenate(
											    [
											        decoder_attention_mask,
											        jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])		),
											    ]				,							axis=-1				,							)

											lowercase_	:     Any													=       model.init_cache(decoder_input_ids.shape[0]				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		)
											lowercase_	:     List[Any]													=       jnp.broadcast_to(
											    jnp.arange(decoder_input_ids.shape[-1] - 1		)[None, :]				,							(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)				,							)

											lowercase_	:     List[Any]													=       model.decode(
											    decoder_input_ids[:, :-1]				,							__SCREAMING_SNAKE_CASE				,							decoder_attention_mask=__SCREAMING_SNAKE_CASE				,							past_key_values=__SCREAMING_SNAKE_CASE				,							decoder_position_ids=__SCREAMING_SNAKE_CASE				,							)
											lowercase_	:     int													=       jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]				,							dtype='''i4'''		)
											lowercase_	:     Dict													=       model.decode(
											    decoder_input_ids[:, -1:]				,							__SCREAMING_SNAKE_CASE				,							past_key_values=outputs_cache.past_key_values				,							decoder_attention_mask=__SCREAMING_SNAKE_CASE				,							decoder_position_ids=__SCREAMING_SNAKE_CASE				,							)

											lowercase_	:     int													=       model.decode(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							decoder_attention_mask=__SCREAMING_SNAKE_CASE		)

											lowercase_	:     Union[str, 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  lowerCAmelCase__   (  unittest.TestCase     ):
					lowerCAmelCase_        =     9_9






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     List[Any]													=       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_	:     str													=       input_ids.shape[0]
											lowercase_	:     Tuple													=       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 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_						, lowercase_						, lowercase_	:     List[Any]													=       self._get_config_and_data()
											lowercase_	:     List[Any]													=       FlaxBlenderbotForConditionalGeneration(__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Union[str, Any]													=       lm_model(input_ids=__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Optional[int]													=       (batch_size, input_ids.shape[1], config.vocab_size)
											self.assertEqual(outputs['''logits'''].shape				,							__SCREAMING_SNAKE_CASE		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     List[str]													=       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_	:     Tuple													=       FlaxBlenderbotForConditionalGeneration(__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Dict													=       np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]				,							dtype=np.intaa		)
											lowercase_	:     Optional[int]													=       np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]				,							dtype=np.intaa		)
											lowercase_	:     Dict													=       lm_model(input_ids=__SCREAMING_SNAKE_CASE				,							decoder_input_ids=__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Optional[Any]													=       (*summary.shape, config.vocab_size)
											self.assertEqual(outputs['''logits'''].shape				,							__SCREAMING_SNAKE_CASE		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





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







@require_flax
class  lowerCAmelCase__   (  lowerCamelCase_    ,						unittest.TestCase    ,						lowerCamelCase_     ):
					lowerCAmelCase_        =     True
					lowerCAmelCase_        =     (
					    (
					        FlaxBlenderbotModel,
					        FlaxBlenderbotForConditionalGeneration,
					    )
					    if is_flax_available()
					    else ()
					)
					lowerCAmelCase_        =     (FlaxBlenderbotForConditionalGeneration,) if is_flax_available() else ()






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Union[str, Any]													=       FlaxBlenderbotModelTester(self		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_						, lowercase_	:     Tuple													=       self.model_tester.prepare_config_and_inputs()
											for model_class in self.all_model_classes:
																	self.model_tester.check_use_cache_forward(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_						, lowercase_	:     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(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_						, lowercase_	:     List[str]													=       self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	with self.subTest(model_class.__name__		):
																							lowercase_	:     Tuple													=       self._prepare_for_class(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		)
																							lowercase_	:     Tuple													=       model_class(__SCREAMING_SNAKE_CASE		)

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

																							with self.subTest('''JIT Enabled'''		):
																													lowercase_	:     Tuple													=       encode_jitted(**__SCREAMING_SNAKE_CASE		).to_tuple()

																							with self.subTest('''JIT Disabled'''		):
																													with jax.disable_jit():
																																			lowercase_	:     str													=       encode_jitted(**__SCREAMING_SNAKE_CASE		).to_tuple()

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






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_						, lowercase_	:     int													=       self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	with self.subTest(model_class.__name__		):
																							lowercase_	:     Dict													=       model_class(__SCREAMING_SNAKE_CASE		)
																							lowercase_	:     Optional[Any]													=       model.encode(inputs_dict['''input_ids''']				,							inputs_dict['''attention_mask''']		)

																							lowercase_	:     Optional[int]													=       {
																							    '''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(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		):
																													return model.decode(
																													    decoder_input_ids=__SCREAMING_SNAKE_CASE				,							decoder_attention_mask=__SCREAMING_SNAKE_CASE				,							encoder_outputs=__SCREAMING_SNAKE_CASE				,							)

																							with self.subTest('''JIT Enabled'''		):
																													lowercase_	:     List[str]													=       decode_jitted(**__SCREAMING_SNAKE_CASE		).to_tuple()

																							with self.subTest('''JIT Disabled'''		):
																													with jax.disable_jit():
																																			lowercase_	:     List[str]													=       decode_jitted(**__SCREAMING_SNAKE_CASE		).to_tuple()

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






					@slow
					def 	_snake_case    (   self		):


											"""simple docstring"""





											for model_class_name in self.all_model_classes:
																	lowercase_	:     List[Any]													=       model_class_name.from_pretrained('''facebook/blenderbot-400M-distill'''		)
																	# FlaxBlenderbotForSequenceClassification expects eos token in input_ids
																	lowercase_	:     Optional[int]													=       np.ones((1, 1)		) * model.config.eos_token_id
																	lowercase_	:     Optional[Any]													=       model(__SCREAMING_SNAKE_CASE		)
																	self.assertIsNotNone(__SCREAMING_SNAKE_CASE		)






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


											"""simple docstring"""





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

											lowercase_	:     Dict													=       FlaxBlenderbotForConditionalGeneration.from_pretrained('''facebook/blenderbot-3B'''				,							from_pt=__SCREAMING_SNAKE_CASE		)
											lowercase_	:     Any													=       BlenderbotTokenizer.from_pretrained('''facebook/blenderbot-3B'''		)

											lowercase_	:     Any													=       ['''Sam''']
											lowercase_	:     Optional[int]													=       tokenizer(__SCREAMING_SNAKE_CASE				,							return_tensors='''jax'''		)

											lowercase_	:     List[str]													=       model.generate(**__SCREAMING_SNAKE_CASE				,							**__SCREAMING_SNAKE_CASE		)
											lowercase_	:     str													=       '''Sam is a great name. It means "sun" in Gaelic.'''

											lowercase_	:     str													=       tokenizer.batch_decode(__SCREAMING_SNAKE_CASE				,							**__SCREAMING_SNAKE_CASE		)
											assert generated_txt[0].strip() == tgt_text





'''simple docstring'''







def 					snake_case_		(				__SCREAMING_SNAKE_CASE					: int							):



						"""simple docstring"""





						lowercase_	:     Optional[int]													=       int(__SCREAMING_SNAKE_CASE							)
						if decimal in (0, 1):  # Exit cases for the recursion
												return str(__SCREAMING_SNAKE_CASE							)
						lowercase_						, lowercase_	:     List[str]													=       divmod(__SCREAMING_SNAKE_CASE      ,      2							)
						return binary_recursive(__SCREAMING_SNAKE_CASE							) + str(__SCREAMING_SNAKE_CASE							)
def 					snake_case_		(				__SCREAMING_SNAKE_CASE					: str							):



						"""simple docstring"""





						lowercase_	:     str													=       str(__SCREAMING_SNAKE_CASE							).strip()
						if not number:
												raise ValueError('''No input value was provided'''							)
						lowercase_	:     Optional[int]													=       '''-''' if number.startswith('''-'''							) else ''''''
						lowercase_	:     Union[str, Any]													=       number.lstrip('''-'''							)
						if not number.isnumeric():
												raise ValueError('''Input value is not an integer'''							)
						return F'''{negative}0b{binary_recursive(int(__SCREAMING_SNAKE_CASE							)							)}'''


if __name__ == "__main__":
							from doctest import testmod

							testmod()



