

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


   def __init__(			self      ,    _UpperCamelCase      ,    _UpperCamelCase=13      ,    _UpperCamelCase=7      ,    _UpperCamelCase=False      ,    _UpperCamelCase=True      ,    _UpperCamelCase=False      ,    _UpperCamelCase=True      ,    _UpperCamelCase=33      ,    _UpperCamelCase=32      ,    _UpperCamelCase=5      ,    _UpperCamelCase=4      ,    _UpperCamelCase=37      ,    _UpperCamelCase="gelu"      ,    _UpperCamelCase=0.1      ,    _UpperCamelCase=0.1      ,    _UpperCamelCase=512      ,    _UpperCamelCase=16      ,    _UpperCamelCase=2      ,    _UpperCamelCase=0.0_2      ,    _UpperCamelCase=3      ,    _UpperCamelCase=4      ,    _UpperCamelCase=None      ,    ):





      """simple docstring"""





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


   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      _lowercase      :				int      					=       ids_tensor([self.batch_size, self.seq_length]      ,    self.vocab_size  )

      _lowercase      :				Tuple      					=       None
      if self.use_input_mask:
         _lowercase      :				List[str]      					=       random_attention_mask([self.batch_size, self.seq_length]  )

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

      _lowercase      :				Tuple      					=       self.get_config()

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


   def 						_lowerCamelCase       (			self  ):





      """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      ,    _UpperCamelCase      ,    _UpperCamelCase      ,    _UpperCamelCase      ,    _UpperCamelCase      ,    _UpperCamelCase      ,    _UpperCamelCase  ):





      """simple docstring"""





      _lowercase      :				Optional[Any]      					=       EsmModel(config=_UpperCamelCase  )
      model.to(_UpperCamelCase  )
      model.eval()
      _lowercase      :				Optional[Any]      					=       model(_UpperCamelCase      ,    attention_mask=_UpperCamelCase  )
      _lowercase      :				Tuple      					=       model(_UpperCamelCase  )
      _lowercase      :				Tuple      					=       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      ,    _UpperCamelCase      ,    _UpperCamelCase      ,    _UpperCamelCase      ,    _UpperCamelCase      ,    _UpperCamelCase      ,    _UpperCamelCase  ):





      """simple docstring"""





      _lowercase      :				str      					=       EsmForMaskedLM(config=_UpperCamelCase  )
      model.to(_UpperCamelCase  )
      model.eval()
      _lowercase      :				str      					=       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      ,    _UpperCamelCase      ,    _UpperCamelCase      ,    _UpperCamelCase      ,    _UpperCamelCase      ,    _UpperCamelCase      ,    _UpperCamelCase  ):





      """simple docstring"""





      _lowercase      :				Any      					=       self.num_labels
      _lowercase      :				Any      					=       EsmForTokenClassification(config=_UpperCamelCase  )
      model.to(_UpperCamelCase  )
      model.eval()
      _lowercase      :				List[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  ):





      """simple docstring"""





      _lowercase      :				List[Any]      					=       self.prepare_config_and_inputs()
      (
          _lowercase				
      )      :				Optional[int]      					=       config_and_inputs
      _lowercase      :				Union[str, Any]      					=       {"input_ids": input_ids, "attention_mask": input_mask}
      return config, inputs_dict




@require_torch
class        a__  (							lowerCamelCase_  ,				lowerCamelCase_  ,				unittest.TestCase  ):
   _SCREAMING_SNAKE_CASE   :					Union[str, Any]     	=							False

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


   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      _lowercase      :				Any      					=       EsmModelTester(self  )
      _lowercase      :				Any      					=       ConfigTester(self      ,    config_class=_UpperCamelCase      ,    hidden_size=37  )


   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      self.config_tester.run_common_tests()


   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      _lowercase      :				List[Any]      					=       self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*_UpperCamelCase  )


   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      _lowercase      :				List[Any]      					=       self.model_tester.prepare_config_and_inputs()
      for type in ["absolute", "relative_key", "relative_key_query"]:
         _lowercase      :				Dict      					=       type
         self.model_tester.create_and_check_model(*_UpperCamelCase  )


   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      _lowercase      :				Any      					=       self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_lm(*_UpperCamelCase  )


   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      _lowercase      :				str      					=       self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_token_classification(*_UpperCamelCase  )


   @slow
   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         _lowercase      :				Optional[int]      					=       EsmModel.from_pretrained(_UpperCamelCase  )
         self.assertIsNotNone(_UpperCamelCase  )


   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      _lowercase      :				str      					=       self.model_tester.prepare_config_and_inputs()[0]
      _lowercase      :				Optional[int]      					=       EsmEmbeddings(config=_UpperCamelCase  )

      _lowercase      :				Dict      					=       torch.as_tensor([[12, 31, 13, model.padding_idx]]  )
      _lowercase      :				Tuple      					=       torch.as_tensor(
          [
              [
                  0 + model.padding_idx + 1,
                  1 + model.padding_idx + 1,
                  2 + model.padding_idx + 1,
                  model.padding_idx,
              ]
          ]  )
      _lowercase      :				Optional[int]      					=       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  ):





      """simple docstring"""





      _lowercase      :				Dict      					=       self.model_tester.prepare_config_and_inputs()[0]
      _lowercase      :				int      					=       EsmEmbeddings(config=_UpperCamelCase  )

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





      """simple docstring"""





      pass


   @unittest.skip("Esm does not support embedding resizing"  )
   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      pass






   @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."  )
   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      pass




@require_torch
class        a__  (							lowerCamelCase_  ):


   @slow
   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      with torch.no_grad():
         _lowercase      :				List[str]      					=       EsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"  )
         model.eval()
         _lowercase      :				Union[str, Any]      					=       torch.tensor([[0, 1, 2, 3, 4, 5]]  )
         _lowercase      :				int      					=       model(_UpperCamelCase  )[0]

         _lowercase      :				Dict      					=       33

         _lowercase      :				Optional[int]      					=       torch.Size((1, 6, vocab_size)  )
         self.assertEqual(output.shape      ,    _UpperCamelCase  )

         _lowercase      :				Dict      					=       torch.tensor(
             [[[8.9_2_1_5, -10.5898, -6.4_6_7_1], [-6.3_9_6_7, -13.9114, -1.1_2_1_2], [-7.7_8_1_2, -13.9516, -3.7_4_0_6]]]  )
         self.assertTrue(torch.allclose(output[:, :3, :3]      ,    _UpperCamelCase      ,    atol=1E-4  )  )






   @slow
   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





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

         _lowercase      :				Tuple      					=       torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]  )
         _lowercase      :				List[Any]      					=       model(_UpperCamelCase  )[0]
         # compare the actual values for a slice.
         _lowercase      :				Union[str, Any]      					=       torch.tensor(
             [[[0.1_4_4_4, 0.5_4_1_3, 0.3_2_4_8], [0.3_0_3_4, 0.0_0_5_3, 0.3_1_0_8], [0.3_2_2_8, -0.2_4_9_9, 0.3_4_1_5]]]  )
         self.assertTrue(torch.allclose(output[:, :3, :3]      ,    _UpperCamelCase      ,    atol=1E-4  )  )









'''simple docstring'''
from __future__ import annotations







def  _A		(     snake_case		,							snake_case = None		,							snake_case = None					)      ->     None:
   if start is None:
      _lowercase      :				Dict      					=       0

   if end is None:
      _lowercase      :				List[Any]      					=       len(snake_case					) - 1

   if start >= end:
      return

   _lowercase      :				int      					=       (start + end) // 2

   slowsort(snake_case		,							snake_case		,							snake_case					)
   slowsort(snake_case		,							mid + 1		,							snake_case					)

   if sequence[end] < sequence[mid]:
      _lowercase					,				_lowercase      :				Any      					=       sequence[mid], sequence[end]

   slowsort(snake_case		,							snake_case		,							end - 1					)


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

    testmod()






