


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						:   Union[str, Any]   ,     __A						:   Optional[int]   ,     __A						:   Optional[int]=13   ,     __A						:   Optional[Any]=7   ,     __A						:   int=False   ,     __A						:   Optional[Any]=True   ,     __A						:   Optional[Any]=False   ,     __A						:   Any=True   ,     __A						:   Optional[Any]=33   ,     __A						:   str=32   ,     __A						:   Dict=5   ,     __A						:   Union[str, Any]=4   ,     __A						:   Any=37   ,     __A						:   Any="gelu"   ,     __A						:   Dict=0.1   ,     __A						:   Union[str, Any]=0.1   ,     __A						:   int=512   ,     __A						:   List[str]=16   ,     __A						:   List[str]=2   ,     __A						:   Dict=0.02   ,     __A						:   int=3   ,     __A						:   Union[str, Any]=4   ,     __A						:   List[Any]=None   ,     )      ->Optional[int]:






									"""simple docstring"""






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




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






									"""simple docstring"""






									a__ :Optional[int]						  =    ids_tensor([self.batch_size, self.seq_length]   ,     self.vocab_size						)

									a__ :Optional[int]						  =    None
									if self.use_input_mask:
												a__ :List[Any]						  =    random_attention_mask([self.batch_size, self.seq_length]						)

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

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

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




						def  _snake_case					(	self						:   Dict						)      ->Optional[Any]:






									"""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  _snake_case					(	self						:   str   ,     __A						:   int   ,     __A						:   Optional[Any]   ,     __A						:   List[str]   ,     __A						:   List[str]   ,     __A						:   Tuple   ,     __A						:   List[str]						)      ->str:






									"""simple docstring"""






									a__ :Optional[int]						  =    EsmModel(config=__A						)
									model.to(__A						)
									model.eval()
									a__ :Any						  =    model(__A   ,     attention_mask=__A						)
									a__ :Dict						  =    model(__A						)
									a__ :List[str]						  =    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  _snake_case					(	self						:   Tuple   ,     __A						:   Optional[int]   ,     __A						:   Optional[Any]   ,     __A						:   str   ,     __A						:   Tuple   ,     __A						:   Any   ,     __A						:   Any						)      ->Union[str, Any]:






									"""simple docstring"""






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




						def  _snake_case					(	self						:   Dict   ,     __A						:   str   ,     __A						:   Union[str, Any]   ,     __A						:   Optional[Any]   ,     __A						:   List[str]   ,     __A						:   Any   ,     __A						:   Any						)      ->Any:






									"""simple docstring"""






									a__ :Optional[Any]						  =    self.num_labels
									a__ :Optional[Any]						  =    EsmForTokenClassification(config=__A						)
									model.to(__A						)
									model.eval()
									a__ :Union[str, Any]						  =    model(__A   ,     attention_mask=__A   ,     labels=__A						)
									self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.seq_length, self.num_labels)						)






						def  _snake_case					(	self						:   str						)      ->Tuple:






									"""simple docstring"""






									a__ :List[str]						  =    self.prepare_config_and_inputs()
									(
									    a__  
									) :Optional[Any]						  =    config_and_inputs
									a__ :Any						  =    {"input_ids": input_ids, "attention_mask": input_mask}
									return config, inputs_dict






@require_torch
class 					lowerCAmelCase_ (       _a			,_a			,unittest.TestCase):
						lowerCamelCase_         =     False

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




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






									"""simple docstring"""






									a__ :Union[str, Any]						  =    EsmModelTester(self						)
									a__ :int						  =    ConfigTester(self   ,     config_class=__A   ,     hidden_size=37						)




						def  _snake_case					(	self						:   List[str]						)      ->Optional[int]:






									"""simple docstring"""






									self.config_tester.run_common_tests()




						def  _snake_case					(	self						:   List[str]						)      ->Optional[Any]:






									"""simple docstring"""






									a__ :Any						  =    self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*__A						)




						def  _snake_case					(	self						:   Optional[int]						)      ->Optional[Any]:






									"""simple docstring"""






									a__ :Union[str, Any]						  =    self.model_tester.prepare_config_and_inputs()
									for type in ["absolute", "relative_key", "relative_key_query"]:
												a__ :Optional[Any]						  =    type
												self.model_tester.create_and_check_model(*__A						)




						def  _snake_case					(	self						:   Optional[Any]						)      ->int:






									"""simple docstring"""






									a__ :List[str]						  =    self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*__A						)




						def  _snake_case					(	self						:   str						)      ->List[str]:






									"""simple docstring"""






									a__ :List[Any]						  =    self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*__A						)




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






									"""simple docstring"""






									for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												a__ :Union[str, Any]						  =    EsmModel.from_pretrained(__A						)
												self.assertIsNotNone(__A						)




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






									"""simple docstring"""






									a__ :Dict						  =    self.model_tester.prepare_config_and_inputs()[0]
									a__ :List[Any]						  =    EsmEmbeddings(config=__A						)

									a__ :Optional[int]						  =    torch.as_tensor([[12, 31, 13, model.padding_idx]]						)
									a__ :Union[str, Any]						  =    torch.as_tensor(
									    [
									        [
									            0 + model.padding_idx + 1,
									            1 + model.padding_idx + 1,
									            2 + model.padding_idx + 1,
									            model.padding_idx,
									        ]
									    ]						)
									a__ :str						  =    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  _snake_case					(	self						:   Tuple						)      ->Union[str, Any]:






									"""simple docstring"""






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

									a__ :List[str]						  =    torch.empty(2   ,     4   ,     30						)
									a__ :Any						  =    [
									    0 + embeddings.padding_idx + 1,
									    1 + embeddings.padding_idx + 1,
									    2 + embeddings.padding_idx + 1,
									    3 + embeddings.padding_idx + 1,
									]
									a__ :Union[str, Any]						  =    torch.as_tensor([expected_single_positions, expected_single_positions]						)
									a__ :Union[str, Any]						  =    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  _snake_case					(	self						:   Union[str, Any]						)      ->Dict:






									"""simple docstring"""






									pass




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






									"""simple docstring"""






									pass






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






									"""simple docstring"""






									pass






@require_torch
class 					lowerCAmelCase_ (       _a):




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






									"""simple docstring"""






									with torch.no_grad():
												a__ :List[str]						  =    EsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"						)
												model.eval()
												a__ :Any						  =    torch.tensor([[0, 1, 2, 3, 4, 5]]						)
												a__ :int						  =    model(__A						)[0]

												a__ :Tuple						  =    33

												a__ :Dict						  =    torch.Size((1, 6, vocab_size)						)
												self.assertEqual(output.shape   ,     __A						)

												a__ :Union[str, Any]						  =    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]   ,     __A   ,     atol=1E-4						)						)






						@slow
						def  _snake_case					(	self						:   Tuple						)      ->Tuple:






									"""simple docstring"""






									with torch.no_grad():
												a__ :Tuple						  =    EsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"						)
												model.eval()

												a__ :List[str]						  =    torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]						)
												a__ :int						  =    model(__A						)[0]
												# compare the actual values for a slice.
												a__ :Union[str, Any]						  =    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]   ,     __A   ,     atol=1E-4						)						)






from .configuration_bert_masked import MaskedBertConfig
from .modeling_bert_masked import (
    MaskedBertForMultipleChoice,
    MaskedBertForQuestionAnswering,
    MaskedBertForSequenceClassification,
    MaskedBertForTokenClassification,
    MaskedBertModel,
)
from .modules import *


