'''simple docstring'''






from typing import Dict, List, Optional, Union

import numpy as np

from .feature_extraction_utils import BatchFeature, FeatureExtractionMixin
from .utils import PaddingStrategy, TensorType, is_tf_tensor, is_torch_tensor, logging, to_numpy


__a										=					logging.get_logger(__name__)





class     A__							(       UpperCamelCase   ):





   """simple docstring"""





   def __init__( self     :		int						,     lowerCAmelCase__     :		List[str]						,     lowerCAmelCase__     :		Dict						,     lowerCAmelCase__     :		Tuple						,     **lowerCAmelCase__     :		Optional[int]			)       ->    Optional[Any]:







      """simple docstring"""




      _UpperCAmelCase							:	Union[str, Any]							=   feature_size
      _UpperCAmelCase							:	int							=   sampling_rate
      _UpperCAmelCase							:	int							=   padding_value

      _UpperCAmelCase							:	int							=   kwargs.pop("padding_side"						,     "right"			)
      _UpperCAmelCase							:	List[Any]							=   kwargs.pop("return_attention_mask"						,     lowerCAmelCase__			)

      super().__init__(**lowerCAmelCase__			)




   def 			_lowerCAmelCase		( self     :		List[Any]						,     lowerCAmelCase__     :		List[Any]						,     lowerCAmelCase__     :		Tuple = True						,     lowerCAmelCase__     :		Dict = None						,     lowerCAmelCase__     :		List[Any] = False						,     lowerCAmelCase__     :		str = None						,     lowerCAmelCase__     :		Tuple = None						,     lowerCAmelCase__     :		Optional[Any] = None						,     )       ->    BatchFeature:







      """simple docstring"""




      if isinstance(lowerCAmelCase__						,     (list, tuple)			) and isinstance(processed_features[0]						,     (dict, BatchFeature)			):
         _UpperCAmelCase							:	List[str]							=   {
             key: [example[key] for example in processed_features] for key in processed_features[0].keys()
         }

      # The model's main input name, usually `input_values`, has be passed for padding
      if self.model_input_names[0] not in processed_features:
         raise ValueError(
             "You should supply an instance of `transformers.BatchFeature` or list of `transformers.BatchFeature`"
             F""" to this method that includes {self.model_input_names[0]}, but you provided"""
             F""" {list(processed_features.keys()			)}"""			)

      _UpperCAmelCase							:	Union[str, Any]							=   processed_features[self.model_input_names[0]]
      _UpperCAmelCase							:	Union[str, Any]							=   (
          return_attention_mask if return_attention_mask is not None else self.return_attention_mask
      )

      if len(lowerCAmelCase__			) == 0:
         if return_attention_mask:
            _UpperCAmelCase							:	Dict							=   []
         return processed_features

      # If we have PyTorch/TF tensors or lists as inputs, we cast them as Numpy arrays
      # and rebuild them afterwards if no return_tensors is specified
      # Note that we lose the specific device the tensor may be on for PyTorch

      _UpperCAmelCase							:	Any							=   required_input[0]
      if isinstance(lowerCAmelCase__						,     (list, tuple)			):
         # first_element might be an empty list/tuple in some edge cases so we grab the first non empty element.
         _UpperCAmelCase							:	List[str]							=   0
         while len(required_input[index]			) == 0:
            index += 1
         if index < len(lowerCAmelCase__			):
            _UpperCAmelCase							:	List[Any]							=   required_input[index][0]

      if return_tensors is None:
         if is_tf_tensor(lowerCAmelCase__			):
            _UpperCAmelCase							:	int							=   "tf"
         elif is_torch_tensor(lowerCAmelCase__			):
            _UpperCAmelCase							:	List[Any]							=   "pt"
         elif isinstance(lowerCAmelCase__						,     (int, float, list, tuple, np.ndarray)			):
            _UpperCAmelCase							:	Union[str, Any]							=   "np"
         else:
            raise ValueError(
                F"""type of {first_element} unknown: {type(lowerCAmelCase__			)}. """
                "Should be one of a python, numpy, pytorch or tensorflow object."			)

      for key, value in processed_features.items():
         if isinstance(value[0]						,     (int, float)			):
            _UpperCAmelCase							:	List[Any]							=   to_numpy(lowerCAmelCase__			)
         else:
            _UpperCAmelCase							:	Dict							=   [to_numpy(lowerCAmelCase__			) for v in value]

            # Convert padding_strategy in PaddingStrategy
      _UpperCAmelCase							:	List[str]							=   self._get_padding_strategies(padding=lowerCAmelCase__						,     max_length=lowerCAmelCase__			)

      _UpperCAmelCase							:	str							=   processed_features[self.model_input_names[0]]

      _UpperCAmelCase							:	Dict							=   len(lowerCAmelCase__			)
      if not all(len(lowerCAmelCase__			) == batch_size for v in processed_features.values()			):
         raise ValueError("Some items in the output dictionary have a different batch size than others."			)

      _UpperCAmelCase							:	List[str]							=   []
      for i in range(lowerCAmelCase__			):
         _UpperCAmelCase							:	Tuple							=   {k: v[i] for k, v in processed_features.items()}
         # truncation
         _UpperCAmelCase							:	List[Any]							=   self._truncate(
             lowerCAmelCase__						,     max_length=lowerCAmelCase__						,     pad_to_multiple_of=lowerCAmelCase__						,     truncation=lowerCAmelCase__						,     )
         truncated_inputs.append(lowerCAmelCase__			)

      if padding_strategy == PaddingStrategy.LONGEST:
         # make sure that `max_length` cannot be longer than the longest truncated length
         _UpperCAmelCase							:	List[str]							=   max(len(input_slice[self.model_input_names[0]]			) for input_slice in truncated_inputs			)
         _UpperCAmelCase							:	Any							=   PaddingStrategy.MAX_LENGTH

      _UpperCAmelCase							:	Union[str, Any]							=   {}
      for i in range(lowerCAmelCase__			):
         # padding
         _UpperCAmelCase							:	int							=   self._pad(
             truncated_inputs[i]						,     max_length=lowerCAmelCase__						,     padding_strategy=lowerCAmelCase__						,     pad_to_multiple_of=lowerCAmelCase__						,     return_attention_mask=lowerCAmelCase__						,     )

         for key, value in outputs.items():
            if key not in batch_outputs:
               _UpperCAmelCase							:	List[Any]							=   []
            if value.dtype is np.dtype(np.floataa			):
               _UpperCAmelCase							:	Optional[Any]							=   value.astype(np.floataa			)
            batch_outputs[key].append(lowerCAmelCase__			)

      return BatchFeature(lowerCAmelCase__						,     tensor_type=lowerCAmelCase__			)




   def 			_lowerCAmelCase		( self     :		int						,     lowerCAmelCase__     :		Optional[Any]						,     lowerCAmelCase__     :		Tuple = None						,     lowerCAmelCase__     :		Optional[Any] = PaddingStrategy.DO_NOT_PAD						,     lowerCAmelCase__     :		int = None						,     lowerCAmelCase__     :		Union[str, Any] = None						,     )       ->    dict:







      """simple docstring"""




      _UpperCAmelCase							:	int							=   processed_features[self.model_input_names[0]]

      if padding_strategy == PaddingStrategy.LONGEST:
         _UpperCAmelCase							:	Optional[int]							=   len(lowerCAmelCase__			)

      if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
         _UpperCAmelCase							:	List[Any]							=   ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of

      _UpperCAmelCase							:	List[str]							=   padding_strategy != PaddingStrategy.DO_NOT_PAD and len(lowerCAmelCase__			) < max_length

      if return_attention_mask and "attention_mask" not in processed_features:
         _UpperCAmelCase							:	Dict							=   np.ones(len(lowerCAmelCase__			)						,     dtype=np.intaa			)

      if needs_to_be_padded:
         _UpperCAmelCase							:	List[Any]							=   max_length - len(lowerCAmelCase__			)
         if self.padding_side == "right":
            if return_attention_mask:
               _UpperCAmelCase							:	Tuple							=   np.pad(
                   processed_features["attention_mask"]						,     (0, difference)			)
            _UpperCAmelCase							:	Dict							=   ((0, difference), (0, 0)) if self.feature_size > 1 else (0, difference)
            _UpperCAmelCase							:	Tuple							=   np.pad(
                lowerCAmelCase__						,     lowerCAmelCase__						,     "constant"						,     constant_values=self.padding_value			)
         elif self.padding_side == "left":
            if return_attention_mask:
               _UpperCAmelCase							:	int							=   np.pad(
                   processed_features["attention_mask"]						,     (difference, 0)			)
            _UpperCAmelCase							:	Any							=   ((difference, 0), (0, 0)) if self.feature_size > 1 else (difference, 0)
            _UpperCAmelCase							:	str							=   np.pad(
                lowerCAmelCase__						,     lowerCAmelCase__						,     "constant"						,     constant_values=self.padding_value			)
         else:
            raise ValueError("Invalid padding strategy:" + str(self.padding_side			)			)

      return processed_features




   def 			_lowerCAmelCase		( self     :		Any						,     lowerCAmelCase__     :		List[str]						,     lowerCAmelCase__     :		int = None						,     lowerCAmelCase__     :		str = None						,     lowerCAmelCase__     :		Dict = None						,     )       ->    List[str]:







      """simple docstring"""




      if not truncation:
         return processed_features
      elif truncation and max_length is None:
         raise ValueError("When setting ``truncation=True``, make sure that ``max_length`` is defined."			)

      _UpperCAmelCase							:	Optional[int]							=   processed_features[self.model_input_names[0]]

      # find `max_length` that fits `pad_to_multiple_of`
      if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
         _UpperCAmelCase							:	List[str]							=   ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of

      _UpperCAmelCase							:	str							=   len(lowerCAmelCase__			) > max_length

      if needs_to_be_truncated:
         _UpperCAmelCase							:	Tuple							=   processed_features[self.model_input_names[0]][:max_length]
         if "attention_mask" in processed_features:
            _UpperCAmelCase							:	Any							=   processed_features["attention_mask"][:max_length]

      return processed_features




   def 			_lowerCAmelCase		( self     :		List[str]						,     lowerCAmelCase__     :		Dict=False						,     lowerCAmelCase__     :		Optional[int]=None			)       ->    Dict:







      """simple docstring"""




      if padding is not False:
         if padding is True:
            _UpperCAmelCase							:	Dict							=   PaddingStrategy.LONGEST  # Default to pad to the longest sequence in the batch
         elif not isinstance(lowerCAmelCase__						,     lowerCAmelCase__			):
            _UpperCAmelCase							:	str							=   PaddingStrategy(lowerCAmelCase__			)
         elif isinstance(lowerCAmelCase__						,     lowerCAmelCase__			):
            _UpperCAmelCase							:	List[Any]							=   padding
      else:
         _UpperCAmelCase							:	Union[str, Any]							=   PaddingStrategy.DO_NOT_PAD

      # Set max length if needed
      if max_length is None:
         if padding_strategy == PaddingStrategy.MAX_LENGTH:
            raise ValueError(
                F"""When setting ``padding={PaddingStrategy.MAX_LENGTH}``, make sure that max_length is defined"""			)

            # Test if we have a padding value
      if padding_strategy != PaddingStrategy.DO_NOT_PAD and (self.padding_value is None):
         raise ValueError(
             "Asking to pad but the feature_extractor does not have a padding value. Please select a value to use"
             " as `padding_value`. For example: `feature_extractor.padding_value = 0.0`."			)

      return padding_strategy



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__   (       lowerCAmelCase):
     SCREAMING_SNAKE_CASE__      			=			'''esm'''





     def __init__(self		,     UpperCAmelCase=None		,     UpperCAmelCase=None		,     UpperCAmelCase=None		,     UpperCAmelCase=7_6_8		,     UpperCAmelCase=1_2		,     UpperCAmelCase=1_2		,     UpperCAmelCase=3_0_7_2		,     UpperCAmelCase=0.1		,     UpperCAmelCase=0.1		,     UpperCAmelCase=1_0_2_6		,     UpperCAmelCase=0.02		,     UpperCAmelCase=1e-12		,     UpperCAmelCase="absolute"		,     UpperCAmelCase=True		,     UpperCAmelCase=None		,     UpperCAmelCase=False		,     UpperCAmelCase=False		,     UpperCAmelCase=None		,     UpperCAmelCase=None		,     **UpperCAmelCase		,     )     ->   Tuple:
      super().__init__(pad_token_id=UpperCAmelCase		,     mask_token_id=UpperCAmelCase		,     **UpperCAmelCase       )

      _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(UpperCAmelCase		,     UpperCAmelCase       ):
        _lowercase							      =EsmFoldConfig(**UpperCAmelCase       )
       _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'''		,     UpperCAmelCase       ):
       raise ValueError('''The HuggingFace port of ESMFold does not support use_esm_attn_map at this time!'''       )





     def      __A					(self       )     ->   List[str]:
      _lowercase							      =super().to_dict()
      if isinstance(self.esmfold_config		,     UpperCAmelCase       ):
       _lowercase							      =self.esmfold_config.to_dict()
      return output







@dataclass
class       lowerCamelCase__   :
     SCREAMING_SNAKE_CASE__      			=			None
     SCREAMING_SNAKE_CASE__      			=			True
     SCREAMING_SNAKE_CASE__      			=			False
     SCREAMING_SNAKE_CASE__      			=			False
     SCREAMING_SNAKE_CASE__      			=			False
     SCREAMING_SNAKE_CASE__      			=			0

     SCREAMING_SNAKE_CASE__      			=			True
     SCREAMING_SNAKE_CASE__      			=			False

     SCREAMING_SNAKE_CASE__      			=			128
     SCREAMING_SNAKE_CASE__      			=			None





     def      __A					(self       )     ->   Union[str, Any]:
      if self.trunk is None:
       _lowercase							      =TrunkConfig()
      elif isinstance(self.trunk		,     UpperCAmelCase       ):
       _lowercase							      =TrunkConfig(**self.trunk       )





     def      __A					(self       )     ->   Tuple:
      _lowercase							      =asdict(self       )
      _lowercase							      =self.trunk.to_dict()
      return output







@dataclass
class       lowerCamelCase__   :
     SCREAMING_SNAKE_CASE__      			=			48
     SCREAMING_SNAKE_CASE__      			=			1024
     SCREAMING_SNAKE_CASE__      			=			128
     SCREAMING_SNAKE_CASE__      			=			32
     SCREAMING_SNAKE_CASE__      			=			32
     SCREAMING_SNAKE_CASE__      			=			32
     SCREAMING_SNAKE_CASE__      			=			0
     SCREAMING_SNAKE_CASE__      			=			0
     SCREAMING_SNAKE_CASE__      			=			False
     SCREAMING_SNAKE_CASE__      			=			4
     SCREAMING_SNAKE_CASE__      			=			128
     SCREAMING_SNAKE_CASE__      			=			None





     def      __A					(self       )     ->   List[str]:
      if self.structure_module is None:
       _lowercase							      =StructureModuleConfig()
      elif isinstance(self.structure_module		,     UpperCAmelCase       ):
       _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      __A					(self       )     ->   Dict:
      _lowercase							      =asdict(self       )
      _lowercase							      =self.structure_module.to_dict()
      return output







@dataclass
class       lowerCamelCase__   :
     SCREAMING_SNAKE_CASE__      			=			384
     SCREAMING_SNAKE_CASE__      			=			128
     SCREAMING_SNAKE_CASE__      			=			16
     SCREAMING_SNAKE_CASE__      			=			128
     SCREAMING_SNAKE_CASE__      			=			12
     SCREAMING_SNAKE_CASE__      			=			4
     SCREAMING_SNAKE_CASE__      			=			8
     SCREAMING_SNAKE_CASE__      			=			0.1
     SCREAMING_SNAKE_CASE__      			=			8
     SCREAMING_SNAKE_CASE__      			=			1
     SCREAMING_SNAKE_CASE__      			=			2
     SCREAMING_SNAKE_CASE__      			=			7
     SCREAMING_SNAKE_CASE__      			=			10
     SCREAMING_SNAKE_CASE__      			=			1E-8
     SCREAMING_SNAKE_CASE__      			=			1E5





     def      __A					(self       )     ->   List[Any]:
      return asdict(self       )



def 				UpperCAmelCase_					(							)     ->			Tuple:







 """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>",
 )

