


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,   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.0_2,   lowerCamelCase=3,   lowerCamelCase=4,   lowerCamelCase=None,   ) ->     Any:




           """simple docstring"""

           _lowercase :							Tuple   	=	parent
           _lowercase :							str   	=	batch_size
           _lowercase :							Optional[Any]   	=	seq_length
           _lowercase :							List[Any]   	=	is_training
           _lowercase :							Optional[Any]   	=	use_input_mask
           _lowercase :							int   	=	use_token_type_ids
           _lowercase :							Tuple   	=	use_labels
           _lowercase :							Union[str, Any]   	=	vocab_size
           _lowercase :							Any   	=	hidden_size
           _lowercase :							Union[str, Any]   	=	num_hidden_layers
           _lowercase :							List[Any]   	=	num_attention_heads
           _lowercase :							List[str]   	=	intermediate_size
           _lowercase :							List[Any]   	=	hidden_act
           _lowercase :							Tuple   	=	hidden_dropout_prob
           _lowercase :							List[Any]   	=	attention_probs_dropout_prob
           _lowercase :							Dict   	=	max_position_embeddings
           _lowercase :							str   	=	type_vocab_size
           _lowercase :							str   	=	type_sequence_label_size
           _lowercase :							Any   	=	initializer_range
           _lowercase :							Dict   	=	num_labels
           _lowercase :							Tuple   	=	num_choices
           _lowercase :							List[Any]   	=	scope
       def     UpperCamelCase  (      self) ->     Tuple:




           """simple docstring"""

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

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

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

           _lowercase :							List[Any]   	=	self.get_config()

           return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
       def     UpperCamelCase  (      self) ->     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     UpperCamelCase  (      self,   lowerCamelCase,   lowerCamelCase,   lowerCamelCase,   lowerCamelCase,   lowerCamelCase,   lowerCamelCase) ->     List[Any]:




           """simple docstring"""

           _lowercase :							Union[str, Any]   	=	EsmModel(config=lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()
           _lowercase :							List[Any]   	=	model(lowerCamelCase,   attention_mask=lowerCamelCase)
           _lowercase :							List[str]   	=	model(lowerCamelCase)
           _lowercase :							str   	=	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     UpperCamelCase  (      self,   lowerCamelCase,   lowerCamelCase,   lowerCamelCase,   lowerCamelCase,   lowerCamelCase,   lowerCamelCase) ->     Dict:




           """simple docstring"""

           _lowercase :							List[str]   	=	EsmForMaskedLM(config=lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()
           _lowercase :							int   	=	model(lowerCamelCase,   attention_mask=lowerCamelCase,   labels=lowerCamelCase)
           self.parent.assertEqual(result.logits.shape,   (self.batch_size, self.seq_length, self.vocab_size))
       def     UpperCamelCase  (      self,   lowerCamelCase,   lowerCamelCase,   lowerCamelCase,   lowerCamelCase,   lowerCamelCase,   lowerCamelCase) ->     List[Any]:




           """simple docstring"""

           _lowercase :							Optional[Any]   	=	self.num_labels
           _lowercase :							int   	=	EsmForTokenClassification(config=lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()
           _lowercase :							List[str]   	=	model(lowerCamelCase,   attention_mask=lowerCamelCase,   labels=lowerCamelCase)
           self.parent.assertEqual(result.logits.shape,   (self.batch_size, self.seq_length, self.num_labels))




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




           """simple docstring"""

           _lowercase :							int   	=	self.prepare_config_and_inputs()
           (
               (
               _lowercase
           )	,       (
               _lowercase
           )	,       (
               _lowercase
           )	,       (
               _lowercase
           )	,       (
               _lowercase
           )	,       (
               _lowercase
           )	,       
           ) :							Optional[Any]   	=	config_and_inputs
           _lowercase :							Union[str, Any]   	=	{'input_ids': input_ids, 'attention_mask': input_mask}
           return config, inputs_dict






@require_torch
class _lowerCamelCase(		_a,				_a,				unittest.TestCase     ):
       lowercase_				:			Union[str, Any]									=    False

       lowercase_				:			str									=    (
           (
               EsmForMaskedLM,
               EsmModel,
               EsmForSequenceClassification,
               EsmForTokenClassification,
           )
           if is_torch_available()
           else ()
       )
       lowercase_				:			List[str]									=    ()
       lowercase_				:			List[Any]									=    (
           {
               """feature-extraction""": EsmModel,
               """fill-mask""": EsmForMaskedLM,
               """text-classification""": EsmForSequenceClassification,
               """token-classification""": EsmForTokenClassification,
               """zero-shot""": EsmForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       lowercase_				:			List[Any]									=    True
       def     UpperCamelCase  (      self) ->     Optional[Any]:




           """simple docstring"""

           _lowercase :							Tuple   	=	EsmModelTester(self)
           _lowercase :							int   	=	ConfigTester(self,   config_class=lowerCamelCase,   hidden_size=37)
       def     UpperCamelCase  (      self) ->     Optional[int]:




           """simple docstring"""

           self.config_tester.run_common_tests()
       def     UpperCamelCase  (      self) ->     int:




           """simple docstring"""

           _lowercase :							Optional[int]   	=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*lowerCamelCase)
       def     UpperCamelCase  (      self) ->     Tuple:




           """simple docstring"""

           _lowercase :							Optional[Any]   	=	self.model_tester.prepare_config_and_inputs()
           for type in ["absolute", "relative_key", "relative_key_query"]:
               _lowercase :							List[str]   	=	type
               self.model_tester.create_and_check_model(*lowerCamelCase)
       def     UpperCamelCase  (      self) ->     Union[str, Any]:




           """simple docstring"""

           _lowercase :							str   	=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase)
       def     UpperCamelCase  (      self) ->     List[str]:




           """simple docstring"""

           _lowercase :							List[str]   	=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(*lowerCamelCase)
       @slow
       def     UpperCamelCase  (      self) ->     List[Any]:




           """simple docstring"""

           for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               _lowercase :							Tuple   	=	EsmModel.from_pretrained(lowerCamelCase)
               self.assertIsNotNone(lowerCamelCase)
       def     UpperCamelCase  (      self) ->     Union[str, Any]:




           """simple docstring"""

           _lowercase :							Tuple   	=	self.model_tester.prepare_config_and_inputs()[0]
           _lowercase :							str   	=	EsmEmbeddings(config=lowerCamelCase)

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




           """simple docstring"""

           _lowercase :							Any   	=	self.model_tester.prepare_config_and_inputs()[0]
           _lowercase :							Dict   	=	EsmEmbeddings(config=lowerCamelCase)

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




           """simple docstring"""

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




           """simple docstring"""

           pass




       @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.')
       def     UpperCamelCase  (      self) ->     str:




           """simple docstring"""

           pass






@require_torch
class _lowerCamelCase(		_a     ):
       @slow
       def     UpperCamelCase  (      self) ->     List[Any]:




           """simple docstring"""

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

               _lowercase :							Union[str, Any]   	=	33

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

               _lowercase :							Optional[int]   	=	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],   lowerCamelCase,   atol=1E-4))




       @slow
       def     UpperCamelCase  (      self) ->     str:




           """simple docstring"""

           with torch.no_grad():
               _lowercase :							int   	=	EsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D')
               model.eval()

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






def        UpperCamelCase_(   lowerCamelCase_							,				lowerCamelCase_				)					->		bool:
    return numa ^ numa < 0


if __name__ == "__main__":
 import doctest

 doctest.testmod()


