



"""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.models.esm.modeling_esmfold import EsmForProteinFolding



class  _lowerCAmelCase     :





		"""simple docstring"""






		def __init__(       self				,     _lowercase				,     _lowercase=1_3				,     _lowercase=7				,     _lowercase=False				,     _lowercase=True				,     _lowercase=False				,     _lowercase=False				,     _lowercase=1_9				,     _lowercase=3_2				,     _lowercase=5				,     _lowercase=4				,     _lowercase=3_7				,     _lowercase="gelu"				,     _lowercase=0.1				,     _lowercase=0.1				,     _lowercase=5_1_2				,     _lowercase=1_6				,     _lowercase=2				,     _lowercase=0.02				,     _lowercase=3				,     _lowercase=4				,     _lowercase=None				,     )			->    Dict:







								'''simple docstring'''

								snake_case_ :			Any            =      parent
								snake_case_ :			str            =      batch_size
								snake_case_ :			Dict            =      seq_length
								snake_case_ :			List[str]            =      is_training
								snake_case_ :			Tuple            =      use_input_mask
								snake_case_ :			str            =      use_token_type_ids
								snake_case_ :			Any            =      use_labels
								snake_case_ :			Any            =      vocab_size
								snake_case_ :			Union[str, Any]            =      hidden_size
								snake_case_ :			Any            =      num_hidden_layers
								snake_case_ :			Dict            =      num_attention_heads
								snake_case_ :			Dict            =      intermediate_size
								snake_case_ :			List[str]            =      hidden_act
								snake_case_ :			Any            =      hidden_dropout_prob
								snake_case_ :			Optional[int]            =      attention_probs_dropout_prob
								snake_case_ :			List[Any]            =      max_position_embeddings
								snake_case_ :			Optional[Any]            =      type_vocab_size
								snake_case_ :			List[str]            =      type_sequence_label_size
								snake_case_ :			str            =      initializer_range
								snake_case_ :			List[str]            =      num_labels
								snake_case_ :			Optional[Any]            =      num_choices
								snake_case_ :			str            =      scope





		def  UpperCAmelCase__				(       self )			->    Optional[Any]:







								'''simple docstring'''

								snake_case_ :			List[Any]            =      ids_tensor([self.batch_size, self.seq_length]				,     self.vocab_size )

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

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

								snake_case_ :			str            =      self.get_config()

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





		def  UpperCAmelCase__				(       self )			->    Any:







								'''simple docstring'''

								snake_case_ :			Dict            =      EsmConfig(
								    vocab_size=3_3				,     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				,     is_folding_model=_lowercase				,     esmfold_config={"""trunk""": {"""num_blocks""": 2}, """fp16_esm""": False}				,     )
								return config





		def  UpperCAmelCase__				(       self				,     _lowercase				,     _lowercase				,     _lowercase				,     _lowercase				,     _lowercase				,     _lowercase )			->    Union[str, Any]:







								'''simple docstring'''

								snake_case_ :			Optional[int]            =      EsmForProteinFolding(config=_lowercase ).float()
								model.to(_lowercase )
								model.eval()
								snake_case_ :			Tuple            =      model(_lowercase				,     attention_mask=_lowercase )
								snake_case_ :			int            =      model(_lowercase )
								snake_case_ :			Any            =      model(_lowercase )

								self.parent.assertEqual(result.positions.shape				,     (8, self.batch_size, self.seq_length, 1_4, 3) )
								self.parent.assertEqual(result.angles.shape				,     (8, self.batch_size, self.seq_length, 7, 2) )





		def  UpperCAmelCase__				(       self )			->    Tuple:







								'''simple docstring'''

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




@require_torch
class  _lowerCAmelCase     (					SCREAMING_SNAKE_CASE__	,					SCREAMING_SNAKE_CASE__	,					unittest.TestCase		):





		"""simple docstring"""






		_lowerCamelCase      	=							False

		_lowerCamelCase      	=							(EsmForProteinFolding,) if is_torch_available() else ()
		_lowerCamelCase      	=							()
		_lowerCamelCase      	=							{} if is_torch_available() else {}
		_lowerCamelCase      	=							False





		def  UpperCAmelCase__				(       self )			->    Any:







								'''simple docstring'''

								snake_case_ :			str            =      EsmFoldModelTester(self )
								snake_case_ :			Union[str, Any]            =      ConfigTester(self				,     config_class=_lowercase				,     hidden_size=3_7 )





		def  UpperCAmelCase__				(       self )			->    int:







								'''simple docstring'''

								self.config_tester.run_common_tests()





		def  UpperCAmelCase__				(       self )			->    Optional[Any]:







								'''simple docstring'''

								snake_case_ :			str            =      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_lowercase )





		@unittest.skip("""Does not support attention outputs""" )
		def  UpperCAmelCase__				(       self )			->    Any:







								'''simple docstring'''

								pass





		@unittest.skip
		def  UpperCAmelCase__				(       self )			->    Dict:







								'''simple docstring'''

								pass





		@unittest.skip("""Esm does not support embedding resizing""" )
		def  UpperCAmelCase__				(       self )			->    Dict:







								'''simple docstring'''

								pass





		@unittest.skip("""Esm does not support embedding resizing""" )
		def  UpperCAmelCase__				(       self )			->    str:







								'''simple docstring'''

								pass





		@unittest.skip("""ESMFold does not support passing input embeds!""" )
		def  UpperCAmelCase__				(       self )			->    List[Any]:







								'''simple docstring'''

								pass





		@unittest.skip("""ESMFold does not support head pruning.""" )
		def  UpperCAmelCase__				(       self )			->    Tuple:







								'''simple docstring'''

								pass





		@unittest.skip("""ESMFold does not support head pruning.""" )
		def  UpperCAmelCase__				(       self )			->    Optional[Any]:







								'''simple docstring'''

								pass





		@unittest.skip("""ESMFold does not support head pruning.""" )
		def  UpperCAmelCase__				(       self )			->    List[str]:







								'''simple docstring'''

								pass





		@unittest.skip("""ESMFold does not support head pruning.""" )
		def  UpperCAmelCase__				(       self )			->    Any:







								'''simple docstring'''

								pass





		@unittest.skip("""ESMFold does not support head pruning.""" )
		def  UpperCAmelCase__				(       self )			->    str:







								'''simple docstring'''

								pass





		@unittest.skip("""ESMFold does not output hidden states in the normal way.""" )
		def  UpperCAmelCase__				(       self )			->    int:







								'''simple docstring'''

								pass





		@unittest.skip("""ESMfold does not output hidden states in the normal way.""" )
		def  UpperCAmelCase__				(       self )			->    Union[str, Any]:







								'''simple docstring'''

								pass





		@unittest.skip("""ESMFold only has one output format.""" )
		def  UpperCAmelCase__				(       self )			->    Dict:







								'''simple docstring'''

								pass





		@unittest.skip("""This test doesn't work for ESMFold and doesn't test core functionality""" )
		def  UpperCAmelCase__				(       self )			->    Tuple:







								'''simple docstring'''

								pass





		@unittest.skip("""ESMFold does not support input chunking.""" )
		def  UpperCAmelCase__				(       self )			->    Optional[int]:







								'''simple docstring'''

								pass





		@unittest.skip("""ESMFold doesn't respect you and it certainly doesn't respect your initialization arguments.""" )
		def  UpperCAmelCase__				(       self )			->    Union[str, Any]:







								'''simple docstring'''

								pass





		@unittest.skip("""ESMFold doesn't support torchscript compilation.""" )
		def  UpperCAmelCase__				(       self )			->    List[str]:







								'''simple docstring'''

								pass





		@unittest.skip("""ESMFold doesn't support torchscript compilation.""" )
		def  UpperCAmelCase__				(       self )			->    int:







								'''simple docstring'''

								pass





		@unittest.skip("""ESMFold doesn't support torchscript compilation.""" )
		def  UpperCAmelCase__				(       self )			->    Any:







								'''simple docstring'''

								pass





		@unittest.skip("""ESMFold doesn't support data parallel.""" )
		def  UpperCAmelCase__				(       self )			->    Union[str, Any]:







								'''simple docstring'''

								pass





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







								'''simple docstring'''

								pass




@require_torch
class  _lowerCAmelCase     (					SCREAMING_SNAKE_CASE__		):





		"""simple docstring"""






		@slow
		def  UpperCAmelCase__				(       self )			->    Union[str, Any]:







								'''simple docstring'''

								snake_case_ :			List[str]            =      EsmForProteinFolding.from_pretrained("""facebook/esmfold_v1""" ).float()
								model.eval()
								snake_case_ :			Any            =      torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]] )
								snake_case_ :			Tuple            =      model(_lowercase )["""positions"""]
								snake_case_ :			Any            =      torch.tensor([2.5828, 0.7993, -10.9334]				,     dtype=torch.floataa )
								self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]				,     _lowercase				,     atol=1E-4 ) )


'''simple docstring'''




from math import ceil
from typing import List, Optional, Union

import numpy as np

from ...audio_utils import mel_filter_bank, spectrogram, window_function
from ...feature_extraction_sequence_utils import BatchFeature, SequenceFeatureExtractor
from ...utils import TensorType, logging


__UpperCAmelCase					=logging.get_logger(__name__)




class        a__ (  UpperCAmelCase__					):
     lowerCamelCase    :	Union[str, Any]				=["audio_values", "audio_mask"]






     def __init__(					self						:						List[str]    ,							a						:						List[Any]=20_48    ,							a						:						Optional[int]=1    ,							a						:						List[str]=[16, 16]    ,							a						:						int=1_28    ,							a						:						Dict=4_41_00    ,							a						:						str=86    ,							a						:						int=20_48    ,							a						:						int=0.0    ,							**a						:						Dict    ,							):





         """simple docstring"""



         super().__init__(
             feature_size=a    ,							sampling_rate=a    ,							padding_value=a    ,							**a    ,							)

         __lowerCamelCase            =    spectrogram_length
         __lowerCamelCase            =    num_channels
         __lowerCamelCase            =    patch_size
         __lowerCamelCase            =    feature_size // self.patch_size[1]
         __lowerCamelCase            =    n_fft
         __lowerCamelCase            =    sampling_rate // hop_length_to_sampling_rate
         __lowerCamelCase            =    sampling_rate
         __lowerCamelCase            =    padding_value
         __lowerCamelCase            =    mel_filter_bank(
             num_frequency_bins=1 + n_fft // 2    ,							num_mel_filters=a    ,							min_frequency=0.0    ,							max_frequency=2_20_50.0    ,							sampling_rate=a    ,							norm='''slaney'''    ,							mel_scale='''slaney'''    ,							).T






     def   SCREAMING_SNAKE_CASE__  (					self						:						Tuple    ,							a						:						np.array       ):





         """simple docstring"""



         __lowerCamelCase            =    spectrogram(
             a    ,							window_function(self.n_fft    ,							'''hann'''       )    ,							frame_length=self.n_fft    ,							hop_length=self.hop_length    ,							power=2.0    ,							mel_filters=self.mel_filters.T    ,							log_mel='''dB'''    ,							db_range=80.0    ,							)
         __lowerCamelCase            =    log_spec[:, :-1]
         __lowerCamelCase            =    log_spec - 20.0
         __lowerCamelCase            =    np.clip(log_spec / 40.0    ,							-2.0    ,							0.0       ) + 1.0
         return log_spec






     def __call__(					self						:						Union[str, Any]    ,							a						:						Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]    ,							a						:						Optional[Union[str, TensorType]] = None    ,							a						:						Optional[bool] = True    ,							a						:						Optional[int] = None    ,							a						:						bool = False    ,							a						:						bool = False    ,							**a						:						int    ,							):





         """simple docstring"""



         if sampling_rate is not None:
             if sampling_rate != self.sampling_rate:
                 raise ValueError(
                     '''This feature extractor is set to support sampling rate'''
                     f""" of {self.sampling_rate}. Please make sure that the provided `raw_speech` input was sampled"""
                     f""" with {self.sampling_rate} and not {sampling_rate}."""       )
         else:
             logger.warning(
                 '''It is strongly recommended to pass the `sampling_rate` argument to this function. '''
                 '''Failing to do so can result in silent errors that might be hard to debug.'''       )

         __lowerCamelCase            =    isinstance(a    ,							np.ndarray       ) and len(raw_speech.shape       ) > 1
         if is_batched_numpy and len(raw_speech.shape       ) > 2:
             raise ValueError(f"""Only mono-channel audio is supported for input to {self}"""       )
         __lowerCamelCase            =    is_batched_numpy or (
             isinstance(a    ,							(list, tuple)       ) and (isinstance(raw_speech[0]    ,							(np.ndarray, tuple, list)       ))
         )
         if is_batched:
             __lowerCamelCase            =    [np.asarray([speech]    ,							dtype=np.floataa       ).T for speech in raw_speech]
         elif not is_batched and not isinstance(a    ,							np.ndarray       ):
             __lowerCamelCase            =    np.asarray(a    ,							dtype=np.floataa       )
         elif isinstance(a    ,							np.ndarray       ) and raw_speech.dtype is np.dtype(np.floataa       ):
             __lowerCamelCase            =    raw_speech.astype(np.floataa       )
         # always return batch
         if not is_batched:
             __lowerCamelCase            =    [np.asarray([raw_speech]       ).T]

         # Convert audio signals to log mel spectrograms, truncate by time axis
         __lowerCamelCase            =    [
             self._np_extract_fbank_features(waveform.squeeze()       ).T[: self.spectrogram_length] for waveform in raw_speech
         ]
         if isinstance(audio_features[0]    ,							a       ):
             __lowerCamelCase            =    [np.asarray(a    ,							dtype=np.floataa       ) for feature in audio_features]

         # Create audio attention mask
         __lowerCamelCase            =    max(
             [ceil(feature.shape[0] / self.patch_size[0]       ) * self.freq_len for feature in audio_features]       )  # The maximum number of audio patches in a batch
         if return_attention_mask:
             __lowerCamelCase            =    [
                 (ceil(feature.shape[0] / self.patch_size[0]       ) * self.freq_len) * [1]
                 + (max_patch_len - ceil(feature.shape[0] / self.patch_size[0]       ) * self.freq_len) * [0]
                 for feature in audio_features
             ]
             __lowerCamelCase            =    np.array(a       ).astype(np.floataa       )

         # convert into correct format for padding
         __lowerCamelCase            =    max_patch_len // self.freq_len * self.patch_size[0]  # The maximum audio size in a batch
         __lowerCamelCase            =    np.ones([len(a       ), 1, max_time_len, self.feature_size]       ).astype(np.floataa       )
         __lowerCamelCase            =    padded_audio_features * self.padding_value
         for i in range(len(a       )       ):
             __lowerCamelCase            =    audio_features[i]
             __lowerCamelCase            =    feature

         # return as BatchFeature
         if return_attention_mask:
             __lowerCamelCase            =    {'''audio_values''': padded_audio_features, '''audio_mask''': audio_mask}
         else:
             __lowerCamelCase            =    {'''audio_values''': padded_audio_features}

         __lowerCamelCase            =    BatchFeature(data=a    ,							tensor_type=a       )
         return encoded_inputs
