
'''simple docstring'''



from sklearn.metrics import mean_squared_error

import datasets


_lowercase          =						"""\
@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}
}
"""

_lowercase          =						"""\
Mean Squared Error(MSE) is the average of the square of difference between the predicted
and actual values.
"""


_lowercase          =						"""
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     UpperCAmelCase_	(						datasets.Metric					):






						'''simple docstring'''






						def    _lowercase      (   self ):

												"""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    _lowercase      (   self ):

												"""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    _lowercase      (   self       ,  _lowercase       ,  _lowercase       ,  _lowercase=None       ,  _lowercase="uniform_average"       ,  _lowercase=True ):

												"""simple docstring"""



												_lowerCAmelCase    		=					mean_squared_error(
												    _lowercase       ,  _lowercase       ,  sample_weight=_lowercase       ,  multioutput=_lowercase       ,  squared=_lowercase )

												return {"mse": mse}





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 	__snake_case   :
		def __init__(							self    ,a_    ,a_=13    ,a_=7    ,a_=False    ,a_=True    ,a_=False    ,a_=True    ,a_=33    ,a_=32    ,a_=5    ,a_=4    ,a_=37    ,a_="gelu"    ,a_=0.1    ,a_=0.1    ,a_=512    ,a_=16    ,a_=2    ,a_=0.02    ,a_=3    ,a_=4    ,a_=None    ,):







						"""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 						SCREAMING_SNAKE_CASE_					(							self						):







						"""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 						SCREAMING_SNAKE_CASE_					(							self						):







						"""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 						SCREAMING_SNAKE_CASE_					(							self    ,a_    ,a_    ,a_    ,a_    ,a_    ,a_						):







						"""simple docstring"""
						lowerCAmelCase__		  =  EsmModel(config=a_						)
						model.to(a_						)
						model.eval()
						lowerCAmelCase__		  =  model(a_    ,attention_mask=a_						)
						lowerCAmelCase__		  =  model(a_						)
						lowerCAmelCase__		  =  model(a_						)

						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 						SCREAMING_SNAKE_CASE_					(							self    ,a_    ,a_    ,a_    ,a_    ,a_    ,a_						):







						"""simple docstring"""
						lowerCAmelCase__		  =  EsmForMaskedLM(config=a_						)
						model.to(a_						)
						model.eval()
						lowerCAmelCase__		  =  model(a_    ,attention_mask=a_    ,labels=a_						)
						self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.seq_length, self.vocab_size)						)
		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_    ,a_    ,a_    ,a_    ,a_    ,a_						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.num_labels
						lowerCAmelCase__		  =  EsmForTokenClassification(config=a_						)
						model.to(a_						)
						model.eval()
						lowerCAmelCase__		  =  model(a_    ,attention_mask=a_    ,labels=a_						)
						self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.seq_length, self.num_labels)						)



		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""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 	__snake_case   ( SCREAMING_SNAKE_CASE							, SCREAMING_SNAKE_CASE							, unittest.TestCase ):
		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 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  EsmModelTester(self						)
						lowerCAmelCase__		  =  ConfigTester(self    ,config_class=a_    ,hidden_size=37						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						self.config_tester.run_common_tests()
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""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(*a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*a_						)
		@slow
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowerCAmelCase__		  =  EsmModel.from_pretrained(a_						)
										self.assertIsNotNone(a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()[0]
						lowerCAmelCase__		  =  EsmEmbeddings(config=a_						)

						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(a_    ,model.padding_idx						)
						self.assertEqual(position_ids.shape    ,expected_positions.shape						)
						self.assertTrue(torch.all(torch.eq(a_    ,a_						)						)						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()[0]
						lowerCAmelCase__		  =  EsmEmbeddings(config=a_						)

						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(a_						)
						self.assertEqual(position_ids.shape    ,expected_positions.shape						)
						self.assertTrue(torch.all(torch.eq(a_    ,a_						)						)						)
		@unittest.skip('Esm does not support embedding resizing'						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						pass
		@unittest.skip('Esm does not support embedding resizing'						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						pass



		@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						pass



@require_torch
class 	__snake_case   ( SCREAMING_SNAKE_CASE ):
		@slow
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""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(a_						)[0]

										lowerCAmelCase__		  =  33

										lowerCAmelCase__		  =  torch.Size((1, 6, vocab_size)						)
										self.assertEqual(output.shape    ,a_						)

										lowerCAmelCase__		  =  torch.tensor(
										    [[[8.9215, -10.5898, -6.4671], [-6.3967, -13.9114, -1.1212], [-7.7812, -13.9516, -3.7406]]]						)
										self.assertTrue(torch.allclose(output[:, :3, :3]    ,a_    ,atol=1e-4						)						)



		@slow
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""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(a_						)[0]
										# compare the actual values for a slice.
										lowerCAmelCase__		  =  torch.tensor(
										    [[[0.1444, 0.5413, 0.3248], [0.3034, 0.0053, 0.3108], [0.3228, -0.2499, 0.3415]]]						)
										self.assertTrue(torch.allclose(output[:, :3, :3]    ,a_    ,atol=1e-4						)						)
