



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  lowerCamelCase__    :




		'''simple docstring'''



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

						"""simple docstring"""






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



		def 	lowerCAmelCase__   (self				)  ->    Dict:

						"""simple docstring"""






						lowerCAmelCase__  :       Union[str, Any]				=  ids_tensor([self.batch_size, self.seq_length]	,self.vocab_size				)

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

						lowerCAmelCase__  :       Optional[int]				=  None
						lowerCAmelCase__  :       Union[str, Any]				=  None
						lowerCAmelCase__  :       List[Any]				=  None
						if self.use_labels:
										lowerCAmelCase__  :       Optional[int]				=  ids_tensor([self.batch_size]	,self.type_sequence_label_size				)
										lowerCAmelCase__  :       str				=  ids_tensor([self.batch_size, self.seq_length]	,self.num_labels				)
										lowerCAmelCase__  :       List[Any]				=  ids_tensor([self.batch_size]	,self.num_choices				)

						lowerCAmelCase__  :       Union[str, Any]				=  self.get_config()

						return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels



		def 	lowerCAmelCase__   (self				)  ->    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 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)  ->    int:

						"""simple docstring"""






						lowerCAmelCase__  :       Dict				=  EsmModel(config=__lowerCamelCase				)
						model.to(__lowerCamelCase				)
						model.eval()
						lowerCAmelCase__  :       str				=  model(__lowerCamelCase	,attention_mask=__lowerCamelCase				)
						lowerCAmelCase__  :       Optional[Any]				=  model(__lowerCamelCase				)
						lowerCAmelCase__  :       List[Any]				=  model(__lowerCamelCase				)

						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 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)  ->    int:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  EsmForMaskedLM(config=__lowerCamelCase				)
						model.to(__lowerCamelCase				)
						model.eval()
						lowerCAmelCase__  :       List[Any]				=  model(__lowerCamelCase	,attention_mask=__lowerCamelCase	,labels=__lowerCamelCase				)
						self.parent.assertEqual(result.logits.shape	,(self.batch_size, self.seq_length, self.vocab_size)				)



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)  ->    List[Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       Dict				=  self.num_labels
						lowerCAmelCase__  :       Optional[int]				=  EsmForTokenClassification(config=__lowerCamelCase				)
						model.to(__lowerCamelCase				)
						model.eval()
						lowerCAmelCase__  :       Optional[Any]				=  model(__lowerCamelCase	,attention_mask=__lowerCamelCase	,labels=__lowerCamelCase				)
						self.parent.assertEqual(result.logits.shape	,(self.batch_size, self.seq_length, self.num_labels)				)



		def 	lowerCAmelCase__   (self				)  ->    Dict:

						"""simple docstring"""






						lowerCAmelCase__  :       Tuple				=  self.prepare_config_and_inputs()
						(
						    (
						    lowerCAmelCase__
						)		,	(
						    lowerCAmelCase__
						)		,	(
						    lowerCAmelCase__
						)		,	(
						    lowerCAmelCase__
						)		,	(
						    lowerCAmelCase__
						)		,	(
						    lowerCAmelCase__
						)		,	
						)  :       List[str]				=  config_and_inputs
						lowerCAmelCase__  :       str				=  {'''input_ids''': input_ids, '''attention_mask''': input_mask}
						return config, inputs_dict



@require_torch
class  lowerCamelCase__    (							lowerCamelCase__   ,       lowerCamelCase__   ,       unittest.TestCase):




		'''simple docstring'''



		snake_case_			=False

		snake_case_			=(
		    (
		        EsmForMaskedLM,
		        EsmModel,
		        EsmForSequenceClassification,
		        EsmForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		snake_case_			=()
		snake_case_			=(
		    {
		        """feature-extraction""": EsmModel,
		        """fill-mask""": EsmForMaskedLM,
		        """text-classification""": EsmForSequenceClassification,
		        """token-classification""": EsmForTokenClassification,
		        """zero-shot""": EsmForSequenceClassification,
		    }
		    if is_torch_available()
		    else {}
		)
		snake_case_			=True



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

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  EsmModelTester(self				)
						lowerCAmelCase__  :       int				=  ConfigTester(self	,config_class=__lowerCamelCase	,hidden_size=37				)



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

						"""simple docstring"""






						self.config_tester.run_common_tests()



		def 	lowerCAmelCase__   (self				)  ->    Any:

						"""simple docstring"""






						lowerCAmelCase__  :       List[str]				=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*__lowerCamelCase				)



		def 	lowerCAmelCase__   (self				)  ->    Any:

						"""simple docstring"""






						lowerCAmelCase__  :       Dict				=  self.model_tester.prepare_config_and_inputs()
						for type in ["absolute", "relative_key", "relative_key_query"]:
										lowerCAmelCase__  :       List[Any]				=  type
										self.model_tester.create_and_check_model(*__lowerCamelCase				)



		def 	lowerCAmelCase__   (self				)  ->    Tuple:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[int]				=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*__lowerCamelCase				)



		def 	lowerCAmelCase__   (self				)  ->    Tuple:

						"""simple docstring"""






						lowerCAmelCase__  :       Union[str, Any]				=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*__lowerCamelCase				)



		@slow
		def 	lowerCAmelCase__   (self				)  ->    Union[str, Any]:

						"""simple docstring"""






						for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowerCAmelCase__  :       int				=  EsmModel.from_pretrained(__lowerCamelCase				)
										self.assertIsNotNone(__lowerCamelCase				)



		def 	lowerCAmelCase__   (self				)  ->    str:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[Any]				=  self.model_tester.prepare_config_and_inputs()[0]
						lowerCAmelCase__  :       int				=  EsmEmbeddings(config=__lowerCamelCase				)

						lowerCAmelCase__  :       Optional[int]				=  torch.as_tensor([[12, 31, 13, model.padding_idx]]				)
						lowerCAmelCase__  :       str				=  torch.as_tensor(
						    [
						        [
						            0 + model.padding_idx + 1,
						            1 + model.padding_idx + 1,
						            2 + model.padding_idx + 1,
						            model.padding_idx,
						        ]
						    ]				)
						lowerCAmelCase__  :       List[Any]				=  create_position_ids_from_input_ids(__lowerCamelCase	,model.padding_idx				)
						self.assertEqual(position_ids.shape	,expected_positions.shape				)
						self.assertTrue(torch.all(torch.eq(__lowerCamelCase	,__lowerCamelCase				)				)				)



		def 	lowerCAmelCase__   (self				)  ->    int:

						"""simple docstring"""






						lowerCAmelCase__  :       Tuple				=  self.model_tester.prepare_config_and_inputs()[0]
						lowerCAmelCase__  :       Union[str, Any]				=  EsmEmbeddings(config=__lowerCamelCase				)

						lowerCAmelCase__  :       Dict				=  torch.empty(2	,4	,30				)
						lowerCAmelCase__  :       Tuple				=  [
						    0 + embeddings.padding_idx + 1,
						    1 + embeddings.padding_idx + 1,
						    2 + embeddings.padding_idx + 1,
						    3 + embeddings.padding_idx + 1,
						]
						lowerCAmelCase__  :       Tuple				=  torch.as_tensor([expected_single_positions, expected_single_positions]				)
						lowerCAmelCase__  :       Optional[int]				=  embeddings.create_position_ids_from_inputs_embeds(__lowerCamelCase				)
						self.assertEqual(position_ids.shape	,expected_positions.shape				)
						self.assertTrue(torch.all(torch.eq(__lowerCamelCase	,__lowerCamelCase				)				)				)



		@unittest.skip('''Esm does not support embedding resizing'''				)
		def 	lowerCAmelCase__   (self				)  ->    str:

						"""simple docstring"""






						pass



		@unittest.skip('''Esm does not support embedding resizing'''				)
		def 	lowerCAmelCase__   (self				)  ->    int:

						"""simple docstring"""






						pass



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

						"""simple docstring"""






						pass



@require_torch
class  lowerCamelCase__    (							lowerCamelCase__):




		'''simple docstring'''



		@slow
		def 	lowerCAmelCase__   (self				)  ->    Optional[int]:

						"""simple docstring"""






						with torch.no_grad():
										lowerCAmelCase__  :       Tuple				=  EsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''				)
										model.eval()
										lowerCAmelCase__  :       Optional[Any]				=  torch.tensor([[0, 1, 2, 3, 4, 5]]				)
										lowerCAmelCase__  :       List[str]				=  model(__lowerCamelCase				)[0]

										lowerCAmelCase__  :       Optional[int]				=  33

										lowerCAmelCase__  :       Tuple				=  torch.Size((1, 6, vocab_size)				)
										self.assertEqual(output.shape	,__lowerCamelCase				)

										lowerCAmelCase__  :       List[Any]				=  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]	,__lowerCamelCase	,atol=1e-4				)				)



		@slow
		def 	lowerCAmelCase__   (self				)  ->    Optional[int]:

						"""simple docstring"""






						with torch.no_grad():
										lowerCAmelCase__  :       Optional[int]				=  EsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''				)
										model.eval()

										lowerCAmelCase__  :       Optional[Any]				=  torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]				)
										lowerCAmelCase__  :       Tuple				=  model(__lowerCamelCase				)[0]
										# compare the actual values for a slice.
										lowerCAmelCase__  :       int				=  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]	,__lowerCamelCase	,atol=1e-4				)				)






from scipy.stats import pearsonr, spearmanr
from sklearn.metrics import fa_score, matthews_corrcoef

import datasets


__snake_case       :	str            ='\\n@inproceedings{wang2019glue,\n  title={{GLUE}: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding},\n  author={Wang, Alex and Singh, Amanpreet and Michael, Julian and Hill, Felix and Levy, Omer and Bowman, Samuel R.},\n  note={In the Proceedings of ICLR.},\n  year={2019}\n}\n'

__snake_case       :	Tuple            ='\\nGLUE, the General Language Understanding Evaluation benchmark\n(https://gluebenchmark.com/) is a collection of resources for training,\nevaluating, and analyzing natural language understanding systems.\n'

__snake_case       :	int            ='\nCompute GLUE evaluation metric associated to each GLUE dataset.\nArgs:\n    predictions: list of predictions to score.\n        Each translation should be tokenized into a list of tokens.\n    references: list of lists of references for each translation.\n        Each reference should be tokenized into a list of tokens.\nReturns: depending on the GLUE subset, one or several of:\n    "accuracy": Accuracy\n    "f1": F1 score\n    "pearson": Pearson Correlation\n    "spearmanr": Spearman Correlation\n    "matthews_correlation": Matthew Correlation\nExamples:\n\n    >>> glue_metric = datasets.load_metric(\'glue\', \'sst2\')  # \'sst2\' or any of ["mnli", "mnli_mismatched", "mnli_matched", "qnli", "rte", "wnli", "hans"]\n    >>> references = [0, 1]\n    >>> predictions = [0, 1]\n    >>> results = glue_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {\'accuracy\': 1.0}\n\n    >>> glue_metric = datasets.load_metric(\'glue\', \'mrpc\')  # \'mrpc\' or \'qqp\'\n    >>> references = [0, 1]\n    >>> predictions = [0, 1]\n    >>> results = glue_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {\'accuracy\': 1.0, \'f1\': 1.0}\n\n    >>> glue_metric = datasets.load_metric(\'glue\', \'stsb\')\n    >>> references = [0., 1., 2., 3., 4., 5.]\n    >>> predictions = [0., 1., 2., 3., 4., 5.]\n    >>> results = glue_metric.compute(predictions=predictions, references=references)\n    >>> print({"pearson": round(results["pearson"], 2), "spearmanr": round(results["spearmanr"], 2)})\n    {\'pearson\': 1.0, \'spearmanr\': 1.0}\n\n    >>> glue_metric = datasets.load_metric(\'glue\', \'cola\')\n    >>> references = [0, 1]\n    >>> predictions = [0, 1]\n    >>> results = glue_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {\'matthews_correlation\': 1.0}\n'






def 			lowerCAmelCase__	(					lowerCamelCase_	:	List[str] ,lowerCamelCase_	:	Dict):





				'''simple docstring'''

				return float((preds == labels).mean())






def 			lowerCAmelCase__	(					lowerCamelCase_	:	List[Any] ,lowerCamelCase_	:	Optional[int]):





				'''simple docstring'''

				lowerCAmelCase__  :       int				=  simple_accuracy(lowerCamelCase_ ,lowerCamelCase_)
				lowerCAmelCase__  :       Any				=  float(fa_score(y_true=lowerCamelCase_ ,y_pred=lowerCamelCase_))
				return {
				    "accuracy": acc,
				    "f1": fa,
				}






def 			lowerCAmelCase__	(					lowerCamelCase_	:	List[Any] ,lowerCamelCase_	:	Optional[Any]):





				'''simple docstring'''

				lowerCAmelCase__  :       Union[str, Any]				=  float(pearsonr(lowerCamelCase_ ,lowerCamelCase_)[0])
				lowerCAmelCase__  :       int				=  float(spearmanr(lowerCamelCase_ ,lowerCamelCase_)[0])
				return {
				    "pearson": pearson_corr,
				    "spearmanr": spearman_corr,
				}



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




		'''simple docstring'''



		def 	lowerCAmelCase__   (self				)  ->    Union[str, Any]:

						"""simple docstring"""






						if self.config_name not in [
						    "sst2",
						    "mnli",
						    "mnli_mismatched",
						    "mnli_matched",
						    "cola",
						    "stsb",
						    "mrpc",
						    "qqp",
						    "qnli",
						    "rte",
						    "wnli",
						    "hans",
						]:
										raise KeyError(
										    '''You should supply a configuration name selected in '''
										    '''["sst2", "mnli", "mnli_mismatched", "mnli_matched", '''
										    '''"cola", "stsb", "mrpc", "qqp", "qnli", "rte", "wnli", "hans"]'''				)
						return datasets.MetricInfo(
						    description=_DESCRIPTION	,citation=_CITATION	,inputs_description=_KWARGS_DESCRIPTION	,features=datasets.Features(
						        {
						            '''predictions''': datasets.Value('''int64''' if self.config_name != '''stsb''' else '''float32'''				),
						            '''references''': datasets.Value('''int64''' if self.config_name != '''stsb''' else '''float32'''				),
						        }				)	,codebase_urls=[]	,reference_urls=[]	,format='''numpy'''	,)



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase				)  ->    Tuple:

						"""simple docstring"""






						if self.config_name == "cola":
										return {"matthews_correlation": matthews_corrcoef(__lowerCamelCase	,__lowerCamelCase				)}
						elif self.config_name == "stsb":
										return pearson_and_spearman(__lowerCamelCase	,__lowerCamelCase				)
						elif self.config_name in ["mrpc", "qqp"]:
										return acc_and_fa(__lowerCamelCase	,__lowerCamelCase				)
						elif self.config_name in ["sst2", "mnli", "mnli_mismatched", "mnli_matched", "qnli", "rte", "wnli", "hans"]:
										return {"accuracy": simple_accuracy(__lowerCamelCase	,__lowerCamelCase				)}
						else:
										raise KeyError(
										    '''You should supply a configuration name selected in '''
										    '''["sst2", "mnli", "mnli_mismatched", "mnli_matched", '''
										    '''"cola", "stsb", "mrpc", "qqp", "qnli", "rte", "wnli", "hans"]'''				)

