





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





      def __init__( self       ,							A_       ,							):
       '''simple docstring'''




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





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					str       	=   ids_tensor([self.batch_size, self.seq_length]       ,							self.vocab_size   )

       UpperCamelCase	:					List[str]       	=   None
       if self.use_input_mask:
        UpperCamelCase	:					int       	=   random_attention_mask([self.batch_size, self.seq_length]   )

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

       UpperCamelCase	:					Dict       	=   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   ):
       '''simple docstring'''




       (
           (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    
       )	:					Optional[int]       	=   self.prepare_config_and_inputs()

       UpperCamelCase	:					Union[str, Any]       	=   True
       UpperCamelCase	:					Any       	=   floats_tensor([self.batch_size, self.seq_length, self.hidden_size]   )
       UpperCamelCase	:					List[Any]       	=   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       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_   ):
       '''simple docstring'''




       UpperCamelCase	:					Optional[int]       	=   TFEsmModel(config=A_   )
       UpperCamelCase	:					Any       	=   {"input_ids": input_ids, "attention_mask": input_mask}
       UpperCamelCase	:					int       	=   model(A_   )

       UpperCamelCase	:					List[Any]       	=   [input_ids, input_mask]
       UpperCamelCase	:					int       	=   model(A_   )

       UpperCamelCase	:					Tuple       	=   model(A_   )

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





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							):
       '''simple docstring'''




       UpperCamelCase	:					Tuple       	=   True

       UpperCamelCase	:					str       	=   TFEsmModel(config=A_   )
       UpperCamelCase	:					Optional[Any]       	=   {
           "input_ids": input_ids,
           "attention_mask": input_mask,
           "encoder_hidden_states": encoder_hidden_states,
           "encoder_attention_mask": encoder_attention_mask,
       }
       UpperCamelCase	:					Dict       	=   model(A_   )

       UpperCamelCase	:					Union[str, Any]       	=   [input_ids, input_mask]
       UpperCamelCase	:					List[Any]       	=   model(A_       ,							encoder_hidden_states=A_   )

       # Also check the case where encoder outputs are not passed
       UpperCamelCase	:					List[str]       	=   model(A_       ,							attention_mask=A_   )

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





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_   ):
       '''simple docstring'''




       UpperCamelCase	:					List[Any]       	=   TFEsmForMaskedLM(config=A_   )
       UpperCamelCase	:					Tuple       	=   model([input_ids, input_mask]   )
       self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.seq_length, self.vocab_size)   )





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_   ):
       '''simple docstring'''




       UpperCamelCase	:					int       	=   self.num_labels
       UpperCamelCase	:					Dict       	=   TFEsmForTokenClassification(config=A_   )
       UpperCamelCase	:					List[str]       	=   {"input_ids": input_ids, "attention_mask": input_mask}
       UpperCamelCase	:					Dict       	=   model(A_   )
       self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.seq_length, self.num_labels)   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					List[str]       	=   self.prepare_config_and_inputs()
       (
           (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    
       )	:					Any       	=   config_and_inputs
       UpperCamelCase	:					str       	=   {"input_ids": input_ids, "attention_mask": input_mask}
       return config, inputs_dict



@require_tf
class 					A__   (       __snake_case       ,   __snake_case       ,   unittest.TestCase			):
      _UpperCAmelCase	:Any						    =   (
          (
              TFEsmModel,
              TFEsmForMaskedLM,
              TFEsmForSequenceClassification,
              TFEsmForTokenClassification,
          )
          if is_tf_available()
          else ()
      )
      _UpperCAmelCase	:List[Any]						    =   (
          {
              'feature-extraction': TFEsmModel,
              'fill-mask': TFEsmForMaskedLM,
              'text-classification': TFEsmForSequenceClassification,
              'token-classification': TFEsmForTokenClassification,
              'zero-shot': TFEsmForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      _UpperCAmelCase	:Tuple						    =   False
      _UpperCAmelCase	:Optional[Any]						    =   False





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Tuple       	=   TFEsmModelTester(self   )
       UpperCamelCase	:					List[Any]       	=   ConfigTester(self       ,							config_class=A_       ,							hidden_size=37   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       self.config_tester.run_common_tests()





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Dict       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Any       	=   self.model_tester.prepare_config_and_inputs_for_decoder()
       self.model_tester.create_and_check_model_as_decoder(*A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Tuple       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					List[Any]       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*A_   )





      @slow
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCamelCase	:					Union[str, Any]       	=   TFEsmModel.from_pretrained(A_   )
        self.assertIsNotNone(A_   )





      @unittest.skip("Protein models do not support embedding resizing."   )
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       pass





      @unittest.skip("Protein models do not support embedding resizing."   )
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       pass





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	,    UpperCamelCase	:					Dict       	=   self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        UpperCamelCase	:					Dict       	=   model_class(A_   )
        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	:					Optional[Any]       	=   model.get_bias()
         assert isinstance(A_       ,							A_   )
         for k, v in name.items():
          assert isinstance(A_       ,							tf.Variable   )
        else:
         UpperCamelCase	:					str       	=   model.get_output_embeddings()
         assert x is None
         UpperCamelCase	:					Dict       	=   model.get_bias()
         assert name is None



@require_tf
class 					A__   (       unittest.TestCase			):





      @slow
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Dict       	=   TFEsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"   )

       UpperCamelCase	:					Optional[int]       	=   tf.constant([[0, 1, 2, 3, 4, 5]]   )
       UpperCamelCase	:					int       	=   model(A_   )[0]
       UpperCamelCase	:					List[str]       	=   [1, 6, 33]
       self.assertEqual(list(output.numpy().shape   )       ,							A_   )
       # compare the actual values for a slice.
       UpperCamelCase	:					str       	=   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        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Optional[int]       	=   TFEsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"   )

       UpperCamelCase	:					Optional[int]       	=   tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]   )
       UpperCamelCase	:					List[Any]       	=   model(A_   )[0]
       # compare the actual values for a slice.
       UpperCamelCase	:					int       	=   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   )   )







import itertools
import random
import unittest

import numpy as np

from transformers import ASTFeatureExtractor
from transformers.testing_utils import require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


__lowerCamelCase       :     str   =     random.Random()

if is_torch_available():
    import torch



def 		A_				(     _lowerCAmelCase							,   _lowerCAmelCase=1.0							,   _lowerCAmelCase=None							,   _lowerCAmelCase=None     )						-> Optional[Any]:
 if rng is None:
  UpperCamelCase	:					Optional[int]       	=   global_rng

 UpperCamelCase	:					Optional[Any]       	=   []
 for batch_idx in range(shape[0]     ):
  values.append([]     )
  for _ in range(shape[1]     ):
   values[-1].append(rng.random() * scale     )

 return values



class 					A__   (       unittest.TestCase			):





      def __init__( self       ,							A_       ,							A_=7       ,							A_=400       ,							A_=2000       ,							A_=1       ,							A_=0.0       ,							A_=1_6000       ,							A_=True       ,							A_=True       ,							):
       '''simple docstring'''




       UpperCamelCase	:					Tuple       	=   parent
       UpperCamelCase	:					List[Any]       	=   batch_size
       UpperCamelCase	:					List[Any]       	=   min_seq_length
       UpperCamelCase	:					List[str]       	=   max_seq_length
       UpperCamelCase	:					int       	=   (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
       UpperCamelCase	:					Union[str, Any]       	=   feature_size
       UpperCamelCase	:					List[str]       	=   padding_value
       UpperCamelCase	:					Optional[Any]       	=   sampling_rate
       UpperCamelCase	:					List[str]       	=   return_attention_mask
       UpperCamelCase	:					List[Any]       	=   do_normalize





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       return {
           "feature_size": self.feature_size,
           "padding_value": self.padding_value,
           "sampling_rate": self.sampling_rate,
           "return_attention_mask": self.return_attention_mask,
           "do_normalize": self.do_normalize,
       }





      def        __UpperCamelCase( self       ,							A_=False       ,							A_=False   ):
       '''simple docstring'''




       def _flatten(A_   ):
        return list(itertools.chain(*A_   )   )

       if equal_length:
        UpperCamelCase	:					List[str]       	=   floats_list((self.batch_size, self.max_seq_length)   )
       else:
        # make sure that inputs increase in size
        UpperCamelCase	:					Dict       	=   [
            _flatten(floats_list((x, self.feature_size)   )   )
            for x in range(self.min_seq_length       ,							self.max_seq_length       ,							self.seq_length_diff   )
        ]

       if numpify:
        UpperCamelCase	:					Union[str, Any]       	=   [np.asarray(A_   ) for x in speech_inputs]

       return speech_inputs



@require_torch
@require_torchaudio
class 					A__   (       __snake_case       ,   unittest.TestCase			):
      _UpperCAmelCase	:Optional[Any]						    =   ASTFeatureExtractor





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Tuple       	=   ASTFeatureExtractionTester(self   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Optional[Any]       	=   self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
       # create three inputs of length 800, 1000, and 1200
       UpperCamelCase	:					Tuple       	=   [floats_list((1, x)   )[0] for x in range(800       ,							1400       ,							200   )]
       UpperCamelCase	:					Dict       	=   [np.asarray(A_   ) for speech_input in speech_inputs]

       # Test not batched input
       UpperCamelCase	:					Dict       	=   feat_extract(speech_inputs[0]       ,							return_tensors="np"   ).input_values
       UpperCamelCase	:					Union[str, Any]       	=   feat_extract(np_speech_inputs[0]       ,							return_tensors="np"   ).input_values
       self.assertTrue(np.allclose(A_       ,							A_       ,							atol=1e-3   )   )

       # Test batched
       UpperCamelCase	:					Any       	=   feat_extract(A_       ,							padding=A_       ,							return_tensors="np"   ).input_values
       UpperCamelCase	:					Any       	=   feat_extract(A_       ,							padding=A_       ,							return_tensors="np"   ).input_values
       for enc_seq_a, enc_seq_a in zip(A_       ,							A_   ):
        self.assertTrue(np.allclose(A_       ,							A_       ,							atol=1e-3   )   )

       # Test 2-D numpy arrays are batched.
       UpperCamelCase	:					Dict       	=   [floats_list((1, x)   )[0] for x in (800, 800, 800)]
       UpperCamelCase	:					int       	=   np.asarray(A_   )
       UpperCamelCase	:					Any       	=   feat_extract(A_       ,							return_tensors="np"   ).input_values
       UpperCamelCase	:					List[str]       	=   feat_extract(A_       ,							return_tensors="np"   ).input_values
       for enc_seq_a, enc_seq_a in zip(A_       ,							A_   ):
        self.assertTrue(np.allclose(A_       ,							A_       ,							atol=1e-3   )   )





      @require_torch
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       import torch

       UpperCamelCase	:					List[Any]       	=   self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
       UpperCamelCase	:					int       	=   np.random.rand(100   ).astype(np.floataa   )
       UpperCamelCase	:					str       	=   np_speech_inputs.tolist()

       for inputs in [py_speech_inputs, np_speech_inputs]:
        UpperCamelCase	:					List[Any]       	=   feature_extractor.pad([{"input_values": inputs}]       ,							return_tensors="np"   )
        self.assertTrue(np_processed.input_values.dtype == np.floataa   )
        UpperCamelCase	:					List[str]       	=   feature_extractor.pad([{"input_values": inputs}]       ,							return_tensors="pt"   )
        self.assertTrue(pt_processed.input_values.dtype == torch.floataa   )





      def        __UpperCamelCase( self       ,							A_   ):
       '''simple docstring'''




       from datasets import load_dataset

       UpperCamelCase	:					Dict       	=   load_dataset("hf-internal-testing/librispeech_asr_dummy"       ,							"clean"       ,							split="validation"   )
       # automatic decoding with librispeech
       UpperCamelCase	:					Any       	=   ds.sort("id"   ).select(range(A_   )   )[:num_samples]["audio"]

       return [x["array"] for x in speech_samples]





      @require_torch
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Any       	=   torch.tensor(
           [-0.98_94, -1.27_76, -0.90_66, -1.27_76, -0.93_49, -1.26_09, -1.03_86, -1.27_76,
            -1.15_61, -1.27_76, -1.20_52, -1.27_23, -1.21_90, -1.21_32, -1.27_76, -1.11_33,
            -1.19_53, -1.13_43, -1.15_84, -1.22_03, -1.17_70, -1.24_74, -1.23_81, -1.19_36,
            -0.92_70, -0.83_17, -0.80_49, -0.77_06, -0.75_65, -0.78_69]   )
       # fmt: on

       UpperCamelCase	:					List[Any]       	=   self._load_datasamples(1   )
       UpperCamelCase	:					Tuple       	=   ASTFeatureExtractor()
       UpperCamelCase	:					str       	=   feature_extractor(A_       ,							return_tensors="pt"   ).input_values
       self.assertEquals(input_values.shape       ,							(1, 1024, 128)   )
       self.assertTrue(torch.allclose(input_values[0, 0, :30]       ,							A_       ,							atol=1e-4   )   )
