



from dataclasses import asdict, dataclass
from typing import Optional

from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCamelCase__ =						logging.get_logger(__name__)

# TODO Update this
UpperCamelCase__ =						{
    "facebook/esm-1b": "https://huggingface.co/facebook/esm-1b/resolve/main/config.json",
    # See all ESM models at https://huggingface.co/models?filter=esm
}
class 			_lowerCAmelCase   ( _UpperCAmelCase			):






			"""simple docstring"""





			lowercase__      :  Union[str, Any]      =			"""esm"""





			def __init__(				self	:       Any			,     lowercase	:       Optional[Any]=None			,     lowercase	:       Optional[int]=None			,     lowercase	:       List[Any]=None			,     lowercase	:       Optional[int]=768			,     lowercase	:       str=12			,     lowercase	:       Union[str, Any]=12			,     lowercase	:       Dict=3_072			,     lowercase	:       Optional[int]=0.1			,     lowercase	:       str=0.1			,     lowercase	:       Dict=1_026			,     lowercase	:       Tuple=0.02			,     lowercase	:       str=1E-1_2			,     lowercase	:       Dict="absolute"			,     lowercase	:       Optional[Any]=True			,     lowercase	:       int=None			,     lowercase	:       int=False			,     lowercase	:       List[str]=False			,     lowercase	:       Tuple=None			,     lowercase	:       Tuple=None			,     **lowercase	:       Union[str, Any]			,     )						->					int:
									"""simple docstring"""







									super().__init__(pad_token_id=lowercase			,     mask_token_id=lowercase			,     **lowercase   )

									__lowercase					=				vocab_size
									__lowercase					=				hidden_size
									__lowercase					=				num_hidden_layers
									__lowercase					=				num_attention_heads
									__lowercase					=				intermediate_size
									__lowercase					=				hidden_dropout_prob
									__lowercase					=				attention_probs_dropout_prob
									__lowercase					=				max_position_embeddings
									__lowercase					=				initializer_range
									__lowercase					=				layer_norm_eps
									__lowercase					=				position_embedding_type
									__lowercase					=				use_cache
									__lowercase					=				emb_layer_norm_before
									__lowercase					=				token_dropout
									__lowercase					=				is_folding_model
									if is_folding_model:
															if esmfold_config is None:
																					logger.info("""No esmfold_config supplied for folding model, using default values."""   )
																					__lowercase					=				EsmFoldConfig()
															elif isinstance(lowercase			,     lowercase   ):
																					__lowercase					=				EsmFoldConfig(**lowercase   )
															__lowercase					=				esmfold_config
															if vocab_list is None:
																					logger.warning("""No vocab_list supplied for folding model, assuming the ESM-2 vocabulary!"""   )
																					__lowercase					=				get_default_vocab_list()
															else:
																					__lowercase					=				vocab_list
									else:
															__lowercase					=				None
															__lowercase					=				None
									if self.esmfold_config is not None and getattr(self.esmfold_config			,     """use_esm_attn_map"""			,     lowercase   ):
															raise ValueError("""The HuggingFace port of ESMFold does not support use_esm_attn_map at this time!"""   )







			def 		snake_case__			(				self	:       Optional[Any]   )						->					Tuple:
									"""simple docstring"""







									__lowercase					=				super().to_dict()
									if isinstance(self.esmfold_config			,     lowercase   ):
															__lowercase					=				self.esmfold_config.to_dict()
									return output





@dataclass
class 			_lowerCAmelCase   :






			"""simple docstring"""





			lowercase__      :  str      =			None
			lowercase__      :  bool      =			True
			lowercase__      :  bool      =			False
			lowercase__      :  bool      =			False
			lowercase__      :  bool      =			False
			lowercase__      :  float      =			0

			lowercase__      :  bool      =			True
			lowercase__      :  bool      =			False

			lowercase__      :  int      =			128
			lowercase__      :  "TrunkConfig"      =			None





			def 		snake_case__			(				self	:       List[str]   )						->					Any:
									"""simple docstring"""







									if self.trunk is None:
															__lowercase					=				TrunkConfig()
									elif isinstance(self.trunk			,     lowercase   ):
															__lowercase					=				TrunkConfig(**self.trunk   )







			def 		snake_case__			(				self	:       Dict   )						->					Any:
									"""simple docstring"""







									__lowercase					=				asdict(self   )
									__lowercase					=				self.trunk.to_dict()
									return output





@dataclass
class 			_lowerCAmelCase   :






			"""simple docstring"""





			lowercase__      :  int      =			48
			lowercase__      :  int      =			1_024
			lowercase__      :  int      =			128
			lowercase__      :  int      =			32
			lowercase__      :  int      =			32
			lowercase__      :  int      =			32
			lowercase__      :  float      =			0
			lowercase__      :  float      =			0
			lowercase__      :  bool      =			False
			lowercase__      :  int      =			4
			lowercase__      :  Optional[int]      =			128
			lowercase__      :  "StructureModuleConfig"      =			None





			def 		snake_case__			(				self	:       Tuple   )						->					str:
									"""simple docstring"""







									if self.structure_module is None:
															__lowercase					=				StructureModuleConfig()
									elif isinstance(self.structure_module			,     lowercase   ):
															__lowercase					=				StructureModuleConfig(**self.structure_module   )

									if self.max_recycles <= 0:
															raise ValueError(F"`max_recycles` should be positive, got {self.max_recycles}."   )
									if self.sequence_state_dim % self.sequence_state_dim != 0:
															raise ValueError(
															    """`sequence_state_dim` should be a round multiple of `sequence_state_dim`, got"""
															    F" {self.sequence_state_dim} and {self.sequence_state_dim}."   )
									if self.pairwise_state_dim % self.pairwise_state_dim != 0:
															raise ValueError(
															    """`pairwise_state_dim` should be a round multiple of `pairwise_state_dim`, got"""
															    F" {self.pairwise_state_dim} and {self.pairwise_state_dim}."   )

									__lowercase					=				self.sequence_state_dim // self.sequence_head_width
									__lowercase					=				self.pairwise_state_dim // self.pairwise_head_width

									if self.sequence_state_dim != sequence_num_heads * self.sequence_head_width:
															raise ValueError(
															    """`sequence_state_dim` should be equal to `sequence_num_heads * sequence_head_width, got"""
															    F" {self.sequence_state_dim} != {sequence_num_heads} * {self.sequence_head_width}."   )
									if self.pairwise_state_dim != pairwise_num_heads * self.pairwise_head_width:
															raise ValueError(
															    """`pairwise_state_dim` should be equal to `pairwise_num_heads * pairwise_head_width, got"""
															    F" {self.pairwise_state_dim} != {pairwise_num_heads} * {self.pairwise_head_width}."   )
									if self.pairwise_state_dim % 2 != 0:
															raise ValueError(F"`pairwise_state_dim` should be even, got {self.pairwise_state_dim}."   )

									if self.dropout >= 0.4:
															raise ValueError(F"`dropout` should not be greater than 0.4, got {self.dropout}."   )







			def 		snake_case__			(				self	:       List[str]   )						->					Union[str, Any]:
									"""simple docstring"""







									__lowercase					=				asdict(self   )
									__lowercase					=				self.structure_module.to_dict()
									return output





@dataclass
class 			_lowerCAmelCase   :






			"""simple docstring"""





			lowercase__      :  int      =			384
			lowercase__      :  int      =			128
			lowercase__      :  int      =			16
			lowercase__      :  int      =			128
			lowercase__      :  int      =			12
			lowercase__      :  int      =			4
			lowercase__      :  int      =			8
			lowercase__      :  float      =			0.1
			lowercase__      :  int      =			8
			lowercase__      :  int      =			1
			lowercase__      :  int      =			2
			lowercase__      :  int      =			7
			lowercase__      :  int      =			10
			lowercase__      :  float      =			1E-8
			lowercase__      :  float      =			1E5





			def 		snake_case__			(				self	:       Tuple   )						->					List[str]:
									"""simple docstring"""







									return asdict(self   )






def  UpperCAmelCase__    (    ) ->      List[Any]:
						return (
						    "<cls>",
						    "<pad>",
						    "<eos>",
						    "<unk>",
						    "L",
						    "A",
						    "G",
						    "V",
						    "S",
						    "E",
						    "R",
						    "T",
						    "I",
						    "D",
						    "P",
						    "K",
						    "Q",
						    "N",
						    "F",
						    "Y",
						    "M",
						    "H",
						    "W",
						    "C",
						    "X",
						    "B",
						    "U",
						    "Z",
						    "O",
						    ".",
						    "-",
						    "<null_1>",
						    "<mask>",
						)










import itertools
import random
import unittest

import numpy as np

from transformers import is_speech_available
from transformers.testing_utils import require_torch, require_torchaudio

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
			from transformers import SpeechaTextFeatureExtractor

UpperCamelCase__ =						random.Random()






def  UpperCAmelCase__    (    lowercase__						,       lowercase__=1.0						,       lowercase__=None						,       lowercase__=None   ) ->      str:
						if rng is None:
												__lowercase					=				global_rng

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

						return values
@require_torch
@require_torchaudio
class 			_lowerCAmelCase   ( unittest.TestCase			):






			"""simple docstring"""





			def __init__(				self	:       List[Any]			,     lowercase	:       Tuple			,     lowercase	:       Union[str, Any]=7			,     lowercase	:       List[Any]=400			,     lowercase	:       Any=2_000			,     lowercase	:       Optional[int]=24			,     lowercase	:       Any=24			,     lowercase	:       List[str]=0.0			,     lowercase	:       Dict=16_000			,     lowercase	:       Union[str, Any]=True			,     lowercase	:       Dict=True			,     )						->					Any:
									"""simple docstring"""







									__lowercase					=				parent
									__lowercase					=				batch_size
									__lowercase					=				min_seq_length
									__lowercase					=				max_seq_length
									__lowercase					=				(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
									__lowercase					=				feature_size
									__lowercase					=				num_mel_bins
									__lowercase					=				padding_value
									__lowercase					=				sampling_rate
									__lowercase					=				return_attention_mask
									__lowercase					=				do_normalize





			def 		snake_case__			(				self	:       Optional[int]   )						->					Any:
									"""simple docstring"""







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







			def 		snake_case__			(				self	:       List[str]			,     lowercase	:       Tuple=False			,     lowercase	:       int=False   )						->					Optional[Any]:
									"""simple docstring"""







									def _flatten(lowercase	:       Optional[Any]   ):
															return list(itertools.chain(*lowercase   )   )

									if equal_length:
															__lowercase					=				[floats_list((self.max_seq_length, self.feature_size)   ) for _ in range(self.batch_size   )]
									else:
															# make sure that inputs increase in size
															__lowercase					=				[
															    floats_list((x, self.feature_size)   )
															    for x in range(self.min_seq_length			,     self.max_seq_length			,     self.seq_length_diff   )
															]
									if numpify:
															__lowercase					=				[np.asarray(lowercase   ) for x in speech_inputs]
									return speech_inputs





@require_torch
@require_torchaudio
class 			_lowerCAmelCase   ( _UpperCAmelCase    ,					unittest.TestCase			):






			"""simple docstring"""





			lowercase__      :  int      =			SpeechaTextFeatureExtractor if is_speech_available() else None





			def 		snake_case__			(				self	:       Dict   )						->					Optional[int]:
									"""simple docstring"""







									__lowercase					=				SpeechaTextFeatureExtractionTester(self   )





			def 		snake_case__			(				self	:       Tuple			,     lowercase	:       Optional[Any]   )						->					Union[str, Any]:
									"""simple docstring"""







									self.assertTrue(np.all(np.mean(lowercase			,     axis=0   ) < 1E-3   )   )
									self.assertTrue(np.all(np.abs(np.var(lowercase			,     axis=0   ) - 1   ) < 1E-3   )   )





			def 		snake_case__			(				self	:       List[Any]   )						->					str:
									"""simple docstring"""







									__lowercase					=				self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
									# create three inputs of length 800, 1000, and 1200
									__lowercase					=				[floats_list((1, x)   )[0] for x in range(800			,     1_400			,     200   )]
									__lowercase					=				[np.asarray(lowercase   ) for speech_input in speech_inputs]

									# Test feature size
									__lowercase					=				feature_extractor(lowercase			,     padding=lowercase			,     return_tensors="""np"""   ).input_features
									self.assertTrue(input_features.ndim == 3   )
									self.assertTrue(input_features.shape[-1] == feature_extractor.feature_size   )

									# Test not batched input
									__lowercase					=				feature_extractor(speech_inputs[0]			,     return_tensors="""np"""   ).input_features
									__lowercase					=				feature_extractor(np_speech_inputs[0]			,     return_tensors="""np"""   ).input_features
									self.assertTrue(np.allclose(lowercase			,     lowercase			,     atol=1E-3   )   )

									# Test batched
									__lowercase					=				feature_extractor(lowercase			,     return_tensors="""np"""   ).input_features
									__lowercase					=				feature_extractor(lowercase			,     return_tensors="""np"""   ).input_features
									for enc_seq_a, enc_seq_a in zip(lowercase			,     lowercase   ):
															self.assertTrue(np.allclose(lowercase			,     lowercase			,     atol=1E-3   )   )

									# Test 2-D numpy arrays are batched.
									__lowercase					=				[floats_list((1, x)   )[0] for x in (800, 800, 800)]
									__lowercase					=				np.asarray(lowercase   )
									__lowercase					=				feature_extractor(lowercase			,     return_tensors="""np"""   ).input_features
									__lowercase					=				feature_extractor(lowercase			,     return_tensors="""np"""   ).input_features
									for enc_seq_a, enc_seq_a in zip(lowercase			,     lowercase   ):
															self.assertTrue(np.allclose(lowercase			,     lowercase			,     atol=1E-3   )   )





			def 		snake_case__			(				self	:       List[str]   )						->					Tuple:
									"""simple docstring"""







									__lowercase					=				self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
									__lowercase					=				[floats_list((1, x)   )[0] for x in range(800			,     1_400			,     200   )]

									__lowercase					=				["""longest""", """max_length""", """do_not_pad"""]
									__lowercase					=				[None, 16, None]
									for max_length, padding in zip(lowercase			,     lowercase   ):
															__lowercase					=				feature_extractor(
															    lowercase			,     padding=lowercase			,     max_length=lowercase			,     return_attention_mask=lowercase   )
															__lowercase					=				inputs.input_features
															__lowercase					=				inputs.attention_mask
															__lowercase					=				[np.sum(lowercase   ) for x in attention_mask]

															self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]   )
															self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]   )
															self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]   )





			def 		snake_case__			(				self	:       Any   )						->					Dict:
									"""simple docstring"""







									__lowercase					=				self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
									__lowercase					=				[floats_list((1, x)   )[0] for x in range(800			,     1_400			,     200   )]

									__lowercase					=				["""longest""", """max_length""", """do_not_pad"""]
									__lowercase					=				[None, 16, None]
									for max_length, padding in zip(lowercase			,     lowercase   ):
															__lowercase					=				feature_extractor(
															    lowercase			,     max_length=lowercase			,     padding=lowercase			,     return_tensors="""np"""			,     return_attention_mask=lowercase   )
															__lowercase					=				inputs.input_features
															__lowercase					=				inputs.attention_mask
															__lowercase					=				[np.sum(lowercase   ) for x in attention_mask]

															self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]   )
															self.assertTrue(input_features[0][fbank_feat_lengths[0] :].sum() < 1E-6   )
															self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]   )
															self.assertTrue(input_features[0][fbank_feat_lengths[1] :].sum() < 1E-6   )
															self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]   )





			def 		snake_case__			(				self	:       str   )						->					List[Any]:
									"""simple docstring"""







									__lowercase					=				self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
									__lowercase					=				[floats_list((1, x)   )[0] for x in range(800			,     1_400			,     200   )]
									__lowercase					=				feature_extractor(
									    lowercase			,     padding="""max_length"""			,     max_length=4			,     truncation=lowercase			,     return_tensors="""np"""			,     return_attention_mask=lowercase			,     )
									__lowercase					=				inputs.input_features
									__lowercase					=				inputs.attention_mask
									__lowercase					=				np.sum(attention_mask == 1			,     axis=1   )

									self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]   )
									self._check_zero_mean_unit_variance(input_features[1]   )
									self._check_zero_mean_unit_variance(input_features[2]   )





			def 		snake_case__			(				self	:       Tuple   )						->					List[str]:
									"""simple docstring"""







									__lowercase					=				self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
									__lowercase					=				[floats_list((1, x)   )[0] for x in range(800			,     1_400			,     200   )]
									__lowercase					=				feature_extractor(
									    lowercase			,     padding="""longest"""			,     max_length=4			,     truncation=lowercase			,     return_tensors="""np"""			,     return_attention_mask=lowercase			,     )
									__lowercase					=				inputs.input_features
									__lowercase					=				inputs.attention_mask
									__lowercase					=				np.sum(attention_mask == 1			,     axis=1   )

									self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]   )
									self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]   )
									self._check_zero_mean_unit_variance(input_features[2]   )

									# make sure that if max_length < longest -> then pad to max_length
									self.assertEqual(input_features.shape			,     (3, 4, 24)   )

									__lowercase					=				[floats_list((1, x)   )[0] for x in range(800			,     1_400			,     200   )]
									__lowercase					=				feature_extractor(
									    lowercase			,     padding="""longest"""			,     max_length=16			,     truncation=lowercase			,     return_tensors="""np"""			,     return_attention_mask=lowercase			,     )
									__lowercase					=				inputs.input_features
									__lowercase					=				inputs.attention_mask
									__lowercase					=				np.sum(attention_mask == 1			,     axis=1   )

									self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]   )
									self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]   )
									self._check_zero_mean_unit_variance(input_features[2]   )

									# make sure that if max_length < longest -> then pad to max_length
									self.assertEqual(input_features.shape			,     (3, 6, 24)   )





			def 		snake_case__			(				self	:       Optional[Any]   )						->					Tuple:
									"""simple docstring"""







									import torch

									__lowercase					=				self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
									__lowercase					=				np.random.rand(100			,     32   ).astype(np.floataa   )
									__lowercase					=				np_speech_inputs.tolist()

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





			def 		snake_case__			(				self	:       Optional[int]			,     lowercase	:       Union[str, Any]   )						->					int:
									"""simple docstring"""







									from datasets import load_dataset

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

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







			def 		snake_case__			(				self	:       str   )						->					Any:
									"""simple docstring"""







									__lowercase					=				np.array([
									    -1.5745, -1.7713, -1.7020, -1.6069, -1.2250, -1.1105, -0.9072, -0.8241,
									    -1.2310, -0.8098, -0.3320, -0.4101, -0.7985, -0.4996, -0.8213, -0.9128,
									    -1.0420, -1.1286, -1.0440, -0.7999, -0.8405, -1.2275, -1.5443, -1.4625,
									]   )
									# fmt: on

									__lowercase					=				self._load_datasamples(1   )
									__lowercase					=				self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
									__lowercase					=				feature_extractor(lowercase			,     return_tensors="""pt"""   ).input_features
									self.assertEquals(input_features.shape			,     (1, 584, 24)   )
									self.assertTrue(np.allclose(input_features[0, 0, :30]			,     lowercase			,     atol=1E-4   )   )





