





import argparse
import pathlib
from pathlib import Path
from tempfile import TemporaryDirectory

import esm as esm_module
import torch
from esm.esmfold.va.misc import batch_encode_sequences as esmfold_encode_sequences
from esm.esmfold.va.pretrained import esmfold_va

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()
UpperCAmelCase_       : List[str]							=    logging.get_logger(__name__)

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

UpperCAmelCase_       : List[Any]							=    {
    """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_va,
}

UpperCAmelCase_       : List[Any]							=    list("""ARNDCQEGHILKMFPSTWYV""")

UpperCAmelCase_       : List[Any]							=    restypes + ["""X"""]
UpperCAmelCase_       : Optional[int]							=    restypes_with_x + ["""<pad>""", """<mask>""", """<cls>""", """<sep>""", """<eos>"""]






def 				_lowerCAmelCase			(	)						->      Tuple:
     with TemporaryDirectory() as tempdir:
          lowerCAmelCase_       :							Optional[int]			 =		"""\n""".join(_a      )
          lowerCAmelCase_       :							Optional[Any]			 =		Path(_a      ) / """vocab.txt"""
          vocab_file.write_text(_a      )
          lowerCAmelCase_       :							List[Any]			 =		EsmTokenizer(vocab_file=str(_a      )      )
     lowerCAmelCase_       :							int			 =		0  # Overlaps with 'A' but that seems to be what they want
     return hf_tokenizer






def 				_lowerCAmelCase			(	_a :       Dict       , _a :       Dict      )						->      List[str]:
     lowerCAmelCase_       :							List[str]			 =		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 				_lowerCAmelCase			(	_a :       str       , _a :       str       , _a :       bool       , _a :       str       , _a :       str      )						->      str:
     if model.startswith("""esmfold"""      ):
          lowerCAmelCase_       :							List[str]			 =		MODEL_MAPPING[model]()
     else:
          lowerCAmelCase_ ,	lowerCAmelCase_       :							Dict			 =		MODEL_MAPPING[model]()
     esm.eval()  # disable dropout

     if model.startswith("""esmfold"""      ):
          lowerCAmelCase_       :							Union[str, Any]			 =		esm.esm.embed_dim
          lowerCAmelCase_       :							Any			 =		esm.esm.num_layers
          lowerCAmelCase_       :							Tuple			 =		esm.esm.attention_heads
          lowerCAmelCase_       :							Dict			 =		4 * embed_dim
          lowerCAmelCase_       :							Dict			 =		esm.esm.token_dropout
          lowerCAmelCase_       :							List[str]			 =		False  # This code path does not exist in ESM-2
          lowerCAmelCase_       :							Any			 =		"""rotary"""
          lowerCAmelCase_       :							str			 =		True
          lowerCAmelCase_       :							Dict			 =		EsmFoldConfig()
          for key, val in esm.cfg.items():
               if hasattr(_a       , _a      ) and key != "trunk":
                    setattr(_a       , _a       , _a      )
          for key, val in esm.cfg.trunk.items():
               if hasattr(esmfold_config.trunk       , _a      ) and key != "structure_module":
                    setattr(esmfold_config.trunk       , _a       , _a      )
          for key, val in esm.cfg.trunk.structure_module.items():
               if hasattr(esmfold_config.trunk.structure_module       , _a      ):
                    setattr(esmfold_config.trunk.structure_module       , _a       , _a      )
     elif hasattr(_a       , """args"""      ):
          # Indicates an ESM-1b or ESM-1v model
          lowerCAmelCase_       :							Any			 =		esm.args.embed_dim
          lowerCAmelCase_       :							Dict			 =		esm.args.layers
          lowerCAmelCase_       :							Optional[int]			 =		esm.args.attention_heads
          lowerCAmelCase_       :							List[str]			 =		esm.args.ffn_embed_dim
          lowerCAmelCase_       :							Union[str, Any]			 =		esm.args.token_dropout
          lowerCAmelCase_       :							Optional[int]			 =		True if esm.emb_layer_norm_before else False
          lowerCAmelCase_       :							Union[str, Any]			 =		"""absolute"""
          lowerCAmelCase_       :							List[str]			 =		False
          lowerCAmelCase_       :							List[str]			 =		None
     else:
          # Indicates an ESM-2 model
          lowerCAmelCase_       :							Any			 =		esm.embed_dim
          lowerCAmelCase_       :							Any			 =		esm.num_layers
          lowerCAmelCase_       :							int			 =		esm.attention_heads
          lowerCAmelCase_       :							Union[str, Any]			 =		4 * embed_dim  # This is hardcoded in ESM-2
          lowerCAmelCase_       :							int			 =		esm.token_dropout
          lowerCAmelCase_       :							Any			 =		False  # This code path does not exist in ESM-2
          lowerCAmelCase_       :							List[str]			 =		"""rotary"""
          lowerCAmelCase_       :							Tuple			 =		False
          lowerCAmelCase_       :							Union[str, Any]			 =		None

     if is_folding_model:
          lowerCAmelCase_       :							Union[str, Any]			 =		esm.esm.alphabet
     lowerCAmelCase_       :							Any			 =		tuple(alphabet.all_toks      )
     lowerCAmelCase_       :							Union[str, Any]			 =		alphabet.mask_idx
     lowerCAmelCase_       :							Optional[Any]			 =		alphabet.padding_idx

     if is_folding_model:
          lowerCAmelCase_       :							Dict			 =		esm.esm
     else:
          lowerCAmelCase_       :							List[str]			 =		esm

     lowerCAmelCase_       :							Any			 =		EsmConfig(
         vocab_size=original_esm_model.embed_tokens.num_embeddings       , mask_token_id=_a       , hidden_size=_a       , num_hidden_layers=_a       , num_attention_heads=_a       , intermediate_size=_a       , max_position_embeddings=10_26       , layer_norm_eps=1E-5       , attention_probs_dropout_prob=0.0       , hidden_dropout_prob=0.0       , pad_token_id=_a       , emb_layer_norm_before=_a       , token_dropout=_a       , position_embedding_type=_a       , is_folding_model=_a       , esmfold_config=_a       , vocab_list=_a       , )
     if classification_head:
          lowerCAmelCase_       :							Tuple			 =		esm.classification_heads["""mnli"""].out_proj.weight.shape[0]
     print("""Our ESM config:"""       , _a      )

     if model.startswith("""esmfold"""      ):
          lowerCAmelCase_       :							int			 =		EsmForProteinFolding
     elif classification_head:
          lowerCAmelCase_       :							str			 =		EsmForSequenceClassification
     else:
          lowerCAmelCase_       :							Optional[int]			 =		EsmForMaskedLM
     lowerCAmelCase_       :							str			 =		model_class(_a      )
     model.eval()

     # Now let's copy all the weights.
     # Embeddings
     lowerCAmelCase_       :							Union[str, Any]			 =		original_esm_model.embed_tokens.weight
     if position_embedding_type == "absolute":
          lowerCAmelCase_       :							int			 =		original_esm_model.embed_positions.weight

     if config.emb_layer_norm_before:
          lowerCAmelCase_       :							List[Any]			 =		original_esm_model.emb_layer_norm_before.weight
          lowerCAmelCase_       :							Dict			 =		original_esm_model.emb_layer_norm_before.bias

     lowerCAmelCase_       :							List[Any]			 =		original_esm_model.emb_layer_norm_after.weight
     lowerCAmelCase_       :							Optional[Any]			 =		original_esm_model.emb_layer_norm_after.bias

     for i in range(config.num_hidden_layers      ):
          # Encoder: start of layer
          lowerCAmelCase_       :							EsmLayer			 =		model.esm.encoder.layer[i]
          # esm_layer: TransformerSentenceEncoderLayer = original_esm_model.layers[i]
          lowerCAmelCase_       :							Optional[int]			 =		original_esm_model.layers[i]

          # self attention
          lowerCAmelCase_       :							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)      )
          )

          lowerCAmelCase_       :							Tuple			 =		esm_layer.self_attn.q_proj.weight
          lowerCAmelCase_       :							List[str]			 =		esm_layer.self_attn.q_proj.bias
          lowerCAmelCase_       :							Optional[int]			 =		esm_layer.self_attn.k_proj.weight
          lowerCAmelCase_       :							List[str]			 =		esm_layer.self_attn.k_proj.bias
          lowerCAmelCase_       :							List[str]			 =		esm_layer.self_attn.v_proj.weight
          lowerCAmelCase_       :							Optional[int]			 =		esm_layer.self_attn.v_proj.bias

          if getattr(esm_layer.self_attn       , """rot_emb"""       , _a      ) is not None:
               # 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.
               lowerCAmelCase_       :							Optional[Any]			 =		esm_layer.self_attn.rot_emb.inv_freq

          # LayerNorm changes for pre-activation
          lowerCAmelCase_       :							int			 =		esm_layer.self_attn_layer_norm.weight
          lowerCAmelCase_       :							List[Any]			 =		esm_layer.self_attn_layer_norm.bias
          lowerCAmelCase_       :							Tuple			 =		esm_layer.final_layer_norm.weight
          lowerCAmelCase_       :							Union[str, Any]			 =		esm_layer.final_layer_norm.bias

          # self-attention output
          lowerCAmelCase_       :							EsmSelfOutput			 =		layer.attention.output
          assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
          lowerCAmelCase_       :							Any			 =		esm_layer.self_attn.out_proj.weight
          lowerCAmelCase_       :							List[Any]			 =		esm_layer.self_attn.out_proj.bias

          # intermediate
          lowerCAmelCase_       :							EsmIntermediate			 =		layer.intermediate
          assert intermediate.dense.weight.shape == esm_layer.fca.weight.shape
          lowerCAmelCase_       :							Optional[int]			 =		esm_layer.fca.weight
          lowerCAmelCase_       :							Tuple			 =		esm_layer.fca.bias

          # output
          lowerCAmelCase_       :							EsmOutput			 =		layer.output
          assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
          lowerCAmelCase_       :							Optional[Any]			 =		esm_layer.fca.weight
          lowerCAmelCase_       :							Union[str, Any]			 =		esm_layer.fca.bias
          # end of layer

     if is_folding_model:
          lowerCAmelCase_       :							str			 =		esm.esm_s_combine.data
          lowerCAmelCase_       :							Optional[int]			 =		esm.afa_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:
          lowerCAmelCase_       :							List[str]			 =		esm.esm.classification_heads["""mnli"""].dense.weight
          lowerCAmelCase_       :							List[Any]			 =		esm.classification_heads["""mnli"""].dense.bias
          lowerCAmelCase_       :							Optional[Any]			 =		esm.classification_heads["""mnli"""].out_proj.weight
          lowerCAmelCase_       :							List[Any]			 =		esm.classification_heads["""mnli"""].out_proj.bias
     else:
          # LM Head
          lowerCAmelCase_       :							List[str]			 =		esm.lm_head.dense.weight
          lowerCAmelCase_       :							Dict			 =		esm.lm_head.dense.bias
          lowerCAmelCase_       :							Tuple			 =		esm.lm_head.layer_norm.weight
          lowerCAmelCase_       :							List[str]			 =		esm.lm_head.layer_norm.bias
          lowerCAmelCase_       :							List[Any]			 =		esm.lm_head.weight
          lowerCAmelCase_       :							int			 =		esm.lm_head.bias

     # Contact prediction head
     transfer_and_check_weights(esm.contact_head       , model.esm.contact_head      )

     # Prepare data (first 2 sequences from ESMStructuralSplitDataset superfamily / 4)
     if is_folding_model:
          # Folding models aren't trained on masked inputs and don't like mask tokens.
          lowerCAmelCase_       :							Union[str, Any]			 =		SAMPLE_DATA[:2]
     else:
          lowerCAmelCase_       :							Optional[int]			 =		SAMPLE_DATA

     if is_folding_model:
          lowerCAmelCase_       :							Optional[Any]			 =		get_esmfold_tokenizer()
          lowerCAmelCase_       :							Dict			 =		hf_tokenizer(
              [row[1] for row in sample_data]       , return_tensors="""pt"""       , padding=_a       , add_special_tokens=_a      )
          lowerCAmelCase_ ,	lowerCAmelCase_ ,	lowerCAmelCase_ ,	lowerCAmelCase_ ,	lowerCAmelCase_       :							Dict			 =		esmfold_encode_sequences([row[1] for row in sample_data]      )
          lowerCAmelCase_       :							Optional[int]			 =		torch.all(hf_tokens["""input_ids"""] == esmfold_aas      ) and torch.all(
              hf_tokens["""attention_mask"""] == esmfold_mask      )
     else:
          # Let's check that we get the same results.
          lowerCAmelCase_       :							Tuple			 =		alphabet.get_batch_converter()
          lowerCAmelCase_ ,	lowerCAmelCase_ ,	lowerCAmelCase_       :							Tuple			 =		batch_converter(_a      )
          # Prepare tokenizer and make sure it matches
          with TemporaryDirectory() as tempdir:
               lowerCAmelCase_       :							Any			 =		"""\n""".join(alphabet.all_toks      )
               lowerCAmelCase_       :							List[str]			 =		Path(_a      ) / """vocab.txt"""
               vocab_file.write_text(_a      )
               lowerCAmelCase_       :							Union[str, Any]			 =		EsmTokenizer(vocab_file=str(_a      )      )

          lowerCAmelCase_       :							Dict			 =		hf_tokenizer([row[1] for row in sample_data]       , return_tensors="""pt"""       , padding=_a      )
          lowerCAmelCase_       :							List[Any]			 =		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:
               # 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.
               lowerCAmelCase_       :							Optional[int]			 =		esm.cuda().infer([row[1] for row in sample_data]      )
               lowerCAmelCase_       :							Union[str, Any]			 =		model.cuda()(
                   input_ids=hf_tokens["""input_ids"""].cuda()       , attention_mask=hf_tokens["""attention_mask"""].cuda()      )
          else:
               lowerCAmelCase_       :							Dict			 =		model(**_a       , output_hidden_states=_a      )
               lowerCAmelCase_       :							List[Any]			 =		our_output["""logits"""]
               if classification_head:
                    lowerCAmelCase_       :							List[str]			 =		esm.model.classification_heads["""mnli"""](esm.extract_features(_a      )      )
               else:
                    lowerCAmelCase_       :							Dict			 =		esm(hf_tokens["""input_ids"""]       , repr_layers=list(range(9_99      )      )      )
                    lowerCAmelCase_       :							Dict			 =		their_output["""logits"""]

          if is_folding_model:
               lowerCAmelCase_       :							Optional[Any]			 =		torch.max(torch.abs(our_output["""positions"""] - their_output["""positions"""]      )      ).item()
               lowerCAmelCase_       :							str			 =		torch.allclose(our_output["""positions"""]       , their_output["""positions"""]       , atol=1E-5      )
          else:
               lowerCAmelCase_       :							List[Any]			 =		torch.max(torch.abs(our_output - their_output      )      ).item()
               lowerCAmelCase_       :							str			 =		torch.allclose(_a       , _a       , atol=1E-5      )

          print(F'max_absolute_diff = {max_absolute_diff}'      )  # ~ 1e-5
          print("""Do both models output the same tensors?"""       , """🔥""" if success else """💩"""      )

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

          if not is_folding_model:
               # Let's check contact prediction too
               lowerCAmelCase_       :							Union[str, Any]			 =		model.predict_contacts(hf_tokens["""input_ids"""]       , hf_tokens["""attention_mask"""]      )
               lowerCAmelCase_       :							Optional[Any]			 =		esm.predict_contacts(hf_tokens["""input_ids"""]      )
               lowerCAmelCase_       :							Dict			 =		torch.max(torch.abs(our_output - their_output      )      ).item()
               lowerCAmelCase_       :							List[str]			 =		torch.allclose(_a       , _a       , atol=1E-5      )

               print("""Contact prediction testing:"""      )
               print(F'max_absolute_diff = {max_absolute_diff}'      )  # ~ 1e-5
               print("""Do both models output the same tensors?"""       , """🔥""" if success else """💩"""      )

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

          pathlib.Path(_a      ).mkdir(parents=_a       , exist_ok=_a      )
          print(F'Saving model to {pytorch_dump_folder_path}'      )
          model.save_pretrained(_a      )

          del esm  # Free up some memory before continuing

     print(F'Saving tokenizer to {pytorch_dump_folder_path}'      )
     hf_tokenizer.save_pretrained(_a      )

     if push_to_repo:
          model.push_to_hub(repo_id=_a       , use_auth_token=_a      )
          hf_tokenizer.push_to_hub(repo_id=_a       , use_auth_token=_a      )


if __name__ == "__main__":
      UpperCAmelCase_       : Optional[int]							=    argparse.ArgumentParser()
      # Required parameters
      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.""")
      UpperCAmelCase_       : List[str]							=    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
      )







def 				_lowerCAmelCase			(	_a :       str      )						->      int:
     lowerCAmelCase_       :							Dict			 =		0
     for plain_chr in plain_text:
          lowerCAmelCase_       :							List[Any]			 =		(
              ord(_a      ) + (hash_value << 6) + (hash_value << 16) - hash_value
          )
     return hash_value
