


"""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 			lowerCamelCase_		:





				"""simple docstring"""







				def __init__(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__=13       ,		UpperCAmelCase__=7       ,		UpperCAmelCase__=False       ,		UpperCAmelCase__=True       ,		UpperCAmelCase__=False       ,		UpperCAmelCase__=True       ,		UpperCAmelCase__=33       ,		UpperCAmelCase__=32       ,		UpperCAmelCase__=5       ,		UpperCAmelCase__=4       ,		UpperCAmelCase__=37       ,		UpperCAmelCase__="gelu"       ,		UpperCAmelCase__=0.1       ,		UpperCAmelCase__=0.1       ,		UpperCAmelCase__=512       ,		UpperCAmelCase__=16       ,		UpperCAmelCase__=2       ,		UpperCAmelCase__=0.02       ,		UpperCAmelCase__=3       ,		UpperCAmelCase__=4       ,		UpperCAmelCase__=None       ,		):
								SCREAMING_SNAKE_CASE__     							=				parent
								SCREAMING_SNAKE_CASE__     							=				batch_size
								SCREAMING_SNAKE_CASE__     							=				seq_length
								SCREAMING_SNAKE_CASE__     							=				is_training
								SCREAMING_SNAKE_CASE__     							=				use_input_mask
								SCREAMING_SNAKE_CASE__     							=				use_token_type_ids
								SCREAMING_SNAKE_CASE__     							=				use_labels
								SCREAMING_SNAKE_CASE__     							=				vocab_size
								SCREAMING_SNAKE_CASE__     							=				hidden_size
								SCREAMING_SNAKE_CASE__     							=				num_hidden_layers
								SCREAMING_SNAKE_CASE__     							=				num_attention_heads
								SCREAMING_SNAKE_CASE__     							=				intermediate_size
								SCREAMING_SNAKE_CASE__     							=				hidden_act
								SCREAMING_SNAKE_CASE__     							=				hidden_dropout_prob
								SCREAMING_SNAKE_CASE__     							=				attention_probs_dropout_prob
								SCREAMING_SNAKE_CASE__     							=				max_position_embeddings
								SCREAMING_SNAKE_CASE__     							=				type_vocab_size
								SCREAMING_SNAKE_CASE__     							=				type_sequence_label_size
								SCREAMING_SNAKE_CASE__     							=				initializer_range
								SCREAMING_SNAKE_CASE__     							=				num_labels
								SCREAMING_SNAKE_CASE__     							=				num_choices
								SCREAMING_SNAKE_CASE__     							=				scope



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				ids_tensor([self.batch_size, self.seq_length]       ,		self.vocab_size     )

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

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

								SCREAMING_SNAKE_CASE__     							=				self.get_config()

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



				def   lowerCAmelCase__		(						self     ):
								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       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__     ):
								SCREAMING_SNAKE_CASE__     							=				EsmModel(config=UpperCAmelCase__     )
								model.to(UpperCAmelCase__     )
								model.eval()
								SCREAMING_SNAKE_CASE__     							=				model(UpperCAmelCase__       ,		attention_mask=UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				model(UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				model(UpperCAmelCase__     )

								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       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__     ):
								SCREAMING_SNAKE_CASE__     							=				EsmForMaskedLM(config=UpperCAmelCase__     )
								model.to(UpperCAmelCase__     )
								model.eval()
								SCREAMING_SNAKE_CASE__     							=				model(UpperCAmelCase__       ,		attention_mask=UpperCAmelCase__       ,		labels=UpperCAmelCase__     )
								self.parent.assertEqual(result.logits.shape       ,		(self.batch_size, self.seq_length, self.vocab_size)     )



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__     ):
								SCREAMING_SNAKE_CASE__     							=				self.num_labels
								SCREAMING_SNAKE_CASE__     							=				EsmForTokenClassification(config=UpperCAmelCase__     )
								model.to(UpperCAmelCase__     )
								model.eval()
								SCREAMING_SNAKE_CASE__     							=				model(UpperCAmelCase__       ,		attention_mask=UpperCAmelCase__       ,		labels=UpperCAmelCase__     )
								self.parent.assertEqual(result.logits.shape       ,		(self.batch_size, self.seq_length, self.num_labels)     )






				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				self.prepare_config_and_inputs()
								(
								    (
								    SCREAMING_SNAKE_CASE__      
								)							,      (
								    SCREAMING_SNAKE_CASE__      
								)							,      (
								    SCREAMING_SNAKE_CASE__      
								)							,      (
								    SCREAMING_SNAKE_CASE__      
								)							,      (
								    SCREAMING_SNAKE_CASE__      
								)							,      (
								    SCREAMING_SNAKE_CASE__      
								)							,            
								)     							=				config_and_inputs
								SCREAMING_SNAKE_CASE__     							=				{"input_ids": input_ids, "attention_mask": input_mask}
								return config, inputs_dict


@require_torch
class 			lowerCamelCase_		(			lowercase       ,  lowercase       ,  unittest.TestCase   ):





				"""simple docstring"""







				_lowerCAmelCase      :      List[Any]      							=					False

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



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				EsmModelTester(self     )
								SCREAMING_SNAKE_CASE__     							=				ConfigTester(self       ,		config_class=UpperCAmelCase__       ,		hidden_size=37     )



				def   lowerCAmelCase__		(						self     ):
								self.config_tester.run_common_tests()



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*UpperCAmelCase__     )



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				self.model_tester.prepare_config_and_inputs()
								for type in ["absolute", "relative_key", "relative_key_query"]:
												SCREAMING_SNAKE_CASE__     							=				type
												self.model_tester.create_and_check_model(*UpperCAmelCase__     )



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*UpperCAmelCase__     )



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*UpperCAmelCase__     )



				@slow
				def   lowerCAmelCase__		(						self     ):
								for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												SCREAMING_SNAKE_CASE__     							=				EsmModel.from_pretrained(UpperCAmelCase__     )
												self.assertIsNotNone(UpperCAmelCase__     )



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				self.model_tester.prepare_config_and_inputs()[0]
								SCREAMING_SNAKE_CASE__     							=				EsmEmbeddings(config=UpperCAmelCase__     )

								SCREAMING_SNAKE_CASE__     							=				torch.as_tensor([[12, 31, 13, model.padding_idx]]     )
								SCREAMING_SNAKE_CASE__     							=				torch.as_tensor(
								    [
								        [
								            0 + model.padding_idx + 1,
								            1 + model.padding_idx + 1,
								            2 + model.padding_idx + 1,
								            model.padding_idx,
								        ]
								    ]     )
								SCREAMING_SNAKE_CASE__     							=				create_position_ids_from_input_ids(UpperCAmelCase__       ,		model.padding_idx     )
								self.assertEqual(position_ids.shape       ,		expected_positions.shape     )
								self.assertTrue(torch.all(torch.eq(UpperCAmelCase__       ,		UpperCAmelCase__     )     )     )



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				self.model_tester.prepare_config_and_inputs()[0]
								SCREAMING_SNAKE_CASE__     							=				EsmEmbeddings(config=UpperCAmelCase__     )

								SCREAMING_SNAKE_CASE__     							=				torch.empty(2       ,		4       ,		30     )
								SCREAMING_SNAKE_CASE__     							=				[
								    0 + embeddings.padding_idx + 1,
								    1 + embeddings.padding_idx + 1,
								    2 + embeddings.padding_idx + 1,
								    3 + embeddings.padding_idx + 1,
								]
								SCREAMING_SNAKE_CASE__     							=				torch.as_tensor([expected_single_positions, expected_single_positions]     )
								SCREAMING_SNAKE_CASE__     							=				embeddings.create_position_ids_from_inputs_embeds(UpperCAmelCase__     )
								self.assertEqual(position_ids.shape       ,		expected_positions.shape     )
								self.assertTrue(torch.all(torch.eq(UpperCAmelCase__       ,		UpperCAmelCase__     )     )     )



				@unittest.skip("Esm does not support embedding resizing"     )
				def   lowerCAmelCase__		(						self     ):
								pass



				@unittest.skip("Esm does not support embedding resizing"     )
				def   lowerCAmelCase__		(						self     ):
								pass






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


@require_torch
class 			lowerCamelCase_		(			lowercase   ):





				"""simple docstring"""







				@slow
				def   lowerCAmelCase__		(						self     ):
								with torch.no_grad():
												SCREAMING_SNAKE_CASE__     							=				EsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"     )
												model.eval()
												SCREAMING_SNAKE_CASE__     							=				torch.tensor([[0, 1, 2, 3, 4, 5]]     )
												SCREAMING_SNAKE_CASE__     							=				model(UpperCAmelCase__     )[0]

												SCREAMING_SNAKE_CASE__     							=				33

												SCREAMING_SNAKE_CASE__     							=				torch.Size((1, 6, vocab_size)     )
												self.assertEqual(output.shape       ,		UpperCAmelCase__     )

												SCREAMING_SNAKE_CASE__     							=				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]       ,		UpperCAmelCase__       ,		atol=1e-4     )     )






				@slow
				def   lowerCAmelCase__		(						self     ):
								with torch.no_grad():
												SCREAMING_SNAKE_CASE__     							=				EsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"     )
												model.eval()

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




"""simple docstring"""






import os

from bleurt import score  # From: git+https://github.com/google-research/bleurt.git

import datasets


_lowerCamelCase     =  datasets.logging.get_logger(__name__)


_lowerCamelCase     =  '\\n@inproceedings{bleurt,\n  title={BLEURT: Learning Robust Metrics for Text Generation},\n  author={Thibault Sellam and Dipanjan Das and Ankur P. Parikh},\n  booktitle={ACL},\n  year={2020},\n  url={https://arxiv.org/abs/2004.04696}\n}\n'

_lowerCamelCase     =  '\\nBLEURT a learnt evaluation metric for Natural Language Generation. It is built using multiple phases of transfer learning starting from a pretrained BERT model (Devlin et al. 2018)\nand then employing another pre-training phrase using synthetic data. Finally it is trained on WMT human annotations. You may run BLEURT out-of-the-box or fine-tune\nit for your specific application (the latter is expected to perform better).\n\nSee the project\'s README at https://github.com/google-research/bleurt#readme for more information.\n'

_lowerCamelCase     =  '\nBLEURT score.\n\nArgs:\n    `predictions` (list of str): prediction/candidate sentences\n    `references` (list of str): reference sentences\n    `checkpoint` BLEURT checkpoint. Will default to BLEURT-tiny if None.\n\nReturns:\n    \'scores\': List of scores.\nExamples:\n\n    >>> predictions = ["hello there", "general kenobi"]\n    >>> references = ["hello there", "general kenobi"]\n    >>> bleurt = datasets.load_metric("bleurt")\n    >>> results = bleurt.compute(predictions=predictions, references=references)\n    >>> print([round(v, 2) for v in results["scores"]])\n    [1.03, 1.04]\n'

_lowerCamelCase     =  {
    'bleurt-tiny-128': 'https://storage.googleapis.com/bleurt-oss/bleurt-tiny-128.zip',
    'bleurt-tiny-512': 'https://storage.googleapis.com/bleurt-oss/bleurt-tiny-512.zip',
    'bleurt-base-128': 'https://storage.googleapis.com/bleurt-oss/bleurt-base-128.zip',
    'bleurt-base-512': 'https://storage.googleapis.com/bleurt-oss/bleurt-base-512.zip',
    'bleurt-large-128': 'https://storage.googleapis.com/bleurt-oss/bleurt-large-128.zip',
    'bleurt-large-512': 'https://storage.googleapis.com/bleurt-oss/bleurt-large-512.zip',
    'BLEURT-20-D3': 'https://storage.googleapis.com/bleurt-oss-21/BLEURT-20-D3.zip',
    'BLEURT-20-D6': 'https://storage.googleapis.com/bleurt-oss-21/BLEURT-20-D6.zip',
    'BLEURT-20-D12': 'https://storage.googleapis.com/bleurt-oss-21/BLEURT-20-D12.zip',
    'BLEURT-20': 'https://storage.googleapis.com/bleurt-oss-21/BLEURT-20.zip',
}


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





				"""simple docstring"""







				def   lowerCAmelCase__		(						self     ):
								return datasets.MetricInfo(
								    description=_DESCRIPTION       ,		citation=_CITATION       ,		homepage="https://github.com/google-research/bleurt"       ,		inputs_description=_KWARGS_DESCRIPTION       ,		features=datasets.Features(
								        {
								            "predictions": datasets.Value("string"       ,		id="sequence"     ),
								            "references": datasets.Value("string"       ,		id="sequence"     ),
								        }     )       ,		codebase_urls=["https://github.com/google-research/bleurt"]       ,		reference_urls=["https://github.com/google-research/bleurt", "https://arxiv.org/abs/2004.04696"]       ,		)



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__     ):
								# check that config name specifies a valid BLEURT model
								if self.config_name == "default":
												logger.warning(
												    "Using default BLEURT-Base checkpoint for sequence maximum length 128. "
												    "You can use a bigger model for better results with e.g.: datasets.load_metric('bleurt', 'bleurt-large-512')."     )
												SCREAMING_SNAKE_CASE__     							=				"bleurt-base-128"

								if self.config_name.lower() in CHECKPOINT_URLS:
												SCREAMING_SNAKE_CASE__     							=				self.config_name.lower()

								elif self.config_name.upper() in CHECKPOINT_URLS:
												SCREAMING_SNAKE_CASE__     							=				self.config_name.upper()

								else:
												raise KeyError(
												    f'''{self.config_name} model not found. You should supply the name of a model checkpoint for bleurt in {CHECKPOINT_URLS.keys()}'''     )

								# download the model checkpoint specified by self.config_name and set up the scorer
								SCREAMING_SNAKE_CASE__     							=				dl_manager.download_and_extract(CHECKPOINT_URLS[checkpoint_name]     )
								SCREAMING_SNAKE_CASE__     							=				score.BleurtScorer(os.path.join(UpperCAmelCase__       ,		UpperCAmelCase__     )     )






				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__     ):
								SCREAMING_SNAKE_CASE__     							=				self.scorer.score(references=UpperCAmelCase__       ,		candidates=UpperCAmelCase__     )
								return {"scores": scores}
