


'''simple docstring'''

from typing import TYPE_CHECKING

from ...utils import _LazyModule


lowerCamelCase				:  List[Any]    		=      {'tokenization_bertweet': ['BertweetTokenizer']}


if TYPE_CHECKING:
 from .tokenization_bertweet import BertweetTokenizer

else:
 import sys

 lowerCamelCase				:  List[str]    		=      _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)








'''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     __lowerCAmelCase		:





      '''simple docstring'''






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

             '''simple docstring'''







             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       :       str					):

             '''simple docstring'''







             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       :       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  UpperCamelCase__						(self       :       Optional[int]  ,		UpperCamelCase       :       Union[str, Any]  ,		UpperCamelCase       :       int  ,		UpperCamelCase       :       List[Any]  ,		UpperCamelCase       :       str  ,		UpperCamelCase       :       str  ,		UpperCamelCase       :       Dict					):

             '''simple docstring'''







             lowercase__      =					EsmModel(config=UpperCamelCase					)
             model.to(UpperCamelCase					)
             model.eval()
             lowercase__      =					model(UpperCamelCase  ,		attention_mask=UpperCamelCase					)
             lowercase__      =					model(UpperCamelCase					)
             lowercase__      =					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  UpperCamelCase__						(self       :       str  ,		UpperCamelCase       :       Any  ,		UpperCamelCase       :       Tuple  ,		UpperCamelCase       :       Any  ,		UpperCamelCase       :       List[Any]  ,		UpperCamelCase       :       Optional[Any]  ,		UpperCamelCase       :       Optional[int]					):

             '''simple docstring'''







             lowercase__      =					EsmForMaskedLM(config=UpperCamelCase					)
             model.to(UpperCamelCase					)
             model.eval()
             lowercase__      =					model(UpperCamelCase  ,		attention_mask=UpperCamelCase  ,		labels=UpperCamelCase					)
             self.parent.assertEqual(result.logits.shape  ,		(self.batch_size, self.seq_length, self.vocab_size)					)



      def  UpperCamelCase__						(self       :       Any  ,		UpperCamelCase       :       str  ,		UpperCamelCase       :       List[Any]  ,		UpperCamelCase       :       Tuple  ,		UpperCamelCase       :       int  ,		UpperCamelCase       :       Union[str, Any]  ,		UpperCamelCase       :       List[Any]					):

             '''simple docstring'''







             lowercase__      =					self.num_labels
             lowercase__      =					EsmForTokenClassification(config=UpperCamelCase					)
             model.to(UpperCamelCase					)
             model.eval()
             lowercase__      =					model(UpperCamelCase  ,		attention_mask=UpperCamelCase  ,		labels=UpperCamelCase					)
             self.parent.assertEqual(result.logits.shape  ,		(self.batch_size, self.seq_length, self.num_labels)					)






      def  UpperCamelCase__						(self       :       Optional[Any]					):

             '''simple docstring'''







             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     __lowerCAmelCase		(lowercase_    ,	lowercase_    ,	unittest.TestCase	):





      '''simple docstring'''






      lowerCAmelCase__   :	Dict       =				False

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



      def  UpperCamelCase__						(self       :       List[Any]					):

             '''simple docstring'''







             lowercase__      =					EsmModelTester(self					)
             lowercase__      =					ConfigTester(self  ,		config_class=UpperCamelCase  ,		hidden_size=37					)



      def  UpperCamelCase__						(self       :       Union[str, Any]					):

             '''simple docstring'''







             self.config_tester.run_common_tests()



      def  UpperCamelCase__						(self       :       Optional[Any]					):

             '''simple docstring'''







             lowercase__      =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*UpperCamelCase					)



      def  UpperCamelCase__						(self       :       Any					):

             '''simple docstring'''







             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(*UpperCamelCase					)



      def  UpperCamelCase__						(self       :       Optional[int]					):

             '''simple docstring'''







             lowercase__      =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase					)



      def  UpperCamelCase__						(self       :       List[Any]					):

             '''simple docstring'''







             lowercase__      =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_token_classification(*UpperCamelCase					)



      @slow
      def  UpperCamelCase__						(self       :       Union[str, Any]					):

             '''simple docstring'''







             for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    lowercase__      =					EsmModel.from_pretrained(UpperCamelCase					)
                    self.assertIsNotNone(UpperCamelCase					)



      def  UpperCamelCase__						(self       :       Optional[int]					):

             '''simple docstring'''







             lowercase__      =					self.model_tester.prepare_config_and_inputs()[0]
             lowercase__      =					EsmEmbeddings(config=UpperCamelCase					)

             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(UpperCamelCase  ,		model.padding_idx					)
             self.assertEqual(position_ids.shape  ,		expected_positions.shape					)
             self.assertTrue(torch.all(torch.eq(UpperCamelCase  ,		UpperCamelCase					)					)					)



      def  UpperCamelCase__						(self       :       Union[str, Any]					):

             '''simple docstring'''







             lowercase__      =					self.model_tester.prepare_config_and_inputs()[0]
             lowercase__      =					EsmEmbeddings(config=UpperCamelCase					)

             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(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  UpperCamelCase__						(self       :       Optional[int]					):

             '''simple docstring'''







             pass



      @unittest.skip('''Esm does not support embedding resizing'''					)
      def  UpperCamelCase__						(self       :       Optional[Any]					):

             '''simple docstring'''







             pass






      @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''					)
      def  UpperCamelCase__						(self       :       Optional[int]					):

             '''simple docstring'''







             pass






@require_torch
class     __lowerCAmelCase		(lowercase_	):





      '''simple docstring'''






      @slow
      def  UpperCamelCase__						(self       :       Optional[Any]					):

             '''simple docstring'''







             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(UpperCamelCase					)[0]

                    lowercase__      =					33

                    lowercase__      =					torch.Size((1, 6, vocab_size)					)
                    self.assertEqual(output.shape  ,		UpperCamelCase					)

                    lowercase__      =					torch.tensor(
                        [[[8.92_15, -10.58_98, -6.46_71], [-6.39_67, -13.91_14, -1.12_12], [-7.78_12, -13.95_16, -3.74_06]]]					)
                    self.assertTrue(torch.allclose(output[:, :3, :3]  ,		UpperCamelCase  ,		atol=1E-4					)					)






      @slow
      def  UpperCamelCase__						(self       :       Optional[Any]					):

             '''simple docstring'''







             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(UpperCamelCase					)[0]
                    # compare the actual values for a slice.
                    lowercase__      =					torch.tensor(
                        [[[0.14_44, 0.54_13, 0.32_48], [0.30_34, 0.00_53, 0.31_08], [0.32_28, -0.24_99, 0.34_15]]]					)
                    self.assertTrue(torch.allclose(output[:, :3, :3]  ,		UpperCamelCase  ,		atol=1E-4					)					)




