
from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import convert_to_rgb, normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
 import PIL


_lowercase											=    logging.get_logger(__name__)


class 			lowercase_		(      A       ):
       __lowerCamelCase				       =						["pixel_values"]


       def __init__(    self						,       __A = True						,       __A = None						,       __A = PILImageResampling.BICUBIC						,       __A = True						,       __A = 1 / 255						,       __A = True						,       __A = None						,       __A = None						,       __A = True						,       **__A						,       )						->    None:
              super().__init__(**__A				)
              SCREAMING_SNAKE_CASE_							:				List[Any]		=size if size is not None else {'''height''': 384, '''width''': 384}
              SCREAMING_SNAKE_CASE_							:				Any		=get_size_dict(__A						,       default_to_square=__A				)

              SCREAMING_SNAKE_CASE_							:				str		=do_resize
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=size
              SCREAMING_SNAKE_CASE_							:				List[Any]		=resample
              SCREAMING_SNAKE_CASE_							:				str		=do_rescale
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=rescale_factor
              SCREAMING_SNAKE_CASE_							:				Dict		=do_normalize
              SCREAMING_SNAKE_CASE_							:				List[Any]		=image_mean if image_mean is not None else OPENAI_CLIP_MEAN
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=image_std if image_std is not None else OPENAI_CLIP_STD
              SCREAMING_SNAKE_CASE_							:				List[Any]		=do_convert_rgb


       def 	_snake_case				(    self						,       __A						,       __A						,       __A = PILImageResampling.BICUBIC						,       __A = None						,       **__A						,       )						->    np.ndarray:
              SCREAMING_SNAKE_CASE_							:				Any		=get_size_dict(__A						,       default_to_square=__A				)
              if "height" not in size or "width" not in size:
                     raise ValueError(F'The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}'				)
              SCREAMING_SNAKE_CASE_							:				List[str]		=(size['''height'''], size['''width'''])
              return resize(__A						,       size=__A						,       resample=__A						,       data_format=__A						,       **__A				)


       def 	_snake_case				(    self						,       __A						,       __A						,       __A = None						,       **__A						,       )						->    str:
              return rescale(__A						,       scale=__A						,       data_format=__A						,       **__A				)


       def 	_snake_case				(    self						,       __A						,       __A						,       __A						,       __A = None						,       **__A						,       )						->    np.ndarray:
              return normalize(__A						,       mean=__A						,       std=__A						,       data_format=__A						,       **__A				)





       def 	_snake_case				(    self						,       __A						,       __A = None						,       __A = None						,       __A = None						,       __A = None						,       __A = None						,       __A = None						,       __A = None						,       __A = None						,       __A = None						,       __A = None						,       __A = ChannelDimension.FIRST						,       **__A						,       )						->    PIL.Image.Image:
              SCREAMING_SNAKE_CASE_							:				Tuple		=do_resize if do_resize is not None else self.do_resize
              SCREAMING_SNAKE_CASE_							:				List[Any]		=resample if resample is not None else self.resample
              SCREAMING_SNAKE_CASE_							:				List[Any]		=do_rescale if do_rescale is not None else self.do_rescale
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=rescale_factor if rescale_factor is not None else self.rescale_factor
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=do_normalize if do_normalize is not None else self.do_normalize
              SCREAMING_SNAKE_CASE_							:				List[str]		=image_mean if image_mean is not None else self.image_mean
              SCREAMING_SNAKE_CASE_							:				Tuple		=image_std if image_std is not None else self.image_std
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

              SCREAMING_SNAKE_CASE_							:				List[Any]		=size if size is not None else self.size
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=get_size_dict(__A						,       default_to_square=__A				)

              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=make_list_of_images(__A				)

              if not valid_images(__A				):
                     raise ValueError(
                         '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
                         '''torch.Tensor, tf.Tensor or jax.ndarray.'''				)

              if do_resize and size is None or resample is None:
                     raise ValueError('''Size and resample must be specified if do_resize is True.'''				)

              if do_rescale and rescale_factor is None:
                     raise ValueError('''Rescale factor must be specified if do_rescale is True.'''				)

              if do_normalize and (image_mean is None or image_std is None):
                     raise ValueError('''Image mean and std must be specified if do_normalize is True.'''				)

              # PIL RGBA images are converted to RGB
              if do_convert_rgb:
                     SCREAMING_SNAKE_CASE_							:				Optional[Any]		=[convert_to_rgb(__A				) for image in images]

              # All transformations expect numpy arrays.
              SCREAMING_SNAKE_CASE_							:				Tuple		=[to_numpy_array(__A				) for image in images]

              if do_resize:
                     SCREAMING_SNAKE_CASE_							:				Optional[int]		=[self.resize(image=__A						,       size=__A						,       resample=__A				) for image in images]

              if do_rescale:
                     SCREAMING_SNAKE_CASE_							:				Optional[int]		=[self.rescale(image=__A						,       scale=__A				) for image in images]

              if do_normalize:
                     SCREAMING_SNAKE_CASE_							:				Tuple		=[self.normalize(image=__A						,       mean=__A						,       std=__A				) for image in images]

              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=[to_channel_dimension_format(__A						,       __A				) for image in images]

              SCREAMING_SNAKE_CASE_							:				List[Any]		=BatchFeature(data={'''pixel_values''': images}						,       tensor_type=__A				)

              return encoded_outputs


from dataclasses import asdict, dataclass
from typing import Optional

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


_lowercase											=    logging.get_logger(__name__)

# TODO Update this
_lowercase											=    {
    """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 			lowercase_		(      A       ):
       __lowerCamelCase				       =						"esm"


       def __init__(    self						,       __A=None						,       __A=None						,       __A=None						,       __A=768						,       __A=12						,       __A=12						,       __A=3_072						,       __A=0.1						,       __A=0.1						,       __A=1_026						,       __A=0.02						,       __A=1e-1_2						,       __A="absolute"						,       __A=True						,       __A=None						,       __A=False						,       __A=False						,       __A=None						,       __A=None						,       **__A						,       )						->    int:
              super().__init__(pad_token_id=__A						,       mask_token_id=__A						,       **__A				)

              SCREAMING_SNAKE_CASE_							:				Any		=vocab_size
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=hidden_size
              SCREAMING_SNAKE_CASE_							:				int		=num_hidden_layers
              SCREAMING_SNAKE_CASE_							:				int		=num_attention_heads
              SCREAMING_SNAKE_CASE_							:				List[str]		=intermediate_size
              SCREAMING_SNAKE_CASE_							:				int		=hidden_dropout_prob
              SCREAMING_SNAKE_CASE_							:				int		=attention_probs_dropout_prob
              SCREAMING_SNAKE_CASE_							:				List[Any]		=max_position_embeddings
              SCREAMING_SNAKE_CASE_							:				List[Any]		=initializer_range
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=layer_norm_eps
              SCREAMING_SNAKE_CASE_							:				str		=position_embedding_type
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=use_cache
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=emb_layer_norm_before
              SCREAMING_SNAKE_CASE_							:				Tuple		=token_dropout
              SCREAMING_SNAKE_CASE_							:				List[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_							:				Optional[int]		=EsmFoldConfig()
                     elif isinstance(__A						,       __A				):
                            SCREAMING_SNAKE_CASE_							:				List[Any]		=EsmFoldConfig(**__A				)
                     SCREAMING_SNAKE_CASE_							:				Optional[int]		=esmfold_config
                     if vocab_list is None:
                            logger.warning('''No vocab_list supplied for folding model, assuming the ESM-2 vocabulary!'''				)
                            SCREAMING_SNAKE_CASE_							:				Tuple		=get_default_vocab_list()
                     else:
                            SCREAMING_SNAKE_CASE_							:				Optional[Any]		=vocab_list
              else:
                     SCREAMING_SNAKE_CASE_							:				List[str]		=None
                     SCREAMING_SNAKE_CASE_							:				Optional[int]		=None
              if self.esmfold_config is not None and getattr(self.esmfold_config						,       '''use_esm_attn_map'''						,       __A				):
                     raise ValueError('''The HuggingFace port of ESMFold does not support use_esm_attn_map at this time!'''				)





       def 	_snake_case				(    self				)						->    List[Any]:
              SCREAMING_SNAKE_CASE_							:				str		=super().to_dict()
              if isinstance(self.esmfold_config						,       __A				):
                     SCREAMING_SNAKE_CASE_							:				Optional[Any]		=self.esmfold_config.to_dict()
              return output






@dataclass
class 			lowercase_		:
       __lowerCamelCase				       =						None
       __lowerCamelCase				       =						True
       __lowerCamelCase				       =						False
       __lowerCamelCase				       =						False
       __lowerCamelCase				       =						False
       __lowerCamelCase				       =						0

       __lowerCamelCase				       =						True
       __lowerCamelCase				       =						False

       __lowerCamelCase				       =						1_2_8
       __lowerCamelCase				       =						None


       def 	_snake_case				(    self				)						->    Optional[Any]:
              if self.trunk is None:
                     SCREAMING_SNAKE_CASE_							:				Tuple		=TrunkConfig()
              elif isinstance(self.trunk						,       __A				):
                     SCREAMING_SNAKE_CASE_							:				List[str]		=TrunkConfig(**self.trunk				)





       def 	_snake_case				(    self				)						->    List[Any]:
              SCREAMING_SNAKE_CASE_							:				int		=asdict(self				)
              SCREAMING_SNAKE_CASE_							:				Tuple		=self.trunk.to_dict()
              return output






@dataclass
class 			lowercase_		:
       __lowerCamelCase				       =						4_8
       __lowerCamelCase				       =						1_0_2_4
       __lowerCamelCase				       =						1_2_8
       __lowerCamelCase				       =						3_2
       __lowerCamelCase				       =						3_2
       __lowerCamelCase				       =						3_2
       __lowerCamelCase				       =						0
       __lowerCamelCase				       =						0
       __lowerCamelCase				       =						False
       __lowerCamelCase				       =						4
       __lowerCamelCase				       =						1_2_8
       __lowerCamelCase				       =						None


       def 	_snake_case				(    self				)						->    int:
              if self.structure_module is None:
                     SCREAMING_SNAKE_CASE_							:				Dict		=StructureModuleConfig()
              elif isinstance(self.structure_module						,       __A				):
                     SCREAMING_SNAKE_CASE_							:				Optional[Any]		=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_							:				Tuple		=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 	_snake_case				(    self				)						->    int:
              SCREAMING_SNAKE_CASE_							:				str		=asdict(self				)
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=self.structure_module.to_dict()
              return output






@dataclass
class 			lowercase_		:
       __lowerCamelCase				       =						3_8_4
       __lowerCamelCase				       =						1_2_8
       __lowerCamelCase				       =						1_6
       __lowerCamelCase				       =						1_2_8
       __lowerCamelCase				       =						1_2
       __lowerCamelCase				       =						4
       __lowerCamelCase				       =						8
       __lowerCamelCase				       =						0.1
       __lowerCamelCase				       =						8
       __lowerCamelCase				       =						1
       __lowerCamelCase				       =						2
       __lowerCamelCase				       =						7
       __lowerCamelCase				       =						1_0
       __lowerCamelCase				       =						1E-8
       __lowerCamelCase				       =						1E5


       def 	_snake_case				(    self				)						->    str:
              return asdict(self				)






def 							SCREAMING_SNAKE_CASE_	(					)			->						int:
       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>",
       )
