codingutf8 2022 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 convert esm checkpoint import argparse import pathlib from pathlib import path from tempfile import temporarydirectory import esm as esmmodule import torch from esm esmfold v1 misc import batchencodesequences as esmfoldencodesequences from esm esmfold v1 pretrained import esmfoldv1 from transformers models esm configurationesm import esmconfig esmfoldconfig from transformers models esm modelingesm import esmformaskedlm esmforsequenceclassification esmintermediate esmlayer esmoutput esmselfattention esmselfoutput from transformers models esm modelingesmfold import esmforproteinfolding from transformers models esm tokenizationesm import esmtokenizer from transformers utils import logging logging setverbosityinfo logger logging getloggername sampledata protein1 mngtegpnfyvpfsnatgvvrspfeypqyylaepwqfsmlaaymfllivlgfpinfltlyvtvqhkklrtplnyillnlavadlfmvlggftstlytslhgyfvfgptgcnlegffatlggeialwslvvlaieryvvvckpmsnfrfgenhaimgvaftwvmalacaapplagwsryipeglqcscgidyytlkpevnnesfviymfvvhftipmiiiffcygqlvftvkeaaaqqqesattqkaekevtrmviimviaflicwvpyasvafyifthqgsnfgpifmtipaffaksaaiynpviyimmnkqfrncmltticcgknplgddeasatvsktetsqvapa protein2 mktvrqerlksivrilerskepvsgaqlaeelsvsrqvivqdiaylrslgynivatprgyvla protein3 mktvrqerlksimaskrilerskepvsgaqlaeelsmasksrqvivqdiaylrslgynmaskvatprgyvlagg protein4 mktvrqerlksimaskrilerskepvsgaqlaeelsmasksrqvivqdiaylrslgynmaskvatprgyvla modelmapping esm1bt33650mur50s esmmodule pretrained esm1bt33650mur50s esm1vt33650mur90s1 esmmodule pretrained esm1vt33650mur90s1 esm1vt33650mur90s2 esmmodule pretrained esm1vt33650mur90s2 esm1vt33650mur90s3 esmmodule pretrained esm1vt33650mur90s3 esm1vt33650mur90s4 esmmodule pretrained esm1vt33650mur90s4 esm1vt33650mur90s5 esmmodule pretrained esm1vt33650mur90s5 esm2t4815bur50d esmmodule pretrained esm2t4815bur50d esm2t363bur50d esmmodule pretrained esm2t363bur50d esm2t33650mur50d esmmodule pretrained esm2t33650mur50d esm2t30150mur50d esmmodule pretrained esm2t30150mur50d esm2t1235mur50d esmmodule pretrained esm2t1235mur50d esm2t68mur50d esmmodule pretrained esm2t68mur50d esmfoldv1 esmfoldv1 restypes listarndcqeghilkmfpstwyv restypeswithx restypes x restypeswithextras restypeswithx pad mask cls sep eos def getesmfoldtokenizer with temporarydirectory as tempdir vocab n joinrestypeswithextras vocabfile pathtempdir vocab txt vocabfile writetextvocab hftokenizer esmtokenizervocabfilestrvocabfile hftokenizer padtokenid 0 overlaps with a but that seems to be what they want return hftokenizer def transferandcheckweightsoriginalmodule ourmodule status ourmodule loadstatedictoriginalmodule statedict if status missingkeys raise valueerrorfmissing keys status missingkeys if status unexpectedkeys raise valueerrorfunexpected keys status unexpectedkeys def convertesmcheckpointtopytorch model str pytorchdumpfolderpath str classificationhead bool pushtorepo str authtoken str if model startswithesmfold esm modelmappingmodel else esm alphabet modelmappingmodel esm eval disable dropout if model startswithesmfold embeddim esm esm embeddim numlayers esm esm numlayers numattentionheads esm esm attentionheads intermediatesize 4 embeddim tokendropout esm esm tokendropout emblayernormbefore false this code path does not exist in esm2 positionembeddingtype rotary isfoldingmodel true esmfoldconfig esmfoldconfig for key val in esm cfg items if hasattresmfoldconfig key and key trunk setattresmfoldconfig key val for key val in esm cfg trunk items if hasattresmfoldconfig trunk key and key structuremodule setattresmfoldconfig trunk key val for key val in esm cfg trunk structuremodule items if hasattresmfoldconfig trunk structuremodule key setattresmfoldconfig trunk structuremodule key val elif hasattresm args indicates an esm1b or esm1v model embeddim esm args embeddim numlayers esm args layers numattentionheads esm args attentionheads intermediatesize esm args ffnembeddim tokendropout esm args tokendropout emblayernormbefore true if esm emblayernormbefore else false positionembeddingtype absolute isfoldingmodel false esmfoldconfig none else indicates an esm2 model embeddim esm embeddim numlayers esm numlayers numattentionheads esm attentionheads intermediatesize 4 embeddim this is hardcoded in esm2 tokendropout esm tokendropout emblayernormbefore false this code path does not exist in esm2 positionembeddingtype rotary isfoldingmodel false esmfoldconfig none if isfoldingmodel alphabet esm esm alphabet vocablist tuplealphabet alltoks masktokenid alphabet maskidx padtokenid alphabet paddingidx if isfoldingmodel originalesmmodel esm esm else originalesmmodel esm config esmconfig vocabsizeoriginalesmmodel embedtokens numembeddings masktokenidmasktokenid hiddensizeembeddim numhiddenlayersnumlayers numattentionheadsnumattentionheads intermediatesizeintermediatesize maxpositionembeddings1026 layernormeps1e5 pytorch default used in fairseq attentionprobsdropoutprob0 0 hiddendropoutprob0 0 padtokenidpadtokenid emblayernormbeforeemblayernormbefore tokendropouttokendropout positionembeddingtypepositionembeddingtype isfoldingmodelisfoldingmodel esmfoldconfigesmfoldconfig vocablistvocablist if classificationhead config numlabels esm classificationheadsmnli outproj weight shape0 printour esm config config if model startswithesmfold modelclass esmforproteinfolding elif classificationhead modelclass esmforsequenceclassification else modelclass esmformaskedlm model modelclassconfig model eval now let s copy all the weights embeddings model esm embeddings wordembeddings weight originalesmmodel embedtokens weight if positionembeddingtype absolute model esm embeddings positionembeddings weight originalesmmodel embedpositions weight if config emblayernormbefore model esm embeddings layernorm weight originalesmmodel emblayernormbefore weight model esm embeddings layernorm bias originalesmmodel emblayernormbefore bias model esm encoder emblayernormafter weight originalesmmodel emblayernormafter weight model esm encoder emblayernormafter bias originalesmmodel emblayernormafter bias for i in rangeconfig numhiddenlayers encoder start of layer layer esmlayer model esm encoder layeri esmlayer transformersentenceencoderlayer originalesmmodel layersi esmlayer originalesmmodel layersi self attention selfattn esmselfattention layer attention self assert esmlayer selfattn kproj weight data shape esmlayer selfattn qproj weight data shape esmlayer selfattn vproj weight data shape torch sizeconfig hiddensize config hiddensize selfattn query weight data esmlayer selfattn qproj weight selfattn query bias data esmlayer selfattn qproj bias selfattn key weight data esmlayer selfattn kproj weight selfattn key bias data esmlayer selfattn kproj bias selfattn value weight data esmlayer selfattn vproj weight selfattn value bias data esmlayer selfattn vproj bias if getattresmlayer selfattn rotemb none is not none matt although invfreq is not a trainable weight it is computed at model init and cached during the training of esm2 the model was converted to float16 precision which also converts the invfreq tensor and the loss of precision remains even if the model is loaded later as float32 if we recompute invfreq without this loss of precision then we will get subtly different rotary embeddings which are enough to cause significant discrepancies in model outputs to avoid this we make sure the new model copies the data from the old invfreq selfattn rotaryembeddings invfreq data esmlayer selfattn rotemb invfreq layernorm changes for preactivation layer attention layernorm weight esmlayer selfattnlayernorm weight layer attention layernorm bias esmlayer selfattnlayernorm bias layer layernorm weight esmlayer finallayernorm weight layer layernorm bias esmlayer finallayernorm bias selfattention output selfoutput esmselfoutput layer attention output assert selfoutput dense weight shape esmlayer selfattn outproj weight shape selfoutput dense weight esmlayer selfattn outproj weight selfoutput dense bias esmlayer selfattn outproj bias intermediate intermediate esmintermediate layer intermediate assert intermediate dense weight shape esmlayer fc1 weight shape intermediate dense weight esmlayer fc1 weight intermediate dense bias esmlayer fc1 bias output bertoutput esmoutput layer output assert bertoutput dense weight shape esmlayer fc2 weight shape bertoutput dense weight esmlayer fc2 weight bertoutput dense bias esmlayer fc2 bias end of layer if isfoldingmodel model esmscombine data esm esmscombine data model af2toesm data esm af2toesm data transferandcheckweightsesm embedding model embedding transferandcheckweightsesm esmsmlp model esmsmlp transferandcheckweightsesm trunk model trunk transferandcheckweightsesm distogramhead model distogramhead transferandcheckweightsesm ptmhead model ptmhead transferandcheckweightsesm lmhead model lmhead transferandcheckweightsesm lddthead model lddthead elif classificationhead model classifier dense weight esm esm classificationheadsmnli dense weight model classifier dense bias esm classificationheadsmnli dense bias model classifier outproj weight esm classificationheadsmnli outproj weight model classifier outproj bias esm classificationheadsmnli outproj bias else lm head model lmhead dense weight esm lmhead dense weight model lmhead dense bias esm lmhead dense bias model lmhead layernorm weight esm lmhead layernorm weight model lmhead layernorm bias esm lmhead layernorm bias model lmhead decoder weight esm lmhead weight model lmhead bias esm lmhead bias contact prediction head transferandcheckweightsesm contacthead model esm contacthead prepare data first 2 sequences from esmstructuralsplitdataset superfamily 4 if isfoldingmodel folding models aren t trained on masked inputs and don t like mask tokens sampledata sampledata 2 else sampledata sampledata if isfoldingmodel hftokenizer getesmfoldtokenizer hftokens hftokenizer row1 for row in sampledata returntensorspt paddingtrue addspecialtokensfalse esmfoldaas esmfoldmask esmfoldencodesequencesrow1 for row in sampledata success torch allhftokensinputids esmfoldaas and torch all hftokensattentionmask esmfoldmask else let s check that we get the same results batchconverter alphabet getbatchconverter batchlabels batchstrs batchtokens batchconvertersampledata prepare tokenizer and make sure it matches with temporarydirectory as tempdir vocab n joinalphabet alltoks vocabfile pathtempdir vocab txt vocabfile writetextvocab hftokenizer esmtokenizervocabfilestrvocabfile hftokens hftokenizerrow1 for row in sampledata returntensorspt paddingtrue success torch allhftokensinputids batchtokens printdo both models tokenizers output the same tokens if success else if not success raise exceptiontokenization does not match with torch nograd if isfoldingmodel let s test the model in parts esmfold always converts the esm stem to float16 which requires float16 ops that don t exist on cpu therefore to test it we need to run it on gpu however esmfold is what we in the community call a big boy and so we desperately avoid putting both the original and the converted model on the gpu at the same time theiroutput esm cuda inferrow1 for row in sampledata ouroutput model cuda inputidshftokensinputids cuda attentionmaskhftokensattentionmask cuda else ouroutput modelhftokens outputhiddenstatestrue ouroutput ouroutputlogits if classificationhead theiroutput esm model classificationheadsmnliesm extractfeaturesbatchtokens else theiroutput esmhftokensinputids reprlayerslistrange999 theiroutput theiroutputlogits if isfoldingmodel maxabsolutediff torch maxtorch absouroutputpositions theiroutputpositions item success torch allcloseouroutputpositions theiroutputpositions atol1e5 else maxabsolutediff torch maxtorch absouroutput theiroutput item success torch allcloseouroutput theiroutput atol1e5 printfmaxabsolutediff maxabsolutediff 1e5 printdo both models output the same tensors if success else if not success raise exceptionsomething went wrong if not isfoldingmodel let s check contact prediction too ouroutput model predictcontactshftokensinputids hftokensattentionmask theiroutput esm predictcontactshftokensinputids maxabsolutediff torch maxtorch absouroutput theiroutput item success torch allcloseouroutput theiroutput atol1e5 printcontact prediction testing printfmaxabsolutediff maxabsolutediff 1e5 printdo both models output the same tensors if success else if not success raise exceptionsomething went wrong pathlib pathpytorchdumpfolderpath mkdirparentstrue existoktrue printfsaving model to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath del esm free up some memory before continuing printfsaving tokenizer to pytorchdumpfolderpath hftokenizer savepretrainedpytorchdumpfolderpath if pushtorepo model pushtohubrepoidpushtorepo tokentokenauthtoken hftokenizer pushtohubrepoidpushtorepo tokentokenauthtoken if name main parser argparse argumentparser required parameters parser addargument pytorchdumpfolderpath typestr requiredtrue helppath to the output pytorch model parser addargument classificationhead actionstoretrue helpwhether to convert a final classification head parser addargumentmodel defaultnone typestr requiredtrue helpname of model to convert parser addargumentpushtorepo typestr helprepo to upload to including username parser addargumentauthtoken typestr helphuggingface auth token args parser parseargs convertesmcheckpointtopytorch args model args pytorchdumpfolderpath args classificationhead args pushtorepo args authtoken coding utf 8 2022 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 convert esm checkpoint overlaps with a but that seems to be what they want copy paste tweak esm s weights to our bert structure disable dropout this code path does not exist in esm 2 indicates an esm 1b or esm 1v model indicates an esm 2 model this is hardcoded in esm 2 this code path does not exist in esm 2 pytorch default used in fairseq now let s copy all the weights embeddings encoder start of layer esm_layer transformersentenceencoderlayer original_esm_model layers i self attention matt although inv_freq is not a trainable weight it is computed at model init and cached during the training of esm 2 the model was converted to float16 precision which also converts the inv_freq tensor and the loss of precision remains even if the model is loaded later as float32 if we recompute inv_freq without this loss of precision then we will get subtly different rotary embeddings which are enough to cause significant discrepancies in model outputs to avoid this we make sure the new model copies the data from the old inv_freq layernorm changes for pre activation self attention output intermediate output end of layer lm head contact prediction head prepare data first 2 sequences from esmstructuralsplitdataset superfamily 4 folding models aren t trained on masked inputs and don t like mask tokens let s check that we get the same results prepare tokenizer and make sure it matches let s test the model in parts esmfold always converts the esm stem to float16 which requires float16 ops that don t exist on cpu therefore to test it we need to run it on gpu however esmfold is what we in the community call a big boy and so we desperately avoid putting both the original and the converted model on the gpu at the same time 1e 5 let s check contact prediction too 1e 5 free up some memory before continuing required parameters
import argparse
import pathlib
from pathlib import Path
from tempfile import TemporaryDirectory

import esm as esm_module
import torch
from esm.esmfold.v1.misc import batch_encode_sequences as esmfold_encode_sequences
from esm.esmfold.v1.pretrained import esmfold_v1

from transformers.models.esm.configuration_esm import EsmConfig, EsmFoldConfig
from transformers.models.esm.modeling_esm import (
    EsmForMaskedLM,
    EsmForSequenceClassification,
    EsmIntermediate,
    EsmLayer,
    EsmOutput,
    EsmSelfAttention,
    EsmSelfOutput,
)
from transformers.models.esm.modeling_esmfold import EsmForProteinFolding
from transformers.models.esm.tokenization_esm import EsmTokenizer
from transformers.utils import logging


logging.set_verbosity_info()
logger = logging.get_logger(__name__)

SAMPLE_DATA = [
    (
        "protein1",
        "MNGTEGPNFYVPFSNATGVVRSPFEYPQYYLAEPWQFSMLAAYMFLLIVLGFPINFLTLYVTVQHKKLRTPLNYILLNLAVADLFMVLGGFTSTLYTSLHGYFVFGPTGCNLEGFFATLGGEIALWSLVVLAIERYVVVCKPMSNFRFGENHAIMGVAFTWVMALACAAPPLAGWSRYIPEGLQCSCGIDYYTLKPEVNNESFVIYMFVVHFTIPMIIIFFCYGQLVFTVKEAAAQQQESATTQKAEKEVTRMVIIMVIAFLICWVPYASVAFYIFTHQGSNFGPIFMTIPAFFAKSAAIYNPVIYIMMNKQFRNCMLTTICCGKNPLGDDEASATVSKTETSQVAPA",
    ),
    ("protein2", "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLA"),
    ("protein3", "MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLAGG"),
    ("protein4", "MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLA"),
]

MODEL_MAPPING = {
    "esm1b_t33_650M_UR50S": esm_module.pretrained.esm1b_t33_650M_UR50S,
    "esm1v_t33_650M_UR90S_1": esm_module.pretrained.esm1v_t33_650M_UR90S_1,
    "esm1v_t33_650M_UR90S_2": esm_module.pretrained.esm1v_t33_650M_UR90S_2,
    "esm1v_t33_650M_UR90S_3": esm_module.pretrained.esm1v_t33_650M_UR90S_3,
    "esm1v_t33_650M_UR90S_4": esm_module.pretrained.esm1v_t33_650M_UR90S_4,
    "esm1v_t33_650M_UR90S_5": esm_module.pretrained.esm1v_t33_650M_UR90S_5,
    "esm2_t48_15B_UR50D": esm_module.pretrained.esm2_t48_15B_UR50D,
    "esm2_t36_3B_UR50D": esm_module.pretrained.esm2_t36_3B_UR50D,
    "esm2_t33_650M_UR50D": esm_module.pretrained.esm2_t33_650M_UR50D,
    "esm2_t30_150M_UR50D": esm_module.pretrained.esm2_t30_150M_UR50D,
    "esm2_t12_35M_UR50D": esm_module.pretrained.esm2_t12_35M_UR50D,
    "esm2_t6_8M_UR50D": esm_module.pretrained.esm2_t6_8M_UR50D,
    "esmfold_v1": esmfold_v1,
}

restypes = list("ARNDCQEGHILKMFPSTWYV")

restypes_with_x = restypes + ["X"]
restypes_with_extras = restypes_with_x + ["<pad>", "<mask>", "<cls>", "<sep>", "<eos>"]


def get_esmfold_tokenizer():
    with TemporaryDirectory() as tempdir:
        vocab = "\n".join(restypes_with_extras)
        vocab_file = Path(tempdir) / "vocab.txt"
        vocab_file.write_text(vocab)
        hf_tokenizer = EsmTokenizer(vocab_file=str(vocab_file))
    hf_tokenizer.pad_token_id = 0  
    return hf_tokenizer


def transfer_and_check_weights(original_module, our_module):
    status = our_module.load_state_dict(original_module.state_dict())
    if status.missing_keys:
        raise ValueError(f"Missing keys: {status.missing_keys}")
    if status.unexpected_keys:
        raise ValueError(f"Unexpected keys: {status.unexpected_keys}")


def convert_esm_checkpoint_to_pytorch(
    model: str, pytorch_dump_folder_path: str, classification_head: bool, push_to_repo: str, auth_token: str
):
    
    if model.startswith("esmfold"):
        esm = MODEL_MAPPING[model]()
    else:
        esm, alphabet = MODEL_MAPPING[model]()
    esm.eval()  

    if model.startswith("esmfold"):
        embed_dim = esm.esm.embed_dim
        num_layers = esm.esm.num_layers
        num_attention_heads = esm.esm.attention_heads
        intermediate_size = 4 * embed_dim
        token_dropout = esm.esm.token_dropout
        emb_layer_norm_before = False  
        position_embedding_type = "rotary"
        is_folding_model = True
        esmfold_config = EsmFoldConfig()
        for key, val in esm.cfg.items():
            if hasattr(esmfold_config, key) and key != "trunk":
                setattr(esmfold_config, key, val)
        for key, val in esm.cfg.trunk.items():
            if hasattr(esmfold_config.trunk, key) and key != "structure_module":
                setattr(esmfold_config.trunk, key, val)
        for key, val in esm.cfg.trunk.structure_module.items():
            if hasattr(esmfold_config.trunk.structure_module, key):
                setattr(esmfold_config.trunk.structure_module, key, val)
    elif hasattr(esm, "args"):
        
        embed_dim = esm.args.embed_dim
        num_layers = esm.args.layers
        num_attention_heads = esm.args.attention_heads
        intermediate_size = esm.args.ffn_embed_dim
        token_dropout = esm.args.token_dropout
        emb_layer_norm_before = True if esm.emb_layer_norm_before else False
        position_embedding_type = "absolute"
        is_folding_model = False
        esmfold_config = None
    else:
        
        embed_dim = esm.embed_dim
        num_layers = esm.num_layers
        num_attention_heads = esm.attention_heads
        intermediate_size = 4 * embed_dim  
        token_dropout = esm.token_dropout
        emb_layer_norm_before = False  
        position_embedding_type = "rotary"
        is_folding_model = False
        esmfold_config = None

    if is_folding_model:
        alphabet = esm.esm.alphabet
    vocab_list = tuple(alphabet.all_toks)
    mask_token_id = alphabet.mask_idx
    pad_token_id = alphabet.padding_idx

    if is_folding_model:
        original_esm_model = esm.esm
    else:
        original_esm_model = esm

    config = EsmConfig(
        vocab_size=original_esm_model.embed_tokens.num_embeddings,
        mask_token_id=mask_token_id,
        hidden_size=embed_dim,
        num_hidden_layers=num_layers,
        num_attention_heads=num_attention_heads,
        intermediate_size=intermediate_size,
        max_position_embeddings=1026,
        layer_norm_eps=1e-5,  
        attention_probs_dropout_prob=0.0,
        hidden_dropout_prob=0.0,
        pad_token_id=pad_token_id,
        emb_layer_norm_before=emb_layer_norm_before,
        token_dropout=token_dropout,
        position_embedding_type=position_embedding_type,
        is_folding_model=is_folding_model,
        esmfold_config=esmfold_config,
        vocab_list=vocab_list,
    )
    if classification_head:
        config.num_labels = esm.classification_heads["mnli"].out_proj.weight.shape[0]
    print("Our ESM config:", config)

    if model.startswith("esmfold"):
        model_class = EsmForProteinFolding
    elif classification_head:
        model_class = EsmForSequenceClassification
    else:
        model_class = EsmForMaskedLM
    model = model_class(config)
    model.eval()

    
    
    model.esm.embeddings.word_embeddings.weight = original_esm_model.embed_tokens.weight
    if position_embedding_type == "absolute":
        model.esm.embeddings.position_embeddings.weight = original_esm_model.embed_positions.weight

    if config.emb_layer_norm_before:
        model.esm.embeddings.layer_norm.weight = original_esm_model.emb_layer_norm_before.weight
        model.esm.embeddings.layer_norm.bias = original_esm_model.emb_layer_norm_before.bias

    model.esm.encoder.emb_layer_norm_after.weight = original_esm_model.emb_layer_norm_after.weight
    model.esm.encoder.emb_layer_norm_after.bias = original_esm_model.emb_layer_norm_after.bias

    for i in range(config.num_hidden_layers):
        
        layer: EsmLayer = model.esm.encoder.layer[i]
        
        esm_layer = original_esm_model.layers[i]

        
        self_attn: EsmSelfAttention = layer.attention.self
        assert (
            esm_layer.self_attn.k_proj.weight.data.shape
            == esm_layer.self_attn.q_proj.weight.data.shape
            == esm_layer.self_attn.v_proj.weight.data.shape
            == torch.Size((config.hidden_size, config.hidden_size))
        )

        self_attn.query.weight.data = esm_layer.self_attn.q_proj.weight
        self_attn.query.bias.data = esm_layer.self_attn.q_proj.bias
        self_attn.key.weight.data = esm_layer.self_attn.k_proj.weight
        self_attn.key.bias.data = esm_layer.self_attn.k_proj.bias
        self_attn.value.weight.data = esm_layer.self_attn.v_proj.weight
        self_attn.value.bias.data = esm_layer.self_attn.v_proj.bias

        if getattr(esm_layer.self_attn, "rot_emb", None) is not None:
            
            
            
            
            
            
            self_attn.rotary_embeddings.inv_freq.data = esm_layer.self_attn.rot_emb.inv_freq

        
        layer.attention.LayerNorm.weight = esm_layer.self_attn_layer_norm.weight
        layer.attention.LayerNorm.bias = esm_layer.self_attn_layer_norm.bias
        layer.LayerNorm.weight = esm_layer.final_layer_norm.weight
        layer.LayerNorm.bias = esm_layer.final_layer_norm.bias

        
        self_output: EsmSelfOutput = layer.attention.output
        assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
        self_output.dense.weight = esm_layer.self_attn.out_proj.weight
        self_output.dense.bias = esm_layer.self_attn.out_proj.bias

        
        intermediate: EsmIntermediate = layer.intermediate
        assert intermediate.dense.weight.shape == esm_layer.fc1.weight.shape
        intermediate.dense.weight = esm_layer.fc1.weight
        intermediate.dense.bias = esm_layer.fc1.bias

        
        bert_output: EsmOutput = layer.output
        assert bert_output.dense.weight.shape == esm_layer.fc2.weight.shape
        bert_output.dense.weight = esm_layer.fc2.weight
        bert_output.dense.bias = esm_layer.fc2.bias
        

    if is_folding_model:
        model.esm_s_combine.data = esm.esm_s_combine.data
        model.af2_to_esm.data = esm.af2_to_esm.data
        transfer_and_check_weights(esm.embedding, model.embedding)
        transfer_and_check_weights(esm.esm_s_mlp, model.esm_s_mlp)
        transfer_and_check_weights(esm.trunk, model.trunk)
        transfer_and_check_weights(esm.distogram_head, model.distogram_head)
        transfer_and_check_weights(esm.ptm_head, model.ptm_head)
        transfer_and_check_weights(esm.lm_head, model.lm_head)
        transfer_and_check_weights(esm.lddt_head, model.lddt_head)

    elif classification_head:
        model.classifier.dense.weight = esm.esm.classification_heads["mnli"].dense.weight
        model.classifier.dense.bias = esm.classification_heads["mnli"].dense.bias
        model.classifier.out_proj.weight = esm.classification_heads["mnli"].out_proj.weight
        model.classifier.out_proj.bias = esm.classification_heads["mnli"].out_proj.bias
    else:
        
        model.lm_head.dense.weight = esm.lm_head.dense.weight
        model.lm_head.dense.bias = esm.lm_head.dense.bias
        model.lm_head.layer_norm.weight = esm.lm_head.layer_norm.weight
        model.lm_head.layer_norm.bias = esm.lm_head.layer_norm.bias
        model.lm_head.decoder.weight = esm.lm_head.weight
        model.lm_head.bias = esm.lm_head.bias

    
    transfer_and_check_weights(esm.contact_head, model.esm.contact_head)

    
    if is_folding_model:
        
        sample_data = SAMPLE_DATA[:2]
    else:
        sample_data = SAMPLE_DATA

    if is_folding_model:
        hf_tokenizer = get_esmfold_tokenizer()
        hf_tokens = hf_tokenizer(
            [row[1] for row in sample_data], return_tensors="pt", padding=True, add_special_tokens=False
        )
        esmfold_aas, esmfold_mask, _, _, _ = esmfold_encode_sequences([row[1] for row in sample_data])
        success = torch.all(hf_tokens["input_ids"] == esmfold_aas) and torch.all(
            hf_tokens["attention_mask"] == esmfold_mask
        )
    else:
        
        batch_converter = alphabet.get_batch_converter()
        batch_labels, batch_strs, batch_tokens = batch_converter(sample_data)
        
        with TemporaryDirectory() as tempdir:
            vocab = "\n".join(alphabet.all_toks)
            vocab_file = Path(tempdir) / "vocab.txt"
            vocab_file.write_text(vocab)
            hf_tokenizer = EsmTokenizer(vocab_file=str(vocab_file))

        hf_tokens = hf_tokenizer([row[1] for row in sample_data], return_tensors="pt", padding=True)
        success = torch.all(hf_tokens["input_ids"] == batch_tokens)

    print("Do both models tokenizers output the same tokens?", "🔥" if success else "💩")
    if not success:
        raise Exception("Tokenization does not match!")

    with torch.no_grad():
        if is_folding_model:
            
            
            
            
            
            their_output = esm.cuda().infer([row[1] for row in sample_data])
            our_output = model.cuda()(
                input_ids=hf_tokens["input_ids"].cuda(), attention_mask=hf_tokens["attention_mask"].cuda()
            )
        else:
            our_output = model(**hf_tokens, output_hidden_states=True)
            our_output = our_output["logits"]
            if classification_head:
                their_output = esm.model.classification_heads["mnli"](esm.extract_features(batch_tokens))
            else:
                their_output = esm(hf_tokens["input_ids"], repr_layers=list(range(999)))
                their_output = their_output["logits"]

        if is_folding_model:
            max_absolute_diff = torch.max(torch.abs(our_output["positions"] - their_output["positions"])).item()
            success = torch.allclose(our_output["positions"], their_output["positions"], atol=1e-5)
        else:
            max_absolute_diff = torch.max(torch.abs(our_output - their_output)).item()
            success = torch.allclose(our_output, their_output, atol=1e-5)

        print(f"max_absolute_diff = {max_absolute_diff}")  
        print("Do both models output the same tensors?", "🔥" if success else "💩")

        if not success:
            raise Exception("Something went wRoNg")

        if not is_folding_model:
            
            our_output = model.predict_contacts(hf_tokens["input_ids"], hf_tokens["attention_mask"])
            their_output = esm.predict_contacts(hf_tokens["input_ids"])
            max_absolute_diff = torch.max(torch.abs(our_output - their_output)).item()
            success = torch.allclose(our_output, their_output, atol=1e-5)

            print("Contact prediction testing:")
            print(f"max_absolute_diff = {max_absolute_diff}")  
            print("Do both models output the same tensors?", "🔥" if success else "💩")

            if not success:
                raise Exception("Something went wRoNg")

        pathlib.Path(pytorch_dump_folder_path).mkdir(parents=True, exist_ok=True)
        print(f"Saving model to {pytorch_dump_folder_path}")
        model.save_pretrained(pytorch_dump_folder_path)

        del esm  

    print(f"Saving tokenizer to {pytorch_dump_folder_path}")
    hf_tokenizer.save_pretrained(pytorch_dump_folder_path)

    if push_to_repo:
        model.push_to_hub(repo_id=push_to_repo, token_token=auth_token)
        hf_tokenizer.push_to_hub(repo_id=push_to_repo, token_token=auth_token)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--pytorch_dump_folder_path", type=str, required=True, help="Path to the output PyTorch model."
    )
    parser.add_argument(
        "--classification_head", action="store_true", help="Whether to convert a final classification head."
    )
    parser.add_argument("--model", default=None, type=str, required=True, help="Name of model to convert.")
    parser.add_argument("--push_to_repo", type=str, help="Repo to upload to (including username!).")
    parser.add_argument("--auth_token", type=str, help="HuggingFace auth token.")
    args = parser.parse_args()
    convert_esm_checkpoint_to_pytorch(
        args.model, args.pytorch_dump_folder_path, args.classification_head, args.push_to_repo, args.auth_token
    )