


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  :






      def __init__(					self						:     List[str]			,   lowercase_						:     str			,   lowercase_						:     Optional[Any]=13			,   lowercase_						:     Tuple=7			,   lowercase_						:     Dict=False			,   lowercase_						:     Any=True			,   lowercase_						:     List[str]=False			,   lowercase_						:     Tuple=True			,   lowercase_						:     Optional[Any]=33			,   lowercase_						:     Tuple=32			,   lowercase_						:     int=5			,   lowercase_						:     str=4			,   lowercase_						:     Dict=37			,   lowercase_						:     List[Any]="gelu"			,   lowercase_						:     List[str]=0.1			,   lowercase_						:     Union[str, Any]=0.1			,   lowercase_						:     Union[str, Any]=512			,   lowercase_						:     Dict=16			,   lowercase_						:     Union[str, Any]=2			,   lowercase_						:     Union[str, Any]=0.02			,   lowercase_						:     Any=3			,   lowercase_						:     List[str]=4			,   lowercase_						:     int=None			,   )					->						str:
             """simple docstring"""




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






      def        lowerCamelCase       (					self						:     Tuple )					->						Optional[int]:
             """simple docstring"""




             _lowerCamelCase       :		Any		=ids_tensor([self.batch_size, self.seq_length]			,   self.vocab_size )

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

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

             _lowerCamelCase       :		Dict		=self.get_config()

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






      def        lowerCamelCase       (					self						:     int )					->						Dict:
             """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        lowerCamelCase       (					self						:     Dict			,   lowercase_						:     List[Any]			,   lowercase_						:     Optional[int]			,   lowercase_						:     List[str]			,   lowercase_						:     List[Any]			,   lowercase_						:     int			,   lowercase_						:     List[Any] )					->						Union[str, Any]:
             """simple docstring"""




             _lowerCamelCase       :		List[str]		=EsmModel(config=lowercase_ )
             model.to(lowercase_ )
             model.eval()
             _lowerCamelCase       :		Tuple		=model(lowercase_			,   attention_mask=lowercase_ )
             _lowerCamelCase       :		Union[str, Any]		=model(lowercase_ )
             _lowerCamelCase       :		List[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        lowerCamelCase       (					self						:     Dict			,   lowercase_						:     Union[str, Any]			,   lowercase_						:     Tuple			,   lowercase_						:     Dict			,   lowercase_						:     List[Any]			,   lowercase_						:     int			,   lowercase_						:     Union[str, Any] )					->						Dict:
             """simple docstring"""




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






      def        lowerCamelCase       (					self						:     str			,   lowercase_						:     int			,   lowercase_						:     Dict			,   lowercase_						:     Any			,   lowercase_						:     Tuple			,   lowercase_						:     List[Any]			,   lowercase_						:     Union[str, Any] )					->						Dict:
             """simple docstring"""




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






      def        lowerCamelCase       (					self						:     Optional[int] )					->						Union[str, Any]:
             """simple docstring"""




             _lowerCamelCase       :		Union[str, Any]		=self.prepare_config_and_inputs()
             (
                 (
                 _lowerCamelCase
             )	,     (
                 _lowerCamelCase
             )	,     (
                 _lowerCamelCase
             )	,     (
                 _lowerCamelCase
             )	,     (
                 _lowerCamelCase
             )	,     (
                 _lowerCamelCase
             )	,     
             )       :		List[str]		=config_and_inputs
             _lowerCamelCase       :		Dict		={'input_ids': input_ids, 'attention_mask': input_mask}
             return config, inputs_dict




@require_torch
class A  (   UpperCamelCase_    ,							UpperCamelCase_    ,							unittest.TestCase     ):
      UpperCamelCase__							:  Dict      =False

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






      def        lowerCamelCase       (					self						:     Optional[int] )					->						Any:
             """simple docstring"""




             _lowerCamelCase       :		int		=EsmModelTester(self )
             _lowerCamelCase       :		Optional[Any]		=ConfigTester(self			,   config_class=lowercase_			,   hidden_size=37 )






      def        lowerCamelCase       (					self						:     Tuple )					->						Union[str, Any]:
             """simple docstring"""




             self.config_tester.run_common_tests()






      def        lowerCamelCase       (					self						:     int )					->						Dict:
             """simple docstring"""




             _lowerCamelCase       :		Union[str, Any]		=self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*lowercase_ )






      def        lowerCamelCase       (					self						:     Optional[int] )					->						Tuple:
             """simple docstring"""




             _lowerCamelCase       :		Optional[int]		=self.model_tester.prepare_config_and_inputs()
             for type in ["absolute", "relative_key", "relative_key_query"]:
                    _lowerCamelCase       :		Dict		=type
                    self.model_tester.create_and_check_model(*lowercase_ )






      def        lowerCamelCase       (					self						:     Dict )					->						Tuple:
             """simple docstring"""




             _lowerCamelCase       :		Dict		=self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_masked_lm(*lowercase_ )






      def        lowerCamelCase       (					self						:     List[str] )					->						List[str]:
             """simple docstring"""




             _lowerCamelCase       :		Dict		=self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_token_classification(*lowercase_ )






      @slow
      def        lowerCamelCase       (					self						:     Optional[int] )					->						str:
             """simple docstring"""




             for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _lowerCamelCase       :		Any		=EsmModel.from_pretrained(lowercase_ )
                    self.assertIsNotNone(lowercase_ )






      def        lowerCamelCase       (					self						:     Tuple )					->						str:
             """simple docstring"""




             _lowerCamelCase       :		List[Any]		=self.model_tester.prepare_config_and_inputs()[0]
             _lowerCamelCase       :		Optional[Any]		=EsmEmbeddings(config=lowercase_ )

             _lowerCamelCase       :		str		=torch.as_tensor([[12, 31, 13, model.padding_idx]] )
             _lowerCamelCase       :		List[Any]		=torch.as_tensor(
                 [
                     [
                         0 + model.padding_idx + 1,
                         1 + model.padding_idx + 1,
                         2 + model.padding_idx + 1,
                         model.padding_idx,
                     ]
                 ] )
             _lowerCamelCase       :		Optional[Any]		=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        lowerCamelCase       (					self						:     Tuple )					->						Any:
             """simple docstring"""




             _lowerCamelCase       :		Union[str, Any]		=self.model_tester.prepare_config_and_inputs()[0]
             _lowerCamelCase       :		int		=EsmEmbeddings(config=lowercase_ )

             _lowerCamelCase       :		str		=torch.empty(2			,   4			,   30 )
             _lowerCamelCase       :		int		=[
                 0 + embeddings.padding_idx + 1,
                 1 + embeddings.padding_idx + 1,
                 2 + embeddings.padding_idx + 1,
                 3 + embeddings.padding_idx + 1,
             ]
             _lowerCamelCase       :		str		=torch.as_tensor([expected_single_positions, expected_single_positions] )
             _lowerCamelCase       :		Dict		=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        lowerCamelCase       (					self						:     Dict )					->						Optional[int]:
             """simple docstring"""




             pass






      @unittest.skip('Esm does not support embedding resizing' )
      def        lowerCamelCase       (					self						:     Union[str, Any] )					->						List[str]:
             """simple docstring"""




             pass






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




             pass




@require_torch
class A  (   UpperCamelCase_     ):






      @slow
      def        lowerCamelCase       (					self						:     Tuple )					->						List[str]:
             """simple docstring"""




             with torch.no_grad():
                    _lowerCamelCase       :		str		=EsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D' )
                    model.eval()
                    _lowerCamelCase       :		Any		=torch.tensor([[0, 1, 2, 3, 4, 5]] )
                    _lowerCamelCase       :		Any		=model(lowercase_ )[0]

                    _lowerCamelCase       :		List[str]		=33

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

                    _lowerCamelCase       :		List[str]		=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]			,   lowercase_			,   atol=1E-4 ) )






      @slow
      def        lowerCamelCase       (					self						:     Tuple )					->						Dict:
             """simple docstring"""




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

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






import sys

from .dependency_versions_table import deps
from .utils.versions import require_version, require_version_core


# define which module versions we always want to check at run time
# (usually the ones defined in `install_requires` in setup.py)
#
# order specific notes:
# - tqdm must be checked before tokenizers

lowerCamelCase											=       'python tqdm regex requests packaging filelock numpy tokenizers'.split()
if sys.version_info < (3, 7):
 pkgs_to_check_at_runtime.append('dataclasses')
if sys.version_info < (3, 8):
 pkgs_to_check_at_runtime.append('importlib_metadata')

for pkg in pkgs_to_check_at_runtime:
 if pkg in deps:
  if pkg == "tokenizers":
   # must be loaded here, or else tqdm check may fail
   from .utils import is_tokenizers_available

   if not is_tokenizers_available():
    continue  # not required, check version only if installed

  require_version_core(deps[pkg])
 else:
  raise ValueError(F"""can't find {pkg} in {deps.keys()}, check dependency_versions_table.py""")




def 	a_      (    SCREAMING_SNAKE_CASE__					:      int		,			SCREAMING_SNAKE_CASE__					:      str=None     ):
       '''simple docstring'''





       require_version(deps[pkg]		,			SCREAMING_SNAKE_CASE__     )


