





"""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   __SCREAMING_SNAKE_CASE     :







  '''simple docstring'''


  def __init__(		self :							Union[str, Any]       , __a :							str       , __a :							Optional[Any]=13       , __a :							Any=7       , __a :							Optional[int]=False       , __a :							Optional[Any]=True       , __a :							Union[str, Any]=False       , __a :							Tuple=True       , __a :							Dict=33       , __a :							int=32       , __a :							Any=5       , __a :							List[Any]=4       , __a :							Any=37       , __a :							Any="gelu"       , __a :							Any=0.1       , __a :							Dict=0.1       , __a :							Tuple=512       , __a :							List[Any]=16       , __a :							Any=2       , __a :							Optional[Any]=0.02       , __a :							str=3       , __a :							Dict=4       , __a :							str=None       , )		->					Tuple:
        _UpperCamelCase   :						Optional[Any]    =							parent
        _UpperCamelCase   :						str    =							batch_size
        _UpperCamelCase   :						Dict    =							seq_length
        _UpperCamelCase   :						Union[str, Any]    =							is_training
        _UpperCamelCase   :						Dict    =							use_input_mask
        _UpperCamelCase   :						List[str]    =							use_token_type_ids
        _UpperCamelCase   :						List[Any]    =							use_labels
        _UpperCamelCase   :						List[Any]    =							vocab_size
        _UpperCamelCase   :						List[Any]    =							hidden_size
        _UpperCamelCase   :						Optional[int]    =							num_hidden_layers
        _UpperCamelCase   :						Optional[int]    =							num_attention_heads
        _UpperCamelCase   :						List[str]    =							intermediate_size
        _UpperCamelCase   :						Tuple    =							hidden_act
        _UpperCamelCase   :						List[Any]    =							hidden_dropout_prob
        _UpperCamelCase   :						Optional[Any]    =							attention_probs_dropout_prob
        _UpperCamelCase   :						int    =							max_position_embeddings
        _UpperCamelCase   :						Optional[Any]    =							type_vocab_size
        _UpperCamelCase   :						List[Any]    =							type_sequence_label_size
        _UpperCamelCase   :						str    =							initializer_range
        _UpperCamelCase   :						Optional[int]    =							num_labels
        _UpperCamelCase   :						Any    =							num_choices
        _UpperCamelCase   :						Tuple    =							scope

  def  __SCREAMING_SNAKE_CASE					(		self :							Union[str, Any]						)		->					Dict:
        _UpperCamelCase   :						List[Any]    =							ids_tensor([self.batch_size, self.seq_length]       , self.vocab_size						)

        _UpperCamelCase   :						str    =							None
        if self.use_input_mask:
              _UpperCamelCase   :						Optional[int]    =							random_attention_mask([self.batch_size, self.seq_length]						)

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

        _UpperCamelCase   :						Dict    =							self.get_config()

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

  def  __SCREAMING_SNAKE_CASE					(		self :							str						)		->					int:
        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 :							Union[str, Any]       , __a :							Dict       , __a :							str       , __a :							int       , __a :							Optional[Any]       , __a :							Optional[int]       , __a :							Optional[int]						)		->					List[Any]:
        _UpperCamelCase   :						int    =							EsmModel(config=__a						)
        model.to(__a						)
        model.eval()
        _UpperCamelCase   :						List[Any]    =							model(__a       , attention_mask=__a						)
        _UpperCamelCase   :						Dict    =							model(__a						)
        _UpperCamelCase   :						Dict    =							model(__a						)

        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 :							Optional[int]       , __a :							Tuple       , __a :							Optional[int]       , __a :							int       , __a :							List[Any]       , __a :							List[str]       , __a :							Optional[int]						)		->					Optional[Any]:
        _UpperCamelCase   :						Dict    =							EsmForMaskedLM(config=__a						)
        model.to(__a						)
        model.eval()
        _UpperCamelCase   :						List[Any]    =							model(__a       , attention_mask=__a       , labels=__a						)
        self.parent.assertEqual(result.logits.shape       , (self.batch_size, self.seq_length, self.vocab_size)						)

  def  __SCREAMING_SNAKE_CASE					(		self :							List[Any]       , __a :							int       , __a :							str       , __a :							int       , __a :							int       , __a :							Any       , __a :							List[str]						)		->					Optional[Any]:
        _UpperCamelCase   :						List[Any]    =							self.num_labels
        _UpperCamelCase   :						Union[str, Any]    =							EsmForTokenClassification(config=__a						)
        model.to(__a						)
        model.eval()
        _UpperCamelCase   :						Union[str, Any]    =							model(__a       , attention_mask=__a       , labels=__a						)
        self.parent.assertEqual(result.logits.shape       , (self.batch_size, self.seq_length, self.num_labels)						)


  def  __SCREAMING_SNAKE_CASE					(		self :							Union[str, Any]						)		->					str:
        _UpperCamelCase   :						Any    =							self.prepare_config_and_inputs()
        (
            (
            _UpperCamelCase
        ),	(
            _UpperCamelCase
        ),	(
            _UpperCamelCase
        ),	(
            _UpperCamelCase
        ),	(
            _UpperCamelCase
        ),	(
            _UpperCamelCase
        ),	
        )   :						List[Any]    =							config_and_inputs
        _UpperCamelCase   :						Any    =							{"input_ids": input_ids, "attention_mask": input_mask}
        return config, inputs_dict





@require_torch
class   __SCREAMING_SNAKE_CASE     (    _UpperCamelCase		,	_UpperCamelCase		,	unittest.TestCase ):







  '''simple docstring'''


  SCREAMING_SNAKE_CASE__				:Union[str, Any]						  =       False

  SCREAMING_SNAKE_CASE__				:Dict						  =       (
      (
          EsmForMaskedLM,
          EsmModel,
          EsmForSequenceClassification,
          EsmForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  SCREAMING_SNAKE_CASE__				:str						  =       ()
  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__				:Dict						  =       True

  def  __SCREAMING_SNAKE_CASE					(		self :							Optional[Any]						)		->					List[str]:
        _UpperCamelCase   :						Optional[int]    =							EsmModelTester(self						)
        _UpperCamelCase   :						int    =							ConfigTester(self       , config_class=__a       , hidden_size=37						)

  def  __SCREAMING_SNAKE_CASE					(		self :							Union[str, Any]						)		->					Optional[int]:
        self.config_tester.run_common_tests()

  def  __SCREAMING_SNAKE_CASE					(		self :							List[str]						)		->					str:
        _UpperCamelCase   :						str    =							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__a						)

  def  __SCREAMING_SNAKE_CASE					(		self :							Union[str, Any]						)		->					Optional[Any]:
        _UpperCamelCase   :						Tuple    =							self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
              _UpperCamelCase   :						List[Any]    =							type
              self.model_tester.create_and_check_model(*__a						)

  def  __SCREAMING_SNAKE_CASE					(		self :							Optional[Any]						)		->					List[str]:
        _UpperCamelCase   :						Tuple    =							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*__a						)

  def  __SCREAMING_SNAKE_CASE					(		self :							Optional[Any]						)		->					Tuple:
        _UpperCamelCase   :						Dict    =							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*__a						)

  @slow
  def  __SCREAMING_SNAKE_CASE					(		self :							int						)		->					Dict:
        for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              _UpperCamelCase   :						List[Any]    =							EsmModel.from_pretrained(__a						)
              self.assertIsNotNone(__a						)

  def  __SCREAMING_SNAKE_CASE					(		self :							List[Any]						)		->					Any:
        _UpperCamelCase   :						Optional[int]    =							self.model_tester.prepare_config_and_inputs()[0]
        _UpperCamelCase   :						Tuple    =							EsmEmbeddings(config=__a						)

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

  def  __SCREAMING_SNAKE_CASE					(		self :							str						)		->					Tuple:
        _UpperCamelCase   :						int    =							self.model_tester.prepare_config_and_inputs()[0]
        _UpperCamelCase   :						Tuple    =							EsmEmbeddings(config=__a						)

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

  @unittest.skip("Esm does not support embedding resizing"						)
  def  __SCREAMING_SNAKE_CASE					(		self :							str						)		->					int:
        pass

  @unittest.skip("Esm does not support embedding resizing"						)
  def  __SCREAMING_SNAKE_CASE					(		self :							str						)		->					int:
        pass


  @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."						)
  def  __SCREAMING_SNAKE_CASE					(		self :							Dict						)		->					Union[str, Any]:
        pass





@require_torch
class   __SCREAMING_SNAKE_CASE     (    _UpperCamelCase ):







  '''simple docstring'''


  @slow
  def  __SCREAMING_SNAKE_CASE					(		self :							int						)		->					Optional[Any]:
        with torch.no_grad():
              _UpperCamelCase   :						int    =							EsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"						)
              model.eval()
              _UpperCamelCase   :						Optional[int]    =							torch.tensor([[0, 1, 2, 3, 4, 5]]						)
              _UpperCamelCase   :						List[str]    =							model(__a						)[0]

              _UpperCamelCase   :						Any    =							33

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

              _UpperCamelCase   :						List[Any]    =							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]       , __a       , atol=1e-4						)						)


  @slow
  def  __SCREAMING_SNAKE_CASE					(		self :							List[Any]						)		->					int:
        with torch.no_grad():
              _UpperCamelCase   :						List[Any]    =							EsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"						)
              model.eval()

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







"""simple docstring"""


import warnings


warnings.warn(
    "memory_utils has been reorganized to utils.memory. Import `find_executable_batchsize` from the main `__init__`: "
    "`from accelerate import find_executable_batch_size` to avoid this warning.",
    FutureWarning,
)
