


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


    def __init__(     self  ,					a  ,					a=13  ,					a=7  ,					a=False  ,					a=True  ,					a=False  ,					a=True  ,					a=33  ,					a=32  ,					a=5  ,					a=4  ,					a=37  ,					a="gelu"  ,					a=0.1  ,					a=0.1  ,					a=5_12  ,					a=16  ,					a=2  ,					a=0.02  ,					a=3  ,					a=4  ,					a=None  ,					):
          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    __a							(     self    ):
          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    __a							(     self    ):
          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    __a							(     self  ,					a  ,					a  ,					a  ,					a  ,					a  ,					a    ):
          UpperCamelCase__										=				EsmModel(config=__snake_case    )
          model.to(__snake_case    )
          model.eval()
          UpperCamelCase__										=				model(__snake_case  ,					attention_mask=__snake_case    )
          UpperCamelCase__										=				model(__snake_case    )
          UpperCamelCase__										=				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    __a							(     self  ,					a  ,					a  ,					a  ,					a  ,					a  ,					a    ):
          UpperCamelCase__										=				EsmForMaskedLM(config=__snake_case    )
          model.to(__snake_case    )
          model.eval()
          UpperCamelCase__										=				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    __a							(     self  ,					a  ,					a  ,					a  ,					a  ,					a  ,					a    ):
          UpperCamelCase__										=				self.num_labels
          UpperCamelCase__										=				EsmForTokenClassification(config=__snake_case    )
          model.to(__snake_case    )
          model.eval()
          UpperCamelCase__										=				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    __a							(     self    ):
          UpperCamelCase__										=				self.prepare_config_and_inputs()
          (
              UpperCamelCase__       
          )										=				config_and_inputs
          UpperCamelCase__										=				{"""input_ids""": input_ids, """attention_mask""": input_mask}
          return config, inputs_dict

@require_torch
class 					lowercase_     (    _snake_case  ,						_snake_case  ,						unittest.TestCase   ):
    __UpperCAmelCase   			=	False

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


    def    __a							(     self    ):
          UpperCamelCase__										=				EsmModelTester(self    )
          UpperCamelCase__										=				ConfigTester(self  ,					config_class=__snake_case  ,					hidden_size=37    )


    def    __a							(     self    ):
          self.config_tester.run_common_tests()


    def    __a							(     self    ):
          UpperCamelCase__										=				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*__snake_case    )


    def    __a							(     self    ):
          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(*__snake_case    )


    def    __a							(     self    ):
          UpperCamelCase__										=				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*__snake_case    )


    def    __a							(     self    ):
          UpperCamelCase__										=				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*__snake_case    )


    @slow
    def    __a							(     self    ):
          for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                UpperCamelCase__										=				EsmModel.from_pretrained(__snake_case    )
                self.assertIsNotNone(__snake_case    )


    def    __a							(     self    ):
          UpperCamelCase__										=				self.model_tester.prepare_config_and_inputs()[0]
          UpperCamelCase__										=				EsmEmbeddings(config=__snake_case    )

          UpperCamelCase__										=				torch.as_tensor([[12, 31, 13, 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(__snake_case  ,					model.padding_idx    )
          self.assertEqual(position_ids.shape  ,					expected_positions.shape    )
          self.assertTrue(torch.all(torch.eq(__snake_case  ,					__snake_case    )    )    )


    def    __a							(     self    ):
          UpperCamelCase__										=				self.model_tester.prepare_config_and_inputs()[0]
          UpperCamelCase__										=				EsmEmbeddings(config=__snake_case    )

          UpperCamelCase__										=				torch.empty(2  ,					4  ,					30    )
          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(__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    __a							(     self    ):
          pass


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


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

@require_torch
class 					lowercase_     (    _snake_case   ):


    @slow
    def    __a							(     self    ):
          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(__snake_case    )[0]

                UpperCamelCase__										=				33

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

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


    @slow
    def    __a							(     self    ):
          with torch.no_grad():
                UpperCamelCase__										=				EsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"    )
                model.eval()

                UpperCamelCase__										=				torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]    )
                UpperCamelCase__										=				model(__snake_case    )[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]  ,					__snake_case  ,					atol=1e-4    )    )









'''simple docstring'''


from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


UpperCAmelCase_		:   Union[str, Any]							   =		logging.get_logger(__name__)

UpperCAmelCase_		:   Optional[int]							   =		{
    'google/bigbird-roberta-base': 'https://huggingface.co/google/bigbird-roberta-base/resolve/main/config.json',
    'google/bigbird-roberta-large': 'https://huggingface.co/google/bigbird-roberta-large/resolve/main/config.json',
    'google/bigbird-base-trivia-itc': 'https://huggingface.co/google/bigbird-base-trivia-itc/resolve/main/config.json',
    # See all BigBird models at https://huggingface.co/models?filter=big_bird
}





class     lowercase__  (						_snake_case					):





							'''simple docstring'''


							A_							:      str      	=  """big_bird"""






							def __init__(    self    ,					__snake_case=5_0358    ,					__snake_case=768    ,					__snake_case=12    ,					__snake_case=12    ,					__snake_case=3072    ,					__snake_case="gelu_new"    ,					__snake_case=0.1    ,					__snake_case=0.1    ,					__snake_case=4096    ,					__snake_case=2    ,					__snake_case=0.02    ,					__snake_case=1e-12    ,					__snake_case=True    ,					__snake_case=0    ,					__snake_case=1    ,					__snake_case=2    ,					__snake_case=66    ,					__snake_case="block_sparse"    ,					__snake_case=True    ,					__snake_case=False    ,					__snake_case=64    ,					__snake_case=3    ,					__snake_case=None    ,					**__snake_case    ,					):
											super().__init__(
											    pad_token_id=__snake_case    ,					bos_token_id=__snake_case    ,					eos_token_id=__snake_case    ,					sep_token_id=__snake_case    ,					**__snake_case    ,					)

											_SCREAMING_SNAKE_CASE :       str     		=      vocab_size
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      max_position_embeddings
											_SCREAMING_SNAKE_CASE :       List[str]     		=      hidden_size
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      num_hidden_layers
											_SCREAMING_SNAKE_CASE :       Any     		=      num_attention_heads
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      intermediate_size
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      hidden_act
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      hidden_dropout_prob
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      attention_probs_dropout_prob
											_SCREAMING_SNAKE_CASE :       Tuple     		=      initializer_range
											_SCREAMING_SNAKE_CASE :       Any     		=      type_vocab_size
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      layer_norm_eps
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      use_cache

											_SCREAMING_SNAKE_CASE :       List[Any]     		=      rescale_embeddings
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      attention_type
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      use_bias
											_SCREAMING_SNAKE_CASE :       int     		=      block_size
											_SCREAMING_SNAKE_CASE :       Any     		=      num_random_blocks
											_SCREAMING_SNAKE_CASE :       List[str]     		=      classifier_dropout





class     lowercase__  (						_snake_case					):





							'''simple docstring'''


							@property
							def UpperCAmelCase_	(    self  ):
											if self.task == "multiple-choice":
															_SCREAMING_SNAKE_CASE :       int     		=      {0: """batch""", 1: """choice""", 2: """sequence"""}
											else:
															_SCREAMING_SNAKE_CASE :       Optional[int]     		=      {0: """batch""", 1: """sequence"""}
											return OrderedDict(
											    [
											        ("""input_ids""", dynamic_axis),
											        ("""attention_mask""", dynamic_axis),
											    ]  )




