





import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotSmallConfig, is_flax_available
from transformers.testing_utils import 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
							UpperCAmelCase_							=				"""platform"""

							import jax
							import jax.numpy as jnp

							from transformers.models.blenderbot_small.modeling_flax_blenderbot_small import (
							    FlaxBlenderbotSmallForConditionalGeneration,
							    FlaxBlenderbotSmallModel,
							    shift_tokens_right,
							)





def  __magic_name__  (				lowercase				,       lowercase				,       lowercase=None				,       lowercase=None				,       lowercase=None				,       lowercase=None				,       lowercase=None				,       lowercase=None				,       )				-> Tuple:

							"""simple docstring"""



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






		'''simple docstring'''






		def __init__(							self,   snake_case__,   snake_case__=13,   snake_case__=7,   snake_case__=True,   snake_case__=False,   snake_case__=99,   snake_case__=16,   snake_case__=2,   snake_case__=4,   snake_case__=4,   snake_case__="gelu",   snake_case__=0.1,   snake_case__=0.1,   snake_case__=32,   snake_case__=2,   snake_case__=1,   snake_case__=0,   snake_case__=0.02,   )       -> Optional[int]:






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






		def 						snake_case__				(							self  )       -> str:






									"""simple docstring"""
									lowercase_	:   Optional[Any]  				=     np.clip(ids_tensor([self.batch_size, self.seq_length - 1],   self.vocab_size  ),   3,   self.vocab_size  )
									lowercase_	:   Optional[int]  				=     np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1),   dtype=np.intaa  )),   -1  )

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

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






		def 						snake_case__				(							self  )       -> int:






									"""simple docstring"""
									lowercase_				,		lowercase_	:   Dict  				=     self.prepare_config_and_inputs()
									return config, inputs_dict






		def 						snake_case__				(							self,   snake_case__,   snake_case__,   snake_case__  )       -> Dict:






									"""simple docstring"""
									lowercase_	:   str  				=     20
									lowercase_	:   int  				=     model_class_name(snake_case__  )

									lowercase_	:   Dict  				=     model.encode(inputs_dict["""input_ids"""]  )

									lowercase_				,		lowercase_	:   int  				=     (
									    inputs_dict["""decoder_input_ids"""],
									    inputs_dict["""decoder_attention_mask"""],
									)

									lowercase_	:   Optional[Any]  				=     model.init_cache(decoder_input_ids.shape[0],   snake_case__,   snake_case__  )
									lowercase_	:   Optional[Any]  				=     jnp.ones((decoder_input_ids.shape[0], max_decoder_length),   dtype="""i4"""  )

									lowercase_	:   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_	:   Dict  				=     model.decode(
									    decoder_input_ids[:, :-1],   snake_case__,   decoder_attention_mask=snake_case__,   past_key_values=snake_case__,   decoder_position_ids=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:],   snake_case__,   decoder_attention_mask=snake_case__,   past_key_values=outputs_cache.past_key_values,   decoder_position_ids=snake_case__,   )

									lowercase_	:   int  				=     model.decode(snake_case__,   snake_case__  )

									lowercase_	:   Dict  				=     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,   snake_case__,   snake_case__,   snake_case__  )       -> List[Any]:






									"""simple docstring"""
									lowercase_	:   Optional[int]  				=     20
									lowercase_	:   List[str]  				=     model_class_name(snake_case__  )

									lowercase_	:   Optional[Any]  				=     model.encode(inputs_dict["""input_ids"""]  )

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

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

									lowercase_	:   Tuple  				=     model.init_cache(decoder_input_ids.shape[0],   snake_case__,   snake_case__  )
									lowercase_	:   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_	:   Union[str, Any]  				=     model.decode(
									    decoder_input_ids[:, :-1],   snake_case__,   decoder_attention_mask=snake_case__,   past_key_values=snake_case__,   decoder_position_ids=snake_case__,   )
									lowercase_	:   Any  				=     jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]],   dtype="""i4"""  )
									lowercase_	:   Optional[int]  				=     model.decode(
									    decoder_input_ids[:, -1:],   snake_case__,   past_key_values=outputs_cache.past_key_values,   decoder_attention_mask=snake_case__,   decoder_position_ids=snake_case__,   )

									lowercase_	:   Any  				=     model.decode(snake_case__,   snake_case__,   decoder_attention_mask=snake_case__  )

									lowercase_	:   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     UpperCamelCase__					(     unittest.TestCase				):






		'''simple docstring'''






		__a       :       List[Any]       =	99






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






									"""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_	:   Optional[Any]  				=     input_ids.shape[0]
									lowercase_	:   str  				=     BlenderbotSmallConfig(
									    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  )       -> Any:






									"""simple docstring"""
									lowercase_				,		lowercase_				,		lowercase_	:   Any  				=     self._get_config_and_data()
									lowercase_	:   List[str]  				=     FlaxBlenderbotSmallForConditionalGeneration(snake_case__  )
									lowercase_	:   Any  				=     lm_model(input_ids=snake_case__  )
									lowercase_	:   Tuple  				=     (batch_size, input_ids.shape[1], config.vocab_size)
									self.assertEqual(outputs["""logits"""].shape,   snake_case__  )






		def 						snake_case__				(							self  )       -> List[Any]:






									"""simple docstring"""
									lowercase_	:   Union[str, Any]  				=     BlenderbotSmallConfig(
									    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  				=     FlaxBlenderbotSmallForConditionalGeneration(snake_case__  )
									lowercase_	:   Tuple  				=     np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]],   dtype=np.intaa  )
									lowercase_	:   Dict  				=     np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]],   dtype=np.intaa  )
									lowercase_	:   Any  				=     lm_model(input_ids=snake_case__,   decoder_input_ids=snake_case__  )
									lowercase_	:   List[str]  				=     (*summary.shape, config.vocab_size)
									self.assertEqual(outputs["""logits"""].shape,   snake_case__  )






		def 						snake_case__				(							self  )       -> Any:






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




@require_flax
class     UpperCamelCase__					(     lowerCamelCase__				, unittest.TestCase				, lowerCamelCase__				):






		'''simple docstring'''






		__a       :       int       =	True
		__a       :       int       =	(
		    (
		        FlaxBlenderbotSmallModel,
		        FlaxBlenderbotSmallForConditionalGeneration,
		    )
		    if is_flax_available()
		    else ()
		)
		__a       :       int       =	(FlaxBlenderbotSmallForConditionalGeneration,) if is_flax_available() else ()






		def 						snake_case__				(							self  )       -> Any:






									"""simple docstring"""
									lowercase_	:   Any  				=     FlaxBlenderbotSmallModelTester(self  )






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






									"""simple docstring"""
									lowercase_				,		lowercase_	:   Union[str, Any]  				=     self.model_tester.prepare_config_and_inputs()
									for model_class in self.all_model_classes:
																self.model_tester.check_use_cache_forward(snake_case__,   snake_case__,   snake_case__  )






		def 						snake_case__				(							self  )       -> str:






									"""simple docstring"""
									lowercase_				,		lowercase_	:   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(snake_case__,   snake_case__,   snake_case__  )






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






									"""simple docstring"""
									lowercase_				,		lowercase_	:   Union[str, Any]  				=     self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
																with self.subTest(model_class.__name__  ):
																							lowercase_	:   List[Any]  				=     self._prepare_for_class(snake_case__,   snake_case__  )
																							lowercase_	:   int  				=     model_class(snake_case__  )

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

																							with self.subTest("""JIT Enabled"""  ):
																														lowercase_	:   List[Any]  				=     encode_jitted(**snake_case__  ).to_tuple()

																							with self.subTest("""JIT Disabled"""  ):
																														with jax.disable_jit():
																																					lowercase_	:   int  				=     encode_jitted(**snake_case__  ).to_tuple()

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






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






									"""simple docstring"""
									lowercase_				,		lowercase_	:   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__  ):
																							lowercase_	:   List[Any]  				=     model_class(snake_case__  )
																							lowercase_	:   List[Any]  				=     model.encode(inputs_dict["""input_ids"""],   inputs_dict["""attention_mask"""]  )

																							lowercase_	:   Union[str, 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(snake_case__,   snake_case__,   snake_case__  ):
																														return model.decode(
																														    decoder_input_ids=snake_case__,   decoder_attention_mask=snake_case__,   encoder_outputs=snake_case__,   )

																							with self.subTest("""JIT Enabled"""  ):
																														lowercase_	:   Union[str, Any]  				=     decode_jitted(**snake_case__  ).to_tuple()

																							with self.subTest("""JIT Disabled"""  ):
																														with jax.disable_jit():
																																					lowercase_	:   Union[str, Any]  				=     decode_jitted(**snake_case__  ).to_tuple()

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






		@slow
		def 						snake_case__				(							self  )       -> List[str]:






									"""simple docstring"""
									for model_class_name in self.all_model_classes:
																lowercase_	:   List[str]  				=     model_class_name.from_pretrained("""facebook/blenderbot_small-90M"""  )
																# FlaxBlenderbotForSequenceClassification expects eos token in input_ids
																lowercase_	:   Union[str, Any]  				=     np.ones((1, 1)  ) * model.config.eos_token_id
																lowercase_	:   Tuple  				=     model(snake_case__  )
																self.assertIsNotNone(snake_case__  )






import pytest
import requests

from datasets.utils.file_utils import http_head

from .utils import OfflineSimulationMode, RequestWouldHangIndefinitelyError, offline





@pytest.mark.integration
def  __magic_name__  (				)				-> Union[str, Any]:

							"""simple docstring"""



							with offline(OfflineSimulationMode.CONNECTION_TIMES_OUT				):
														with pytest.raises(lowercase				):
																					requests.request("""GET"""				,       """https://huggingface.co"""				)
														with pytest.raises(requests.exceptions.ConnectTimeout				):
																					requests.request("""GET"""				,       """https://huggingface.co"""				,       timeout=1.0				)





@pytest.mark.integration
def  __magic_name__  (				)				-> Union[str, Any]:

							"""simple docstring"""



							with offline(OfflineSimulationMode.CONNECTION_FAILS				):
														with pytest.raises(requests.exceptions.ConnectionError				):
																					requests.request("""GET"""				,       """https://huggingface.co"""				)





def  __magic_name__  (				)				-> Dict:

							"""simple docstring"""



							with offline(OfflineSimulationMode.HF_DATASETS_OFFLINE_SET_TO_1				):
														with pytest.raises(lowercase				):
																					http_head("""https://huggingface.co"""				)