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


  """simple docstring"""

  def __init__(		self      :       Any   ,				lowercase_      :       Optional[Any]   ,				lowercase_      :       str=13   ,				lowercase_      :       str=7   ,				lowercase_      :       Dict=False   ,				lowercase_      :       Union[str, Any]=True   ,				lowercase_      :       Dict=False   ,				lowercase_      :       Optional[int]=True   ,				lowercase_      :       List[Any]=33   ,				lowercase_      :       int=32   ,				lowercase_      :       str=5   ,				lowercase_      :       List[str]=4   ,				lowercase_      :       str=37   ,				lowercase_      :       str="gelu"   ,				lowercase_      :       Tuple=0.1   ,				lowercase_      :       str=0.1   ,				lowercase_      :       Any=512   ,				lowercase_      :       Optional[int]=16   ,				lowercase_      :       Optional[Any]=2   ,				lowercase_      :       List[Any]=0.0_2   ,				lowercase_      :       str=3   ,				lowercase_      :       Dict=4   ,				lowercase_      :       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 		lowerCamelCase__  (		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 		lowerCamelCase__  (		self      :       Dict      ):







    '''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      :       Tuple   ,				lowercase_      :       Optional[int]   ,				lowercase_      :       List[Any]   ,				lowercase_      :       Optional[Any]   ,				lowercase_      :       Optional[int]   ,				lowercase_      :       Any   ,				lowercase_      :       List[Any]      ):







    '''simple docstring'''


    lowercase_     =       EsmModel(config=lowercase_      )
    model.to(lowercase_      )
    model.eval()
    lowercase_     =       model(lowercase_   ,				attention_mask=lowercase_      )
    lowercase_     =       model(lowercase_      )
    lowercase_     =       model(lowercase_      )

    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      :       str   ,				lowercase_      :       Optional[Any]   ,				lowercase_      :       str   ,				lowercase_      :       List[str]   ,				lowercase_      :       Optional[int]   ,				lowercase_      :       str   ,				lowercase_      :       Any      ):







    '''simple docstring'''


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




  def 		lowerCamelCase__  (		self      :       List[str]   ,				lowercase_      :       Any   ,				lowercase_      :       List[str]   ,				lowercase_      :       str   ,				lowercase_      :       Dict   ,				lowercase_      :       Dict   ,				lowercase_      :       Optional[Any]      ):







    '''simple docstring'''


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





  def 		lowerCamelCase__  (		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       _a    (		__a ,				__a ,				unittest.TestCase      ):


  """simple docstring"""

  A_    							=						False

  A_    							=						(
      (
          EsmForMaskedLM,
          EsmModel,
          EsmForSequenceClassification,
          EsmForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  A_    							=						()
  A_    							=						(
      {
          '''feature-extraction''': EsmModel,
          '''fill-mask''': EsmForMaskedLM,
          '''text-classification''': EsmForSequenceClassification,
          '''token-classification''': EsmForTokenClassification,
          '''zero-shot''': EsmForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  A_    							=						True




  def 		lowerCamelCase__  (		self      :       Dict      ):







    '''simple docstring'''


    lowercase_     =       EsmModelTester(self      )
    lowercase_     =       ConfigTester(self   ,				config_class=lowercase_   ,				hidden_size=37      )




  def 		lowerCamelCase__  (		self      :       Tuple      ):







    '''simple docstring'''


    self.config_tester.run_common_tests()




  def 		lowerCamelCase__  (		self      :       List[str]      ):







    '''simple docstring'''


    lowercase_     =       self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*lowercase_      )




  def 		lowerCamelCase__  (		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(*lowercase_      )




  def 		lowerCamelCase__  (		self      :       str      ):







    '''simple docstring'''


    lowercase_     =       self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_masked_lm(*lowercase_      )




  def 		lowerCamelCase__  (		self      :       Dict      ):







    '''simple docstring'''


    lowercase_     =       self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_token_classification(*lowercase_      )




  @slow
  def 		lowerCamelCase__  (		self      :       List[str]      ):







    '''simple docstring'''


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




  def 		lowerCamelCase__  (		self      :       Any      ):







    '''simple docstring'''


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

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




  def 		lowerCamelCase__  (		self      :       Dict      ):







    '''simple docstring'''


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

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




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







    '''simple docstring'''


    pass




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







    '''simple docstring'''


    pass





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







    '''simple docstring'''


    pass






@require_torch
class       _a    (		__a      ):


  """simple docstring"""

  @slow
  def 		lowerCamelCase__  (		self      :       List[str]      ):







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

      lowercase_     =       33

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

      lowercase_     =       torch.tensor(
          [[[8.9_2_1_5, -1_0.5_8_9_8, -6.4_6_7_1], [-6.3_9_6_7, -1_3.9_1_1_4, -1.1_2_1_2], [-7.7_8_1_2, -1_3.9_5_1_6, -3.7_4_0_6]]]      )
      self.assertTrue(torch.allclose(output[:, :3, :3]   ,				lowercase_   ,				atol=1e-4      )      )





  @slow
  def 		lowerCamelCase__  (		self      :       Union[str, 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(lowercase_      )[0]
      # compare the actual values for a slice.
      lowercase_     =       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]   ,				lowercase_   ,				atol=1e-4      )      )


'''simple docstring'''


import math
import sys

def   A_      (       SCREAMING_SNAKE_CASE_						)			->int:
  if number != int(SCREAMING_SNAKE_CASE_						):
    raise ValueError("""the value of input must be a natural number"""						)
  if number < 0:
    raise ValueError("""the value of input must not be a negative number"""						)
  if number == 0:
    return 1
  lowercase_     =       [-1] * (number + 1)
  lowercase_     =       0
  for i in range(1 ,	number + 1						):
    lowercase_     =       sys.maxsize
    lowercase_     =       int(math.sqrt(SCREAMING_SNAKE_CASE_						)						)
    for j in range(1 ,	root + 1						):
      lowercase_     =       1 + answers[i - (j**2)]
      lowercase_     =       min(SCREAMING_SNAKE_CASE_ ,	SCREAMING_SNAKE_CASE_						)
    lowercase_     =       answer
  return answers[number]


if __name__ == "__main__":
       import doctest

       doctest.testmod()

