


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()
_snake_case				=			logging.get_logger(__name__)

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

_snake_case				=			{
    "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,
}

_snake_case				=			list("ARNDCQEGHILKMFPSTWYV")

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







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







def 	A (		A__		,		A__			):
     UpperCAmelCase    :					Any			=      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 	A (		A__		,		A__		,		A__		,		A__		,		A__			):
     if model.startswith("esmfold"			):
          UpperCAmelCase    :					str			=      MODEL_MAPPING[model]()
     else:
          UpperCAmelCase	,			UpperCAmelCase    :					Any			=      MODEL_MAPPING[model]()
     esm.eval()  # disable dropout

     if model.startswith("esmfold"			):
          UpperCAmelCase    :					List[str]			=      esm.esm.embed_dim
          UpperCAmelCase    :					List[Any]			=      esm.esm.num_layers
          UpperCAmelCase    :					List[str]			=      esm.esm.attention_heads
          UpperCAmelCase    :					Optional[int]			=      4 * embed_dim
          UpperCAmelCase    :					Optional[int]			=      esm.esm.token_dropout
          UpperCAmelCase    :					List[str]			=      False  # This code path does not exist in ESM-2
          UpperCAmelCase    :					Tuple			=      "rotary"
          UpperCAmelCase    :					List[str]			=      True
          UpperCAmelCase    :					str			=      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
          UpperCAmelCase    :					Tuple			=      esm.args.embed_dim
          UpperCAmelCase    :					List[Any]			=      esm.args.layers
          UpperCAmelCase    :					str			=      esm.args.attention_heads
          UpperCAmelCase    :					List[str]			=      esm.args.ffn_embed_dim
          UpperCAmelCase    :					int			=      esm.args.token_dropout
          UpperCAmelCase    :					int			=      True if esm.emb_layer_norm_before else False
          UpperCAmelCase    :					Optional[int]			=      "absolute"
          UpperCAmelCase    :					Union[str, Any]			=      False
          UpperCAmelCase    :					Any			=      None
     else:
          # Indicates an ESM-2 model
          UpperCAmelCase    :					Optional[Any]			=      esm.embed_dim
          UpperCAmelCase    :					Dict			=      esm.num_layers
          UpperCAmelCase    :					Union[str, Any]			=      esm.attention_heads
          UpperCAmelCase    :					Any			=      4 * embed_dim  # This is hardcoded in ESM-2
          UpperCAmelCase    :					Tuple			=      esm.token_dropout
          UpperCAmelCase    :					Dict			=      False  # This code path does not exist in ESM-2
          UpperCAmelCase    :					str			=      "rotary"
          UpperCAmelCase    :					List[str]			=      False
          UpperCAmelCase    :					str			=      None

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

     if is_folding_model:
          UpperCAmelCase    :					Dict			=      esm.esm
     else:
          UpperCAmelCase    :					Dict			=      esm

     UpperCAmelCase    :					List[str]			=      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=1_0_2_6		,		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:
          UpperCAmelCase    :					Union[str, Any]			=      esm.classification_heads["mnli"].out_proj.weight.shape[0]
     print("Our ESM config:"		,		A__			)

     if model.startswith("esmfold"			):
          UpperCAmelCase    :					Optional[Any]			=      EsmForProteinFolding
     elif classification_head:
          UpperCAmelCase    :					Any			=      EsmForSequenceClassification
     else:
          UpperCAmelCase    :					Optional[int]			=      EsmForMaskedLM
     UpperCAmelCase    :					Dict			=      model_class(A__			)
     model.eval()

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

     if config.emb_layer_norm_before:
          UpperCAmelCase    :					Tuple			=      original_esm_model.emb_layer_norm_before.weight
          UpperCAmelCase    :					List[str]			=      original_esm_model.emb_layer_norm_before.bias

     UpperCAmelCase    :					Optional[Any]			=      original_esm_model.emb_layer_norm_after.weight
     UpperCAmelCase    :					Dict			=      original_esm_model.emb_layer_norm_after.bias

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

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

          UpperCAmelCase    :					Any			=      esm_layer.self_attn.q_proj.weight
          UpperCAmelCase    :					int			=      esm_layer.self_attn.q_proj.bias
          UpperCAmelCase    :					Any			=      esm_layer.self_attn.k_proj.weight
          UpperCAmelCase    :					Tuple			=      esm_layer.self_attn.k_proj.bias
          UpperCAmelCase    :					Any			=      esm_layer.self_attn.v_proj.weight
          UpperCAmelCase    :					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.
               UpperCAmelCase    :					str			=      esm_layer.self_attn.rot_emb.inv_freq

          # LayerNorm changes for pre-activation
          UpperCAmelCase    :					Tuple			=      esm_layer.self_attn_layer_norm.weight
          UpperCAmelCase    :					Dict			=      esm_layer.self_attn_layer_norm.bias
          UpperCAmelCase    :					List[str]			=      esm_layer.final_layer_norm.weight
          UpperCAmelCase    :					List[Any]			=      esm_layer.final_layer_norm.bias

          # self-attention output
          UpperCAmelCase    :					EsmSelfOutput			=      layer.attention.output
          assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
          UpperCAmelCase    :					Optional[Any]			=      esm_layer.self_attn.out_proj.weight
          UpperCAmelCase    :					Union[str, Any]			=      esm_layer.self_attn.out_proj.bias

          # intermediate
          UpperCAmelCase    :					EsmIntermediate			=      layer.intermediate
          assert intermediate.dense.weight.shape == esm_layer.fca.weight.shape
          UpperCAmelCase    :					List[Any]			=      esm_layer.fca.weight
          UpperCAmelCase    :					List[str]			=      esm_layer.fca.bias

          # output
          UpperCAmelCase    :					EsmOutput			=      layer.output
          assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
          UpperCAmelCase    :					str			=      esm_layer.fca.weight
          UpperCAmelCase    :					int			=      esm_layer.fca.bias
          # end of layer

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

     if is_folding_model:
          UpperCAmelCase    :					Optional[int]			=      get_esmfold_tokenizer()
          UpperCAmelCase    :					int			=      hf_tokenizer(
              [row[1] for row in sample_data]		,		return_tensors="pt"		,		padding=A__		,		add_special_tokens=A__			)
          UpperCAmelCase	,			UpperCAmelCase	,			UpperCAmelCase	,			UpperCAmelCase	,			UpperCAmelCase    :					List[str]			=      esmfold_encode_sequences([row[1] for row in sample_data]			)
          UpperCAmelCase    :					Dict			=      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.
          UpperCAmelCase    :					Tuple			=      alphabet.get_batch_converter()
          UpperCAmelCase	,			UpperCAmelCase	,			UpperCAmelCase    :					int			=      batch_converter(A__			)
          # Prepare tokenizer and make sure it matches
          with TemporaryDirectory() as tempdir:
               UpperCAmelCase    :					Dict			=      "\n".join(alphabet.all_toks			)
               UpperCAmelCase    :					Tuple			=      Path(A__			) / "vocab.txt"
               vocab_file.write_text(A__			)
               UpperCAmelCase    :					str			=      EsmTokenizer(vocab_file=str(A__			)			)

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

          if is_folding_model:
               UpperCAmelCase    :					List[str]			=      torch.max(torch.abs(our_output["positions"] - their_output["positions"]			)			).item()
               UpperCAmelCase    :					Any			=      torch.allclose(our_output["positions"]		,		their_output["positions"]		,		atol=1E-5			)
          else:
               UpperCAmelCase    :					int			=      torch.max(torch.abs(our_output - their_output			)			).item()
               UpperCAmelCase    :					List[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
               UpperCAmelCase    :					List[str]			=      model.predict_contacts(hf_tokens["input_ids"]		,		hf_tokens["attention_mask"]			)
               UpperCAmelCase    :					List[str]			=      esm.predict_contacts(hf_tokens["input_ids"]			)
               UpperCAmelCase    :					List[Any]			=      torch.max(torch.abs(our_output - their_output			)			).item()
               UpperCAmelCase    :					Optional[Any]			=      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__":
      _snake_case				=			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.")
      _snake_case				=			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
      )







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


_snake_case				=			{
    "configuration_nystromformer": ["NYSTROMFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP", "NystromformerConfig"],
}

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      _snake_case				=			[
          "NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
          "NystromformerForMaskedLM",
          "NystromformerForMultipleChoice",
          "NystromformerForQuestionAnswering",
          "NystromformerForSequenceClassification",
          "NystromformerForTokenClassification",
          "NystromformerLayer",
          "NystromformerModel",
          "NystromformerPreTrainedModel",
      ]


if TYPE_CHECKING:
      from .configuration_nystromformer import NYSTROMFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, NystromformerConfig

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_nystromformer import (
                NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
                NystromformerForMaskedLM,
                NystromformerForMultipleChoice,
                NystromformerForQuestionAnswering,
                NystromformerForSequenceClassification,
                NystromformerForTokenClassification,
                NystromformerLayer,
                NystromformerModel,
                NystromformerPreTrainedModel,
            )


else:
      import sys

      _snake_case				=			_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



