





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
						_SCREAMING_SNAKE_CASE 		=						"""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 							SCREAMING_SNAKE_CASE__  (      __a				,	__a				,	__a=None				,	__a=None				,	__a=None				,	__a=None				,	__a=None				,	__a=None				,	):
	if attention_mask is None:
		snake_case_						:     Dict     =		np.where(input_ids != config.pad_token_id				,	1				,	0				)
	if decoder_attention_mask is None:
		snake_case_						:     List[str]     =		np.where(decoder_input_ids != config.pad_token_id				,	1				,	0				)
	if head_mask is None:
		snake_case_						:     List[str]     =		np.ones((config.encoder_layers, config.encoder_attention_heads)				)
	if decoder_head_mask is None:
		snake_case_						:     str     =		np.ones((config.decoder_layers, config.decoder_attention_heads)				)
	if cross_attn_head_mask is None:
		snake_case_						:     Optional[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 SCREAMING_SNAKE_CASE_  :






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







							"""simple docstring"""


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






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







							"""simple docstring"""


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

							snake_case_						:     Any     =		shift_tokens_right(_A						,				1						,				2		)

							snake_case_						:     int     =		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						,				)
							snake_case_						:     List[str]     =		prepare_blenderbot_inputs_dict(_A						,				_A						,				_A		)
							return config, inputs_dict






						def UpperCAmelCase_			( self : Any		)							->      int:







							"""simple docstring"""


							snake_case_  ,snake_case_						:     Dict     =		self.prepare_config_and_inputs()
							return config, inputs_dict






						def UpperCAmelCase_			( self : int						,				_A : Optional[Any]						,				_A : List[Any]						,				_A : Optional[Any]		)							->      Optional[int]:







							"""simple docstring"""


							snake_case_						:     Optional[Any]     =		20
							snake_case_						:     List[str]     =		model_class_name(_A		)

							snake_case_						:     List[Any]     =		model.encode(inputs_dict['input_ids']		)

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

							snake_case_						:     List[Any]     =		model.init_cache(decoder_input_ids.shape[0]						,				_A						,				_A		)
							snake_case_						:     List[Any]     =		jnp.ones((decoder_input_ids.shape[0], max_decoder_length)						,				dtype='i4'		)

							snake_case_						:     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)						,				)
							snake_case_						:     Any     =		model.decode(
							    decoder_input_ids[:, :-1]						,				_A						,				decoder_attention_mask=_A						,				past_key_values=_A						,				decoder_position_ids=_A						,				)

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

							snake_case_						:     Optional[Any]     =		model.decode(_A						,				_A		)

							snake_case_						:     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}"""		)






						def UpperCAmelCase_			( self : Any						,				_A : Any						,				_A : Tuple						,				_A : Any		)							->      Tuple:







							"""simple docstring"""


							snake_case_						:     Dict     =		20
							snake_case_						:     Dict     =		model_class_name(_A		)

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

							snake_case_  ,snake_case_						:     List[str]     =		(
							    inputs_dict['decoder_input_ids'],
							    inputs_dict['decoder_attention_mask'],
							)

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

							snake_case_						:     Optional[Any]     =		model.init_cache(decoder_input_ids.shape[0]						,				_A						,				_A		)
							snake_case_						:     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)						,				)

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

							snake_case_						:     Dict     =		model.decode(_A						,				_A						,				decoder_attention_mask=_A		)

							snake_case_						:     str     =		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 SCREAMING_SNAKE_CASE_  (   unittest.TestCase   ):
						__magic_name__:     List[str]	 =    99






						def UpperCAmelCase_			( self : int		)							->      Optional[int]:







							"""simple docstring"""


							snake_case_						:     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						,				)

							snake_case_						:     Tuple     =		input_ids.shape[0]
							snake_case_						:     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 UpperCAmelCase_			( self : Tuple		)							->      List[str]:







							"""simple docstring"""


							snake_case_  ,snake_case_  ,snake_case_						:     List[Any]     =		self._get_config_and_data()
							snake_case_						:     str     =		FlaxBlenderbotForConditionalGeneration(_A		)
							snake_case_						:     int     =		lm_model(input_ids=_A		)
							snake_case_						:     Tuple     =		(batch_size, input_ids.shape[1], config.vocab_size)
							self.assertEqual(outputs['logits'].shape						,				_A		)






						def UpperCAmelCase_			( self : str		)							->      Union[str, Any]:







							"""simple docstring"""


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






						def UpperCAmelCase_			( self : List[str]		)							->      Dict:







							"""simple docstring"""


							snake_case_						:     Union[str, Any]     =		np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]						,				dtype=np.intaa		)
							snake_case_						:     Dict     =		shift_tokens_right(_A						,				1						,				2		)
							snake_case_						:     Tuple     =		np.equal(_A						,				1		).astype(np.floataa		).sum()
							snake_case_						:     List[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 SCREAMING_SNAKE_CASE_  (   snake_case_    ,			unittest.TestCase    ,			snake_case_   ):
						__magic_name__:     List[str]	 =    True
						__magic_name__:     int	 =    (
						    (
						        FlaxBlenderbotModel,
						        FlaxBlenderbotForConditionalGeneration,
						    )
						    if is_flax_available()
						    else ()
						)
						__magic_name__:     Optional[int]	 =    (FlaxBlenderbotForConditionalGeneration,) if is_flax_available() else ()






						def UpperCAmelCase_			( self : Any		)							->      int:







							"""simple docstring"""


							snake_case_						:     Optional[Any]     =		FlaxBlenderbotModelTester(self		)






						def UpperCAmelCase_			( self : List[str]		)							->      Tuple:







							"""simple docstring"""


							snake_case_  ,snake_case_						:     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 UpperCAmelCase_			( self : Tuple		)							->      List[str]:







							"""simple docstring"""


							snake_case_  ,snake_case_						:     Optional[int]     =		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 UpperCAmelCase_			( self : Optional[Any]		)							->      str:







							"""simple docstring"""


							snake_case_  ,snake_case_						:     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__		):
									snake_case_						:     Dict     =		self._prepare_for_class(_A						,				_A		)
									snake_case_						:     Dict     =		model_class(_A		)

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

									with self.subTest('JIT Enabled'		):
										snake_case_						:     Dict     =		encode_jitted(**_A		).to_tuple()

									with self.subTest('JIT Disabled'		):
										with jax.disable_jit():
											snake_case_						:     int     =		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 UpperCAmelCase_			( self : Optional[int]		)							->      Union[str, Any]:







							"""simple docstring"""


							snake_case_  ,snake_case_						:     Tuple     =		self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
								with self.subTest(model_class.__name__		):
									snake_case_						:     Tuple     =		model_class(_A		)
									snake_case_						:     int     =		model.encode(inputs_dict['input_ids']						,				inputs_dict['attention_mask']		)

									snake_case_						:     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 : Optional[int]						,				_A : Tuple						,				_A : int		):
										return model.decode(
										    decoder_input_ids=_A						,				decoder_attention_mask=_A						,				encoder_outputs=_A						,				)

									with self.subTest('JIT Enabled'		):
										snake_case_						:     Union[str, Any]     =		decode_jitted(**_A		).to_tuple()

									with self.subTest('JIT Disabled'		):
										with jax.disable_jit():
											snake_case_						:     Dict     =		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 UpperCAmelCase_			( self : str		)							->      Union[str, Any]:







							"""simple docstring"""


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






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







							"""simple docstring"""


							snake_case_						:     Optional[int]     =		{'num_beams': 1, 'early_stopping': True, 'min_length': 15, 'max_length': 25}
							snake_case_						:     int     =		{'skip_special_tokens': True, 'clean_up_tokenization_spaces': True}

							snake_case_						:     Any     =		FlaxBlenderbotForConditionalGeneration.from_pretrained('facebook/blenderbot-3B'						,				from_pt=_A		)
							snake_case_						:     Optional[Any]     =		BlenderbotTokenizer.from_pretrained('facebook/blenderbot-3B'		)

							snake_case_						:     Optional[int]     =		['Sam']
							snake_case_						:     Optional[Any]     =		tokenizer(_A						,				return_tensors='jax'		)

							snake_case_						:     int     =		model.generate(**_A						,				**_A		)
							snake_case_						:     Tuple     =		'Sam is a great name. It means "sun" in Gaelic.'

							snake_case_						:     Union[str, Any]     =		tokenizer.batch_decode(_A						,				**_A		)
							assert generated_txt[0].strip() == tgt_text








from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


_SCREAMING_SNAKE_CASE 		=						{
    """configuration_xlm""": ["""XLM_PRETRAINED_CONFIG_ARCHIVE_MAP""", """XLMConfig""", """XLMOnnxConfig"""],
    """tokenization_xlm""": ["""XLMTokenizer"""],
}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_SCREAMING_SNAKE_CASE 		=						[
						    """XLM_PRETRAINED_MODEL_ARCHIVE_LIST""",
						    """XLMForMultipleChoice""",
						    """XLMForQuestionAnswering""",
						    """XLMForQuestionAnsweringSimple""",
						    """XLMForSequenceClassification""",
						    """XLMForTokenClassification""",
						    """XLMModel""",
						    """XLMPreTrainedModel""",
						    """XLMWithLMHeadModel""",
						]

try:
						if not is_tf_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_SCREAMING_SNAKE_CASE 		=						[
						    """TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST""",
						    """TFXLMForMultipleChoice""",
						    """TFXLMForQuestionAnsweringSimple""",
						    """TFXLMForSequenceClassification""",
						    """TFXLMForTokenClassification""",
						    """TFXLMMainLayer""",
						    """TFXLMModel""",
						    """TFXLMPreTrainedModel""",
						    """TFXLMWithLMHeadModel""",
						]


if TYPE_CHECKING:
						from .configuration_xlm import XLM_PRETRAINED_CONFIG_ARCHIVE_MAP, XLMConfig, XLMOnnxConfig
						from .tokenization_xlm import XLMTokenizer

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_xlm import (
												    XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
												    XLMForMultipleChoice,
												    XLMForQuestionAnswering,
												    XLMForQuestionAnsweringSimple,
												    XLMForSequenceClassification,
												    XLMForTokenClassification,
												    XLMModel,
												    XLMPreTrainedModel,
												    XLMWithLMHeadModel,
												)

						try:
												if not is_tf_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_tf_xlm import (
												    TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
												    TFXLMForMultipleChoice,
												    TFXLMForQuestionAnsweringSimple,
												    TFXLMForSequenceClassification,
												    TFXLMForTokenClassification,
												    TFXLMMainLayer,
												    TFXLMModel,
												    TFXLMPreTrainedModel,
												    TFXLMWithLMHeadModel,
												)

else:
						import sys

						_SCREAMING_SNAKE_CASE 		=						_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

