



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


	def __init__(self       ,      lowercase       ,      lowercase=13       ,      lowercase=7       ,      lowercase=False       ,      lowercase=True       ,      lowercase=False       ,      lowercase=True       ,      lowercase=33       ,      lowercase=32       ,      lowercase=5       ,      lowercase=4       ,      lowercase=37       ,      lowercase="gelu"       ,      lowercase=0.1       ,      lowercase=0.1       ,      lowercase=512       ,      lowercase=16       ,      lowercase=2       ,      lowercase=0.02       ,      lowercase=3       ,      lowercase=4       ,      lowercase=None       ,      ):
				A_	:						Dict	=     parent
				A_	:						Optional[Any]	=     batch_size
				A_	:						Any	=     seq_length
				A_	:						Optional[Any]	=     is_training
				A_	:						Union[str, Any]	=     use_input_mask
				A_	:						Union[str, Any]	=     use_token_type_ids
				A_	:						Tuple	=     use_labels
				A_	:						Any	=     vocab_size
				A_	:						Tuple	=     hidden_size
				A_	:						Optional[Any]	=     num_hidden_layers
				A_	:						Tuple	=     num_attention_heads
				A_	:						List[str]	=     intermediate_size
				A_	:						Union[str, Any]	=     hidden_act
				A_	:						Optional[int]	=     hidden_dropout_prob
				A_	:						Optional[int]	=     attention_probs_dropout_prob
				A_	:						str	=     max_position_embeddings
				A_	:						Dict	=     type_vocab_size
				A_	:						Dict	=     type_sequence_label_size
				A_	:						Dict	=     initializer_range
				A_	:						Tuple	=     num_labels
				A_	:						Union[str, Any]	=     num_choices
				A_	:						Dict	=     scope


	def     _a	(self							):
				A_	:						List[Any]	=     ids_tensor([self.batch_size, self.seq_length]       ,      self.vocab_size							)

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

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

				A_	:						Tuple	=     self.get_config()

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


	def     _a	(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     _a	(self       ,      lowercase       ,      lowercase       ,      lowercase       ,      lowercase       ,      lowercase       ,      lowercase							):
				A_	:						Optional[Any]	=     EsmModel(config=lowercase							)
				model.to(lowercase							)
				model.eval()
				A_	:						Optional[int]	=     model(lowercase       ,      attention_mask=lowercase							)
				A_	:						Union[str, Any]	=     model(lowercase							)
				A_	:						str	=     model(lowercase							)

				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     _a	(self       ,      lowercase       ,      lowercase       ,      lowercase       ,      lowercase       ,      lowercase       ,      lowercase							):
				A_	:						int	=     EsmForMaskedLM(config=lowercase							)
				model.to(lowercase							)
				model.eval()
				A_	:						List[str]	=     model(lowercase       ,      attention_mask=lowercase       ,      labels=lowercase							)
				self.parent.assertEqual(result.logits.shape       ,      (self.batch_size, self.seq_length, self.vocab_size)							)


	def     _a	(self       ,      lowercase       ,      lowercase       ,      lowercase       ,      lowercase       ,      lowercase       ,      lowercase							):
				A_	:						str	=     self.num_labels
				A_	:						Dict	=     EsmForTokenClassification(config=lowercase							)
				model.to(lowercase							)
				model.eval()
				A_	:						List[str]	=     model(lowercase       ,      attention_mask=lowercase       ,      labels=lowercase							)
				self.parent.assertEqual(result.logits.shape       ,      (self.batch_size, self.seq_length, self.num_labels)							)


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







@require_torch
class 	_lowerCAmelCase							(			__UpperCAmelCase      ,			__UpperCAmelCase      ,			unittest.TestCase ):
	__SCREAMING_SNAKE_CASE   :		Union[str, Any]					=					False

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


	def     _a	(self							):
				A_	:						str	=     EsmModelTester(self							)
				A_	:						Optional[int]	=     ConfigTester(self       ,      config_class=lowercase       ,      hidden_size=37							)


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


	def     _a	(self							):
				A_	:						List[str]	=     self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*lowercase							)


	def     _a	(self							):
				A_	:						List[str]	=     self.model_tester.prepare_config_and_inputs()
				for type in ["absolute", "relative_key", "relative_key_query"]:
							A_	:						str	=     type
							self.model_tester.create_and_check_model(*lowercase							)


	def     _a	(self							):
				A_	:						str	=     self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_masked_lm(*lowercase							)


	def     _a	(self							):
				A_	:						Any	=     self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_token_classification(*lowercase							)


	@slow
	def     _a	(self							):
				for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							A_	:						List[Any]	=     EsmModel.from_pretrained(lowercase							)
							self.assertIsNotNone(lowercase							)


	def     _a	(self							):
				A_	:						str	=     self.model_tester.prepare_config_and_inputs()[0]
				A_	:						List[str]	=     EsmEmbeddings(config=lowercase							)

				A_	:						int	=     torch.as_tensor([[12, 31, 13, model.padding_idx]]							)
				A_	:						Any	=     torch.as_tensor(
				    [
				        [
				            0 + model.padding_idx + 1,
				            1 + model.padding_idx + 1,
				            2 + model.padding_idx + 1,
				            model.padding_idx,
				        ]
				    ]							)
				A_	:						Tuple	=     create_position_ids_from_input_ids(lowercase       ,      model.padding_idx							)
				self.assertEqual(position_ids.shape       ,      expected_positions.shape							)
				self.assertTrue(torch.all(torch.eq(lowercase       ,      lowercase							)							)							)


	def     _a	(self							):
				A_	:						str	=     self.model_tester.prepare_config_and_inputs()[0]
				A_	:						Optional[int]	=     EsmEmbeddings(config=lowercase							)

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


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


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


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







@require_torch
class 	_lowerCAmelCase							(			__UpperCAmelCase ):


	@slow
	def     _a	(self							):
				with torch.no_grad():
							A_	:						Optional[Any]	=     EsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""							)
							model.eval()
							A_	:						str	=     torch.tensor([[0, 1, 2, 3, 4, 5]]							)
							A_	:						List[Any]	=     model(lowercase							)[0]

							A_	:						int	=     33

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

							A_	:						int	=     torch.tensor(
							    [[[8.92_15, -10.58_98, -6.46_71], [-6.39_67, -13.91_14, -1.12_12], [-7.78_12, -13.95_16, -3.74_06]]]							)
							self.assertTrue(torch.allclose(output[:, :3, :3]       ,      lowercase       ,      atol=1E-4							)							)


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

							A_	:						Tuple	=     torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]							)
							A_	:						int	=     model(lowercase							)[0]
							# compare the actual values for a slice.
							A_	:						Tuple	=     torch.tensor(
							    [[[0.14_44, 0.54_13, 0.32_48], [0.30_34, 0.00_53, 0.31_08], [0.32_28, -0.24_99, 0.34_15]]]							)
							self.assertTrue(torch.allclose(output[:, :3, :3]       ,      lowercase       ,      atol=1E-4							)							)




'''simple docstring'''




import unittest

import numpy as np
import torch

from .utils_summarization import build_mask, compute_token_type_ids, process_story, truncate_or_pad







class 	_lowerCAmelCase							(			unittest.TestCase ):


	def     _a	(self							):
				A_	:						Optional[Any]	=     10


	def     _a	(self							):
				A_	:						Dict	=     [1, 2, 3, 4]
				A_	:						List[Any]	=     [1, 2, 3, 4, 0, 0, 0, 0, 0, 0]
				self.assertEqual(truncate_or_pad(lowercase       ,      self.block_size       ,      0							)       ,      lowercase							)


	def     _a	(self							):
				A_	:						List[Any]	=     [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
				A_	:						Tuple	=     [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
				self.assertEqual(truncate_or_pad(lowercase       ,      self.block_size       ,      0							)       ,      lowercase							)


	def     _a	(self							):
				A_	:						Optional[Any]	=     [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]
				A_	:						Any	=     [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
				self.assertEqual(truncate_or_pad(lowercase       ,      self.block_size       ,      0							)       ,      lowercase							)


	def     _a	(self							):
				A_	:						List[str]	=     """It was the year of Our Lord one thousand seven hundred and
        seventy-five.\n\nSpiritual revelations were conceded to England at that
        favoured period, as at this."""
				A_,    A_	:						Dict	=     process_story(lowercase							)
				self.assertEqual(lowercase       ,      []							)


	def     _a	(self							):
				A_	:						Optional[int]	=     """"""
				A_,    A_	:						List[str]	=     process_story(lowercase							)
				self.assertEqual(lowercase       ,      []							)
				self.assertEqual(lowercase       ,      []							)


	def     _a	(self							):
				A_	:						Optional[Any]	=     (
				    """It was the year of Our Lord one thousand seven hundred and """
				    """seventy-five\n\nSpiritual revelations were conceded to England """
				    """at that favoured period, as at this.\n@highlight\n\nIt was the best of times"""
				)
				A_,    A_	:						int	=     process_story(lowercase							)

				A_	:						Optional[Any]	=     [
				    """It was the year of Our Lord one thousand seven hundred and seventy-five.""",
				    """Spiritual revelations were conceded to England at that favoured period, as at this.""",
				]
				self.assertEqual(lowercase       ,      lowercase							)

				A_	:						Dict	=     ["""It was the best of times."""]
				self.assertEqual(lowercase       ,      lowercase							)


	def     _a	(self							):
				A_	:						Optional[int]	=     torch.tensor([1, 2, 3, 4]							)
				A_	:						Dict	=     torch.tensor([1, 1, 1, 1]							)
				np.testing.assert_array_equal(build_mask(lowercase       ,      0							).numpy()       ,      expected.numpy()							)


	def     _a	(self							):
				A_	:						str	=     torch.tensor([1, 2, 3, 4, 23, 23, 23]							)
				A_	:						str	=     torch.tensor([1, 1, 1, 1, 0, 0, 0]							)
				np.testing.assert_array_equal(build_mask(lowercase       ,      23							).numpy()       ,      expected.numpy()							)


	def     _a	(self							):
				A_	:						Any	=     torch.tensor([8, 2, 3, 4, 1, 1, 1]							)
				A_	:						List[str]	=     torch.tensor([1, 1, 1, 1, 0, 0, 0]							)
				np.testing.assert_array_equal(build_mask(lowercase       ,      1							).numpy()       ,      expected.numpy()							)


	def     _a	(self							):
				A_	:						List[Any]	=     101
				A_	:						List[Any]	=     torch.tensor([[1, 2, 3, 4, 5, 6], [1, 2, 3, 101, 5, 6], [1, 101, 3, 4, 101, 6]]							)
				A_	:						List[str]	=     torch.tensor([[1, 1, 1, 1, 1, 1], [1, 1, 1, 0, 0, 0], [1, 0, 0, 0, 1, 1]]							)

				A_	:						Dict	=     compute_token_type_ids(lowercase       ,      lowercase							)
				np.testing.assert_array_equal(lowercase       ,      lowercase							)