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					:						Dict	,  __UpperCAmelCase					:						Optional[int]	,  __UpperCAmelCase					:						Optional[Any]=1_3	,  __UpperCAmelCase					:						Union[str, Any]=7	,  __UpperCAmelCase					:						Union[str, Any]=False	,  __UpperCAmelCase					:						List[str]=True	,  __UpperCAmelCase					:						Optional[Any]=False	,  __UpperCAmelCase					:						int=True	,  __UpperCAmelCase					:						Any=3_3	,  __UpperCAmelCase					:						List[Any]=3_2	,  __UpperCAmelCase					:						Union[str, Any]=5	,  __UpperCAmelCase					:						Optional[int]=4	,  __UpperCAmelCase					:						int=3_7	,  __UpperCAmelCase					:						Any="gelu"	,  __UpperCAmelCase					:						Union[str, Any]=0.1	,  __UpperCAmelCase					:						Optional[int]=0.1	,  __UpperCAmelCase					:						Optional[int]=5_1_2	,  __UpperCAmelCase					:						int=1_6	,  __UpperCAmelCase					:						List[Any]=2	,  __UpperCAmelCase					:						int=0.02	,  __UpperCAmelCase					:						Dict=3	,  __UpperCAmelCase					:						List[str]=4	,  __UpperCAmelCase					:						Dict=None	,  )						->    Dict:
         """simple docstring"""



         UpperCAmelCase__								=  parent
         UpperCAmelCase__								=  batch_size
         UpperCAmelCase__								=  seq_length
         UpperCAmelCase__								=  is_training
         UpperCAmelCase__								=  use_input_mask
         UpperCAmelCase__								=  use_token_type_ids
         UpperCAmelCase__								=  use_labels
         UpperCAmelCase__								=  vocab_size
         UpperCAmelCase__								=  hidden_size
         UpperCAmelCase__								=  num_hidden_layers
         UpperCAmelCase__								=  num_attention_heads
         UpperCAmelCase__								=  intermediate_size
         UpperCAmelCase__								=  hidden_act
         UpperCAmelCase__								=  hidden_dropout_prob
         UpperCAmelCase__								=  attention_probs_dropout_prob
         UpperCAmelCase__								=  max_position_embeddings
         UpperCAmelCase__								=  type_vocab_size
         UpperCAmelCase__								=  type_sequence_label_size
         UpperCAmelCase__								=  initializer_range
         UpperCAmelCase__								=  num_labels
         UpperCAmelCase__								=  num_choices
         UpperCAmelCase__								=  scope


   def       lowercase_				(self					:						Tuple		)						->    str:
         """simple docstring"""



         UpperCAmelCase__								=  ids_tensor([self.batch_size, self.seq_length]	,  self.vocab_size		)

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

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

         UpperCAmelCase__								=  self.get_config()

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


   def       lowercase_				(self					:						Union[str, Any]		)						->    Optional[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       lowercase_				(self					:						Optional[int]	,  __UpperCAmelCase					:						List[str]	,  __UpperCAmelCase					:						Union[str, Any]	,  __UpperCAmelCase					:						Union[str, Any]	,  __UpperCAmelCase					:						Union[str, Any]	,  __UpperCAmelCase					:						Optional[int]	,  __UpperCAmelCase					:						List[str]		)						->    List[str]:
         """simple docstring"""



         UpperCAmelCase__								=  EsmModel(config=__UpperCAmelCase		)
         model.to(__UpperCAmelCase		)
         model.eval()
         UpperCAmelCase__								=  model(__UpperCAmelCase	,  attention_mask=__UpperCAmelCase		)
         UpperCAmelCase__								=  model(__UpperCAmelCase		)
         UpperCAmelCase__								=  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       lowercase_				(self					:						List[str]	,  __UpperCAmelCase					:						Union[str, Any]	,  __UpperCAmelCase					:						str	,  __UpperCAmelCase					:						Optional[int]	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						List[Any]	,  __UpperCAmelCase					:						Dict		)						->    List[str]:
         """simple docstring"""



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


   def       lowercase_				(self					:						Dict	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						Union[str, Any]	,  __UpperCAmelCase					:						Any	,  __UpperCAmelCase					:						List[str]	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						Optional[Any]		)						->    Optional[Any]:
         """simple docstring"""



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





   def       lowercase_				(self					:						Optional[int]		)						->    List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  self.prepare_config_and_inputs()
         (
             (
             UpperCAmelCase__
         )		,				(
             UpperCAmelCase__
         )		,				(
             UpperCAmelCase__
         )		,				(
             UpperCAmelCase__
         )		,				(
             UpperCAmelCase__
         )		,				(
             UpperCAmelCase__
         )		,				
         )								=  config_and_inputs
         UpperCAmelCase__								=  {"input_ids": input_ids, "attention_mask": input_mask}
         return config, inputs_dict





@require_torch
class 					A							(    UpperCAmelCase_		,    UpperCAmelCase_		,    unittest.TestCase					):
   __UpperCAmelCase			:			str              =      False

   __UpperCAmelCase			:			Dict              =      (
       (
           EsmForMaskedLM,
           EsmModel,
           EsmForSequenceClassification,
           EsmForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   __UpperCAmelCase			:			List[str]              =      ()
   __UpperCAmelCase			:			List[str]              =      (
       {
           'feature-extraction': EsmModel,
           'fill-mask': EsmForMaskedLM,
           'text-classification': EsmForSequenceClassification,
           'token-classification': EsmForTokenClassification,
           'zero-shot': EsmForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   __UpperCAmelCase			:			str              =      True


   def       lowercase_				(self					:						str		)						->    Optional[int]:
         """simple docstring"""



         UpperCAmelCase__								=  EsmModelTester(self		)
         UpperCAmelCase__								=  ConfigTester(self	,  config_class=__UpperCAmelCase	,  hidden_size=3_7		)


   def       lowercase_				(self					:						int		)						->    Union[str, Any]:
         """simple docstring"""



         self.config_tester.run_common_tests()


   def       lowercase_				(self					:						List[Any]		)						->    Optional[int]:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__UpperCAmelCase		)


   def       lowercase_				(self					:						List[Any]		)						->    Tuple:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         for type in ["absolute", "relative_key", "relative_key_query"]:
               UpperCAmelCase__								=  type
               self.model_tester.create_and_check_model(*__UpperCAmelCase		)


   def       lowercase_				(self					:						str		)						->    int:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*__UpperCAmelCase		)


   def       lowercase_				(self					:						int		)						->    Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*__UpperCAmelCase		)


   @slow
   def       lowercase_				(self					:						List[Any]		)						->    str:
         """simple docstring"""



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


   def       lowercase_				(self					:						Union[str, Any]		)						->    List[str]:
         """simple docstring"""



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

         UpperCAmelCase__								=  torch.as_tensor([[1_2, 3_1, 1_3, model.padding_idx]]		)
         UpperCAmelCase__								=  torch.as_tensor(
             [
                 [
                     0 + model.padding_idx + 1,
                     1 + model.padding_idx + 1,
                     2 + model.padding_idx + 1,
                     model.padding_idx,
                 ]
             ]		)
         UpperCAmelCase__								=  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       lowercase_				(self					:						Optional[Any]		)						->    Union[str, Any]:
         """simple docstring"""



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

         UpperCAmelCase__								=  torch.empty(2	,  4	,  3_0		)
         UpperCAmelCase__								=  [
             0 + embeddings.padding_idx + 1,
             1 + embeddings.padding_idx + 1,
             2 + embeddings.padding_idx + 1,
             3 + embeddings.padding_idx + 1,
         ]
         UpperCAmelCase__								=  torch.as_tensor([expected_single_positions, expected_single_positions]		)
         UpperCAmelCase__								=  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       lowercase_				(self					:						Tuple		)						->    Any:
         """simple docstring"""



         pass


   @unittest.skip("Esm does not support embedding resizing"		)
   def       lowercase_				(self					:						List[Any]		)						->    List[Any]:
         """simple docstring"""



         pass





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



         pass





@require_torch
class 					A							(    UpperCAmelCase_					):


   @slow
   def       lowercase_				(self					:						str		)						->    int:
         """simple docstring"""



         with torch.no_grad():
               UpperCAmelCase__								=  EsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"		)
               model.eval()
               UpperCAmelCase__								=  torch.tensor([[0, 1, 2, 3, 4, 5]]		)
               UpperCAmelCase__								=  model(__UpperCAmelCase		)[0]

               UpperCAmelCase__								=  3_3

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

               UpperCAmelCase__								=  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       lowercase_				(self					:						Dict		)						->    List[Any]:
         """simple docstring"""



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

               UpperCAmelCase__								=  torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]		)
               UpperCAmelCase__								=  model(__UpperCAmelCase		)[0]
               # compare the actual values for a slice.
               UpperCAmelCase__								=  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		)		)




# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

from .config import config_command_parser
from .config_args import default_config_file, load_config_from_file  # noqa: F401
from .default import default_command_parser
from .update import update_command_parser




def        lowerCAmelCase_			(					__A=None							)					->    str:







      '''simple docstring'''


      UpperCAmelCase__								=  argparse.ArgumentParser(add_help=__A,  allow_abbrev=__A							)
      # The main config parser
      UpperCAmelCase__								=  config_command_parser(__A							)
      # The subparser to add commands to
      UpperCAmelCase__								=  config_parser.add_subparsers(title="subcommands",  dest="subcommand"							)

      # Then add other parsers with the parent parser
      default_command_parser(__A,  parents=[parent_parser]							)
      update_command_parser(__A,  parents=[parent_parser]							)

      return config_parser




def        lowerCAmelCase_			(					)					->    Any:







      '''simple docstring'''


      UpperCAmelCase__								=  get_config_parser()
      UpperCAmelCase__								=  config_parser.parse_args()

      if not hasattr(__A,  "func"							):
            config_parser.print_help()
            exit(1							)

      # Run
      args.func(__A							)


if __name__ == "__main__":
      main()



