

from __future__ import annotations

import unittest

from transformers import EsmConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
     import numpy
     import tensorflow as tf

     from transformers.models.esm.modeling_tf_esm import (
         TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
         TFEsmForMaskedLM,
         TFEsmForSequenceClassification,
         TFEsmForTokenClassification,
         TFEsmModel,
     )



class 						_A     :





  '''simple docstring'''



  def __init__(  self   ,SCREAMING_SNAKE_CASE_   ,):

      '''simple docstring'''
      snake_case							:		str        =		parent
      snake_case							:		Optional[Any]        =		13
      snake_case							:		Tuple        =		7
      snake_case							:		Union[str, Any]        =		True
      snake_case							:		Union[str, Any]        =		True
      snake_case							:		Optional[Any]        =		True
      snake_case							:		Any        =		99
      snake_case							:		List[str]        =		32
      snake_case							:		Optional[Any]        =		2
      snake_case							:		Optional[Any]        =		4
      snake_case							:		Any        =		37
      snake_case							:		List[str]        =		"gelu"
      snake_case							:		Optional[int]        =		0.1
      snake_case							:		Dict        =		0.1
      snake_case							:		Tuple        =		512
      snake_case							:		Optional[int]        =		16
      snake_case							:		Tuple        =		2
      snake_case							:		Optional[int]        =		0.02
      snake_case							:		str        =		3
      snake_case							:		Dict        =		4
      snake_case							:		List[Any]        =		None





  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case							:		List[str]        =		ids_tensor([self.batch_size, self.seq_length]   ,self.vocab_size    )

      snake_case							:		str        =		None
      if self.use_input_mask:
          snake_case							:		Dict        =		random_attention_mask([self.batch_size, self.seq_length]    )

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

      snake_case							:		Optional[Any]        =		EsmConfig(
          vocab_size=self.vocab_size   ,hidden_size=self.hidden_size   ,num_hidden_layers=self.num_hidden_layers   ,pad_token_id=1   ,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   ,)

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





  def    snake_case_       (  self    ):

      '''simple docstring'''
      (
          snake_case  
      )							:		Tuple        =		self.prepare_config_and_inputs()

      snake_case							:		Optional[Any]        =		True
      snake_case							:		List[Any]        =		floats_tensor([self.batch_size, self.seq_length, self.hidden_size]    )
      snake_case							:		Dict        =		ids_tensor([self.batch_size, self.seq_length]   ,vocab_size=2    )

      return (
          config,
          input_ids,
          input_mask,
          sequence_labels,
          token_labels,
          choice_labels,
          encoder_hidden_states,
          encoder_attention_mask,
      )





  def    snake_case_       (  self   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_    ):

      '''simple docstring'''
      snake_case							:		Tuple        =		TFEsmModel(config=SCREAMING_SNAKE_CASE_    )
      snake_case							:		List[Any]        =		{"input_ids": input_ids, "attention_mask": input_mask}
      snake_case							:		str        =		model(SCREAMING_SNAKE_CASE_    )

      snake_case							:		Any        =		[input_ids, input_mask]
      snake_case							:		Any        =		model(SCREAMING_SNAKE_CASE_    )

      snake_case							:		List[Any]        =		model(SCREAMING_SNAKE_CASE_    )

      self.parent.assertEqual(result.last_hidden_state.shape   ,(self.batch_size, self.seq_length, self.hidden_size)    )





  def    snake_case_       (  self   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,):

      '''simple docstring'''
      snake_case							:		List[Any]        =		True

      snake_case							:		List[str]        =		TFEsmModel(config=SCREAMING_SNAKE_CASE_    )
      snake_case							:		Optional[int]        =		{
          "input_ids": input_ids,
          "attention_mask": input_mask,
          "encoder_hidden_states": encoder_hidden_states,
          "encoder_attention_mask": encoder_attention_mask,
      }
      snake_case							:		int        =		model(SCREAMING_SNAKE_CASE_    )

      snake_case							:		List[Any]        =		[input_ids, input_mask]
      snake_case							:		Tuple        =		model(SCREAMING_SNAKE_CASE_   ,encoder_hidden_states=SCREAMING_SNAKE_CASE_    )

      # Also check the case where encoder outputs are not passed
      snake_case							:		Tuple        =		model(SCREAMING_SNAKE_CASE_   ,attention_mask=SCREAMING_SNAKE_CASE_    )

      self.parent.assertEqual(result.last_hidden_state.shape   ,(self.batch_size, self.seq_length, self.hidden_size)    )





  def    snake_case_       (  self   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_    ):

      '''simple docstring'''
      snake_case							:		List[Any]        =		TFEsmForMaskedLM(config=SCREAMING_SNAKE_CASE_    )
      snake_case							:		Dict        =		model([input_ids, input_mask]    )
      self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.seq_length, self.vocab_size)    )





  def    snake_case_       (  self   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_    ):

      '''simple docstring'''
      snake_case							:		List[Any]        =		self.num_labels
      snake_case							:		str        =		TFEsmForTokenClassification(config=SCREAMING_SNAKE_CASE_    )
      snake_case							:		int        =		{"input_ids": input_ids, "attention_mask": input_mask}
      snake_case							:		Union[str, Any]        =		model(SCREAMING_SNAKE_CASE_    )
      self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.seq_length, self.num_labels)    )





  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case							:		Optional[int]        =		self.prepare_config_and_inputs()
      (
          snake_case  
      )							:		int        =		config_and_inputs
      snake_case							:		int        =		{"input_ids": input_ids, "attention_mask": input_mask}
      return config, inputs_dict







@require_tf
class 						_A     (				_UpperCamelCase						,    _UpperCamelCase						,    unittest.TestCase				):





  '''simple docstring'''



  __lowerCamelCase			:      Union[str, Any]					=   (
      (
          TFEsmModel,
          TFEsmForMaskedLM,
          TFEsmForSequenceClassification,
          TFEsmForTokenClassification,
      )
      if is_tf_available()
      else ()
  )
  __lowerCamelCase			:      int					=   (
      {
          """feature-extraction""": TFEsmModel,
          """fill-mask""": TFEsmForMaskedLM,
          """text-classification""": TFEsmForSequenceClassification,
          """token-classification""": TFEsmForTokenClassification,
          """zero-shot""": TFEsmForSequenceClassification,
      }
      if is_tf_available()
      else {}
  )
  __lowerCamelCase			:      Optional[int]					=   False
  __lowerCamelCase			:      Dict					=   False





  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case							:		int        =		TFEsmModelTester(self    )
      snake_case							:		str        =		ConfigTester(self   ,config_class=SCREAMING_SNAKE_CASE_   ,hidden_size=37    )





  def    snake_case_       (  self    ):

      '''simple docstring'''
      self.config_tester.run_common_tests()





  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case							:		Union[str, Any]        =		self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_    )





  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case							:		Tuple        =		self.model_tester.prepare_config_and_inputs_for_decoder()
      self.model_tester.create_and_check_model_as_decoder(*SCREAMING_SNAKE_CASE_    )





  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case							:		Dict        =		self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE_    )





  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case							:		Union[str, Any]        =		self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE_    )





  @slow
  def    snake_case_       (  self    ):

      '''simple docstring'''





      for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          snake_case							:		Dict        =		TFEsmModel.from_pretrained(SCREAMING_SNAKE_CASE_    )
          self.assertIsNotNone(SCREAMING_SNAKE_CASE_    )





  @unittest.skip("""Protein models do not support embedding resizing."""    )
  def    snake_case_       (  self    ):

      '''simple docstring'''
      pass





  @unittest.skip("""Protein models do not support embedding resizing."""    )
  def    snake_case_       (  self    ):

      '''simple docstring'''
      pass





  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case							:		Optional[int]        =		self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
          snake_case							:		List[Any]        =		model_class(SCREAMING_SNAKE_CASE_    )
          assert isinstance(model.get_input_embeddings()   ,tf.keras.layers.Layer    )
          if model_class is TFEsmForMaskedLM:
              # Output embedding test differs from the main test because they're a matrix, not a layer
              snake_case							:		str        =		model.get_bias()
              assert isinstance(SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_    )
              for k, v in name.items():
                  assert isinstance(SCREAMING_SNAKE_CASE_   ,tf.Variable    )
          else:
              snake_case							:		Union[str, Any]        =		model.get_output_embeddings()
              assert x is None
              snake_case							:		Union[str, Any]        =		model.get_bias()
              assert name is None







@require_tf
class 						_A     (				unittest.TestCase				):





  '''simple docstring'''



  @slow
  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case							:		List[Any]        =		TFEsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""    )

      snake_case							:		Any        =		tf.constant([[0, 1, 2, 3, 4, 5]]    )
      snake_case							:		Dict        =		model(SCREAMING_SNAKE_CASE_    )[0]
      snake_case							:		Optional[Any]        =		[1, 6, 33]
      self.assertEqual(list(output.numpy().shape    )   ,SCREAMING_SNAKE_CASE_    )
      # compare the actual values for a slice.
      snake_case							:		Tuple        =		tf.constant(
          [
              [
                  [8.92_15_18, -10.58_98_14, -6.4_67_13_07],
                  [-6.3_96_71_56, -13.91_13_77, -1.1_21_19_15],
                  [-7.78_12_47, -13.95_15_57, -3.74_05_92],
              ]
          ]    )
      self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()   ,expected_slice.numpy()   ,atol=1E-2    )    )





  @slow
  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case							:		str        =		TFEsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""    )

      snake_case							:		str        =		tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]    )
      snake_case							:		Optional[Any]        =		model(SCREAMING_SNAKE_CASE_    )[0]
      # compare the actual values for a slice.
      snake_case							:		Any        =		tf.constant(
          [
              [
                  [0.14_44_30_92, 0.54_12_53_27, 0.3_24_77_39],
                  [0.30_34_04_84, 0.00_52_66_76, 0.31_07_77_22],
                  [0.32_27_80_43, -0.24_98_70_96, 0.3_41_46_28],
              ]
          ]    )
      self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()   ,expected_slice.numpy()   ,atol=1E-4    )    )










from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Mapping, Optional, Union

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


if TYPE_CHECKING:
  from ...feature_extraction_utils import FeatureExtractionMixin
  from ...tokenization_utils_base import PreTrainedTokenizerBase
  from ...utils import TensorType

SCREAMING_SNAKE_CASE__							:      List[str]         =    logging.get_logger(__name__)

SCREAMING_SNAKE_CASE__							:      int         =    {
    """openai/whisper-base""": """https://huggingface.co/openai/whisper-base/resolve/main/config.json""",
}

# fmt: off
SCREAMING_SNAKE_CASE__							:      int         =    [
    1, 2, 7, 8, 9, 1_0, 1_4, 2_5,
    2_6, 2_7, 2_8, 2_9, 3_1, 5_8, 5_9, 6_0, 6_1, 6_2,
    6_3, 9_0, 9_1, 9_2, 9_3, 3_5_7, 3_6_6, 4_3_8, 5_3_2, 6_8_5,
    7_0_5, 7_9_6, 9_3_0, 1_0_5_8, 1_2_2_0, 1_2_6_7, 1_2_7_9, 1_3_0_3, 1_3_4_3, 1_3_7_7,
    1_3_9_1, 1_6_3_5, 1_7_8_2, 1_8_7_5, 2_1_6_2, 2_3_6_1, 2_4_8_8, 3_4_6_7, 4_0_0_8, 4_2_1_1,
    4_6_0_0, 4_8_0_8, 5_2_9_9, 5_8_5_5, 6_3_2_9, 7_2_0_3, 9_6_0_9, 9_9_5_9, 1_0_5_6_3, 1_0_7_8_6,
    1_1_4_2_0, 1_1_7_0_9, 1_1_9_0_7, 1_3_1_6_3, 1_3_6_9_7, 1_3_7_0_0, 1_4_8_0_8, 1_5_3_0_6, 1_6_4_1_0, 1_6_7_9_1,
    1_7_9_9_2, 1_9_2_0_3, 1_9_5_1_0, 2_0_7_2_4, 2_2_3_0_5, 2_2_9_3_5, 2_7_0_0_7, 3_0_1_0_9, 3_0_4_2_0, 3_3_4_0_9,
    3_4_9_4_9, 4_0_2_8_3, 4_0_4_9_3, 4_0_5_4_9, 4_7_2_8_2, 4_9_1_4_6, 5_0_2_5_7, 5_0_3_5_9, 5_0_3_6_0, 5_0_3_6_1
]
SCREAMING_SNAKE_CASE__							:      List[str]         =    [
    1, 2, 7, 8, 9, 1_0, 1_4, 2_5,
    2_6, 2_7, 2_8, 2_9, 3_1, 5_8, 5_9, 6_0, 6_1, 6_2,
    6_3, 9_0, 9_1, 9_2, 9_3, 3_5_9, 5_0_3, 5_2_2, 5_4_2, 8_7_3,
    8_9_3, 9_0_2, 9_1_8, 9_2_2, 9_3_1, 1_3_5_0, 1_8_5_3, 1_9_8_2, 2_4_6_0, 2_6_2_7,
    3_2_4_6, 3_2_5_3, 3_2_6_8, 3_5_3_6, 3_8_4_6, 3_9_6_1, 4_1_8_3, 4_6_6_7, 6_5_8_5, 6_6_4_7,
    7_2_7_3, 9_0_6_1, 9_3_8_3, 1_0_4_2_8, 1_0_9_2_9, 1_1_9_3_8, 1_2_0_3_3, 1_2_3_3_1, 1_2_5_6_2, 1_3_7_9_3,
    1_4_1_5_7, 1_4_6_3_5, 1_5_2_6_5, 1_5_6_1_8, 1_6_5_5_3, 1_6_6_0_4, 1_8_3_6_2, 1_8_9_5_6, 2_0_0_7_5, 2_1_6_7_5,
    2_2_5_2_0, 2_6_1_3_0, 2_6_1_6_1, 2_6_4_3_5, 2_8_2_7_9, 2_9_4_6_4, 3_1_6_5_0, 3_2_3_0_2, 3_2_4_7_0, 3_6_8_6_5,
    4_2_8_6_3, 4_7_4_2_5, 4_9_8_7_0, 5_0_2_5_4, 5_0_2_5_8, 5_0_3_6_0, 5_0_3_6_1, 5_0_3_6_2
]







class 	__lowerCAmelCase      (       _UpperCamelCase	):
   _UpperCamelCase			:    Any  		=					"""whisper"""
   _UpperCamelCase			:    Union[str, Any]  		=					["""past_key_values"""]
   _UpperCamelCase			:    Union[str, Any]  		=					{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}





   def __init__(			self       ,      snake_case=51_865       ,      snake_case=80       ,      snake_case=6       ,      snake_case=4       ,      snake_case=6       ,      snake_case=4       ,      snake_case=1_536       ,      snake_case=1_536       ,      snake_case=0.0       ,      snake_case=0.0       ,      snake_case=50_257       ,      snake_case=True       ,      snake_case=True       ,      snake_case="gelu"       ,      snake_case=256       ,      snake_case=0.0       ,      snake_case=0.0       ,      snake_case=0.0       ,      snake_case=0.02       ,      snake_case=False       ,      snake_case=1_500       ,      snake_case=448       ,      snake_case=50_256       ,      snake_case=50_256       ,      snake_case=50_256       ,      snake_case=None       ,      snake_case=[220, 50_256]       ,      snake_case=False       ,      snake_case=256       ,      snake_case=False       ,      snake_case=0.05       ,      snake_case=10       ,      snake_case=2       ,      snake_case=0.0       ,      snake_case=10       ,      snake_case=0       ,      snake_case=7       ,      **snake_case       ,      )					->	Dict:
      """simple docstring"""





      a__						:    Optional[Any]   		=	vocab_size
      a__						:    int   		=	num_mel_bins
      a__						:    Dict   		=	d_model
      a__						:    List[Any]   		=	encoder_layers
      a__						:    List[Any]   		=	encoder_attention_heads
      a__						:    Optional[int]   		=	decoder_layers
      a__						:    int   		=	decoder_attention_heads
      a__						:    Optional[Any]   		=	decoder_ffn_dim
      a__						:    List[Any]   		=	encoder_ffn_dim
      a__						:    int   		=	dropout
      a__						:    Optional[int]   		=	attention_dropout
      a__						:    Tuple   		=	activation_dropout
      a__						:    Optional[Any]   		=	activation_function
      a__						:    List[Any]   		=	init_std
      a__						:    List[Any]   		=	encoder_layerdrop
      a__						:    Dict   		=	decoder_layerdrop
      a__						:    List[Any]   		=	use_cache
      a__						:    Union[str, Any]   		=	encoder_layers
      a__						:    Union[str, Any]   		=	scale_embedding  # scale factor will be sqrt(d_model) if True
      a__						:    Tuple   		=	max_source_positions
      a__						:    Optional[Any]   		=	max_target_positions

      # Audio Classification-specific parameters. Feel free to ignore for other classes.
      a__						:    Optional[Any]   		=	classifier_proj_size
      a__						:    int   		=	use_weighted_layer_sum

      # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
      a__						:    Tuple   		=	apply_spec_augment
      a__						:    int   		=	mask_time_prob
      a__						:    Optional[Any]   		=	mask_time_length
      a__						:    List[str]   		=	mask_time_min_masks
      a__						:    List[str]   		=	mask_feature_prob
      a__						:    Dict   		=	mask_feature_length
      a__						:    Any   		=	mask_feature_min_masks

      a__						:    List[str]   		=	median_filter_width

      super().__init__(
          pad_token_id=snake_case       ,      bos_token_id=snake_case       ,      eos_token_id=snake_case       ,      is_encoder_decoder=snake_case       ,      decoder_start_token_id=snake_case       ,      suppress_tokens=snake_case       ,      begin_suppress_tokens=snake_case       ,      **snake_case       ,      )







class 	__lowerCAmelCase      (       _UpperCamelCase	):





   @property
   def 						_snake_case			(			self   )					->	Mapping[str, Mapping[int, str]]:
      """simple docstring"""





      a__						:    Dict   		=	OrderedDict(
          [
              ("input_features", {0: "batch", 1: "feature_size", 2: "encoder_sequence"}),
          ]   )
      if self.use_past:
         a__						:    Tuple   		=	{0: "batch"}
      else:
         a__						:    Union[str, Any]   		=	{0: "batch", 1: "decoder_sequence"}

      if self.use_past:
         self.fill_with_past_key_values_(snake_case       ,      direction="inputs"   )

      return common_inputs





   def 						_snake_case			(			self       ,      snake_case       ,      snake_case = -1       ,      snake_case = -1       ,      snake_case = False       ,      snake_case = None       ,      snake_case = 22_050       ,      snake_case = 5.0       ,      snake_case = 220       ,      )					->	Mapping[str, Any]:
      """simple docstring"""





      a__						:    int   		=	OrderedDict()
      a__						:    List[str]   		=	OnnxConfig.generate_dummy_inputs(
          self       ,      preprocessor=preprocessor.feature_extractor       ,      batch_size=snake_case       ,      framework=snake_case       ,      sampling_rate=snake_case       ,      time_duration=snake_case       ,      frequency=snake_case       ,      )
      a__						:    Optional[int]   		=	encoder_inputs["input_features"].shape[2]
      a__						:    str   		=	encoder_sequence_length // 2 if self.use_past else seq_length

      a__						:    Optional[int]   		=	super().generate_dummy_inputs(
          preprocessor.tokenizer       ,      snake_case       ,      snake_case       ,      snake_case       ,      snake_case   )

      a__						:    Any   		=	encoder_inputs.pop("input_features"   )
      a__						:    Dict   		=	decoder_inputs.pop("decoder_input_ids"   )

      if "past_key_values" in decoder_inputs:
         a__						:    Tuple   		=	decoder_inputs.pop("past_key_values"   )

      return dummy_inputs





   @property
   def 						_snake_case			(			self   )					->	float:
      """simple docstring"""





      return 1E-3

