




'''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
				_lowerCAmelCase      	=					'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					(    snake_case__					,      snake_case__					,      snake_case__=None					,      snake_case__=None					,      snake_case__=None					,      snake_case__=None					,      snake_case__=None					,      snake_case__=None					,      ):
			if attention_mask is None:
						__UpperCamelCase      :			Any									=   np.where(input_ids != config.pad_token_id					,      1					,      0						)
			if decoder_attention_mask is None:
						__UpperCamelCase      :			int									=   np.where(decoder_input_ids != config.pad_token_id					,      1					,      0						)
			if head_mask is None:
						__UpperCamelCase      :			List[str]									=   np.ones((config.encoder_layers, config.encoder_attention_heads)						)
			if decoder_head_mask is None:
						__UpperCamelCase      :			List[str]									=   np.ones((config.decoder_layers, config.decoder_attention_heads)						)
			if cross_attn_head_mask is None:
						__UpperCamelCase      :			Tuple									=   np.ones((config.decoder_layers, config.decoder_attention_heads)						)
			return {
			    "input_ids": input_ids,
			    "decoder_input_ids": decoder_input_ids,
			    "attention_mask": attention_mask,
			    "decoder_attention_mask": attention_mask,
			}



class 					A		:




				'''simple docstring'''





				def __init__(self ,  _UpperCAmelCase ,  _UpperCAmelCase=1_3 ,  _UpperCAmelCase=7 ,  _UpperCAmelCase=True ,  _UpperCAmelCase=False ,  _UpperCAmelCase=9_9 ,  _UpperCAmelCase=1_6 ,  _UpperCAmelCase=2 ,  _UpperCAmelCase=4 ,  _UpperCAmelCase=4 ,  _UpperCAmelCase="gelu" ,  _UpperCAmelCase=0.1 ,  _UpperCAmelCase=0.1 ,  _UpperCAmelCase=3_2 ,  _UpperCAmelCase=2 ,  _UpperCAmelCase=1 ,  _UpperCAmelCase=0 ,  _UpperCAmelCase=0.02 ,  )	->       int:
							__UpperCamelCase      :			Dict									=   parent
							__UpperCamelCase      :			Tuple									=   batch_size
							__UpperCamelCase      :			List[Any]									=   seq_length
							__UpperCamelCase      :			Union[str, Any]									=   is_training
							__UpperCamelCase      :			Dict									=   use_labels
							__UpperCamelCase      :			Optional[int]									=   vocab_size
							__UpperCamelCase      :			Optional[Any]									=   hidden_size
							__UpperCamelCase      :			Any									=   num_hidden_layers
							__UpperCamelCase      :			List[str]									=   num_attention_heads
							__UpperCamelCase      :			Optional[Any]									=   intermediate_size
							__UpperCamelCase      :			int									=   hidden_act
							__UpperCamelCase      :			Any									=   hidden_dropout_prob
							__UpperCamelCase      :			List[Any]									=   attention_probs_dropout_prob
							__UpperCamelCase      :			Dict									=   max_position_embeddings
							__UpperCamelCase      :			Dict									=   eos_token_id
							__UpperCamelCase      :			Union[str, Any]									=   pad_token_id
							__UpperCamelCase      :			Optional[Any]									=   bos_token_id
							__UpperCamelCase      :			Tuple									=   initializer_range
				def        a_       (self			)	->       Optional[int]:
							__UpperCamelCase      :			Optional[int]									=   np.clip(ids_tensor([self.batch_size, self.seq_length - 1] ,  self.vocab_size			) ,  3 ,  self.vocab_size			)
							__UpperCamelCase      :			Dict									=   np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1) ,  dtype=np.intaa			)) ,  -1			)

							__UpperCamelCase      :			str									=   shift_tokens_right(__a ,  1 ,  2			)

							__UpperCamelCase      :			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 ,  )
							__UpperCamelCase      :			Tuple									=   prepare_blenderbot_inputs_dict(__a ,  __a ,  __a			)
							return config, inputs_dict
				def        a_       (self			)	->       List[Any]:
							__UpperCamelCase		,					__UpperCamelCase      :			Tuple									=   self.prepare_config_and_inputs()
							return config, inputs_dict
				def        a_       (self ,  _UpperCAmelCase ,  _UpperCAmelCase ,  _UpperCAmelCase			)	->       int:
							__UpperCamelCase      :			List[Any]									=   2_0
							__UpperCamelCase      :			Optional[int]									=   model_class_name(__a			)

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

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

							__UpperCamelCase      :			Optional[int]									=   model.init_cache(decoder_input_ids.shape[0] ,  __a ,  __a			)
							__UpperCamelCase      :			int									=   jnp.ones((decoder_input_ids.shape[0], max_decoder_length) ,  dtype="i4"			)

							__UpperCamelCase      :			List[str]									=   jnp.broadcast_to(
							    jnp.arange(decoder_input_ids.shape[-1] - 1			)[None, :] ,  (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1) ,  )
							__UpperCamelCase      :			List[str]									=   model.decode(
							    decoder_input_ids[:, :-1] ,  __a ,  decoder_attention_mask=__a ,  past_key_values=__a ,  decoder_position_ids=__a ,  )

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

							__UpperCamelCase      :			List[Any]									=   model.decode(__a ,  __a			)

							__UpperCamelCase      :			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        a_       (self ,  _UpperCAmelCase ,  _UpperCAmelCase ,  _UpperCAmelCase			)	->       Union[str, Any]:
							__UpperCamelCase      :			Union[str, Any]									=   2_0
							__UpperCamelCase      :			List[str]									=   model_class_name(__a			)

							__UpperCamelCase      :			List[Any]									=   model.encode(inputs_dict["input_ids"]			)

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

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

							__UpperCamelCase      :			str									=   model.init_cache(decoder_input_ids.shape[0] ,  __a ,  __a			)
							__UpperCamelCase      :			Any									=   jnp.broadcast_to(
							    jnp.arange(decoder_input_ids.shape[-1] - 1			)[None, :] ,  (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1) ,  )

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

							__UpperCamelCase      :			List[Any]									=   model.decode(__a ,  __a ,  decoder_attention_mask=__a			)

							__UpperCamelCase      :			List[Any]									=   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])			)			)
							self.parent.assertTrue(diff < 1E-3 ,  msg=f"Max diff is {diff}"			)







@require_flax
class 					A		( unittest.TestCase		):




				'''simple docstring'''





				A      	=				9_9
				def        a_       (self			)	->       Dict:
							__UpperCamelCase      :			Union[str, Any]									=   np.array(
							    [
							        [7_1, 8_2, 1_8, 3_3, 4_6, 9_1, 2],
							        [6_8, 3_4, 2_6, 5_8, 3_0, 8_2, 2],
							        [5, 9_7, 1_7, 3_9, 9_4, 4_0, 2],
							        [7_6, 8_3, 9_4, 2_5, 7_0, 7_8, 2],
							        [8_7, 5_9, 4_1, 3_5, 4_8, 6_6, 2],
							        [5_5, 1_3, 1_6, 5_8, 5, 2, 1],  # note padding
							        [6_4, 2_7, 3_1, 5_1, 1_2, 7_5, 2],
							        [5_2, 6_4, 8_6, 1_7, 8_3, 3_9, 2],
							        [4_8, 6_1, 9, 2_4, 7_1, 8_2, 2],
							        [2_6, 1, 6_0, 4_8, 2_2, 1_3, 2],
							        [2_1, 5, 6_2, 2_8, 1_4, 7_6, 2],
							        [4_5, 9_8, 3_7, 8_6, 5_9, 4_8, 2],
							        [7_0, 7_0, 5_0, 9, 2_8, 0, 2],
							    ] ,  dtype=np.intaa ,  )

							__UpperCamelCase      :			Optional[Any]									=   input_ids.shape[0]
							__UpperCamelCase      :			Any									=   BlenderbotConfig(
							    vocab_size=self.vocab_size ,  d_model=2_4 ,  encoder_layers=2 ,  decoder_layers=2 ,  encoder_attention_heads=2 ,  decoder_attention_heads=2 ,  encoder_ffn_dim=3_2 ,  decoder_ffn_dim=3_2 ,  max_position_embeddings=4_8 ,  eos_token_id=2 ,  pad_token_id=1 ,  bos_token_id=0 ,  )
							return config, input_ids, batch_size
				def        a_       (self			)	->       List[Any]:
							__UpperCamelCase		,					__UpperCamelCase		,					__UpperCamelCase      :			Dict									=   self._get_config_and_data()
							__UpperCamelCase      :			List[Any]									=   FlaxBlenderbotForConditionalGeneration(__a			)
							__UpperCamelCase      :			str									=   lm_model(input_ids=__a			)
							__UpperCamelCase      :			Any									=   (batch_size, input_ids.shape[1], config.vocab_size)
							self.assertEqual(outputs["logits"].shape ,  __a			)
				def        a_       (self			)	->       Optional[Any]:
							__UpperCamelCase      :			Optional[Any]									=   BlenderbotConfig(
							    vocab_size=self.vocab_size ,  d_model=1_4 ,  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=4_8 ,  )
							__UpperCamelCase      :			List[str]									=   FlaxBlenderbotForConditionalGeneration(__a			)
							__UpperCamelCase      :			str									=   np.array([[7_1, 8_2, 1_8, 3_3, 4_6, 9_1, 2], [6_8, 3_4, 2_6, 5_8, 3_0, 2, 1]] ,  dtype=np.intaa			)
							__UpperCamelCase      :			Optional[int]									=   np.array([[8_2, 7_1, 8_2, 1_8, 2], [5_8, 6_8, 2, 1, 1]] ,  dtype=np.intaa			)
							__UpperCamelCase      :			Optional[Any]									=   lm_model(input_ids=__a ,  decoder_input_ids=__a			)
							__UpperCamelCase      :			Any									=   (*summary.shape, config.vocab_size)
							self.assertEqual(outputs["logits"].shape ,  __a			)





				def        a_       (self			)	->       Union[str, Any]:
							__UpperCamelCase      :			str									=   np.array([[7_1, 8_2, 1_8, 3_3, 2, 1, 1], [6_8, 3_4, 2_6, 5_8, 3_0, 8_2, 2]] ,  dtype=np.intaa			)
							__UpperCamelCase      :			Union[str, Any]									=   shift_tokens_right(__a ,  1 ,  2			)
							__UpperCamelCase      :			List[Any]									=   np.equal(__a ,  1			).astype(np.floataa			).sum()
							__UpperCamelCase      :			int									=   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 					A		( lowerCamelCase_			,							unittest.TestCase			,							lowerCamelCase_		):




				'''simple docstring'''





				A      	=				True
				A      	=				(
				    (
				        FlaxBlenderbotModel,
				        FlaxBlenderbotForConditionalGeneration,
				    )
				    if is_flax_available()
				    else ()
				)
				A      	=				(FlaxBlenderbotForConditionalGeneration,) if is_flax_available() else ()
				def        a_       (self			)	->       int:
							__UpperCamelCase      :			Optional[int]									=   FlaxBlenderbotModelTester(self			)
				def        a_       (self			)	->       List[Any]:
							__UpperCamelCase		,					__UpperCamelCase      :			int									=   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        a_       (self			)	->       int:
							__UpperCamelCase		,					__UpperCamelCase      :			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        a_       (self			)	->       Optional[Any]:
							__UpperCamelCase		,					__UpperCamelCase      :			int									=   self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
										with self.subTest(model_class.__name__			):
													__UpperCamelCase      :			List[Any]									=   self._prepare_for_class(__a ,  __a			)
													__UpperCamelCase      :			Union[str, Any]									=   model_class(__a			)

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

													with self.subTest("JIT Enabled"			):
																__UpperCamelCase      :			Tuple									=   encode_jitted(**__a			).to_tuple()

													with self.subTest("JIT Disabled"			):
																with jax.disable_jit():
																			__UpperCamelCase      :			str									=   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        a_       (self			)	->       Optional[int]:
							__UpperCamelCase		,					__UpperCamelCase      :			Dict									=   self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
										with self.subTest(model_class.__name__			):
													__UpperCamelCase      :			Optional[Any]									=   model_class(__a			)
													__UpperCamelCase      :			Any									=   model.encode(inputs_dict["input_ids"] ,  inputs_dict["attention_mask"]			)

													__UpperCamelCase      :			Tuple									=   {
													    "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(_UpperCAmelCase ,  _UpperCAmelCase ,  _UpperCAmelCase			):
																return model.decode(
																    decoder_input_ids=__a ,  decoder_attention_mask=__a ,  encoder_outputs=__a ,  )

													with self.subTest("JIT Enabled"			):
																__UpperCamelCase      :			Any									=   decode_jitted(**__a			).to_tuple()

													with self.subTest("JIT Disabled"			):
																with jax.disable_jit():
																			__UpperCamelCase      :			List[str]									=   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        a_       (self			)	->       Optional[Any]:
							for model_class_name in self.all_model_classes:
										__UpperCamelCase      :			Any									=   model_class_name.from_pretrained("facebook/blenderbot-400M-distill"			)
										# FlaxBlenderbotForSequenceClassification expects eos token in input_ids
										__UpperCamelCase      :			List[str]									=   np.ones((1, 1)			) * model.config.eos_token_id
										__UpperCamelCase      :			Optional[int]									=   model(__a			)
										self.assertIsNotNone(__a			)





				@unittest.skipUnless(jax_device != "cpu" ,  "3B test too slow on CPU."			)
				@slow
				def        a_       (self			)	->       Optional[Any]:
							__UpperCamelCase      :			Dict									=   {"num_beams": 1, "early_stopping": True, "min_length": 1_5, "max_length": 2_5}
							__UpperCamelCase      :			List[Any]									=   {"skip_special_tokens": True, "clean_up_tokenization_spaces": True}

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

							__UpperCamelCase      :			Dict									=   ["Sam"]
							__UpperCamelCase      :			Optional[Any]									=   tokenizer(__a ,  return_tensors="jax"			)

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

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










'''simple docstring'''







from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


lowerCAmelCase_					:    int      	=			logging.get_logger(__name__)

lowerCAmelCase_					:    Tuple      	=			{
    'google/bigbird-roberta-base': 'https://huggingface.co/google/bigbird-roberta-base/resolve/main/config.json',
    'google/bigbird-roberta-large': 'https://huggingface.co/google/bigbird-roberta-large/resolve/main/config.json',
    'google/bigbird-base-trivia-itc': 'https://huggingface.co/google/bigbird-base-trivia-itc/resolve/main/config.json',
    # See all BigBird models at https://huggingface.co/models?filter=big_bird
}
class 				__SCREAMING_SNAKE_CASE  (lowerCamelCase_      ):






				"""simple docstring"""







				__a            ='big_bird'




				def __init__(  self     :					Optional[int]					,				__a     :					Dict=5_03_58					,				__a     :					str=7_68					,				__a     :					List[Any]=12					,				__a     :					List[str]=12					,				__a     :					Union[str, Any]=30_72					,				__a     :					str="gelu_new"					,				__a     :					Dict=0.1					,				__a     :					Union[str, Any]=0.1					,				__a     :					Any=40_96					,				__a     :					int=2					,				__a     :					Tuple=0.02					,				__a     :					List[Any]=1e-1_2					,				__a     :					int=True					,				__a     :					List[str]=0					,				__a     :					Tuple=1					,				__a     :					Optional[Any]=2					,				__a     :					Tuple=66					,				__a     :					str="block_sparse"					,				__a     :					Tuple=True					,				__a     :					Optional[int]=False					,				__a     :					str=64					,				__a     :					Tuple=3					,				__a     :					Any=None					,				**__a     :					Dict					,				):
								super().__init__(
								    pad_token_id=__a					,				bos_token_id=__a					,				eos_token_id=__a					,				sep_token_id=__a					,				**__a					,				)

								_a								= vocab_size
								_a								= max_position_embeddings
								_a								= hidden_size
								_a								= num_hidden_layers
								_a								= num_attention_heads
								_a								= intermediate_size
								_a								= hidden_act
								_a								= hidden_dropout_prob
								_a								= attention_probs_dropout_prob
								_a								= initializer_range
								_a								= type_vocab_size
								_a								= layer_norm_eps
								_a								= use_cache

								_a								= rescale_embeddings
								_a								= attention_type
								_a								= use_bias
								_a								= block_size
								_a								= num_random_blocks
								_a								= classifier_dropout


class 				__SCREAMING_SNAKE_CASE  (lowerCamelCase_      ):






				"""simple docstring"""







				@property
				def   UpperCamelCase__   (  self     :					Optional[int]					):
								if self.task == "multiple-choice":
												_a								= {0: "batch", 1: "choice", 2: "sequence"}
								else:
												_a								= {0: "batch", 1: "sequence"}
								return OrderedDict(
								    [
								        ("input_ids", dynamic_axis),
								        ("attention_mask", dynamic_axis),
								    ]					)
