




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: Dict    ,       UpperCamelCase_: Optional[int]    ,       UpperCamelCase_: str=13    ,       UpperCamelCase_: Any=7    ,       UpperCamelCase_: str=False    ,       UpperCamelCase_: List[Any]=True    ,       UpperCamelCase_: Optional[int]=False    ,       UpperCamelCase_: Union[str, Any]=True    ,       UpperCamelCase_: Dict=33    ,       UpperCamelCase_: Union[str, Any]=32    ,       UpperCamelCase_: Optional[int]=5    ,       UpperCamelCase_: Optional[Any]=4    ,       UpperCamelCase_: Union[str, Any]=37    ,       UpperCamelCase_: List[str]="gelu"    ,       UpperCamelCase_: int=0.1    ,       UpperCamelCase_: Any=0.1    ,       UpperCamelCase_: str=5_12    ,       UpperCamelCase_: Dict=16    ,       UpperCamelCase_: str=2    ,       UpperCamelCase_: List[Any]=0.02    ,       UpperCamelCase_: Any=3    ,       UpperCamelCase_: int=4    ,       UpperCamelCase_: List[str]=None    ,       ):
             __lowerCamelCase   =						parent
             __lowerCamelCase   =						batch_size
             __lowerCamelCase   =						seq_length
             __lowerCamelCase   =						is_training
             __lowerCamelCase   =						use_input_mask
             __lowerCamelCase   =						use_token_type_ids
             __lowerCamelCase   =						use_labels
             __lowerCamelCase   =						vocab_size
             __lowerCamelCase   =						hidden_size
             __lowerCamelCase   =						num_hidden_layers
             __lowerCamelCase   =						num_attention_heads
             __lowerCamelCase   =						intermediate_size
             __lowerCamelCase   =						hidden_act
             __lowerCamelCase   =						hidden_dropout_prob
             __lowerCamelCase   =						attention_probs_dropout_prob
             __lowerCamelCase   =						max_position_embeddings
             __lowerCamelCase   =						type_vocab_size
             __lowerCamelCase   =						type_sequence_label_size
             __lowerCamelCase   =						initializer_range
             __lowerCamelCase   =						num_labels
             __lowerCamelCase   =						num_choices
             __lowerCamelCase   =						scope
      def 			lowerCAmelCase__						(				self: str   ):
             __lowerCamelCase   =						ids_tensor([self.batch_size, self.seq_length]    ,       self.vocab_size   )

             __lowerCamelCase   =						None
             if self.use_input_mask:
                    __lowerCamelCase   =						random_attention_mask([self.batch_size, self.seq_length]   )

             __lowerCamelCase   =						None
             __lowerCamelCase   =						None
             __lowerCamelCase   =						None
             if self.use_labels:
                    __lowerCamelCase   =						ids_tensor([self.batch_size]    ,       self.type_sequence_label_size   )
                    __lowerCamelCase   =						ids_tensor([self.batch_size, self.seq_length]    ,       self.num_labels   )
                    __lowerCamelCase   =						ids_tensor([self.batch_size]    ,       self.num_choices   )

             __lowerCamelCase   =						self.get_config()

             return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
      def 			lowerCAmelCase__						(				self: Any   ):
             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: Optional[Any]    ,       UpperCamelCase_: List[Any]    ,       UpperCamelCase_: Optional[int]    ,       UpperCamelCase_: Optional[Any]    ,       UpperCamelCase_: Any    ,       UpperCamelCase_: Union[str, Any]    ,       UpperCamelCase_: Union[str, Any]   ):
             __lowerCamelCase   =						EsmModel(config=UpperCamelCase_   )
             model.to(UpperCamelCase_   )
             model.eval()
             __lowerCamelCase   =						model(UpperCamelCase_    ,       attention_mask=UpperCamelCase_   )
             __lowerCamelCase   =						model(UpperCamelCase_   )
             __lowerCamelCase   =						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_: List[str]    ,       UpperCamelCase_: List[Any]    ,       UpperCamelCase_: List[Any]    ,       UpperCamelCase_: str    ,       UpperCamelCase_: List[str]    ,       UpperCamelCase_: int   ):
             __lowerCamelCase   =						EsmForMaskedLM(config=UpperCamelCase_   )
             model.to(UpperCamelCase_   )
             model.eval()
             __lowerCamelCase   =						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: int    ,       UpperCamelCase_: Optional[int]    ,       UpperCamelCase_: List[str]    ,       UpperCamelCase_: Any    ,       UpperCamelCase_: Dict    ,       UpperCamelCase_: Optional[Any]    ,       UpperCamelCase_: int   ):
             __lowerCamelCase   =						self.num_labels
             __lowerCamelCase   =						EsmForTokenClassification(config=UpperCamelCase_   )
             model.to(UpperCamelCase_   )
             model.eval()
             __lowerCamelCase   =						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: int   ):
             __lowerCamelCase   =						self.prepare_config_and_inputs()
             (
                 (
                 __lowerCamelCase
             ),   (
                 __lowerCamelCase
             ),   (
                 __lowerCamelCase
             ),   (
                 __lowerCamelCase
             ),   (
                 __lowerCamelCase
             ),   (
                 __lowerCamelCase
             ),   
             )   =						config_and_inputs
             __lowerCamelCase   =						{"""input_ids""": input_ids, """attention_mask""": input_mask}
             return config, inputs_dict





@require_torch
class     lowerCamelCase__(					__lowerCamelCase ,     __lowerCamelCase ,     unittest.TestCase):
      UpperCAmelCase__      :	Union[str, Any]    		=   False

      UpperCAmelCase__      :	Tuple    		=   (
          (
              EsmForMaskedLM,
              EsmModel,
              EsmForSequenceClassification,
              EsmForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      UpperCAmelCase__      :	Optional[int]    		=   ()
      UpperCAmelCase__      :	str    		=   (
          {
              'feature-extraction': EsmModel,
              'fill-mask': EsmForMaskedLM,
              'text-classification': EsmForSequenceClassification,
              'token-classification': EsmForTokenClassification,
              'zero-shot': EsmForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      UpperCAmelCase__      :	Any    		=   True
      def 			lowerCAmelCase__						(				self: List[str]   ):
             __lowerCamelCase   =						EsmModelTester(self   )
             __lowerCamelCase   =						ConfigTester(self    ,       config_class=UpperCamelCase_    ,       hidden_size=37   )
      def 			lowerCAmelCase__						(				self: List[Any]   ):
             self.config_tester.run_common_tests()
      def 			lowerCAmelCase__						(				self: Optional[Any]   ):
             __lowerCamelCase   =						self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*UpperCamelCase_   )
      def 			lowerCAmelCase__						(				self: Dict   ):
             __lowerCamelCase   =						self.model_tester.prepare_config_and_inputs()
             for type in ["absolute", "relative_key", "relative_key_query"]:
                    __lowerCamelCase   =						type
                    self.model_tester.create_and_check_model(*UpperCamelCase_   )
      def 			lowerCAmelCase__						(				self: Optional[int]   ):
             __lowerCamelCase   =						self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase_   )
      def 			lowerCAmelCase__						(				self: str   ):
             __lowerCamelCase   =						self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_token_classification(*UpperCamelCase_   )
      @slow
      def 			lowerCAmelCase__						(				self: Tuple   ):
             for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    __lowerCamelCase   =						EsmModel.from_pretrained(UpperCamelCase_   )
                    self.assertIsNotNone(UpperCamelCase_   )
      def 			lowerCAmelCase__						(				self: int   ):
             __lowerCamelCase   =						self.model_tester.prepare_config_and_inputs()[0]
             __lowerCamelCase   =						EsmEmbeddings(config=UpperCamelCase_   )

             __lowerCamelCase   =						torch.as_tensor([[12, 31, 13, model.padding_idx]]   )
             __lowerCamelCase   =						torch.as_tensor(
                 [
                     [
                         0 + model.padding_idx + 1,
                         1 + model.padding_idx + 1,
                         2 + model.padding_idx + 1,
                         model.padding_idx,
                     ]
                 ]   )
             __lowerCamelCase   =						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: Tuple   ):
             __lowerCamelCase   =						self.model_tester.prepare_config_and_inputs()[0]
             __lowerCamelCase   =						EsmEmbeddings(config=UpperCamelCase_   )

             __lowerCamelCase   =						torch.empty(2    ,       4    ,       30   )
             __lowerCamelCase   =						[
                 0 + embeddings.padding_idx + 1,
                 1 + embeddings.padding_idx + 1,
                 2 + embeddings.padding_idx + 1,
                 3 + embeddings.padding_idx + 1,
             ]
             __lowerCamelCase   =						torch.as_tensor([expected_single_positions, expected_single_positions]   )
             __lowerCamelCase   =						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: int   ):
             pass
      @unittest.skip("""Esm does not support embedding resizing"""   )
      def 			lowerCAmelCase__						(				self: int   ):
             pass




      @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""   )
      def 			lowerCAmelCase__						(				self: List[str]   ):
             pass





@require_torch
class     lowerCamelCase__(					__lowerCamelCase):
      @slow
      def 			lowerCAmelCase__						(				self: Tuple   ):
             with torch.no_grad():
                    __lowerCamelCase   =						EsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""   )
                    model.eval()
                    __lowerCamelCase   =						torch.tensor([[0, 1, 2, 3, 4, 5]]   )
                    __lowerCamelCase   =						model(UpperCamelCase_   )[0]

                    __lowerCamelCase   =						33

                    __lowerCamelCase   =						torch.Size((1, 6, vocab_size)   )
                    self.assertEqual(output.shape    ,       UpperCamelCase_   )

                    __lowerCamelCase   =						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]    ,       UpperCamelCase_    ,       atol=1E-4   )   )




      @slow
      def 			lowerCAmelCase__						(				self: List[Any]   ):
             with torch.no_grad():
                    __lowerCamelCase   =						EsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""   )
                    model.eval()

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





"""simple docstring"""



from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__snake_case   			=    logging.get_logger(__name__)

__snake_case   			=    {
    """facebook/xmod-base""": """https://huggingface.co/facebook/xmod-base/resolve/main/config.json""",
    """facebook/xmod-large-prenorm""": """https://huggingface.co/facebook/xmod-large-prenorm/resolve/main/config.json""",
    """facebook/xmod-base-13-125k""": """https://huggingface.co/facebook/xmod-base-13-125k/resolve/main/config.json""",
    """facebook/xmod-base-30-125k""": """https://huggingface.co/facebook/xmod-base-30-125k/resolve/main/config.json""",
    """facebook/xmod-base-30-195k""": """https://huggingface.co/facebook/xmod-base-30-195k/resolve/main/config.json""",
    """facebook/xmod-base-60-125k""": """https://huggingface.co/facebook/xmod-base-60-125k/resolve/main/config.json""",
    """facebook/xmod-base-60-265k""": """https://huggingface.co/facebook/xmod-base-60-265k/resolve/main/config.json""",
    """facebook/xmod-base-75-125k""": """https://huggingface.co/facebook/xmod-base-75-125k/resolve/main/config.json""",
    """facebook/xmod-base-75-269k""": """https://huggingface.co/facebook/xmod-base-75-269k/resolve/main/config.json""",
}







class _lowerCAmelCase				(      snake_case_	):
   __UpperCAmelCase :       Optional[Any]							  =  '''xmod'''







   def __init__(	self						,	UpperCamelCase__=3_0522						,	UpperCamelCase__=768						,	UpperCamelCase__=12						,	UpperCamelCase__=12						,	UpperCamelCase__=3072						,	UpperCamelCase__="gelu"						,	UpperCamelCase__=0.1						,	UpperCamelCase__=0.1						,	UpperCamelCase__=512						,	UpperCamelCase__=2						,	UpperCamelCase__=0.02						,	UpperCamelCase__=1e-12						,	UpperCamelCase__=1						,	UpperCamelCase__=0						,	UpperCamelCase__=2						,	UpperCamelCase__="absolute"						,	UpperCamelCase__=True						,	UpperCamelCase__=None						,	UpperCamelCase__=False						,	UpperCamelCase__=2						,	UpperCamelCase__=False						,	UpperCamelCase__=True						,	UpperCamelCase__=True						,	UpperCamelCase__=("en_XX",)						,	UpperCamelCase__=None						,	**UpperCamelCase__						,	)	->		int:

       '''simple docstring'''




       super().__init__(pad_token_id=UpperCamelCase__						,	bos_token_id=UpperCamelCase__						,	eos_token_id=UpperCamelCase__						,	**UpperCamelCase__						)

       snake_case    :      List[Any]      					=    vocab_size
       snake_case    :      List[Any]      					=    hidden_size
       snake_case    :      Union[str, Any]      					=    num_hidden_layers
       snake_case    :      Any      					=    num_attention_heads
       snake_case    :      List[str]      					=    hidden_act
       snake_case    :      Union[str, Any]      					=    intermediate_size
       snake_case    :      int      					=    hidden_dropout_prob
       snake_case    :      Any      					=    attention_probs_dropout_prob
       snake_case    :      Optional[int]      					=    max_position_embeddings
       snake_case    :      Tuple      					=    type_vocab_size
       snake_case    :      List[str]      					=    initializer_range
       snake_case    :      int      					=    layer_norm_eps
       snake_case    :      Optional[Any]      					=    position_embedding_type
       snake_case    :      int      					=    use_cache
       snake_case    :      Dict      					=    classifier_dropout
       snake_case    :      Dict      					=    pre_norm
       snake_case    :      Union[str, Any]      					=    adapter_reduction_factor
       snake_case    :      Any      					=    adapter_layer_norm
       snake_case    :      Optional[int]      					=    adapter_reuse_layer_norm
       snake_case    :      List[Any]      					=    ln_before_adapter
       snake_case    :      str      					=    list(UpperCamelCase__						)
       snake_case    :      int      					=    default_language







class _lowerCAmelCase				(      snake_case_	):







   @property
   def  lowerCamelCase		(	self						)	->		Mapping[str, Mapping[int, str]]:

       '''simple docstring'''




       if self.task == "multiple-choice":
           snake_case    :      Optional[int]      					=    {0: "batch", 1: "choice", 2: "sequence"}
       else:
           snake_case    :      Optional[Any]      					=    {0: "batch", 1: "sequence"}
       return OrderedDict(
           [
               ("input_ids", dynamic_axis),
               ("attention_mask", dynamic_axis),
           ]						)





