

'''simple docstring'''





import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
				import torch

				from transformers import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
				from transformers.models.esm.modeling_esm import (
				    ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
				    EsmEmbeddings,
				    create_position_ids_from_input_ids,
				)


class  UpperCAmelCase :

		'''simple docstring'''






		def __init__(					self   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_=13   ,			SCREAMING_SNAKE_CASE_=7   ,			SCREAMING_SNAKE_CASE_=False   ,			SCREAMING_SNAKE_CASE_=True   ,			SCREAMING_SNAKE_CASE_=False   ,			SCREAMING_SNAKE_CASE_=True   ,			SCREAMING_SNAKE_CASE_=33   ,			SCREAMING_SNAKE_CASE_=32   ,			SCREAMING_SNAKE_CASE_=5   ,			SCREAMING_SNAKE_CASE_=4   ,			SCREAMING_SNAKE_CASE_=37   ,			SCREAMING_SNAKE_CASE_="gelu"   ,			SCREAMING_SNAKE_CASE_=0.1   ,			SCREAMING_SNAKE_CASE_=0.1   ,			SCREAMING_SNAKE_CASE_=512   ,			SCREAMING_SNAKE_CASE_=16   ,			SCREAMING_SNAKE_CASE_=2   ,			SCREAMING_SNAKE_CASE_=0.02   ,			SCREAMING_SNAKE_CASE_=3   ,			SCREAMING_SNAKE_CASE_=4   ,			SCREAMING_SNAKE_CASE_=None   ,			)  ->     int:






							'''simple docstring'''



							lowerCamelCase_  				=      parent
							lowerCamelCase_  				=      batch_size
							lowerCamelCase_  				=      seq_length
							lowerCamelCase_  				=      is_training
							lowerCamelCase_  				=      use_input_mask
							lowerCamelCase_  				=      use_token_type_ids
							lowerCamelCase_  				=      use_labels
							lowerCamelCase_  				=      vocab_size
							lowerCamelCase_  				=      hidden_size
							lowerCamelCase_  				=      num_hidden_layers
							lowerCamelCase_  				=      num_attention_heads
							lowerCamelCase_  				=      intermediate_size
							lowerCamelCase_  				=      hidden_act
							lowerCamelCase_  				=      hidden_dropout_prob
							lowerCamelCase_  				=      attention_probs_dropout_prob
							lowerCamelCase_  				=      max_position_embeddings
							lowerCamelCase_  				=      type_vocab_size
							lowerCamelCase_  				=      type_sequence_label_size
							lowerCamelCase_  				=      initializer_range
							lowerCamelCase_  				=      num_labels
							lowerCamelCase_  				=      num_choices
							lowerCamelCase_  				=      scope
		def 			UpperCamelCase(					self		)  ->     str:






							'''simple docstring'''



							lowerCamelCase_  				=      ids_tensor([self.batch_size, self.seq_length]   ,			self.vocab_size		)

							lowerCamelCase_  				=      None
							if self.use_input_mask:
												lowerCamelCase_  				=      random_attention_mask([self.batch_size, self.seq_length]		)

							lowerCamelCase_  				=      None
							lowerCamelCase_  				=      None
							lowerCamelCase_  				=      None
							if self.use_labels:
												lowerCamelCase_  				=      ids_tensor([self.batch_size]   ,			self.type_sequence_label_size		)
												lowerCamelCase_  				=      ids_tensor([self.batch_size, self.seq_length]   ,			self.num_labels		)
												lowerCamelCase_  				=      ids_tensor([self.batch_size]   ,			self.num_choices		)

							lowerCamelCase_  				=      self.get_config()

							return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
		def 			UpperCamelCase(					self		)  ->     List[str]:






							'''simple docstring'''



							return EsmConfig(
							    vocab_size=self.vocab_size   ,			hidden_size=self.hidden_size   ,			pad_token_id=1   ,			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   ,			initializer_range=self.initializer_range   ,			)
		def 			UpperCamelCase(					self   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)  ->     Optional[int]:






							'''simple docstring'''



							lowerCamelCase_  				=      EsmModel(config=SCREAMING_SNAKE_CASE_		)
							model.to(SCREAMING_SNAKE_CASE_		)
							model.eval()
							lowerCamelCase_  				=      model(SCREAMING_SNAKE_CASE_   ,			attention_mask=SCREAMING_SNAKE_CASE_		)
							lowerCamelCase_  				=      model(SCREAMING_SNAKE_CASE_		)
							lowerCamelCase_  				=      model(SCREAMING_SNAKE_CASE_		)

							self.parent.assertEqual(result.last_hidden_state.shape   ,			(self.batch_size, self.seq_length, self.hidden_size)		)
							self.parent.assertEqual(result.pooler_output.shape   ,			(self.batch_size, self.hidden_size)		)
		def 			UpperCamelCase(					self   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)  ->     Tuple:






							'''simple docstring'''



							lowerCamelCase_  				=      EsmForMaskedLM(config=SCREAMING_SNAKE_CASE_		)
							model.to(SCREAMING_SNAKE_CASE_		)
							model.eval()
							lowerCamelCase_  				=      model(SCREAMING_SNAKE_CASE_   ,			attention_mask=SCREAMING_SNAKE_CASE_   ,			labels=SCREAMING_SNAKE_CASE_		)
							self.parent.assertEqual(result.logits.shape   ,			(self.batch_size, self.seq_length, self.vocab_size)		)
		def 			UpperCamelCase(					self   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)  ->     Any:






							'''simple docstring'''



							lowerCamelCase_  				=      self.num_labels
							lowerCamelCase_  				=      EsmForTokenClassification(config=SCREAMING_SNAKE_CASE_		)
							model.to(SCREAMING_SNAKE_CASE_		)
							model.eval()
							lowerCamelCase_  				=      model(SCREAMING_SNAKE_CASE_   ,			attention_mask=SCREAMING_SNAKE_CASE_   ,			labels=SCREAMING_SNAKE_CASE_		)
							self.parent.assertEqual(result.logits.shape   ,			(self.batch_size, self.seq_length, self.num_labels)		)






		def 			UpperCamelCase(					self		)  ->     Any:






							'''simple docstring'''



							lowerCamelCase_  				=      self.prepare_config_and_inputs()
							(
							    (
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,
							)  				=      config_and_inputs
							lowerCamelCase_  				=      {'input_ids': input_ids, 'attention_mask': input_mask}
							return config, inputs_dict


@require_torch
class  UpperCAmelCase (							UpperCAmelCase__  , UpperCAmelCase__  , unittest.TestCase    ):

		'''simple docstring'''






		SCREAMING_SNAKE_CASE_           =					False

		SCREAMING_SNAKE_CASE_           =					(
		    (
		        EsmForMaskedLM,
		        EsmModel,
		        EsmForSequenceClassification,
		        EsmForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		SCREAMING_SNAKE_CASE_           =					()
		SCREAMING_SNAKE_CASE_           =					(
		    {
		        'feature-extraction': EsmModel,
		        'fill-mask': EsmForMaskedLM,
		        'text-classification': EsmForSequenceClassification,
		        'token-classification': EsmForTokenClassification,
		        'zero-shot': EsmForSequenceClassification,
		    }
		    if is_torch_available()
		    else {}
		)
		SCREAMING_SNAKE_CASE_           =					True
		def 			UpperCamelCase(					self		)  ->     List[Any]:






							'''simple docstring'''



							lowerCamelCase_  				=      EsmModelTester(self		)
							lowerCamelCase_  				=      ConfigTester(self   ,			config_class=SCREAMING_SNAKE_CASE_   ,			hidden_size=37		)
		def 			UpperCamelCase(					self		)  ->     List[Any]:






							'''simple docstring'''



							self.config_tester.run_common_tests()
		def 			UpperCamelCase(					self		)  ->     Any:






							'''simple docstring'''



							lowerCamelCase_  				=      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_		)
		def 			UpperCamelCase(					self		)  ->     str:






							'''simple docstring'''



							lowerCamelCase_  				=      self.model_tester.prepare_config_and_inputs()
							for type in ["absolute", "relative_key", "relative_key_query"]:
												lowerCamelCase_  				=      type
												self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_		)
		def 			UpperCamelCase(					self		)  ->     Dict:






							'''simple docstring'''



							lowerCamelCase_  				=      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE_		)
		def 			UpperCamelCase(					self		)  ->     int:






							'''simple docstring'''



							lowerCamelCase_  				=      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE_		)
		@slow
		def 			UpperCamelCase(					self		)  ->     Any:






							'''simple docstring'''



							for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												lowerCamelCase_  				=      EsmModel.from_pretrained(SCREAMING_SNAKE_CASE_		)
												self.assertIsNotNone(SCREAMING_SNAKE_CASE_		)
		def 			UpperCamelCase(					self		)  ->     Any:






							'''simple docstring'''



							lowerCamelCase_  				=      self.model_tester.prepare_config_and_inputs()[0]
							lowerCamelCase_  				=      EsmEmbeddings(config=SCREAMING_SNAKE_CASE_		)

							lowerCamelCase_  				=      torch.as_tensor([[12, 31, 13, model.padding_idx]]		)
							lowerCamelCase_  				=      torch.as_tensor(
							    [
							        [
							            0 + model.padding_idx + 1,
							            1 + model.padding_idx + 1,
							            2 + model.padding_idx + 1,
							            model.padding_idx,
							        ]
							    ]		)
							lowerCamelCase_  				=      create_position_ids_from_input_ids(SCREAMING_SNAKE_CASE_   ,			model.padding_idx		)
							self.assertEqual(position_ids.shape   ,			expected_positions.shape		)
							self.assertTrue(torch.all(torch.eq(SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)		)		)
		def 			UpperCamelCase(					self		)  ->     List[Any]:






							'''simple docstring'''



							lowerCamelCase_  				=      self.model_tester.prepare_config_and_inputs()[0]
							lowerCamelCase_  				=      EsmEmbeddings(config=SCREAMING_SNAKE_CASE_		)

							lowerCamelCase_  				=      torch.empty(2   ,			4   ,			30		)
							lowerCamelCase_  				=      [
							    0 + embeddings.padding_idx + 1,
							    1 + embeddings.padding_idx + 1,
							    2 + embeddings.padding_idx + 1,
							    3 + embeddings.padding_idx + 1,
							]
							lowerCamelCase_  				=      torch.as_tensor([expected_single_positions, expected_single_positions]		)
							lowerCamelCase_  				=      embeddings.create_position_ids_from_inputs_embeds(SCREAMING_SNAKE_CASE_		)
							self.assertEqual(position_ids.shape   ,			expected_positions.shape		)
							self.assertTrue(torch.all(torch.eq(SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)		)		)
		@unittest.skip('Esm does not support embedding resizing'		)
		def 			UpperCamelCase(					self		)  ->     Dict:






							'''simple docstring'''



							pass
		@unittest.skip('Esm does not support embedding resizing'		)
		def 			UpperCamelCase(					self		)  ->     Any:






							'''simple docstring'''



							pass






		@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'		)
		def 			UpperCamelCase(					self		)  ->     List[Any]:






							'''simple docstring'''



							pass


@require_torch
class  UpperCAmelCase (							UpperCAmelCase__    ):

		'''simple docstring'''






		@slow
		def 			UpperCamelCase(					self		)  ->     Any:






							'''simple docstring'''



							with torch.no_grad():
												lowerCamelCase_  				=      EsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D'		)
												model.eval()
												lowerCamelCase_  				=      torch.tensor([[0, 1, 2, 3, 4, 5]]		)
												lowerCamelCase_  				=      model(SCREAMING_SNAKE_CASE_		)[0]

												lowerCamelCase_  				=      33

												lowerCamelCase_  				=      torch.Size((1, 6, vocab_size)		)
												self.assertEqual(output.shape   ,			SCREAMING_SNAKE_CASE_		)

												lowerCamelCase_  				=      torch.tensor(
												    [[[8.9_215, -10.5_898, -6.4_671], [-6.3_967, -13.9_114, -1.1_212], [-7.7_812, -13.9_516, -3.7_406]]]		)
												self.assertTrue(torch.allclose(output[:, :3, :3]   ,			SCREAMING_SNAKE_CASE_   ,			atol=1E-4		)		)






		@slow
		def 			UpperCamelCase(					self		)  ->     Tuple:






							'''simple docstring'''



							with torch.no_grad():
												lowerCamelCase_  				=      EsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D'		)
												model.eval()

												lowerCamelCase_  				=      torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]		)
												lowerCamelCase_  				=      model(SCREAMING_SNAKE_CASE_		)[0]
												# compare the actual values for a slice.
												lowerCamelCase_  				=      torch.tensor(
												    [[[0.1_444, 0.5_413, 0.3_248], [0.3_034, 0.0_053, 0.3_108], [0.3_228, -0.2_499, 0.3_415]]]		)
												self.assertTrue(torch.allclose(output[:, :3, :3]   ,			SCREAMING_SNAKE_CASE_   ,			atol=1E-4		)		)





from sklearn.metrics import mean_squared_error

import datasets


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

A		: str								=     '''\
Mean Squared Error(MSE) is the average of the square of difference between the predicted
and actual values.
'''


A		: Any								=     '''
Args:
    predictions: array-like of shape (n_samples,) or (n_samples, n_outputs)
        Estimated target values.
    references: array-like of shape (n_samples,) or (n_samples, n_outputs)
        Ground truth (correct) target values.
    sample_weight: array-like of shape (n_samples,), default=None
        Sample weights.
    multioutput: {"raw_values", "uniform_average"} or array-like of shape (n_outputs,), default="uniform_average"
        Defines aggregating of multiple output values. Array-like value defines weights used to average errors.

        "raw_values" : Returns a full set of errors in case of multioutput input.

        "uniform_average" : Errors of all outputs are averaged with uniform weight.

    squared : bool, default=True
        If True returns MSE value, if False returns RMSE (Root Mean Squared Error) value.

Returns:
    mse : mean squared error.
Examples:

    >>> mse_metric = datasets.load_metric("mse")
    >>> predictions = [2.5, 0.0, 2, 8]
    >>> references = [3, -0.5, 2, 7]
    >>> results = mse_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'mse\': 0.375}
    >>> rmse_result = mse_metric.compute(predictions=predictions, references=references, squared=False)
    >>> print(rmse_result)
    {\'mse\': 0.6123724356957945}

    If you\'re using multi-dimensional lists, then set the config as follows :

    >>> mse_metric = datasets.load_metric("mse", "multilist")
    >>> predictions = [[0.5, 1], [-1, 1], [7, -6]]
    >>> references = [[0, 2], [-1, 2], [8, -5]]
    >>> results = mse_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'mse\': 0.7083333333333334}
    >>> results = mse_metric.compute(predictions=predictions, references=references, multioutput=\'raw_values\')
    >>> print(results) # doctest: +NORMALIZE_WHITESPACE
    {\'mse\': array([0.41666667, 1. ])}
'''



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






			def    UpperCamelCase_				(      self  ):
									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    UpperCamelCase_				(      self  ):
									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    UpperCamelCase_				(      self				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase=None				,			__UpperCamelCase="uniform_average"				,			__UpperCamelCase=True  ):
									_lowercase					       =					mean_squared_error(
									    __UpperCamelCase				,			__UpperCamelCase				,			sample_weight=__UpperCamelCase				,			multioutput=__UpperCamelCase				,			squared=__UpperCamelCase  )

									return {"mse": mse}