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



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


          '''simple docstring'''


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



   def 				_SCREAMING_SNAKE_CASE   (self		:					int )		->	int:


          '''simple docstring'''


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

          snake_case       :   Union[str, Any]							=      None
          if self.use_input_mask:
                 snake_case       :   List[Any]							=      random_attention_mask([self.batch_size, self.seq_length] )

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

          snake_case       :   List[Any]							=      self.get_config()

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



   def 				_SCREAMING_SNAKE_CASE   (self		:					List[Any] )		->	List[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 				_SCREAMING_SNAKE_CASE   (self		:					Optional[int]    ,						snake_case__		:					List[Any]    ,						snake_case__		:					int    ,						snake_case__		:					Union[str, Any]    ,						snake_case__		:					Optional[Any]    ,						snake_case__		:					List[str]    ,						snake_case__		:					List[Any] )		->	Union[str, Any]:


          '''simple docstring'''


          snake_case       :   Tuple							=      EsmModel(config=snake_case__ )
          model.to(snake_case__ )
          model.eval()
          snake_case       :   List[str]							=      model(snake_case__    ,						attention_mask=snake_case__ )
          snake_case       :   List[str]							=      model(snake_case__ )
          snake_case       :   List[Any]							=      model(snake_case__ )

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


          '''simple docstring'''


          snake_case       :   Optional[int]							=      EsmForMaskedLM(config=snake_case__ )
          model.to(snake_case__ )
          model.eval()
          snake_case       :   int							=      model(snake_case__    ,						attention_mask=snake_case__    ,						labels=snake_case__ )
          self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length, self.vocab_size) )



   def 				_SCREAMING_SNAKE_CASE   (self		:					List[str]    ,						snake_case__		:					str    ,						snake_case__		:					Union[str, Any]    ,						snake_case__		:					str    ,						snake_case__		:					Optional[int]    ,						snake_case__		:					Optional[Any]    ,						snake_case__		:					List[str] )		->	Optional[Any]:


          '''simple docstring'''


          snake_case       :   Any							=      self.num_labels
          snake_case       :   int							=      EsmForTokenClassification(config=snake_case__ )
          model.to(snake_case__ )
          model.eval()
          snake_case       :   Dict							=      model(snake_case__    ,						attention_mask=snake_case__    ,						labels=snake_case__ )
          self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length, self.num_labels) )



   def 				_SCREAMING_SNAKE_CASE   (self		:					Optional[Any] )		->	Union[str, Any]:


          '''simple docstring'''


          snake_case       :   Tuple							=      self.prepare_config_and_inputs()
          (
              (
              snake_case
          )				,					(
              snake_case
          )				,					(
              snake_case
          )				,					(
              snake_case
          )				,					(
              snake_case
          )				,					(
              snake_case
          )				,					
          )       :   Tuple							=      config_and_inputs
          snake_case       :   Optional[Any]							=      {"input_ids": input_ids, "attention_mask": input_mask}
          return config, inputs_dict






@require_torch
class   UpperCAmelCase		(						A_       ,A_       ,unittest.TestCase ):
   A__					:		List[Any]						     =     False

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



   def 				_SCREAMING_SNAKE_CASE   (self		:					Any )		->	Optional[Any]:


          '''simple docstring'''


          snake_case       :   int							=      EsmModelTester(self )
          snake_case       :   Tuple							=      ConfigTester(self    ,						config_class=snake_case__    ,						hidden_size=37 )



   def 				_SCREAMING_SNAKE_CASE   (self		:					int )		->	Any:


          '''simple docstring'''


          self.config_tester.run_common_tests()



   def 				_SCREAMING_SNAKE_CASE   (self		:					Optional[int] )		->	List[Any]:


          '''simple docstring'''


          snake_case       :   Dict							=      self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*snake_case__ )



   def 				_SCREAMING_SNAKE_CASE   (self		:					int )		->	Optional[Any]:


          '''simple docstring'''


          snake_case       :   List[str]							=      self.model_tester.prepare_config_and_inputs()
          for type in ["absolute", "relative_key", "relative_key_query"]:
                 snake_case       :   Optional[int]							=      type
                 self.model_tester.create_and_check_model(*snake_case__ )



   def 				_SCREAMING_SNAKE_CASE   (self		:					int )		->	List[str]:


          '''simple docstring'''


          snake_case       :   Tuple							=      self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*snake_case__ )



   def 				_SCREAMING_SNAKE_CASE   (self		:					Dict )		->	Union[str, Any]:


          '''simple docstring'''


          snake_case       :   Any							=      self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*snake_case__ )



   @slow
   def 				_SCREAMING_SNAKE_CASE   (self		:					Optional[int] )		->	str:


          '''simple docstring'''


          for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 snake_case       :   int							=      EsmModel.from_pretrained(snake_case__ )
                 self.assertIsNotNone(snake_case__ )



   def 				_SCREAMING_SNAKE_CASE   (self		:					str )		->	List[Any]:


          '''simple docstring'''


          snake_case       :   Dict							=      self.model_tester.prepare_config_and_inputs()[0]
          snake_case       :   int							=      EsmEmbeddings(config=snake_case__ )

          snake_case       :   List[Any]							=      torch.as_tensor([[12, 31, 13, model.padding_idx]] )
          snake_case       :   List[Any]							=      torch.as_tensor(
              [
                  [
                      0 + model.padding_idx + 1,
                      1 + model.padding_idx + 1,
                      2 + model.padding_idx + 1,
                      model.padding_idx,
                  ]
              ] )
          snake_case       :   Optional[int]							=      create_position_ids_from_input_ids(snake_case__    ,						model.padding_idx )
          self.assertEqual(position_ids.shape    ,						expected_positions.shape )
          self.assertTrue(torch.all(torch.eq(snake_case__    ,						snake_case__ ) ) )



   def 				_SCREAMING_SNAKE_CASE   (self		:					Dict )		->	List[Any]:


          '''simple docstring'''


          snake_case       :   List[Any]							=      self.model_tester.prepare_config_and_inputs()[0]
          snake_case       :   Any							=      EsmEmbeddings(config=snake_case__ )

          snake_case       :   Optional[int]							=      torch.empty(2    ,						4    ,						30 )
          snake_case       :   Optional[Any]							=      [
              0 + embeddings.padding_idx + 1,
              1 + embeddings.padding_idx + 1,
              2 + embeddings.padding_idx + 1,
              3 + embeddings.padding_idx + 1,
          ]
          snake_case       :   Dict							=      torch.as_tensor([expected_single_positions, expected_single_positions] )
          snake_case       :   int							=      embeddings.create_position_ids_from_inputs_embeds(snake_case__ )
          self.assertEqual(position_ids.shape    ,						expected_positions.shape )
          self.assertTrue(torch.all(torch.eq(snake_case__    ,						snake_case__ ) ) )



   @unittest.skip("Esm does not support embedding resizing" )
   def 				_SCREAMING_SNAKE_CASE   (self		:					Optional[Any] )		->	Tuple:


          '''simple docstring'''


          pass



   @unittest.skip("Esm does not support embedding resizing" )
   def 				_SCREAMING_SNAKE_CASE   (self		:					Optional[Any] )		->	Union[str, Any]:


          '''simple docstring'''


          pass



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


          '''simple docstring'''


          pass






@require_torch
class   UpperCAmelCase		(						A_ ):



   @slow
   def 				_SCREAMING_SNAKE_CASE   (self		:					str )		->	str:


          '''simple docstring'''


          with torch.no_grad():
                 snake_case       :   Tuple							=      EsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D" )
                 model.eval()
                 snake_case       :   str							=      torch.tensor([[0, 1, 2, 3, 4, 5]] )
                 snake_case       :   Optional[Any]							=      model(snake_case__ )[0]

                 snake_case       :   List[str]							=      33

                 snake_case       :   Tuple							=      torch.Size((1, 6, vocab_size) )
                 self.assertEqual(output.shape    ,						snake_case__ )

                 snake_case       :   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]    ,						snake_case__    ,						atol=1e-4 ) )



   @slow
   def 				_SCREAMING_SNAKE_CASE   (self		:					Tuple )		->	Tuple:


          '''simple docstring'''


          with torch.no_grad():
                 snake_case       :   Dict							=      EsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D" )
                 model.eval()

                 snake_case       :   str							=      torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]] )
                 snake_case       :   int							=      model(snake_case__ )[0]
                 # compare the actual values for a slice.
                 snake_case       :   Optional[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]    ,						snake_case__    ,						atol=1e-4 ) )













"""simple docstring"""



_UpperCamelCase		:  List[str]						=		{
    'meter': 'm',
    'kilometer': 'km',
    'megametre': 'Mm',
    'gigametre': 'Gm',
    'terametre': 'Tm',
    'petametre': 'Pm',
    'exametre': 'Em',
    'zettametre': 'Zm',
    'yottametre': 'Ym',
}
# Exponent of the factor(meter)
_UpperCamelCase		:  str						=		{
    'm': 0,
    'km': 3,
    'Mm': 6,
    'Gm': 9,
    'Tm': 1_2,
    'Pm': 1_5,
    'Em': 1_8,
    'Zm': 2_1,
    'Ym': 2_4,
}
def 						_SCREAMING_SNAKE_CASE (    __snake_case			:  float		,					__snake_case			:  str		,					__snake_case			:  str    ):





  '''simple docstring'''





  lowercase							    =       from_type.lower().strip('s'    )
  lowercase							    =       to_type.lower().strip('s'    )

  lowercase							    =       UNIT_SYMBOL.get(__snake_case		,					__snake_case    )
  lowercase							    =       UNIT_SYMBOL.get(__snake_case		,					__snake_case    )

  if from_sanitized not in METRIC_CONVERSION:
    lowercase							    =       (
        f'Invalid \'from_type\' value: {from_type!r}.\n'
        f'Conversion abbreviations are: {", ".join(__snake_case    )}'
    )
    raise ValueError(__snake_case    )
  if to_sanitized not in METRIC_CONVERSION:
    lowercase							    =       (
        f'Invalid \'to_type\' value: {to_type!r}.\n'
        f'Conversion abbreviations are: {", ".join(__snake_case    )}'
    )
    raise ValueError(__snake_case    )
  lowercase							    =       METRIC_CONVERSION[from_sanitized]
  lowercase							    =       METRIC_CONVERSION[to_sanitized]
  lowercase							    =       1

  if from_exponent > to_exponent:
    lowercase							    =       from_exponent - to_exponent
  else:
    lowercase							    =       -(to_exponent - from_exponent)

  return value * pow(10		,					__snake_case    )


if __name__ == "__main__":
    from doctest import testmod

    testmod()


