




"""simple docstring"""




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_						=			'platform'

 import jax
 import jax.numpy as jnp

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






def  snake_case(    __magic_name__			,		__magic_name__			,		__magic_name__=None			,		__magic_name__=None			,		__magic_name__=None			,		__magic_name__=None			,		__magic_name__=None			,		__magic_name__=None			,		)			->							str:





     '''simple docstring'''


     if attention_mask is None:
          lowercase :       Union[str, Any]       =			np.where(input_ids != config.pad_token_id			,		1			,		0   )
     if decoder_attention_mask is None:
          lowercase :       List[Any]       =			np.where(decoder_input_ids != config.pad_token_id			,		1			,		0   )
     if head_mask is None:
          lowercase :       str       =			np.ones((config.encoder_layers, config.encoder_attention_heads)   )
     if decoder_head_mask is None:
          lowercase :       Tuple       =			np.ones((config.decoder_layers, config.decoder_attention_heads)   )
     if cross_attn_head_mask is None:
          lowercase :       Optional[int]       =			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		:     List[Any]				,						_A		:     List[Any]				,						_A		:     Any=13				,						_A		:     Optional[int]=7				,						_A		:     List[str]=True				,						_A		:     Optional[Any]=False				,						_A		:     Optional[int]=99				,						_A		:     List[str]=16				,						_A		:     Any=2				,						_A		:     Tuple=4				,						_A		:     Dict=4				,						_A		:     List[str]="gelu"				,						_A		:     str=0.1				,						_A		:     List[str]=0.1				,						_A		:     Dict=32				,						_A		:     List[Any]=2				,						_A		:     Any=1				,						_A		:     Union[str, Any]=0				,						_A		:     Dict=0.02				,						)		->   int:






       """simple docstring"""







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






  def  __a      (  self		:     Dict				)		->   List[Any]:






       """simple docstring"""







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

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

       lowercase :       Dict       =			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				,						)
       lowercase :       str       =			prepare_blenderbot_inputs_dict(_A				,						_A				,						_A				)
       return config, inputs_dict






  def  __a      (  self		:     List[Any]				)		->   Dict:






       """simple docstring"""







       lowercase :       Optional[int]       =			self.prepare_config_and_inputs()
       return config, inputs_dict






  def  __a      (  self		:     Optional[Any]				,						_A		:     str				,						_A		:     Dict				,						_A		:     int				)		->   Optional[Any]:






       """simple docstring"""







       lowercase :       Tuple       =			20
       lowercase :       List[Any]       =			model_class_name(_A				)

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

       lowercase :       Union[str, Any]       =			(
           inputs_dict['''decoder_input_ids'''],
           inputs_dict['''decoder_attention_mask'''],
       )

       lowercase :       Union[str, Any]       =			model.init_cache(decoder_input_ids.shape[0]				,						_A				,						_A				)
       lowercase :       Tuple       =			jnp.ones((decoder_input_ids.shape[0], max_decoder_length)				,						dtype='''i4'''				)

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

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

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

       lowercase :       Tuple       =			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		:     Dict				,						_A		:     List[str]				,						_A		:     str				,						_A		:     Any				)		->   List[str]:






       """simple docstring"""







       lowercase :       List[str]       =			20
       lowercase :       Union[str, Any]       =			model_class_name(_A				)

       lowercase :       Dict       =			model.encode(inputs_dict['''input_ids''']				)

       lowercase :       Union[str, Any]       =			(
           inputs_dict['''decoder_input_ids'''],
           inputs_dict['''decoder_attention_mask'''],
       )

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

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

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

       lowercase :       int       =			model.decode(_A				,						_A				,						decoder_attention_mask=_A				)

       lowercase :       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 						_A (  unittest.TestCase				):
  _UpperCamelCase       :      Any			=      9_9






  def  __a      (  self		:     Optional[Any]				)		->   str:






       """simple docstring"""







       lowercase :       Union[str, 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 :       int       =			input_ids.shape[0]
       lowercase :       int       =			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  __a      (  self		:     Union[str, Any]				)		->   List[str]:






       """simple docstring"""







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






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






       """simple docstring"""







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






  def  __a      (  self		:     Dict				)		->   Dict:






       """simple docstring"""







       lowercase :       Union[str, 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(_A				,						1				,						2				)
       lowercase :       Any       =			np.equal(_A				,						1				).astype(np.floataa				).sum()
       lowercase :       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 						_A (  _lowerCamelCase      ,       unittest.TestCase      ,       _lowerCamelCase				):
  _UpperCamelCase       :      Optional[int]			=      True
  _UpperCamelCase       :      Tuple			=      (
      (
          FlaxBlenderbotSmallModel,
          FlaxBlenderbotSmallForConditionalGeneration,
      )
      if is_flax_available()
      else ()
  )
  _UpperCamelCase       :      int			=      (FlaxBlenderbotSmallForConditionalGeneration,) if is_flax_available() else ()






  def  __a      (  self		:     Any				)		->   str:






       """simple docstring"""







       lowercase :       Tuple       =			FlaxBlenderbotSmallModelTester(self				)






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






       """simple docstring"""







       lowercase :       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  __a      (  self		:     str				)		->   Optional[int]:






       """simple docstring"""







       lowercase :       List[str]       =			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		:     List[Any]				)		->   Optional[Any]:






       """simple docstring"""







       lowercase :       Optional[int]       =			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(_A				,						_A				)
                 lowercase :       Dict       =			model_class(_A				)

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

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

                 with self.subTest('''JIT Disabled'''				):
                      with jax.disable_jit():
                           lowercase :       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  __a      (  self		:     List[Any]				)		->   Optional[Any]:






       """simple docstring"""







       lowercase :       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[str]       =			model_class(_A				)
                 lowercase :       Dict       =			model.encode(inputs_dict['''input_ids''']				,						inputs_dict['''attention_mask''']				)

                 lowercase :       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		:     str				,						_A		:     Union[str, Any]				,						_A		:     Union[str, Any]				):
                      return model.decode(
                          decoder_input_ids=_A				,						decoder_attention_mask=_A				,						encoder_outputs=_A				,						)

                 with self.subTest('''JIT Enabled'''				):
                      lowercase :       List[Any]       =			decode_jitted(**_A				).to_tuple()

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






       """simple docstring"""







       for model_class_name in self.all_model_classes:
            lowercase :       Dict       =			model_class_name.from_pretrained('''facebook/blenderbot_small-90M'''				)
            # FlaxBlenderbotForSequenceClassification expects eos token in input_ids
            lowercase :       Dict       =			np.ones((1, 1)				) * model.config.eos_token_id
            lowercase :       str       =			model(_A				)
            self.assertIsNotNone(_A				)





from sklearn.metrics import mean_squared_error

import datasets


lowerCAmelCase_						=			'\\n@article{scikit-learn,\n title={Scikit-learn: Machine Learning in {P}ython},\n author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V.\n         and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P.\n         and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and\n         Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},\n journal={Journal of Machine Learning Research},\n volume={12},\n pages={2825--2830},\n year={2011}\n}\n'

lowerCAmelCase_						=			'\\nMean Squared Error(MSE) is the average of the square of difference between the predicted\nand actual values.\n'


lowerCAmelCase_						=			'\nArgs:\n    predictions: array-like of shape (n_samples,) or (n_samples, n_outputs)\n        Estimated target values.\n    references: array-like of shape (n_samples,) or (n_samples, n_outputs)\n        Ground truth (correct) target values.\n    sample_weight: array-like of shape (n_samples,), default=None\n        Sample weights.\n    multioutput: {"raw_values", "uniform_average"} or array-like of shape (n_outputs,), default="uniform_average"\n        Defines aggregating of multiple output values. Array-like value defines weights used to average errors.\n\n        "raw_values" : Returns a full set of errors in case of multioutput input.\n\n        "uniform_average" : Errors of all outputs are averaged with uniform weight.\n\n    squared : bool, default=True\n        If True returns MSE value, if False returns RMSE (Root Mean Squared Error) value.\n\nReturns:\n    mse : mean squared error.\nExamples:\n\n    >>> mse_metric = datasets.load_metric("mse")\n    >>> predictions = [2.5, 0.0, 2, 8]\n    >>> references = [3, -0.5, 2, 7]\n    >>> results = mse_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {\'mse\': 0.375}\n    >>> rmse_result = mse_metric.compute(predictions=predictions, references=references, squared=False)\n    >>> print(rmse_result)\n    {\'mse\': 0.6123724356957945}\n\n    If you\'re using multi-dimensional lists, then set the config as follows :\n\n    >>> mse_metric = datasets.load_metric("mse", "multilist")\n    >>> predictions = [[0.5, 1], [-1, 1], [7, -6]]\n    >>> references = [[0, 2], [-1, 2], [8, -5]]\n    >>> results = mse_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {\'mse\': 0.7083333333333334}\n    >>> results = mse_metric.compute(predictions=predictions, references=references, multioutput=\'raw_values\')\n    >>> print(results) # doctest: +NORMALIZE_WHITESPACE\n    {\'mse\': array([0.41666667, 1. ])}\n'
@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION      ,       _KWARGS_DESCRIPTION				)
class 						_A (  datasets.Metric				):






  def  __a      (  self		:     List[Any]				)		->   Optional[Any]:






       """simple docstring"""







       return datasets.MetricInfo(
           description=_DESCRIPTION				,						citation=_CITATION				,						inputs_description=_KWARGS_DESCRIPTION				,						features=datasets.Features(self._get_feature_types()				)				,						reference_urls=[
               '''https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_squared_error.html'''
           ]				,						)






  def  __a      (  self		:     List[Any]				)		->   int:






       """simple docstring"""







       if self.config_name == "multilist":
            return {
                "predictions": datasets.Sequence(datasets.Value('''float'''				)				),
                "references": datasets.Sequence(datasets.Value('''float'''				)				),
            }
       else:
            return {
                "predictions": datasets.Value('''float'''				),
                "references": datasets.Value('''float'''				),
            }






  def  __a      (  self		:     Any				,						_A		:     Dict				,						_A		:     Any				,						_A		:     Any=None				,						_A		:     Any="uniform_average"				,						_A		:     Optional[Any]=True				)		->   Dict:






       """simple docstring"""







       lowercase :       Any       =			mean_squared_error(
           _A				,						_A				,						sample_weight=_A				,						multioutput=_A				,						squared=_A				)

       return {"mse": mse}