





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







       """simple docstring"""





       def __init__(    self	:	Any						,__A	:	List[Any]						,__A	:	Union[str, Any]=13						,__A	:	List[Any]=7						,__A	:	Optional[Any]=False						,__A	:	List[Any]=True						,__A	:	List[Any]=False						,__A	:	Union[str, Any]=True						,__A	:	str=33						,__A	:	int=32						,__A	:	Tuple=5						,__A	:	List[Any]=4						,__A	:	Any=37						,__A	:	int="gelu"						,__A	:	Optional[int]=0.1						,__A	:	Optional[int]=0.1						,__A	:	int=512						,__A	:	Dict=16						,__A	:	List[Any]=2						,__A	:	Union[str, Any]=0.02						,__A	:	Dict=3						,__A	:	Optional[int]=4						,__A	:	Optional[int]=None						,)						->					Dict:
         _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	:	Optional[Any]  )						->					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	:	str  )						->					List[str]:
         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	:	str						,__A	:	Any						,__A	:	Union[str, Any]						,__A	:	List[Any]						,__A	:	Dict						,__A	:	Dict						,__A	:	List[str]  )						->					str:
         _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[int]						,__A	:	Optional[Any]						,__A	:	Optional[Any]						,__A	:	Dict						,__A	:	Optional[Any]						,__A	:	Union[str, Any]						,__A	:	List[str]  )						->					List[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	:	Any						,__A	:	Optional[int]						,__A	:	str						,__A	:	str						,__A	:	str						,__A	:	Tuple  )						->					List[str]:
         _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[int]  )						->					List[str]:
         _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 				A_				(      UpperCAmelCase					,	UpperCAmelCase					,	unittest.TestCase	):







       """simple docstring"""





       SCREAMING_SNAKE_CASE_      : Optional[int]      =  False

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







       def    __UpperCAmelCase (    self	:	str  )						->					List[str]:
         _lowercase							= EsmModelTester(self  )
         _lowercase							= ConfigTester(self						,config_class=__A						,hidden_size=37  )







       def    __UpperCAmelCase (    self	:	Dict  )						->					List[str]:
         self.config_tester.run_common_tests()







       def    __UpperCAmelCase (    self	:	int  )						->					Optional[int]:
         _lowercase							= self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__A  )







       def    __UpperCAmelCase (    self	:	int  )						->					Optional[Any]:
         _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	:	Tuple  )						->					str:
         _lowercase							= self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*__A  )







       def    __UpperCAmelCase (    self	:	str  )						->					List[str]:
         _lowercase							= self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*__A  )







       @slow
       def    __UpperCAmelCase (    self	:	List[str]  )						->					Optional[int]:
         for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           _lowercase							= EsmModel.from_pretrained(__A  )
           self.assertIsNotNone(__A  )







       def    __UpperCAmelCase (    self	:	Union[str, Any]  )						->					Union[str, Any]:
         _lowercase							= self.model_tester.prepare_config_and_inputs()[0]
         _lowercase							= EsmEmbeddings(config=__A  )

         _lowercase							= torch.as_tensor([[12, 31, 13, 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	:	Tuple  )						->					Tuple:
         _lowercase							= self.model_tester.prepare_config_and_inputs()[0]
         _lowercase							= EsmEmbeddings(config=__A  )

         _lowercase							= torch.empty(2						,4						,30  )
         _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	:	Dict  )						->					str:
         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	:	Tuple  )						->					Union[str, Any]:
         pass





@require_torch
class 				A_				(      UpperCAmelCase	):







       """simple docstring"""





       @slow
       def    __UpperCAmelCase (    self	:	int  )						->					Dict:
         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							= 33

           _lowercase							= torch.Size((1, 6, vocab_size)  )
           self.assertEqual(output.shape						,__A  )

           _lowercase							= torch.tensor(
               [[[8.9215, -10.5898, -6.4671], [-6.3967, -13.9114, -1.1212], [-7.7812, -13.9516, -3.7406]]]  )
           self.assertTrue(torch.allclose(output[:, :3, :3]						,__A						,atol=1e-4  )  )







       @slow
       def    __UpperCAmelCase (    self	:	str  )						->					Any:
         with torch.no_grad():
           _lowercase							= EsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D'  )
           model.eval()

           _lowercase							= torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]  )
           _lowercase							= model(__A  )[0]
           # compare the actual values for a slice.
           _lowercase							= torch.tensor(
               [[[0.1444, 0.5413, 0.3248], [0.3034, 0.0053, 0.3108], [0.3228, -0.2499, 0.3415]]]  )
           self.assertTrue(torch.allclose(output[:, :3, :3]						,__A						,atol=1e-4  )  )






from ....utils import logging


snake_case 				=       logging.get_logger(__name__)





class 				A_				(      UpperCAmelCase	):







       """simple docstring"""





       def __init__(    self	:	Union[str, Any]						,__A	:	int						,__A	:	List[Any]=None						,__A	:	List[str]=2048  )						->					List[Any]:
         _lowercase							= config.__dict__
         _lowercase							= modal_hidden_size
         if num_labels:
           _lowercase							= num_labels