




"""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
					a_				:    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 						UpperCAmelCase				(			A__:    List[Any]							,				A__:    Any							,				A__:    Tuple=None							,				A__:    List[str]=None							,				A__:    Dict=None							,				A__:    Dict=None							,				A__:    Dict=None							,				A__:    Any=None							,				)  ->     Optional[Any]:
			if attention_mask is None:
						__lowerCamelCase					:	Tuple	       =				np.where(input_ids != config.pad_token_id							,				1							,				0     )
			if decoder_attention_mask is None:
						__lowerCamelCase					:	int	       =				np.where(decoder_input_ids != config.pad_token_id							,				1							,				0     )
			if head_mask is None:
						__lowerCamelCase					:	Any	       =				np.ones((config.encoder_layers, config.encoder_attention_heads)     )
			if decoder_head_mask is None:
						__lowerCamelCase					:	Tuple	       =				np.ones((config.decoder_layers, config.decoder_attention_heads)     )
			if cross_attn_head_mask is None:
						__lowerCamelCase					:	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 __lowercase:




					'''simple docstring'''


					def __init__(    self						,	__a						,	__a=13						,	__a=7						,	__a=True						,	__a=False						,	__a=99						,	__a=16						,	__a=2						,	__a=4						,	__a=4						,	__a="gelu"						,	__a=0.1						,	__a=0.1						,	__a=32						,	__a=2						,	__a=1						,	__a=0						,	__a=0.02						,	):
								__lowerCamelCase					:	List[str]	       =				parent
								__lowerCamelCase					:	List[str]	       =				batch_size
								__lowerCamelCase					:	int	       =				seq_length
								__lowerCamelCase					:	Dict	       =				is_training
								__lowerCamelCase					:	List[Any]	       =				use_labels
								__lowerCamelCase					:	Optional[int]	       =				vocab_size
								__lowerCamelCase					:	str	       =				hidden_size
								__lowerCamelCase					:	Dict	       =				num_hidden_layers
								__lowerCamelCase					:	Union[str, Any]	       =				num_attention_heads
								__lowerCamelCase					:	List[Any]	       =				intermediate_size
								__lowerCamelCase					:	int	       =				hidden_act
								__lowerCamelCase					:	List[Any]	       =				hidden_dropout_prob
								__lowerCamelCase					:	Tuple	       =				attention_probs_dropout_prob
								__lowerCamelCase					:	Dict	       =				max_position_embeddings
								__lowerCamelCase					:	Any	       =				eos_token_id
								__lowerCamelCase					:	List[str]	       =				pad_token_id
								__lowerCamelCase					:	int	       =				bos_token_id
								__lowerCamelCase					:	Tuple	       =				initializer_range
					def     snake_case_				(    self     ):
								__lowerCamelCase					:	Dict	       =				np.clip(ids_tensor([self.batch_size, self.seq_length - 1]						,	self.vocab_size     )						,	3						,	self.vocab_size     )
								__lowerCamelCase					:	Any	       =				np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)						,	dtype=np.intaa     ))						,	-1     )

								__lowerCamelCase					:	Optional[Any]	       =				shift_tokens_right(__a						,	1						,	2     )

								__lowerCamelCase					:	List[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=__a						,	)
								__lowerCamelCase					:	int	       =				prepare_blenderbot_inputs_dict(__a						,	__a						,	__a     )
								return config, inputs_dict
					def     snake_case_				(    self     ):
								__lowerCamelCase       ,      __lowerCamelCase					:	str	       =				self.prepare_config_and_inputs()
								return config, inputs_dict
					def     snake_case_				(    self						,	__a						,	__a						,	__a     ):
								__lowerCamelCase					:	List[str]	       =				20
								__lowerCamelCase					:	List[str]	       =				model_class_name(__a     )

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

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

								__lowerCamelCase					:	Union[str, Any]	       =				model.init_cache(decoder_input_ids.shape[0]						,	__a						,	__a     )
								__lowerCamelCase					:	List[str]	       =				jnp.ones((decoder_input_ids.shape[0], max_decoder_length)						,	dtype='i4'     )

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

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

								__lowerCamelCase					:	Tuple	       =				model.decode(__a						,	__a     )

								__lowerCamelCase					:	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}'''     )





					def     snake_case_				(    self						,	__a						,	__a						,	__a     ):
								__lowerCamelCase					:	List[Any]	       =				20
								__lowerCamelCase					:	str	       =				model_class_name(__a     )

								__lowerCamelCase					:	Union[str, Any]	       =				model.encode(inputs_dict['input_ids']     )

								__lowerCamelCase       ,      __lowerCamelCase					:	Optional[Any]	       =				(
								    inputs_dict['decoder_input_ids'],
								    inputs_dict['decoder_attention_mask'],
								)

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

								__lowerCamelCase					:	List[Any]	       =				model.init_cache(decoder_input_ids.shape[0]						,	__a						,	__a     )
								__lowerCamelCase					:	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)						,	)

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

								__lowerCamelCase					:	str	       =				model.decode(__a						,	__a						,	decoder_attention_mask=__a     )

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




					'''simple docstring'''


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





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



@require_flax
class __lowercase(      lowercase__	,	unittest.TestCase	,	lowercase__				):




					'''simple docstring'''


					__a						:						List[Any]     		= True
					__a						:						Union[str, Any]     		= (
					    (
					        FlaxBlenderbotModel,
					        FlaxBlenderbotForConditionalGeneration,
					    )
					    if is_flax_available()
					    else ()
					)
					__a						:						Optional[int]     		= (FlaxBlenderbotForConditionalGeneration,) if is_flax_available() else ()
					def     snake_case_				(    self     ):
								__lowerCamelCase					:	Union[str, Any]	       =				FlaxBlenderbotModelTester(self     )
					def     snake_case_				(    self     ):
								__lowerCamelCase       ,      __lowerCamelCase					:	str	       =				self.model_tester.prepare_config_and_inputs()
								for model_class in self.all_model_classes:
											self.model_tester.check_use_cache_forward(__a						,	__a						,	__a     )
					def     snake_case_				(    self     ):
								__lowerCamelCase       ,      __lowerCamelCase					:	List[Any]	       =				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(__a						,	__a						,	__a     )
					def     snake_case_				(    self     ):
								__lowerCamelCase       ,      __lowerCamelCase					:	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__     ):
														__lowerCamelCase					:	str	       =				self._prepare_for_class(__a						,	__a     )
														__lowerCamelCase					:	Optional[Any]	       =				model_class(__a     )

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

														with self.subTest('JIT Enabled'     ):
																	__lowerCamelCase					:	List[str]	       =				encode_jitted(**__a     ).to_tuple()

														with self.subTest('JIT Disabled'     ):
																	with jax.disable_jit():
																				__lowerCamelCase					:	Optional[Any]	       =				encode_jitted(**__a     ).to_tuple()

														self.assertEqual(len(__a     )						,	len(__a     )     )
														for jitted_output, output in zip(__a						,	__a     ):
																	self.assertEqual(jitted_output.shape						,	output.shape     )
					def     snake_case_				(    self     ):
								__lowerCamelCase       ,      __lowerCamelCase					:	str	       =				self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
											with self.subTest(model_class.__name__     ):
														__lowerCamelCase					:	Optional[int]	       =				model_class(__a     )
														__lowerCamelCase					:	Union[str, Any]	       =				model.encode(inputs_dict['input_ids']						,	inputs_dict['attention_mask']     )

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

														with self.subTest('JIT Enabled'     ):
																	__lowerCamelCase					:	Optional[Any]	       =				decode_jitted(**__a     ).to_tuple()

														with self.subTest('JIT Disabled'     ):
																	with jax.disable_jit():
																				__lowerCamelCase					:	Tuple	       =				decode_jitted(**__a     ).to_tuple()

														self.assertEqual(len(__a     )						,	len(__a     )     )
														for jitted_output, output in zip(__a						,	__a     ):
																	self.assertEqual(jitted_output.shape						,	output.shape     )
					@slow
					def     snake_case_				(    self     ):
								for model_class_name in self.all_model_classes:
											__lowerCamelCase					:	Union[str, Any]	       =				model_class_name.from_pretrained('facebook/blenderbot-400M-distill'     )
											# FlaxBlenderbotForSequenceClassification expects eos token in input_ids
											__lowerCamelCase					:	List[Any]	       =				np.ones((1, 1)     ) * model.config.eos_token_id
											__lowerCamelCase					:	Dict	       =				model(__a     )
											self.assertIsNotNone(__a     )





					@unittest.skipUnless(jax_device != 'cpu'						,	'3B test too slow on CPU.'     )
					@slow
					def     snake_case_				(    self     ):
								__lowerCamelCase					:	Optional[int]	       =				{'num_beams': 1, 'early_stopping': True, 'min_length': 15, 'max_length': 25}
								__lowerCamelCase					:	Tuple	       =				{'skip_special_tokens': True, 'clean_up_tokenization_spaces': True}

								__lowerCamelCase					:	str	       =				FlaxBlenderbotForConditionalGeneration.from_pretrained('facebook/blenderbot-3B'						,	from_pt=__a     )
								__lowerCamelCase					:	List[str]	       =				BlenderbotTokenizer.from_pretrained('facebook/blenderbot-3B'     )

								__lowerCamelCase					:	Union[str, Any]	       =				['Sam']
								__lowerCamelCase					:	Tuple	       =				tokenizer(__a						,	return_tensors='jax'     )

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

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









"""simple docstring"""
import pytest

import datasets


# Import fixture modules as plugins
a_				:    Tuple										=     ['''tests.fixtures.files''', '''tests.fixtures.hub''', '''tests.fixtures.fsspec''']





def 						UpperCAmelCase				(			A__:    List[Any]							,				A__:    Optional[Any]     )  ->     List[str]:
			# Mark tests as "unit" by default if not marked as "integration" (or already marked as "unit")
			for item in items:
						if any(marker in item.keywords for marker in ['integration', 'unit']     ):
									continue
						item.add_marker(pytest.mark.unit     )





def 						UpperCAmelCase				(			A__:    int     )  ->     Optional[int]:
			config.addinivalue_line('markers'							,				'torchaudio_latest: mark test to run with torchaudio>=0.12'     )





@pytest.fixture(autouse=A__     )
def 						UpperCAmelCase				(			A__:    Optional[Any]							,				A__:    List[Any]     )  ->     Any:
			# test_hf_cache_home = tmp_path_factory.mktemp("cache")  # TODO: why a cache dir per test function does not work?
			__lowerCamelCase					:	str	       =				tmp_path_factory.getbasetemp() / 'cache'
			__lowerCamelCase					:	Dict	       =				test_hf_cache_home / 'datasets'
			__lowerCamelCase					:	int	       =				test_hf_cache_home / 'metrics'
			__lowerCamelCase					:	List[str]	       =				test_hf_cache_home / 'modules'
			monkeypatch.setattr('datasets.config.HF_DATASETS_CACHE'							,				str(A__     )     )
			monkeypatch.setattr('datasets.config.HF_METRICS_CACHE'							,				str(A__     )     )
			monkeypatch.setattr('datasets.config.HF_MODULES_CACHE'							,				str(A__     )     )
			__lowerCamelCase					:	List[Any]	       =				test_hf_datasets_cache / 'downloads'
			monkeypatch.setattr('datasets.config.DOWNLOADED_DATASETS_PATH'							,				str(A__     )     )
			__lowerCamelCase					:	Optional[int]	       =				test_hf_datasets_cache / 'downloads' / 'extracted'
			monkeypatch.setattr('datasets.config.EXTRACTED_DATASETS_PATH'							,				str(A__     )     )





@pytest.fixture(autouse=A__							,				scope='session'     )
def 						UpperCAmelCase				(			)  ->     Any:
			datasets.disable_progress_bar()





@pytest.fixture(autouse=A__     )
def 						UpperCAmelCase				(			A__:    Optional[Any]     )  ->     Dict:
			# don't take tests into account when counting downloads
			monkeypatch.setattr('datasets.config.HF_UPDATE_DOWNLOAD_COUNTS'							,				A__     )






@pytest.fixture
def 						UpperCAmelCase				(			A__:    Union[str, Any]     )  ->     Optional[Any]:
			# Required to suppress RemovedIn20Warning when feature(s) are not compatible with SQLAlchemy 2.0
			# To be removed once SQLAlchemy 2.0 supported
			monkeypatch.setattr('sqlalchemy.util.deprecations.SILENCE_UBER_WARNING'							,				A__     )



