




import unittest

import numpy as np

from transformers import AlbertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
 import jax.numpy as jnp

 from transformers.models.albert.modeling_flax_albert import (
     FlaxAlbertForMaskedLM,
     FlaxAlbertForMultipleChoice,
     FlaxAlbertForPreTraining,
     FlaxAlbertForQuestionAnswering,
     FlaxAlbertForSequenceClassification,
     FlaxAlbertForTokenClassification,
     FlaxAlbertModel,
 )
class 						_A (  unittest.TestCase				):






  def __init__(  self		:     List[Any]				,						_A		:     str				,						_A		:     Optional[int]=13				,						_A		:     Dict=7				,						_A		:     Tuple=True				,						_A		:     Any=True				,						_A		:     List[str]=True				,						_A		:     int=True				,						_A		:     List[Any]=99				,						_A		:     int=32				,						_A		:     Any=5				,						_A		:     List[str]=4				,						_A		:     int=37				,						_A		:     Optional[Any]="gelu"				,						_A		:     str=0.1				,						_A		:     Any=0.1				,						_A		:     List[str]=512				,						_A		:     str=16				,						_A		:     str=2				,						_A		:     List[str]=0.02				,						_A		:     List[Any]=4				,						)		->   Dict:






       """simple docstring"""







       lowercase :       Union[str, Any]       =			parent
       lowercase :       List[Any]       =			batch_size
       lowercase :       List[str]       =			seq_length
       lowercase :       Optional[int]       =			is_training
       lowercase :       int       =			use_attention_mask
       lowercase :       Any       =			use_token_type_ids
       lowercase :       List[str]       =			use_labels
       lowercase :       List[Any]       =			vocab_size
       lowercase :       List[str]       =			hidden_size
       lowercase :       Dict       =			num_hidden_layers
       lowercase :       int       =			num_attention_heads
       lowercase :       List[Any]       =			intermediate_size
       lowercase :       List[str]       =			hidden_act
       lowercase :       Dict       =			hidden_dropout_prob
       lowercase :       Tuple       =			attention_probs_dropout_prob
       lowercase :       Dict       =			max_position_embeddings
       lowercase :       Union[str, Any]       =			type_vocab_size
       lowercase :       Optional[int]       =			type_sequence_label_size
       lowercase :       Optional[Any]       =			initializer_range
       lowercase :       str       =			num_choices






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






       """simple docstring"""







       lowercase :       List[str]       =			ids_tensor([self.batch_size, self.seq_length]				,						self.vocab_size				)

       lowercase :       List[str]       =			None
       if self.use_attention_mask:
            lowercase :       Optional[int]       =			random_attention_mask([self.batch_size, self.seq_length]				)

       lowercase :       Union[str, Any]       =			None
       if self.use_token_type_ids:
            lowercase :       Union[str, Any]       =			ids_tensor([self.batch_size, self.seq_length]				,						self.type_vocab_size				)

       lowercase :       Union[str, Any]       =			AlbertConfig(
           vocab_size=self.vocab_size				,						hidden_size=self.hidden_size				,						num_hidden_layers=self.num_hidden_layers				,						num_attention_heads=self.num_attention_heads				,						intermediate_size=self.intermediate_size				,						hidden_act=self.hidden_act				,						hidden_dropout_prob=self.hidden_dropout_prob				,						attention_probs_dropout_prob=self.attention_probs_dropout_prob				,						max_position_embeddings=self.max_position_embeddings				,						type_vocab_size=self.type_vocab_size				,						is_decoder=_A				,						initializer_range=self.initializer_range				,						)

       return config, input_ids, token_type_ids, attention_mask






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






       """simple docstring"""







       lowercase :       Optional[Any]       =			self.prepare_config_and_inputs()
       lowercase      ,	lowercase      ,	lowercase      ,	lowercase :       int       =			config_and_inputs
       lowercase :       List[str]       =			{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
       return config, inputs_dict


@require_flax
class 						_A (  _lowerCamelCase      ,       unittest.TestCase				):
  _UpperCamelCase       :      Optional[int]			=      (
      (
          FlaxAlbertModel,
          FlaxAlbertForPreTraining,
          FlaxAlbertForMaskedLM,
          FlaxAlbertForMultipleChoice,
          FlaxAlbertForQuestionAnswering,
          FlaxAlbertForSequenceClassification,
          FlaxAlbertForTokenClassification,
          FlaxAlbertForQuestionAnswering,
      )
      if is_flax_available()
      else ()
  )






  def  __a      (  self		:     str				)		->   Tuple:






       """simple docstring"""







       lowercase :       List[str]       =			FlaxAlbertModelTester(self				)






  @slow
  def  __a      (  self		:     int				)		->   Dict:






       """simple docstring"""







       for model_class_name in self.all_model_classes:
            lowercase :       Tuple       =			model_class_name.from_pretrained('''albert-base-v2'''				)
            lowercase :       Dict       =			model(np.ones((1, 1)				)				)
            self.assertIsNotNone(_A				)


@require_flax
class 						_A (  unittest.TestCase				):






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






       """simple docstring"""







       lowercase :       Optional[int]       =			FlaxAlbertModel.from_pretrained('''albert-base-v2'''				)
       lowercase :       Union[str, Any]       =			np.array([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]				)
       lowercase :       str       =			np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]				)
       lowercase :       Union[str, Any]       =			model(_A				,						attention_mask=_A				)[0]
       lowercase :       List[str]       =			(1, 11, 768)
       self.assertEqual(output.shape				,						_A				)
       lowercase :       List[str]       =			np.array(
           [[[-0.6_513, 1.5_035, -0.2_766], [-0.6_515, 1.5_046, -0.2_780], [-0.6_512, 1.5_049, -0.2_784]]]				)

       self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]				,						_A				,						atol=1E-4				)				)





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}