





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
    __lowerCamelCase	:					int			  =					'platform'

    import jax
    import jax.numpy as jnp

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



def 		A_    (					_lowerCAmelCase							,       _lowerCAmelCase							,       _lowerCAmelCase=None							,       _lowerCAmelCase=None							,       _lowerCAmelCase=None							,       _lowerCAmelCase=None							,       _lowerCAmelCase=None							,       _lowerCAmelCase=None							,       )						->			Dict:


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





      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	,							):
       '''simple docstring'''




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





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




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

       UpperCamelCase     : int	     =       shift_tokens_right(A_	,							1	,							2       )

       UpperCamelCase     : 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=A_	,							)
       UpperCamelCase     : Dict	     =       prepare_blenderbot_inputs_dict(A_	,							A_	,							A_       )
       return config, inputs_dict





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




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





      def 	__UpperCamelCase(			self	,							A_	,							A_	,							A_       ):
       '''simple docstring'''




       UpperCamelCase     : Optional[int]	     =       20
       UpperCamelCase     : Union[str, Any]	     =       model_class_name(A_       )

       UpperCamelCase     : List[str]	     =       model.encode(inputs_dict["input_ids"]       )

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

       UpperCamelCase     : Union[str, Any]	     =       model.init_cache(decoder_input_ids.shape[0]	,							A_	,							A_       )
       UpperCamelCase     : Optional[int]	     =       jnp.ones((decoder_input_ids.shape[0], max_decoder_length)	,							dtype="i4"       )

       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[str]	     =       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     : Union[str, 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     : Any	     =       model.decode(A_	,							A_       )

       UpperCamelCase     : 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	,							A_	,							A_	,							A_       ):
       '''simple docstring'''




       UpperCamelCase     : Dict	     =       20
       UpperCamelCase     : Optional[int]	     =       model_class_name(A_       )

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

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

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

       UpperCamelCase     : Any	     =       model.init_cache(decoder_input_ids.shape[0]	,							A_	,							A_       )
       UpperCamelCase     : Dict	     =       jnp.broadcast_to(
           jnp.arange(decoder_input_ids.shape[-1] - 1       )[None, :]	,							(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)	,							)

       UpperCamelCase     : Union[str, Any]	     =       model.decode(
           decoder_input_ids[:, :-1]	,							A_	,							decoder_attention_mask=A_	,							past_key_values=A_	,							decoder_position_ids=A_	,							)
       UpperCamelCase     : List[str]	     =       jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]	,							dtype="i4"       )
       UpperCamelCase     : str	     =       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     : Optional[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      ):
      _UpperCAmelCase     :int						   =       9_9





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




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

       UpperCamelCase     : List[Any]	     =       input_ids.shape[0]
       UpperCamelCase     : 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 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : Optional[Any]	     =       self._get_config_and_data()
       UpperCamelCase     : Any	     =       FlaxBlenderbotSmallForConditionalGeneration(A_       )
       UpperCamelCase     : Dict	     =       lm_model(input_ids=A_       )
       UpperCamelCase     : Dict	     =       (batch_size, input_ids.shape[1], config.vocab_size)
       self.assertEqual(outputs["logits"].shape	,							A_       )





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




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





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : List[Any]	     =       np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]	,							dtype=np.intaa       )
       UpperCamelCase     : List[Any]	     =       shift_tokens_right(A_	,							1	,							2       )
       UpperCamelCase     : int	     =       np.equal(A_	,							1       ).astype(np.floataa       ).sum()
       UpperCamelCase     : Dict	     =       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__       (	SCREAMING_SNAKE_CASE__	,       unittest.TestCase	,       SCREAMING_SNAKE_CASE__      ):
      _UpperCAmelCase     :List[Any]						   =       True
      _UpperCAmelCase     :List[str]						   =       (
          (
              FlaxBlenderbotSmallModel,
              FlaxBlenderbotSmallForConditionalGeneration,
          )
          if is_flax_available()
          else ()
      )
      _UpperCAmelCase     :Tuple						   =       (FlaxBlenderbotSmallForConditionalGeneration,) if is_flax_available() else ()





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : str	     =       FlaxBlenderbotSmallModelTester(self       )





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : Tuple	     =       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       ):
       '''simple docstring'''




       UpperCamelCase     : 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 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       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     : Any	     =       self._prepare_for_class(A_	,							A_       )
         UpperCamelCase     : str	     =       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"       ):
          UpperCamelCase     : Any	     =       encode_jitted(**A_       ).to_tuple()

         with self.subTest("JIT Disabled"       ):
          with jax.disable_jit():
           UpperCamelCase     : List[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 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : Any	     =       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[str]	     =       model_class(A_       )
         UpperCamelCase     : Optional[int]	     =       model.encode(inputs_dict["input_ids"]	,							inputs_dict["attention_mask"]       )

         UpperCamelCase     : 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(A_	,							A_	,							A_       ):
          return model.decode(
              decoder_input_ids=A_	,							decoder_attention_mask=A_	,							encoder_outputs=A_	,							)

         with self.subTest("JIT Enabled"       ):
          UpperCamelCase     : Optional[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 	__UpperCamelCase(			self       ):
       '''simple docstring'''




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









import sys

from .dependency_versions_table import deps
from .utils.versions import require_version, require_version_core


# define which module versions we always want to check at run time
# (usually the ones defined in `install_requires` in setup.py)
#
# order specific notes:
# - tqdm must be checked before tokenizers

__lowerCamelCase	:					Optional[int]			  =					"""python tqdm regex requests packaging filelock numpy tokenizers""".split()
if sys.version_info < (3, 7):
    pkgs_to_check_at_runtime.append("""dataclasses""")
if sys.version_info < (3, 8):
    pkgs_to_check_at_runtime.append("""importlib_metadata""")

for pkg in pkgs_to_check_at_runtime:
    if pkg in deps:
        if pkg == "tokenizers":
            # must be loaded here, or else tqdm check may fail
            from .utils import is_tokenizers_available

            if not is_tokenizers_available():
                continue  # not required, check version only if installed

        require_version_core(deps[pkg])
    else:
        raise ValueError(f"""can't find {pkg} in {deps.keys()}, check dependency_versions_table.py""")



def 		A_    (					_lowerCAmelCase							,       _lowerCAmelCase=None					)						->			Optional[Any]:
 require_version(deps[pkg]							,       _lowerCAmelCase					)
