codingutf8 2022 meta and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license esm model configuration from dataclasses import asdict dataclass from typing import optional from configurationutils import pretrainedconfig from utils import logging logger logging getloggername todo update this esmpretrainedconfigarchivemap facebookesm1b https huggingface cofacebookesm1bresolvemainconfig json see all esm models at https huggingface comodels filteresm class esmconfigpretrainedconfig r this is the configuration class to store the configuration of a esmmodel it is used to instantiate a esm model according to the specified arguments defining the model architecture instantiating a configuration with the defaults will yield a similar configuration to that of the esm facebookesm1bhttps huggingface cofacebookesm1b architecture configuration objects inherit from pretrainedconfig and can be used to control the model outputs read the documentation from pretrainedconfig for more information args vocabsize int optional vocabulary size of the esm model defines the number of different tokens that can be represented by the inputsids passed when calling esmmodel masktokenid int optional the index of the mask token in the vocabulary this must be included in the config because of the maskdropout scaling trick which will scale the inputs depending on the number of masked tokens padtokenid int optional the index of the padding token in the vocabulary this must be included in the config because certain parts of the esm code use this instead of the attention mask hiddensize int optional defaults to 768 dimensionality of the encoder layers and the pooler layer numhiddenlayers int optional defaults to 12 number of hidden layers in the transformer encoder numattentionheads int optional defaults to 12 number of attention heads for each attention layer in the transformer encoder intermediatesize int optional defaults to 3072 dimensionality of the intermediate often named feedforward layer in the transformer encoder hiddendropoutprob float optional defaults to 0 1 the dropout probability for all fully connected layers in the embeddings encoder and pooler attentionprobsdropoutprob float optional defaults to 0 1 the dropout ratio for the attention probabilities maxpositionembeddings int optional defaults to 1026 the maximum sequence length that this model might ever be used with typically set this to something large just in case e g 512 or 1024 or 2048 initializerrange float optional defaults to 0 02 the standard deviation of the truncatednormalinitializer for initializing all weight matrices layernormeps float optional defaults to 1e12 the epsilon used by the layer normalization layers positionembeddingtype str optional defaults to absolute type of position embedding choose one of absolute relativekey relativekeyquery rotary for positional embeddings use absolute for more information on relativekey please refer to selfattention with relative position representations shaw et al https arxiv orgabs1803 02155 for more information on relativekeyquery please refer to method 4 in improve transformer models with better relative position embeddings huang et al https arxiv orgabs2009 13658 isdecoder bool optional defaults to false whether the model is used as a decoder or not if false the model is used as an encoder usecache bool optional defaults to true whether or not the model should return the last keyvalues attentions not used by all models only relevant if config isdecodertrue emblayernormbefore bool optional whether to apply layer normalization after embeddings but before the main stem of the network tokendropout bool defaults to false when this is enabled masked tokens are treated as if they had been dropped out by input dropout examples python from transformers import esmmodel esmconfig initializing a esm facebookesm1b style configuration configuration esmconfig initializing a model from the configuration model esmmodelconfiguration accessing the model configuration configuration model config modeltype esm def init self vocabsizenone masktokenidnone padtokenidnone hiddensize768 numhiddenlayers12 numattentionheads12 intermediatesize3072 hiddendropoutprob0 1 attentionprobsdropoutprob0 1 maxpositionembeddings1026 initializerrange0 02 layernormeps1e12 positionembeddingtypeabsolute usecachetrue emblayernormbeforenone tokendropoutfalse isfoldingmodelfalse esmfoldconfignone vocablistnone kwargs super initpadtokenidpadtokenid masktokenidmasktokenid kwargs self vocabsize vocabsize self hiddensize hiddensize self numhiddenlayers numhiddenlayers self numattentionheads numattentionheads self intermediatesize intermediatesize self hiddendropoutprob hiddendropoutprob self attentionprobsdropoutprob attentionprobsdropoutprob self maxpositionembeddings maxpositionembeddings self initializerrange initializerrange self layernormeps layernormeps self positionembeddingtype positionembeddingtype self usecache usecache self emblayernormbefore emblayernormbefore self tokendropout tokendropout self isfoldingmodel isfoldingmodel if isfoldingmodel if esmfoldconfig is none logger infono esmfoldconfig supplied for folding model using default values esmfoldconfig esmfoldconfig elif isinstanceesmfoldconfig dict esmfoldconfig esmfoldconfigesmfoldconfig self esmfoldconfig esmfoldconfig if vocablist is none logger warningno vocablist supplied for folding model assuming the esm2 vocabulary self vocablist getdefaultvocablist else self vocablist vocablist else self esmfoldconfig none self vocablist none if self esmfoldconfig is not none and getattrself esmfoldconfig useesmattnmap false raise valueerrorthe huggingface port of esmfold does not support useesmattnmap at this time def todictself output super todict if isinstanceself esmfoldconfig esmfoldconfig outputesmfoldconfig self esmfoldconfig todict return output dataclass class esmfoldconfig esmtype str none fp16esm bool true useesmattnmap bool false esmablatepairwise bool false esmablatesequence bool false esminputdropout float 0 embedaa bool true bypasslm bool false lddtheadhiddim int 128 trunk trunkconfig none def postinitself if self trunk is none self trunk trunkconfig elif isinstanceself trunk dict self trunk trunkconfigself trunk def todictself output asdictself outputtrunk self trunk todict return output dataclass class trunkconfig numblocks int 48 sequencestatedim int 1024 pairwisestatedim int 128 sequenceheadwidth int 32 pairwiseheadwidth int 32 positionbins int 32 dropout float 0 layerdrop float 0 cpugradcheckpoint bool false maxrecycles int 4 chunksize optionalint 128 structuremodule structuremoduleconfig none def postinitself if self structuremodule is none self structuremodule structuremoduleconfig elif isinstanceself structuremodule dict self structuremodule structuremoduleconfigself structuremodule if self maxrecycles 0 raise valueerrorfmaxrecycles should be positive got self maxrecycles if self sequencestatedim self sequencestatedim 0 raise valueerror sequencestatedim should be a round multiple of sequencestatedim got f self sequencestatedim and self sequencestatedim if self pairwisestatedim self pairwisestatedim 0 raise valueerror pairwisestatedim should be a round multiple of pairwisestatedim got f self pairwisestatedim and self pairwisestatedim sequencenumheads self sequencestatedim self sequenceheadwidth pairwisenumheads self pairwisestatedim self pairwiseheadwidth if self sequencestatedim sequencenumheads self sequenceheadwidth raise valueerror sequencestatedim should be equal to sequencenumheads sequenceheadwidth got f self sequencestatedim sequencenumheads self sequenceheadwidth if self pairwisestatedim pairwisenumheads self pairwiseheadwidth raise valueerror pairwisestatedim should be equal to pairwisenumheads pairwiseheadwidth got f self pairwisestatedim pairwisenumheads self pairwiseheadwidth if self pairwisestatedim 2 0 raise valueerrorfpairwisestatedim should be even got self pairwisestatedim if self dropout 0 4 raise valueerrorfdropout should not be greater than 0 4 got self dropout def todictself output asdictself outputstructuremodule self structuremodule todict return output dataclass class structuremoduleconfig sequencedim int 384 pairwisedim int 128 ipadim int 16 resnetdim int 128 numheadsipa int 12 numqkpoints int 4 numvpoints int 8 dropoutrate float 0 1 numblocks int 8 numtransitionlayers int 1 numresnetblocks int 2 numangles int 7 transscalefactor int 10 epsilon float 1e8 inf float 1e5 def todictself return asdictself def getdefaultvocablist 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 null1 mask coding utf 8 2022 meta and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license esm model configuration todo update this see all esm models at https huggingface co models filter esm this is the configuration class to store the configuration of a esmmodel it is used to instantiate a esm model according to the specified arguments defining the model architecture instantiating a configuration with the defaults will yield a similar configuration to that of the esm facebook esm 1b https huggingface co facebook esm 1b architecture configuration objects inherit from pretrainedconfig and can be used to control the model outputs read the documentation from pretrainedconfig for more information args vocab_size int optional vocabulary size of the esm model defines the number of different tokens that can be represented by the inputs_ids passed when calling esmmodel mask_token_id int optional the index of the mask token in the vocabulary this must be included in the config because of the mask dropout scaling trick which will scale the inputs depending on the number of masked tokens pad_token_id int optional the index of the padding token in the vocabulary this must be included in the config because certain parts of the esm code use this instead of the attention mask hidden_size int optional defaults to 768 dimensionality of the encoder layers and the pooler layer num_hidden_layers int optional defaults to 12 number of hidden layers in the transformer encoder num_attention_heads int optional defaults to 12 number of attention heads for each attention layer in the transformer encoder intermediate_size int optional defaults to 3072 dimensionality of the intermediate often named feed forward layer in the transformer encoder hidden_dropout_prob float optional defaults to 0 1 the dropout probability for all fully connected layers in the embeddings encoder and pooler attention_probs_dropout_prob float optional defaults to 0 1 the dropout ratio for the attention probabilities max_position_embeddings int optional defaults to 1026 the maximum sequence length that this model might ever be used with typically set this to something large just in case e g 512 or 1024 or 2048 initializer_range float optional defaults to 0 02 the standard deviation of the truncated_normal_initializer for initializing all weight matrices layer_norm_eps float optional defaults to 1e 12 the epsilon used by the layer normalization layers position_embedding_type str optional defaults to absolute type of position embedding choose one of absolute relative_key relative_key_query rotary for positional embeddings use absolute for more information on relative_key please refer to self attention with relative position representations shaw et al https arxiv org abs 1803 02155 for more information on relative_key_query please refer to method 4 in improve transformer models with better relative position embeddings huang et al https arxiv org abs 2009 13658 is_decoder bool optional defaults to false whether the model is used as a decoder or not if false the model is used as an encoder use_cache bool optional defaults to true whether or not the model should return the last key values attentions not used by all models only relevant if config is_decoder true emb_layer_norm_before bool optional whether to apply layer normalization after embeddings but before the main stem of the network token_dropout bool defaults to false when this is enabled masked tokens are treated as if they had been dropped out by input dropout examples python from transformers import esmmodel esmconfig initializing a esm facebook esm 1b style configuration configuration esmconfig initializing a model from the configuration model esmmodel configuration accessing the model configuration configuration model config serializes this instance to a python dictionary override the default pretrainedconfig to_dict returns dict str any dictionary of all the attributes that make up this configuration instance serializes this instance to a python dictionary override the default pretrainedconfig to_dict returns dict str any dictionary of all the attributes that make up this configuration instance serializes this instance to a python dictionary override the default pretrainedconfig to_dict returns dict str any dictionary of all the attributes that make up this configuration instance args sequence_dim single representation channel dimension pairwise_dim pair representation channel dimension ipa_dim ipa hidden channel dimension resnet_dim angle resnet alg 23 lines 11 14 hidden channel dimension num_heads_ipa number of ipa heads num_qk_points number of query key points to generate during ipa num_v_points number of value points to generate during ipa dropout_rate dropout rate used throughout the layer num_blocks number of structure module blocks num_transition_layers number of layers in the single representation transition alg 23 lines 8 9 num_resnet_blocks number of blocks in the angle resnet num_angles number of angles to generate in the angle resnet trans_scale_factor scale of single representation transition hidden dimension epsilon small number used in angle resnet normalization inf large number used for attention masking
from dataclasses import asdict, dataclass
from typing import Optional

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


logger = logging.get_logger(__name__)


ESM_PRETRAINED_CONFIG_ARCHIVE_MAP = {
    "facebook/esm-1b": "https://huggingface.co/facebook/esm-1b/resolve/main/config.json",
    
}


class EsmConfig(PretrainedConfig):
    r

    model_type = "esm"

    def __init__(
        self,
        vocab_size=None,
        mask_token_id=None,
        pad_token_id=None,
        hidden_size=768,
        num_hidden_layers=12,
        num_attention_heads=12,
        intermediate_size=3072,
        hidden_dropout_prob=0.1,
        attention_probs_dropout_prob=0.1,
        max_position_embeddings=1026,
        initializer_range=0.02,
        layer_norm_eps=1e-12,
        position_embedding_type="absolute",
        use_cache=True,
        emb_layer_norm_before=None,
        token_dropout=False,
        is_folding_model=False,
        esmfold_config=None,
        vocab_list=None,
        **kwargs,
    ):
        super().__init__(pad_token_id=pad_token_id, mask_token_id=mask_token_id, **kwargs)

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

    def to_dict(self):
        
        output = super().to_dict()
        if isinstance(self.esmfold_config, EsmFoldConfig):
            output["esmfold_config"] = self.esmfold_config.to_dict()
        return output


@dataclass
class EsmFoldConfig:
    esm_type: str = None
    fp16_esm: bool = True
    use_esm_attn_map: bool = False
    esm_ablate_pairwise: bool = False
    esm_ablate_sequence: bool = False
    esm_input_dropout: float = 0

    embed_aa: bool = True
    bypass_lm: bool = False

    lddt_head_hid_dim: int = 128
    trunk: "TrunkConfig" = None

    def __post_init__(self):
        if self.trunk is None:
            self.trunk = TrunkConfig()
        elif isinstance(self.trunk, dict):
            self.trunk = TrunkConfig(**self.trunk)

    def to_dict(self):
        
        output = asdict(self)
        output["trunk"] = self.trunk.to_dict()
        return output


@dataclass
class TrunkConfig:
    num_blocks: int = 48
    sequence_state_dim: int = 1024
    pairwise_state_dim: int = 128
    sequence_head_width: int = 32
    pairwise_head_width: int = 32
    position_bins: int = 32
    dropout: float = 0
    layer_drop: float = 0
    cpu_grad_checkpoint: bool = False
    max_recycles: int = 4
    chunk_size: Optional[int] = 128
    structure_module: "StructureModuleConfig" = None

    def __post_init__(self):
        if self.structure_module is None:
            self.structure_module = StructureModuleConfig()
        elif isinstance(self.structure_module, dict):
            self.structure_module = 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}."
            )

        sequence_num_heads = self.sequence_state_dim // self.sequence_head_width
        pairwise_num_heads = 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 to_dict(self):
        
        output = asdict(self)
        output["structure_module"] = self.structure_module.to_dict()
        return output


@dataclass
class StructureModuleConfig:
    

    sequence_dim: int = 384
    pairwise_dim: int = 128
    ipa_dim: int = 16
    resnet_dim: int = 128
    num_heads_ipa: int = 12
    num_qk_points: int = 4
    num_v_points: int = 8
    dropout_rate: float = 0.1
    num_blocks: int = 8
    num_transition_layers: int = 1
    num_resnet_blocks: int = 2
    num_angles: int = 7
    trans_scale_factor: int = 10
    epsilon: float = 1e-8
    inf: float = 1e5

    def to_dict(self):
        return asdict(self)


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