





'''simple docstring'''






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


  """simple docstring"""







  def __init__(				self    :					Tuple							,	UpperCamelCase__    :					int							,	)     ->						Any:
    _UpperCamelCase       			=parent
    _UpperCamelCase       			=13
    _UpperCamelCase       			=7
    _UpperCamelCase       			=True
    _UpperCamelCase       			=True
    _UpperCamelCase       			=True
    _UpperCamelCase       			=99
    _UpperCamelCase       			=32
    _UpperCamelCase       			=2
    _UpperCamelCase       			=4
    _UpperCamelCase       			=37
    _UpperCamelCase       			='''gelu'''
    _UpperCamelCase       			=0.1
    _UpperCamelCase       			=0.1
    _UpperCamelCase       			=512
    _UpperCamelCase       			=16
    _UpperCamelCase       			=2
    _UpperCamelCase       			=0.02
    _UpperCamelCase       			=3
    _UpperCamelCase       			=4
    _UpperCamelCase       			=None


  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						int:
    _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       			=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   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						Tuple:
    (
        (
        _UpperCamelCase
    )			,     (
        _UpperCamelCase
    )			,     (
        _UpperCamelCase
    )			,     (
        _UpperCamelCase
    )			,     (
        _UpperCamelCase
    )			,     (
        _UpperCamelCase
    )			,     
    )       			=self.prepare_config_and_inputs()

    _UpperCamelCase       			=True
    _UpperCamelCase       			=floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
    _UpperCamelCase       			=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   UpperCamelCase__      (				self    :					List[str]							,	UpperCamelCase__    :					str							,	UpperCamelCase__    :					Tuple							,	UpperCamelCase__    :					str							,	UpperCamelCase__    :					List[Any]							,	UpperCamelCase__    :					List[str]							,	UpperCamelCase__    :					Optional[int]  )     ->						Union[str, Any]:
    _UpperCamelCase       			=TFEsmModel(config=UpperCamelCase__  )
    _UpperCamelCase       			={'''input_ids''': input_ids, '''attention_mask''': input_mask}
    _UpperCamelCase       			=model(UpperCamelCase__  )

    _UpperCamelCase       			=[input_ids, input_mask]
    _UpperCamelCase       			=model(UpperCamelCase__  )

    _UpperCamelCase       			=model(UpperCamelCase__  )

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


  def   UpperCamelCase__      (				self    :					List[str]							,	UpperCamelCase__    :					int							,	UpperCamelCase__    :					int							,	UpperCamelCase__    :					Union[str, Any]							,	UpperCamelCase__    :					str							,	UpperCamelCase__    :					Optional[Any]							,	UpperCamelCase__    :					Optional[int]							,	UpperCamelCase__    :					Tuple							,	UpperCamelCase__    :					int							,	)     ->						List[str]:
    _UpperCamelCase       			=True

    _UpperCamelCase       			=TFEsmModel(config=UpperCamelCase__  )
    _UpperCamelCase       			={
        '''input_ids''': input_ids,
        '''attention_mask''': input_mask,
        '''encoder_hidden_states''': encoder_hidden_states,
        '''encoder_attention_mask''': encoder_attention_mask,
    }
    _UpperCamelCase       			=model(UpperCamelCase__  )

    _UpperCamelCase       			=[input_ids, input_mask]
    _UpperCamelCase       			=model(UpperCamelCase__							,	encoder_hidden_states=UpperCamelCase__  )

    # Also check the case where encoder outputs are not passed
    _UpperCamelCase       			=model(UpperCamelCase__							,	attention_mask=UpperCamelCase__  )

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


  def   UpperCamelCase__      (				self    :					List[Any]							,	UpperCamelCase__    :					int							,	UpperCamelCase__    :					str							,	UpperCamelCase__    :					Dict							,	UpperCamelCase__    :					str							,	UpperCamelCase__    :					Dict							,	UpperCamelCase__    :					Optional[Any]  )     ->						Union[str, Any]:
    _UpperCamelCase       			=TFEsmForMaskedLM(config=UpperCamelCase__  )
    _UpperCamelCase       			=model([input_ids, input_mask]  )
    self.parent.assertEqual(result.logits.shape							,	(self.batch_size, self.seq_length, self.vocab_size)  )


  def   UpperCamelCase__      (				self    :					str							,	UpperCamelCase__    :					Any							,	UpperCamelCase__    :					Tuple							,	UpperCamelCase__    :					Union[str, Any]							,	UpperCamelCase__    :					Any							,	UpperCamelCase__    :					Union[str, Any]							,	UpperCamelCase__    :					str  )     ->						Any:
    _UpperCamelCase       			=self.num_labels
    _UpperCamelCase       			=TFEsmForTokenClassification(config=UpperCamelCase__  )
    _UpperCamelCase       			={'''input_ids''': input_ids, '''attention_mask''': input_mask}
    _UpperCamelCase       			=model(UpperCamelCase__  )
    self.parent.assertEqual(result.logits.shape							,	(self.batch_size, self.seq_length, self.num_labels)  )




  def   UpperCamelCase__      (				self    :					str  )     ->						str:
    _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_tf
class 		UpperCAmelCase   (				lowercase_   ,     lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				(
      (
          TFEsmModel,
          TFEsmForMaskedLM,
          TFEsmForSequenceClassification,
          TFEsmForTokenClassification,
      )
      if is_tf_available()
      else ()
  )
  lowerCAmelCase_					=				(
      {
          """feature-extraction""": TFEsmModel,
          """fill-mask""": TFEsmForMaskedLM,
          """text-classification""": TFEsmForSequenceClassification,
          """token-classification""": TFEsmForTokenClassification,
          """zero-shot""": TFEsmForSequenceClassification,
      }
      if is_tf_available()
      else {}
  )
  lowerCAmelCase_					=				False
  lowerCAmelCase_					=				False


  def   UpperCamelCase__      (				self    :					List[Any]  )     ->						Any:
    _UpperCamelCase       			=TFEsmModelTester(self  )
    _UpperCamelCase       			=ConfigTester(self							,	config_class=UpperCamelCase__							,	hidden_size=37  )


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


  def   UpperCamelCase__      (				self    :					Tuple  )     ->						Optional[int]:
    _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*UpperCamelCase__  )


  def   UpperCamelCase__      (				self    :					Dict  )     ->						Any:
    _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs_for_decoder()
    self.model_tester.create_and_check_model_as_decoder(*UpperCamelCase__  )


  def   UpperCamelCase__      (				self    :					Dict  )     ->						Any:
    _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase__  )


  def   UpperCamelCase__      (				self    :					int  )     ->						Union[str, Any]:
    _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_token_classification(*UpperCamelCase__  )


  @slow
  def   UpperCamelCase__      (				self    :					List[Any]  )     ->						int:
    for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      _UpperCamelCase       			=TFEsmModel.from_pretrained(UpperCamelCase__  )
      self.assertIsNotNone(UpperCamelCase__  )


  @unittest.skip('''Protein models do not support embedding resizing.'''  )
  def   UpperCamelCase__      (				self    :					List[str]  )     ->						str:
    pass


  @unittest.skip('''Protein models do not support embedding resizing.'''  )
  def   UpperCamelCase__      (				self    :					Optional[Any]  )     ->						List[str]:
    pass




  def   UpperCamelCase__      (				self    :					Optional[int]  )     ->						List[str]:
    _UpperCamelCase			,     _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs_for_common()

    for model_class in self.all_model_classes:
      _UpperCamelCase       			=model_class(UpperCamelCase__  )
      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
        _UpperCamelCase       			=model.get_bias()
        assert isinstance(UpperCamelCase__							,	UpperCamelCase__  )
        for k, v in name.items():
          assert isinstance(UpperCamelCase__							,	tf.Variable  )
      else:
        _UpperCamelCase       			=model.get_output_embeddings()
        assert x is None
        _UpperCamelCase       			=model.get_bias()
        assert name is None







@require_tf
class 		UpperCAmelCase   (				unittest.TestCase):


  """simple docstring"""







  @slow
  def   UpperCamelCase__      (				self    :					Optional[Any]  )     ->						List[str]:
    _UpperCamelCase       			=TFEsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''  )

    _UpperCamelCase       			=tf.constant([[0, 1, 2, 3, 4, 5]]  )
    _UpperCamelCase       			=model(UpperCamelCase__  )[0]
    _UpperCamelCase       			=[1, 6, 33]
    self.assertEqual(list(output.numpy().shape  )							,	UpperCamelCase__  )
    # compare the actual values for a slice.
    _UpperCamelCase       			=tf.constant(
        [
            [
                [8.921518, -10.589814, -6.4671307],
                [-6.3967156, -13.911377, -1.1211915],
                [-7.781247, -13.951557, -3.740592],
            ]
        ]  )
    self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()							,	expected_slice.numpy()							,	atol=1E-2  )  )




  @slow
  def   UpperCamelCase__      (				self    :					List[Any]  )     ->						int:
    _UpperCamelCase       			=TFEsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''  )

    _UpperCamelCase       			=tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]  )
    _UpperCamelCase       			=model(UpperCamelCase__  )[0]
    # compare the actual values for a slice.
    _UpperCamelCase       			=tf.constant(
        [
            [
                [0.14443092, 0.54125327, 0.3247739],
                [0.30340484, 0.00526676, 0.31077722],
                [0.32278043, -0.24987096, 0.3414628],
            ]
        ]  )
    self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()							,	expected_slice.numpy()							,	atol=1E-4  )  )













'''simple docstring'''






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_speech_available,
    is_tf_available,
    is_torch_available,
)


__lowerCamelCase						: str		=   {
    'configuration_speech_to_text': ['SPEECH_TO_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'Speech2TextConfig'],
    'processing_speech_to_text': ['Speech2TextProcessor'],
}

try:
      if not is_sentencepiece_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __lowerCamelCase						: List[str]		=   ['Speech2TextTokenizer']

try:
      if not is_speech_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __lowerCamelCase						: str		=   ['Speech2TextFeatureExtractor']

try:
      if not is_tf_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __lowerCamelCase						: Optional[Any]		=   [
          'TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST',
          'TFSpeech2TextForConditionalGeneration',
          'TFSpeech2TextModel',
          'TFSpeech2TextPreTrainedModel',
      ]

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __lowerCamelCase						: Union[str, Any]		=   [
          'SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST',
          'Speech2TextForConditionalGeneration',
          'Speech2TextModel',
          'Speech2TextPreTrainedModel',
      ]


if TYPE_CHECKING:
      from .configuration_speech_to_text import SPEECH_TO_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP, SpeechaTextConfig
      from .processing_speech_to_text import SpeechaTextProcessor

      try:
            if not is_sentencepiece_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_speech_to_text import SpeechaTextTokenizer

      try:
            if not is_speech_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .feature_extraction_speech_to_text import SpeechaTextFeatureExtractor

      try:
            if not is_tf_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_tf_speech_to_text import (
                TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST,
                TFSpeechaTextForConditionalGeneration,
                TFSpeechaTextModel,
                TFSpeechaTextPreTrainedModel,
            )

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_speech_to_text import (
                SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST,
                SpeechaTextForConditionalGeneration,
                SpeechaTextModel,
                SpeechaTextPreTrainedModel,
            )

else:
      import sys

      __lowerCamelCase						: Optional[int]		=   _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)






