





from dataclasses import asdict, dataclass
from typing import Optional

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


_lowerCamelCase :    Any  							= logging.get_logger(__name__)

# TODO Update this
_lowerCamelCase :    Tuple  							= {
    '''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       (__lowerCamelCase      ):







							"""simple docstring"""




							UpperCAmelCase_      						=   "esm"






							def __init__( self     :					Union[str, Any],			_UpperCAmelCase     :					str=None,			_UpperCAmelCase     :					Dict=None,			_UpperCAmelCase     :					List[Any]=None,			_UpperCAmelCase     :					Optional[Any]=7_6_8,			_UpperCAmelCase     :					Optional[Any]=1_2,			_UpperCAmelCase     :					List[Any]=1_2,			_UpperCAmelCase     :					str=3_0_7_2,			_UpperCAmelCase     :					Optional[int]=0.1,			_UpperCAmelCase     :					Optional[Any]=0.1,			_UpperCAmelCase     :					Any=1_0_2_6,			_UpperCAmelCase     :					Optional[int]=0.02,			_UpperCAmelCase     :					Union[str, Any]=1E-12,			_UpperCAmelCase     :					Optional[int]="absolute",			_UpperCAmelCase     :					List[Any]=True,			_UpperCAmelCase     :					Tuple=None,			_UpperCAmelCase     :					str=False,			_UpperCAmelCase     :					int=False,			_UpperCAmelCase     :					str=None,			_UpperCAmelCase     :					Tuple=None,			**_UpperCAmelCase     :					str,			)				->			List[Any]:





														"""simple docstring"""
														super().__init__(pad_token_id=_UpperCAmelCase,			mask_token_id=_UpperCAmelCase,			**_UpperCAmelCase						)

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






							def  A_     ( self     :					Any						)				->			List[Any]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      List[str]     =						super().to_dict()
														if isinstance(self.esmfold_config,			_UpperCAmelCase						):
																					SCREAMING_SNAKE_CASE__       :      Optional[int]     =						self.esmfold_config.to_dict()
														return output


@dataclass
class 						lowerCamelCase       :







							"""simple docstring"""




							UpperCAmelCase_      						=   None
							UpperCAmelCase_      						=   True
							UpperCAmelCase_      						=   False
							UpperCAmelCase_      						=   False
							UpperCAmelCase_      						=   False
							UpperCAmelCase_      						=   0

							UpperCAmelCase_      						=   True
							UpperCAmelCase_      						=   False

							UpperCAmelCase_      						=   128
							UpperCAmelCase_      						=   None






							def  A_     ( self     :					Optional[int]						)				->			Tuple:





														"""simple docstring"""
														if self.trunk is None:
																					SCREAMING_SNAKE_CASE__       :      Dict     =						TrunkConfig()
														elif isinstance(self.trunk,			_UpperCAmelCase						):
																					SCREAMING_SNAKE_CASE__       :      str     =						TrunkConfig(**self.trunk						)






							def  A_     ( self     :					List[str]						)				->			int:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      str     =						asdict(self						)
														SCREAMING_SNAKE_CASE__       :      int     =						self.trunk.to_dict()
														return output


@dataclass
class 						lowerCamelCase       :







							"""simple docstring"""




							UpperCAmelCase_      						=   48
							UpperCAmelCase_      						=   1024
							UpperCAmelCase_      						=   128
							UpperCAmelCase_      						=   32
							UpperCAmelCase_      						=   32
							UpperCAmelCase_      						=   32
							UpperCAmelCase_      						=   0
							UpperCAmelCase_      						=   0
							UpperCAmelCase_      						=   False
							UpperCAmelCase_      						=   4
							UpperCAmelCase_      						=   128
							UpperCAmelCase_      						=   None






							def  A_     ( self     :					Dict						)				->			Dict:





														"""simple docstring"""
														if self.structure_module is None:
																					SCREAMING_SNAKE_CASE__       :      Dict     =						StructureModuleConfig()
														elif isinstance(self.structure_module,			_UpperCAmelCase						):
																					SCREAMING_SNAKE_CASE__       :      List[str]     =						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}.'''						)

														SCREAMING_SNAKE_CASE__       :      List[Any]     =						self.sequence_state_dim // self.sequence_head_width
														SCREAMING_SNAKE_CASE__       :      str     =						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  A_     ( self     :					int						)				->			str:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      str     =						asdict(self						)
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						self.structure_module.to_dict()
														return output


@dataclass
class 						lowerCamelCase       :







							"""simple docstring"""




							UpperCAmelCase_      						=   384
							UpperCAmelCase_      						=   128
							UpperCAmelCase_      						=   16
							UpperCAmelCase_      						=   128
							UpperCAmelCase_      						=   12
							UpperCAmelCase_      						=   4
							UpperCAmelCase_      						=   8
							UpperCAmelCase_      						=   0.1
							UpperCAmelCase_      						=   8
							UpperCAmelCase_      						=   1
							UpperCAmelCase_      						=   2
							UpperCAmelCase_      						=   7
							UpperCAmelCase_      						=   10
							UpperCAmelCase_      						=   1E-8
							UpperCAmelCase_      						=   1E5






							def  A_     ( self     :					List[Any]						)				->			str:





														"""simple docstring"""
														return asdict(self						)




def _a (  )	->    List[str]:

							'''simple docstring'''






							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 json
import os
from typing import Optional

import numpy as np

from ...feature_extraction_utils import BatchFeature
from ...processing_utils import ProcessorMixin
from ...utils import logging
from ...utils.hub import get_file_from_repo
from ..auto import AutoTokenizer


_lowerCamelCase :    Optional[int]  							= logging.get_logger(__name__)


class 						lowerCamelCase       (__lowerCamelCase      ):







							"""simple docstring"""




							UpperCAmelCase_      						=   "AutoTokenizer"
							UpperCAmelCase_      						=   ["tokenizer"]

							UpperCAmelCase_      						=   {
							    "semantic_prompt": 1,
							    "coarse_prompt": 2,
							    "fine_prompt": 2,
							}






							def __init__( self     :					Union[str, Any],			_UpperCAmelCase     :					Optional[int],			_UpperCAmelCase     :					Union[str, Any]=None						)				->			Union[str, Any]:





														"""simple docstring"""
														super().__init__(_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						speaker_embeddings






							@classmethod
							def  A_     ( cls     :					Any,			_UpperCAmelCase     :					List[str],			_UpperCAmelCase     :					Dict="speaker_embeddings_path.json",			**_UpperCAmelCase     :					Optional[Any]						)				->			List[Any]:





														"""simple docstring"""

														if speaker_embeddings_dict_path is not None:
																					SCREAMING_SNAKE_CASE__       :      Any     =						get_file_from_repo(
																					    _UpperCAmelCase,			_UpperCAmelCase,			subfolder=kwargs.pop("subfolder",			_UpperCAmelCase						),			cache_dir=kwargs.pop("cache_dir",			_UpperCAmelCase						),			force_download=kwargs.pop("force_download",			_UpperCAmelCase						),			proxies=kwargs.pop("proxies",			_UpperCAmelCase						),			resume_download=kwargs.pop("resume_download",			_UpperCAmelCase						),			local_files_only=kwargs.pop("local_files_only",			_UpperCAmelCase						),			use_auth_token=kwargs.pop("use_auth_token",			_UpperCAmelCase						),			revision=kwargs.pop("revision",			_UpperCAmelCase						),			)
																					if speaker_embeddings_path is None:
																												logger.warning(
																												    F'''`{os.path.join(_UpperCAmelCase,			_UpperCAmelCase						)}` does not exists
                    , no preloaded speaker embeddings will be used - Make sure to provide a correct path to the json
                    dictionnary if wanted, otherwise set `speaker_embeddings_dict_path=None`.'''						)
																												SCREAMING_SNAKE_CASE__       :      Dict     =						None
																					else:
																												with open(_UpperCAmelCase						) as speaker_embeddings_json:
																																			SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						json.load(_UpperCAmelCase						)
														else:
																					SCREAMING_SNAKE_CASE__       :      List[Any]     =						None

														SCREAMING_SNAKE_CASE__       :      List[str]     =						AutoTokenizer.from_pretrained(_UpperCAmelCase,			**_UpperCAmelCase						)

														return cls(tokenizer=_UpperCAmelCase,			speaker_embeddings=_UpperCAmelCase						)






							def  A_     ( self     :					str,			_UpperCAmelCase     :					Optional[int],			_UpperCAmelCase     :					List[str]="speaker_embeddings_path.json",			_UpperCAmelCase     :					Optional[Any]="speaker_embeddings",			_UpperCAmelCase     :					bool = False,			**_UpperCAmelCase     :					List[str],			)				->			Union[str, Any]:





														"""simple docstring"""
														if self.speaker_embeddings is not None:
																					os.makedirs(os.path.join(_UpperCAmelCase,			_UpperCAmelCase,			"v2"						),			exist_ok=_UpperCAmelCase						)

																					SCREAMING_SNAKE_CASE__       :      Optional[int]     =						{}

																					SCREAMING_SNAKE_CASE__       :      Any     =						save_directory

																					for prompt_key in self.speaker_embeddings:
																												if prompt_key != "repo_or_path":
																																			SCREAMING_SNAKE_CASE__       :      List[Any]     =						self._load_voice_preset(_UpperCAmelCase						)

																																			SCREAMING_SNAKE_CASE__       :      int     =						{}
																																			for key in self.speaker_embeddings[prompt_key]:
																																										np.save(
																																										    os.path.join(
																																										        embeddings_dict["repo_or_path"],			_UpperCAmelCase,			F'''{prompt_key}_{key}'''						),			voice_preset[key],			allow_pickle=_UpperCAmelCase,			)
																																										SCREAMING_SNAKE_CASE__       :      List[str]     =						os.path.join(_UpperCAmelCase,			F'''{prompt_key}_{key}.npy'''						)

																																			SCREAMING_SNAKE_CASE__       :      Optional[int]     =						tmp_dict

																					with open(os.path.join(_UpperCAmelCase,			_UpperCAmelCase						),			"w"						) as fp:
																												json.dump(_UpperCAmelCase,			_UpperCAmelCase						)

														super().save_pretrained(_UpperCAmelCase,			_UpperCAmelCase,			**_UpperCAmelCase						)






							def  A_     ( self     :					List[Any],			_UpperCAmelCase     :					str = None,			**_UpperCAmelCase     :					List[Any]						)				->			List[str]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      str     =						self.speaker_embeddings[voice_preset]

														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						{}
														for key in ["semantic_prompt", "coarse_prompt", "fine_prompt"]:
																					if key not in voice_preset_paths:
																												raise ValueError(
																												    F'''Voice preset unrecognized, missing {key} as a key in self.speaker_embeddings[{voice_preset}].'''						)

																					SCREAMING_SNAKE_CASE__       :      List[Any]     =						get_file_from_repo(
																					    self.speaker_embeddings.get("repo_or_path",			"/"						),			voice_preset_paths[key],			subfolder=kwargs.pop("subfolder",			_UpperCAmelCase						),			cache_dir=kwargs.pop("cache_dir",			_UpperCAmelCase						),			force_download=kwargs.pop("force_download",			_UpperCAmelCase						),			proxies=kwargs.pop("proxies",			_UpperCAmelCase						),			resume_download=kwargs.pop("resume_download",			_UpperCAmelCase						),			local_files_only=kwargs.pop("local_files_only",			_UpperCAmelCase						),			use_auth_token=kwargs.pop("use_auth_token",			_UpperCAmelCase						),			revision=kwargs.pop("revision",			_UpperCAmelCase						),			)
																					if path is None:
																												raise ValueError(
																												    F'''`{os.path.join(self.speaker_embeddings.get("repo_or_path",			"/"						),			voice_preset_paths[key]						)}` does not exists
                    , no preloaded voice preset will be used - Make sure to provide correct paths to the {voice_preset}
                    embeddings.'''						)

																					SCREAMING_SNAKE_CASE__       :      int     =						np.load(_UpperCAmelCase						)

														return voice_preset_dict






							def  A_     ( self     :					int,			_UpperCAmelCase     :					Optional[dict] = None						)				->			Optional[int]:





														"""simple docstring"""
														for key in ["semantic_prompt", "coarse_prompt", "fine_prompt"]:
																					if key not in voice_preset:
																												raise ValueError(F'''Voice preset unrecognized, missing {key} as a key.'''						)

																					if not isinstance(voice_preset[key],			np.ndarray						):
																												raise ValueError(F'''{key} voice preset must be a {str(self.preset_shape[key]						)}D ndarray.'''						)

																					if len(voice_preset[key].shape						) != self.preset_shape[key]:
																												raise ValueError(F'''{key} voice preset must be a {str(self.preset_shape[key]						)}D ndarray.'''						)






							def __call__( self     :					List[Any],			_UpperCAmelCase     :					Optional[Any]=None,			_UpperCAmelCase     :					Union[str, Any]=None,			_UpperCAmelCase     :					Optional[int]="pt",			_UpperCAmelCase     :					List[str]=2_5_6,			_UpperCAmelCase     :					int=False,			_UpperCAmelCase     :					Optional[int]=True,			_UpperCAmelCase     :					Any=False,			**_UpperCAmelCase     :					List[str],			)				->			List[Any]:





														"""simple docstring"""
														if voice_preset is not None and not isinstance(_UpperCAmelCase,			_UpperCAmelCase						):
																					if (
																					    isinstance(_UpperCAmelCase,			_UpperCAmelCase						)
																					    and self.speaker_embeddings is not None
																					    and voice_preset in self.speaker_embeddings
																					):
																												SCREAMING_SNAKE_CASE__       :      List[str]     =						self._load_voice_preset(_UpperCAmelCase						)

																					else:
																												if isinstance(_UpperCAmelCase,			_UpperCAmelCase						) and not voice_preset.endswith(".npz"						):
																																			SCREAMING_SNAKE_CASE__       :      Optional[int]     =						voice_preset + ".npz"

																												SCREAMING_SNAKE_CASE__       :      List[str]     =						np.load(_UpperCAmelCase						)

														if voice_preset is not None:
																					self._validate_voice_preset_dict(_UpperCAmelCase,			**_UpperCAmelCase						)
																					SCREAMING_SNAKE_CASE__       :      Any     =						BatchFeature(data=_UpperCAmelCase,			tensor_type=_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      Dict     =						self.tokenizer(
														    _UpperCAmelCase,			return_tensors=_UpperCAmelCase,			padding="max_length",			max_length=_UpperCAmelCase,			return_attention_mask=_UpperCAmelCase,			return_token_type_ids=_UpperCAmelCase,			add_special_tokens=_UpperCAmelCase,			**_UpperCAmelCase,			)

														if voice_preset is not None:
																					SCREAMING_SNAKE_CASE__       :      str     =						voice_preset

														return encoded_text






