


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

	'''simple docstring'''

	def __init__(     self		,				_UpperCAmelCase		,				_UpperCAmelCase=13		,				_UpperCAmelCase=7		,				_UpperCAmelCase=True		,				_UpperCAmelCase=True		,				_UpperCAmelCase=True		,				_UpperCAmelCase=True		,				_UpperCAmelCase=99		,				_UpperCAmelCase=32		,				_UpperCAmelCase=5		,				_UpperCAmelCase=4		,				_UpperCAmelCase=37		,				_UpperCAmelCase="gelu"		,				_UpperCAmelCase=0.1		,				_UpperCAmelCase=0.1		,				_UpperCAmelCase=5_12		,				_UpperCAmelCase=16		,				_UpperCAmelCase=2		,				_UpperCAmelCase=0.02		,				_UpperCAmelCase=4		,				):
							snake_case_	=				parent
							snake_case_	=				batch_size
							snake_case_	=				seq_length
							snake_case_	=				is_training
							snake_case_	=				use_attention_mask
							snake_case_	=				use_token_type_ids
							snake_case_	=				use_labels
							snake_case_	=				vocab_size
							snake_case_	=				hidden_size
							snake_case_	=				num_hidden_layers
							snake_case_	=				num_attention_heads
							snake_case_	=				intermediate_size
							snake_case_	=				hidden_act
							snake_case_	=				hidden_dropout_prob
							snake_case_	=				attention_probs_dropout_prob
							snake_case_	=				max_position_embeddings
							snake_case_	=				type_vocab_size
							snake_case_	=				type_sequence_label_size
							snake_case_	=				initializer_range
							snake_case_	=				num_choices



	def 						UpperCamelCase__					(     self ):
							snake_case_	=				ids_tensor([self.batch_size, self.seq_length]		,				self.vocab_size )

							snake_case_	=				None
							if self.use_attention_mask:
													snake_case_	=				random_attention_mask([self.batch_size, self.seq_length] )

							snake_case_	=				None
							if self.use_token_type_ids:
													snake_case_	=				ids_tensor([self.batch_size, self.seq_length]		,				self.type_vocab_size )

							snake_case_	=				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 						UpperCamelCase__					(     self ):
							snake_case_	=				self.prepare_config_and_inputs()
							snake_case_	,   snake_case_	,   snake_case_	,   snake_case_	=				config_and_inputs
							snake_case_	=				{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
							return config, inputs_dict







@require_flax
class lowerCAmelCase_						(			__a ,						unittest.TestCase	):

	'''simple docstring'''

	__snake_case												=  (
	    (
	        FlaxAlbertModel,
	        FlaxAlbertForPreTraining,
	        FlaxAlbertForMaskedLM,
	        FlaxAlbertForMultipleChoice,
	        FlaxAlbertForQuestionAnswering,
	        FlaxAlbertForSequenceClassification,
	        FlaxAlbertForTokenClassification,
	        FlaxAlbertForQuestionAnswering,
	    )
	    if is_flax_available()
	    else ()
	)



	def 						UpperCamelCase__					(     self ):
							snake_case_	=				FlaxAlbertModelTester(self )





	@slow
	def 						UpperCamelCase__					(     self ):
							for model_class_name in self.all_model_classes:
													snake_case_	=				model_class_name.from_pretrained('''albert-base-v2''' )
													snake_case_	=				model(np.ones((1, 1) ) )
													self.assertIsNotNone(a__ )







@require_flax
class lowerCAmelCase_						(			unittest.TestCase	):

	'''simple docstring'''

	@slow
	def 						UpperCamelCase__					(     self ):
							snake_case_	=				FlaxAlbertModel.from_pretrained('''albert-base-v2''' )
							snake_case_	=				np.array([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]] )
							snake_case_	=				np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]] )
							snake_case_	=				model(a__		,				attention_mask=a__ )[0]
							snake_case_	=				(1, 11, 7_68)
							self.assertEqual(output.shape		,				a__ )
							snake_case_	=				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 ) )



import numpy as np

import datasets


UpperCAmelCase      =   """
Compute the Mahalanobis Distance

Mahalonobis distance is the distance between a point and a distribution.
And not between two distinct points. It is effectively a multivariate equivalent of the Euclidean distance.
It was introduced by Prof. P. C. Mahalanobis in 1936
and has been used in various statistical applications ever since
[source: https://www.machinelearningplus.com/statistics/mahalanobis-distance/]
"""

UpperCAmelCase      =   """\
@article{de2000mahalanobis,
  title={The mahalanobis distance},
  author={De Maesschalck, Roy and Jouan-Rimbaud, Delphine and Massart, D{\'e}sir{\'e} L},
  journal={Chemometrics and intelligent laboratory systems},
  volume={50},
  number={1},
  pages={1--18},
  year={2000},
  publisher={Elsevier}
}
"""

UpperCAmelCase      =   """
Args:
    X: List of datapoints to be compared with the `reference_distribution`.
    reference_distribution: List of datapoints from the reference distribution we want to compare to.
Returns:
    mahalanobis: The Mahalonobis distance for each datapoint in `X`.
Examples:

    >>> mahalanobis_metric = datasets.load_metric(\"mahalanobis\")
    >>> results = mahalanobis_metric.compute(reference_distribution=[[0, 1], [1, 0]], X=[[0, 1]])
    >>> print(results)
    {'mahalanobis': array([0.5])}
"""







@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION ,						_KWARGS_DESCRIPTION	)
class lowerCAmelCase_						(			datasets.Metric	):

 '''simple docstring'''

 def 						UpperCamelCase__					(     self ):
       return datasets.MetricInfo(
           description=_DESCRIPTION		,				citation=_CITATION		,				inputs_description=_KWARGS_DESCRIPTION		,				features=datasets.Features(
               {
                   '''X''': datasets.Sequence(datasets.Value('''float'''		,				id='''sequence''' )		,				id='''X''' ),
               } )		,				)





 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase		,				_UpperCAmelCase ):
       # convert to numpy arrays
       snake_case_	=				np.array(_UpperCAmelCase )
       snake_case_	=				np.array(_UpperCAmelCase )

       # Assert that arrays are 2D
       if len(X.shape ) != 2:
             raise ValueError('''Expected `X` to be a 2D vector''' )
       if len(reference_distribution.shape ) != 2:
             raise ValueError('''Expected `reference_distribution` to be a 2D vector''' )
       if reference_distribution.shape[0] < 2:
             raise ValueError(
                 '''Expected `reference_distribution` to be a 2D vector with more than one element in the first dimension''' )

       # Get mahalanobis distance for each prediction
       snake_case_	=				X - np.mean(_UpperCAmelCase )
       snake_case_	=				np.cov(reference_distribution.T )
       try:
             snake_case_	=				np.linalg.inv(_UpperCAmelCase )
       except np.linalg.LinAlgError:
             snake_case_	=				np.linalg.pinv(_UpperCAmelCase )
       snake_case_	=				np.dot(_UpperCAmelCase		,				_UpperCAmelCase )
       snake_case_	=				np.dot(_UpperCAmelCase		,				X_minus_mu.T ).diagonal()

       return {"mahalanobis": mahal_dist}