


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






     def __init__(		self     :    str    , UpperCamelCase_     :    Any    , UpperCamelCase_     :    List[Any]=1_3    , UpperCamelCase_     :    Dict=7    , UpperCamelCase_     :    Optional[Any]=False    , UpperCamelCase_     :    Tuple=True    , UpperCamelCase_     :    Dict=False    , UpperCamelCase_     :    Union[str, Any]=True    , UpperCamelCase_     :    int=3_3    , UpperCamelCase_     :    Optional[int]=3_2    , UpperCamelCase_     :    Optional[Any]=5    , UpperCamelCase_     :    Optional[Any]=4    , UpperCamelCase_     :    str=3_7    , UpperCamelCase_     :    Dict="gelu"    , UpperCamelCase_     :    Tuple=0.1    , UpperCamelCase_     :    int=0.1    , UpperCamelCase_     :    int=5_1_2    , UpperCamelCase_     :    Tuple=1_6    , UpperCamelCase_     :    Optional[Any]=2    , UpperCamelCase_     :    Tuple=0.02    , UpperCamelCase_     :    List[Any]=3    , UpperCamelCase_     :    Any=4    , UpperCamelCase_     :    Optional[Any]=None    , ):
           lowerCAmelCase    :		Tuple				=      parent
           lowerCAmelCase    :		Any				=      batch_size
           lowerCAmelCase    :		Optional[int]				=      seq_length
           lowerCAmelCase    :		Any				=      is_training
           lowerCAmelCase    :		Dict				=      use_input_mask
           lowerCAmelCase    :		Union[str, Any]				=      use_token_type_ids
           lowerCAmelCase    :		Optional[Any]				=      use_labels
           lowerCAmelCase    :		int				=      vocab_size
           lowerCAmelCase    :		Union[str, Any]				=      hidden_size
           lowerCAmelCase    :		List[str]				=      num_hidden_layers
           lowerCAmelCase    :		int				=      num_attention_heads
           lowerCAmelCase    :		Dict				=      intermediate_size
           lowerCAmelCase    :		Dict				=      hidden_act
           lowerCAmelCase    :		int				=      hidden_dropout_prob
           lowerCAmelCase    :		Union[str, Any]				=      attention_probs_dropout_prob
           lowerCAmelCase    :		Tuple				=      max_position_embeddings
           lowerCAmelCase    :		Dict				=      type_vocab_size
           lowerCAmelCase    :		Optional[Any]				=      type_sequence_label_size
           lowerCAmelCase    :		int				=      initializer_range
           lowerCAmelCase    :		Tuple				=      num_labels
           lowerCAmelCase    :		Optional[int]				=      num_choices
           lowerCAmelCase    :		Dict				=      scope






     def    lowerCamelCase__   (		self     :    Tuple ):
           lowerCAmelCase    :		Dict				=      ids_tensor([self.batch_size, self.seq_length]    , self.vocab_size )

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

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

           lowerCAmelCase    :		Optional[Any]				=      self.get_config()

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






     def    lowerCamelCase__   (		self     :    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    lowerCamelCase__   (		self     :    int    , UpperCamelCase_     :    Optional[int]    , UpperCamelCase_     :    List[Any]    , UpperCamelCase_     :    Optional[int]    , UpperCamelCase_     :    List[Any]    , UpperCamelCase_     :    List[str]    , UpperCamelCase_     :    List[str] ):
           lowerCAmelCase    :		Dict				=      EsmModel(config=UpperCamelCase_ )
           model.to(UpperCamelCase_ )
           model.eval()
           lowerCAmelCase    :		Dict				=      model(UpperCamelCase_    , attention_mask=UpperCamelCase_ )
           lowerCAmelCase    :		Dict				=      model(UpperCamelCase_ )
           lowerCAmelCase    :		Union[str, Any]				=      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     :    Any    , UpperCamelCase_     :    str    , UpperCamelCase_     :    Tuple    , UpperCamelCase_     :    int    , UpperCamelCase_     :    Union[str, Any]    , UpperCamelCase_     :    Optional[int]    , UpperCamelCase_     :    Union[str, Any] ):
           lowerCAmelCase    :		Optional[int]				=      EsmForMaskedLM(config=UpperCamelCase_ )
           model.to(UpperCamelCase_ )
           model.eval()
           lowerCAmelCase    :		Dict				=      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     :    Union[str, Any]    , UpperCamelCase_     :    Union[str, Any]    , UpperCamelCase_     :    Optional[int]    , UpperCamelCase_     :    List[Any]    , UpperCamelCase_     :    Optional[int]    , UpperCamelCase_     :    List[Any]    , UpperCamelCase_     :    List[str] ):
           lowerCAmelCase    :		Tuple				=      self.num_labels
           lowerCAmelCase    :		int				=      EsmForTokenClassification(config=UpperCamelCase_ )
           model.to(UpperCamelCase_ )
           model.eval()
           lowerCAmelCase    :		Optional[Any]				=      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     :    List[Any] ):
           lowerCAmelCase    :		List[Any]				=      self.prepare_config_and_inputs()
           (
               (
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					
           )    :		List[str]				=      config_and_inputs
           lowerCAmelCase    :		Union[str, Any]				=      {'''input_ids''': input_ids, '''attention_mask''': input_mask}
           return config, inputs_dict



@require_torch
class 						snake_case_(				a__							,    a__							,    unittest.TestCase	):
     __UpperCamelCase     =    False

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






     def    lowerCamelCase__   (		self     :    List[str] ):
           lowerCAmelCase    :		List[Any]				=      EsmModelTester(self )
           lowerCAmelCase    :		Dict				=      ConfigTester(self    , config_class=UpperCamelCase_    , hidden_size=3_7 )






     def    lowerCamelCase__   (		self     :    Optional[int] ):
           self.config_tester.run_common_tests()






     def    lowerCamelCase__   (		self     :    Any ):
           lowerCAmelCase    :		Union[str, Any]				=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCamelCase_ )






     def    lowerCamelCase__   (		self     :    Dict ):
           lowerCAmelCase    :		Optional[Any]				=      self.model_tester.prepare_config_and_inputs()
           for type in ["absolute", "relative_key", "relative_key_query"]:
                 lowerCAmelCase    :		int				=      type
                 self.model_tester.create_and_check_model(*UpperCamelCase_ )






     def    lowerCamelCase__   (		self     :    Union[str, Any] ):
           lowerCAmelCase    :		str				=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase_ )






     def    lowerCamelCase__   (		self     :    int ):
           lowerCAmelCase    :		List[str]				=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(*UpperCamelCase_ )






     @slow
     def    lowerCamelCase__   (		self     :    Union[str, Any] ):
           for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 lowerCAmelCase    :		int				=      EsmModel.from_pretrained(UpperCamelCase_ )
                 self.assertIsNotNone(UpperCamelCase_ )






     def    lowerCamelCase__   (		self     :    List[Any] ):
           lowerCAmelCase    :		Tuple				=      self.model_tester.prepare_config_and_inputs()[0]
           lowerCAmelCase    :		Tuple				=      EsmEmbeddings(config=UpperCamelCase_ )

           lowerCAmelCase    :		Optional[Any]				=      torch.as_tensor([[1_2, 3_1, 1_3, model.padding_idx]] )
           lowerCAmelCase    :		Optional[int]				=      torch.as_tensor(
               [
                   [
                       0 + model.padding_idx + 1,
                       1 + model.padding_idx + 1,
                       2 + model.padding_idx + 1,
                       model.padding_idx,
                   ]
               ] )
           lowerCAmelCase    :		Any				=      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     :    List[Any] ):
           lowerCAmelCase    :		List[str]				=      self.model_tester.prepare_config_and_inputs()[0]
           lowerCAmelCase    :		Tuple				=      EsmEmbeddings(config=UpperCamelCase_ )

           lowerCAmelCase    :		Any				=      torch.empty(2    , 4    , 3_0 )
           lowerCAmelCase    :		Union[str, Any]				=      [
               0 + embeddings.padding_idx + 1,
               1 + embeddings.padding_idx + 1,
               2 + embeddings.padding_idx + 1,
               3 + embeddings.padding_idx + 1,
           ]
           lowerCAmelCase    :		Union[str, Any]				=      torch.as_tensor([expected_single_positions, expected_single_positions] )
           lowerCAmelCase    :		Tuple				=      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     :    Tuple ):
           pass






     @unittest.skip('''Esm does not support embedding resizing''' )
     def    lowerCamelCase__   (		self     :    Dict ):
           pass






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



@require_torch
class 						snake_case_(				a__	):






     @slow
     def    lowerCamelCase__   (		self     :    Union[str, Any] ):
           with torch.no_grad():
                 lowerCAmelCase    :		int				=      EsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D''' )
                 model.eval()
                 lowerCAmelCase    :		Any				=      torch.tensor([[0, 1, 2, 3, 4, 5]] )
                 lowerCAmelCase    :		Any				=      model(UpperCamelCase_ )[0]

                 lowerCAmelCase    :		Any				=      3_3

                 lowerCAmelCase    :		str				=      torch.Size((1, 6, vocab_size) )
                 self.assertEqual(output.shape    , UpperCamelCase_ )

                 lowerCAmelCase    :		int				=      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     :    Union[str, Any] ):
           with torch.no_grad():
                 lowerCAmelCase    :		Optional[Any]				=      EsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D''' )
                 model.eval()

                 lowerCAmelCase    :		int				=      torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]] )
                 lowerCAmelCase    :		int				=      model(UpperCamelCase_ )[0]
                 # compare the actual values for a slice.
                 lowerCAmelCase    :		Optional[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]    , UpperCamelCase_    , atol=1E-4 ) )






"""simple docstring"""





import numpy as np


def    _snake_case  (  _snake_case :	np.array	):
      return 1 / (1 + np.exp(-vector	))


if __name__ == "__main__":
 import doctest

 doctest.testmod()


