





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




		def __init__(self      :	str							,			A__      :	Any							,			A__      :	str=1_3							,			A__      :	str=7							,			A__      :	int=False							,			A__      :	Tuple=True							,			A__      :	Tuple=False							,			A__      :	Optional[int]=True							,			A__      :	Union[str, Any]=3_3							,			A__      :	Tuple=3_2							,			A__      :	Optional[int]=5							,			A__      :	Union[str, Any]=4							,			A__      :	Tuple=3_7							,			A__      :	int="gelu"							,			A__      :	Union[str, Any]=0.1							,			A__      :	Any=0.1							,			A__      :	Tuple=5_1_2							,			A__      :	Dict=1_6							,			A__      :	str=2							,			A__      :	Union[str, Any]=0.0_2							,			A__      :	str=3							,			A__      :	str=4							,			A__      :	int=None							,			)					-> int:
						lowercase         = parent
						lowercase         = batch_size
						lowercase         = seq_length
						lowercase         = is_training
						lowercase         = use_input_mask
						lowercase         = use_token_type_ids
						lowercase         = use_labels
						lowercase         = vocab_size
						lowercase         = hidden_size
						lowercase         = num_hidden_layers
						lowercase         = num_attention_heads
						lowercase         = intermediate_size
						lowercase         = hidden_act
						lowercase         = hidden_dropout_prob
						lowercase         = attention_probs_dropout_prob
						lowercase         = max_position_embeddings
						lowercase         = type_vocab_size
						lowercase         = type_sequence_label_size
						lowercase         = initializer_range
						lowercase         = num_labels
						lowercase         = num_choices
						lowercase         = scope




		def 				UpperCAmelCase__   (self      :	Any				)					-> Union[str, Any]:
						lowercase         = ids_tensor([self.batch_size, self.seq_length]							,			self.vocab_size				)

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

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

						lowercase         = self.get_config()

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




		def 				UpperCAmelCase__   (self      :	Dict				)					-> Tuple:
						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      :	List[Any]							,			A__      :	Optional[Any]							,			A__      :	Optional[int]							,			A__      :	Union[str, Any]							,			A__      :	int							,			A__      :	List[str]							,			A__      :	List[Any]				)					-> Union[str, Any]:
						lowercase         = EsmModel(config=A__				)
						model.to(A__				)
						model.eval()
						lowercase         = model(A__							,			attention_mask=A__				)
						lowercase         = model(A__				)
						lowercase         = 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 				UpperCAmelCase__   (self      :	Optional[Any]							,			A__      :	Dict							,			A__      :	Optional[Any]							,			A__      :	str							,			A__      :	List[Any]							,			A__      :	int							,			A__      :	List[Any]				)					-> Optional[Any]:
						lowercase         = EsmForMaskedLM(config=A__				)
						model.to(A__				)
						model.eval()
						lowercase         = model(A__							,			attention_mask=A__							,			labels=A__				)
						self.parent.assertEqual(result.logits.shape							,			(self.batch_size, self.seq_length, self.vocab_size)				)




		def 				UpperCAmelCase__   (self      :	Optional[Any]							,			A__      :	Tuple							,			A__      :	List[Any]							,			A__      :	str							,			A__      :	Optional[Any]							,			A__      :	Optional[int]							,			A__      :	Optional[int]				)					-> Any:
						lowercase         = self.num_labels
						lowercase         = EsmForTokenClassification(config=A__				)
						model.to(A__				)
						model.eval()
						lowercase         = model(A__							,			attention_mask=A__							,			labels=A__				)
						self.parent.assertEqual(result.logits.shape							,			(self.batch_size, self.seq_length, self.num_labels)				)





		def 				UpperCAmelCase__   (self      :	Optional[Any]				)					-> int:
						lowercase         = self.prepare_config_and_inputs()
						(
						    (
						    lowercase
						)   ,							(
						    lowercase
						)   ,							(
						    lowercase
						)   ,							(
						    lowercase
						)   ,							(
						    lowercase
						)   ,							(
						    lowercase
						)   ,							
						)         = config_and_inputs
						lowercase         = {"input_ids": input_ids, "attention_mask": input_mask}
						return config, inputs_dict




@require_torch
class       UpperCAmelCase (						_lowercase			,       _lowercase			,       unittest.TestCase  ):
		UpperCAmelCase					:	int 			=		False

		UpperCAmelCase					:	Optional[int] 			=		(
		    (
		        EsmForMaskedLM,
		        EsmModel,
		        EsmForSequenceClassification,
		        EsmForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		UpperCAmelCase					:	Dict 			=		()
		UpperCAmelCase					:	Dict 			=		(
		    {
		        '''feature-extraction''': EsmModel,
		        '''fill-mask''': EsmForMaskedLM,
		        '''text-classification''': EsmForSequenceClassification,
		        '''token-classification''': EsmForTokenClassification,
		        '''zero-shot''': EsmForSequenceClassification,
		    }
		    if is_torch_available()
		    else {}
		)
		UpperCAmelCase					:	List[Any] 			=		True




		def 				UpperCAmelCase__   (self      :	List[Any]				)					-> Dict:
						lowercase         = EsmModelTester(self				)
						lowercase         = ConfigTester(self							,			config_class=A__							,			hidden_size=3_7				)




		def 				UpperCAmelCase__   (self      :	Union[str, Any]				)					-> Tuple:
						self.config_tester.run_common_tests()




		def 				UpperCAmelCase__   (self      :	str				)					-> List[Any]:
						lowercase         = self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*A__				)




		def 				UpperCAmelCase__   (self      :	Optional[Any]				)					-> List[str]:
						lowercase         = self.model_tester.prepare_config_and_inputs()
						for type in ["absolute", "relative_key", "relative_key_query"]:
										lowercase         = type
										self.model_tester.create_and_check_model(*A__				)




		def 				UpperCAmelCase__   (self      :	int				)					-> List[Any]:
						lowercase         = self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*A__				)




		def 				UpperCAmelCase__   (self      :	Optional[int]				)					-> Union[str, Any]:
						lowercase         = self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*A__				)




		@slow
		def 				UpperCAmelCase__   (self      :	int				)					-> Optional[int]:
						for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowercase         = EsmModel.from_pretrained(A__				)
										self.assertIsNotNone(A__				)




		def 				UpperCAmelCase__   (self      :	Tuple				)					-> Any:
						lowercase         = self.model_tester.prepare_config_and_inputs()[0]
						lowercase         = EsmEmbeddings(config=A__				)

						lowercase         = torch.as_tensor([[1_2, 3_1, 1_3, model.padding_idx]]				)
						lowercase         = torch.as_tensor(
						    [
						        [
						            0 + model.padding_idx + 1,
						            1 + model.padding_idx + 1,
						            2 + model.padding_idx + 1,
						            model.padding_idx,
						        ]
						    ]				)
						lowercase         = 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 				UpperCAmelCase__   (self      :	Optional[int]				)					-> Any:
						lowercase         = self.model_tester.prepare_config_and_inputs()[0]
						lowercase         = EsmEmbeddings(config=A__				)

						lowercase         = torch.empty(2							,			4							,			3_0				)
						lowercase         = [
						    0 + embeddings.padding_idx + 1,
						    1 + embeddings.padding_idx + 1,
						    2 + embeddings.padding_idx + 1,
						    3 + embeddings.padding_idx + 1,
						]
						lowercase         = torch.as_tensor([expected_single_positions, expected_single_positions]				)
						lowercase         = 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 				UpperCAmelCase__   (self      :	Optional[Any]				)					-> List[Any]:
						pass




		@unittest.skip("Esm does not support embedding resizing"				)
		def 				UpperCAmelCase__   (self      :	List[str]				)					-> Tuple:
						pass





		@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."				)
		def 				UpperCAmelCase__   (self      :	Dict				)					-> Tuple:
						pass




@require_torch
class       UpperCAmelCase (						_lowercase  ):




		@slow
		def 				UpperCAmelCase__   (self      :	Dict				)					-> Optional[Any]:
						with torch.no_grad():
										lowercase         = EsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"				)
										model.eval()
										lowercase         = torch.tensor([[0, 1, 2, 3, 4, 5]]				)
										lowercase         = model(A__				)[0]

										lowercase         = 3_3

										lowercase         = torch.Size((1, 6, vocab_size)				)
										self.assertEqual(output.shape							,			A__				)

										lowercase         = torch.tensor(
										    [[[8.9_2_1_5, -1_0.5_8_9_8, -6.4_6_7_1], [-6.3_9_6_7, -1_3.9_1_1_4, -1.1_2_1_2], [-7.7_8_1_2, -1_3.9_5_1_6, -3.7_4_0_6]]]				)
										self.assertTrue(torch.allclose(output[:, :3, :3]							,			A__							,			atol=1e-4				)				)





		@slow
		def 				UpperCAmelCase__   (self      :	Tuple				)					-> Optional[Any]:
						with torch.no_grad():
										lowercase         = EsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"				)
										model.eval()

										lowercase         = torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]				)
										lowercase         = model(A__				)[0]
										# compare the actual values for a slice.
										lowercase         = torch.tensor(
										    [[[0.1_4_4_4, 0.5_4_1_3, 0.3_2_4_8], [0.3_0_3_4, 0.0_0_5_3, 0.3_1_0_8], [0.3_2_2_8, -0.2_4_9_9, 0.3_4_1_5]]]				)
										self.assertTrue(torch.allclose(output[:, :3, :3]							,			A__							,			atol=1e-4				)				)











'''simple docstring'''

from sklearn.metrics import fa_score, matthews_corrcoef

import datasets

from .record_evaluation import evaluate as evaluate_record


__lowerCamelCase : str											=			"\\n@article{wang2019superglue,\n  title={SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systems},\n  author={Wang, Alex and Pruksachatkun, Yada and Nangia, Nikita and Singh, Amanpreet and Michael, Julian and Hill, Felix and Levy, Omer and Bowman, Samuel R},\n  journal={arXiv preprint arXiv:1905.00537},\n  year={2019}\n}\n"

__lowerCamelCase : Optional[int]											=			"\\nSuperGLUE (https://super.gluebenchmark.com/) is a new benchmark styled after\nGLUE with a new set of more difficult language understanding tasks, improved\nresources, and a new public leaderboard.\n"

__lowerCamelCase : int											=			"\nCompute SuperGLUE evaluation metric associated to each SuperGLUE dataset.\nArgs:\n    predictions: list of predictions to score. Depending on the SuperGlUE subset:\n        - for 'record': list of question-answer dictionaries with the following keys:\n            - 'idx': index of the question as specified by the dataset\n            - 'prediction_text': the predicted answer text\n        - for 'multirc': list of question-answer dictionaries with the following keys:\n            - 'idx': index of the question-answer pair as specified by the dataset\n            - 'prediction': the predicted answer label\n        - otherwise: list of predicted labels\n    references: list of reference labels. Depending on the SuperGLUE subset:\n        - for 'record': list of question-answers dictionaries with the following keys:\n            - 'idx': index of the question as specified by the dataset\n            - 'answers': list of possible answers\n        - otherwise: list of reference labels\nReturns: depending on the SuperGLUE subset:\n    - for 'record':\n        - 'exact_match': Exact match between answer and gold answer\n        - 'f1': F1 score\n    - for 'multirc':\n        - 'exact_match': Exact match between answer and gold answer\n        - 'f1_m': Per-question macro-F1 score\n        - 'f1_a': Average F1 score over all answers\n    - for 'axb':\n        'matthews_correlation': Matthew Correlation\n    - for 'cb':\n        - 'accuracy': Accuracy\n        - 'f1': F1 score\n    - for all others:\n        - 'accuracy': Accuracy\nExamples:\n\n    >>> super_glue_metric = datasets.load_metric('super_glue', 'copa')  # any of [\"copa\", \"rte\", \"wic\", \"wsc\", \"wsc.fixed\", \"boolq\", \"axg\"]\n    >>> predictions = [0, 1]\n    >>> references = [0, 1]\n    >>> results = super_glue_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'accuracy': 1.0}\n\n    >>> super_glue_metric = datasets.load_metric('super_glue', 'cb')\n    >>> predictions = [0, 1]\n    >>> references = [0, 1]\n    >>> results = super_glue_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'accuracy': 1.0, 'f1': 1.0}\n\n    >>> super_glue_metric = datasets.load_metric('super_glue', 'record')\n    >>> predictions = [{'idx': {'passage': 0, 'query': 0}, 'prediction_text': 'answer'}]\n    >>> references = [{'idx': {'passage': 0, 'query': 0}, 'answers': ['answer', 'another_answer']}]\n    >>> results = super_glue_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'exact_match': 1.0, 'f1': 1.0}\n\n    >>> super_glue_metric = datasets.load_metric('super_glue', 'multirc')\n    >>> predictions = [{'idx': {'answer': 0, 'paragraph': 0, 'question': 0}, 'prediction': 0}, {'idx': {'answer': 1, 'paragraph': 2, 'question': 3}, 'prediction': 1}]\n    >>> references = [0, 1]\n    >>> results = super_glue_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'exact_match': 1.0, 'f1_m': 1.0, 'f1_a': 1.0}\n\n    >>> super_glue_metric = datasets.load_metric('super_glue', 'axb')\n    >>> references = [0, 1]\n    >>> predictions = [0, 1]\n    >>> results = super_glue_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'matthews_correlation': 1.0}\n"







def 							UpperCAmelCase_    (	lowerCAmelCase_		,							lowerCAmelCase_							):

				"""simple docstring"""






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







def 							UpperCAmelCase_    (	lowerCAmelCase_		,							lowerCAmelCase_		,							lowerCAmelCase_="binary"							):

				"""simple docstring"""






				lowercase         = simple_accuracy(lowerCAmelCase_		,							lowerCAmelCase_							)
				lowercase         = float(fa_score(y_true=lowerCAmelCase_		,							y_pred=lowerCAmelCase_		,							average=lowerCAmelCase_							)							)
				return {
				    "accuracy": acc,
				    "f1": fa,
				}







def 							UpperCAmelCase_    (	lowerCAmelCase_		,							lowerCAmelCase_							):

				"""simple docstring"""






				lowercase         = {}
				for id_pred, label in zip(lowerCAmelCase_		,							lowerCAmelCase_							):
								lowercase         = f'{id_pred["idx"]["paragraph"]}-{id_pred["idx"]["question"]}'
								lowercase         = id_pred["prediction"]
								if question_id in question_map:
												question_map[question_id].append((pred, label)							)
								else:
												lowercase         = [(pred, label)]
				lowercase   ,							lowercase         = [], []
				for question, preds_labels in question_map.items():
								lowercase   ,							lowercase         = zip(*lowerCAmelCase_							)
								lowercase         = fa_score(y_true=lowerCAmelCase_		,							y_pred=lowerCAmelCase_		,							average="macro"							)
								fas.append(lowerCAmelCase_							)
								lowercase         = int(sum(pred == label for pred, label in preds_labels							) == len(lowerCAmelCase_							)							)
								ems.append(lowerCAmelCase_							)
				lowercase         = float(sum(lowerCAmelCase_							) / len(lowerCAmelCase_							)							)
				lowercase         = sum(lowerCAmelCase_							) / len(lowerCAmelCase_							)
				lowercase         = float(fa_score(y_true=lowerCAmelCase_		,							y_pred=[id_pred["prediction"] for id_pred in ids_preds]							)							)
				return {"exact_match": em, "f1_m": fa_m, "f1_a": fa_a}


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




		def 				UpperCAmelCase__   (self      :	Any				)					-> str:
						if self.config_name not in [
						    "boolq",
						    "cb",
						    "copa",
						    "multirc",
						    "record",
						    "rte",
						    "wic",
						    "wsc",
						    "wsc.fixed",
						    "axb",
						    "axg",
						]:
										raise KeyError(
										    "You should supply a configuration name selected in "
										    "[\"boolq\", \"cb\", \"copa\", \"multirc\", \"record\", \"rte\", \"wic\", \"wsc\", \"wsc.fixed\", \"axb\", \"axg\",]"				)
						return datasets.MetricInfo(
						    description=_DESCRIPTION							,			citation=_CITATION							,			inputs_description=_KWARGS_DESCRIPTION							,			features=datasets.Features(self._get_feature_types()				)							,			codebase_urls=[]							,			reference_urls=[]							,			format="numpy" if not self.config_name == "record" and not self.config_name == "multirc" else None							,			)




		def 				UpperCAmelCase__   (self      :	int				)					-> Optional[int]:
						if self.config_name == "record":
										return {
										    "predictions": {
										        "idx": {
										            "passage": datasets.Value("int64"				),
										            "query": datasets.Value("int64"				),
										        },
										        "prediction_text": datasets.Value("string"				),
										    },
										    "references": {
										        "idx": {
										            "passage": datasets.Value("int64"				),
										            "query": datasets.Value("int64"				),
										        },
										        "answers": datasets.Sequence(datasets.Value("string"				)				),
										    },
										}
						elif self.config_name == "multirc":
										return {
										    "predictions": {
										        "idx": {
										            "answer": datasets.Value("int64"				),
										            "paragraph": datasets.Value("int64"				),
										            "question": datasets.Value("int64"				),
										        },
										        "prediction": datasets.Value("int64"				),
										    },
										    "references": datasets.Value("int64"				),
										}
						else:
										return {
										    "predictions": datasets.Value("int64"				),
										    "references": datasets.Value("int64"				),
										}





		def 				UpperCAmelCase__   (self      :	Any							,			A__      :	int							,			A__      :	List[str]				)					-> Dict:
						if self.config_name == "axb":
										return {"matthews_correlation": matthews_corrcoef(A__							,			A__				)}
						elif self.config_name == "cb":
										return acc_and_fa(A__							,			A__							,			fa_avg="macro"				)
						elif self.config_name == "record":
										lowercase         = [
										    {
										        "qas": [
										            {"id": ref["idx"]["query"], "answers": [{"text": ans} for ans in ref["answers"]]}
										            for ref in references
										        ]
										    }
										]
										lowercase         = {pred["idx"]["query"]: pred["prediction_text"] for pred in predictions}
										return evaluate_record(A__							,			A__				)[0]
						elif self.config_name == "multirc":
										return evaluate_multirc(A__							,			A__				)
						elif self.config_name in ["copa", "rte", "wic", "wsc", "wsc.fixed", "boolq", "axg"]:
										return {"accuracy": simple_accuracy(A__							,			A__				)}
						else:
										raise KeyError(
										    "You should supply a configuration name selected in "
										    "[\"boolq\", \"cb\", \"copa\", \"multirc\", \"record\", \"rte\", \"wic\", \"wsc\", \"wsc.fixed\", \"axb\", \"axg\",]"				)




