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

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

lowerCamelCase__   = {
    '''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,
}

lowerCamelCase__   = list('''ARNDCQEGHILKMFPSTWYV''')

lowerCamelCase__   = restypes + ['''X''']
lowerCamelCase__   = restypes_with_x + ['''<pad>''', '''<mask>''', '''<cls>''', '''<sep>''', '''<eos>''']


def 							lowerCAmelCase__    (       )	->List[str]:




  '''simple docstring'''

  with TemporaryDirectory() as tempdir:
    _UpperCamelCase				=  "\n".join(a__      )
    _UpperCamelCase				=  Path(a__      ) / "vocab.txt"
    vocab_file.write_text(a__      )
    _UpperCamelCase				=  EsmTokenizer(vocab_file=str(a__      )      )
  _UpperCamelCase				=  0  # Overlaps with 'A' but that seems to be what they want
  return hf_tokenizer


def 							lowerCAmelCase__    (       a__  ,			a__      )	->str:




  '''simple docstring'''

  _UpperCamelCase				=  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__  ,			a__  ,			a__  ,			a__  ,			a__      )	->Tuple:




  '''simple docstring'''

  if model.startswith("esmfold"      ):
    _UpperCamelCase				=  MODEL_MAPPING[model]()
  else:
    _UpperCamelCase	,		_UpperCamelCase				=  MODEL_MAPPING[model]()
  esm.eval()  # disable dropout

  if model.startswith("esmfold"      ):
    _UpperCamelCase				=  esm.esm.embed_dim
    _UpperCamelCase				=  esm.esm.num_layers
    _UpperCamelCase				=  esm.esm.attention_heads
    _UpperCamelCase				=  4 * embed_dim
    _UpperCamelCase				=  esm.esm.token_dropout
    _UpperCamelCase				=  False  # This code path does not exist in ESM-2
    _UpperCamelCase				=  "rotary"
    _UpperCamelCase				=  True
    _UpperCamelCase				=  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				=  esm.args.embed_dim
    _UpperCamelCase				=  esm.args.layers
    _UpperCamelCase				=  esm.args.attention_heads
    _UpperCamelCase				=  esm.args.ffn_embed_dim
    _UpperCamelCase				=  esm.args.token_dropout
    _UpperCamelCase				=  True if esm.emb_layer_norm_before else False
    _UpperCamelCase				=  "absolute"
    _UpperCamelCase				=  False
    _UpperCamelCase				=  None
  else:
    # Indicates an ESM-2 model
    _UpperCamelCase				=  esm.embed_dim
    _UpperCamelCase				=  esm.num_layers
    _UpperCamelCase				=  esm.attention_heads
    _UpperCamelCase				=  4 * embed_dim  # This is hardcoded in ESM-2
    _UpperCamelCase				=  esm.token_dropout
    _UpperCamelCase				=  False  # This code path does not exist in ESM-2
    _UpperCamelCase				=  "rotary"
    _UpperCamelCase				=  False
    _UpperCamelCase				=  None

  if is_folding_model:
    _UpperCamelCase				=  esm.esm.alphabet
  _UpperCamelCase				=  tuple(alphabet.all_toks      )
  _UpperCamelCase				=  alphabet.mask_idx
  _UpperCamelCase				=  alphabet.padding_idx

  if is_folding_model:
    _UpperCamelCase				=  esm.esm
  else:
    _UpperCamelCase				=  esm

  _UpperCamelCase				=  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_026  ,			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				=  esm.classification_heads["mnli"].out_proj.weight.shape[0]
  print("Our ESM config:"  ,			a__      )

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

  # Now let's copy all the weights.
  # Embeddings
  _UpperCamelCase				=  original_esm_model.embed_tokens.weight
  if position_embedding_type == "absolute":
    _UpperCamelCase				=  original_esm_model.embed_positions.weight

  if config.emb_layer_norm_before:
    _UpperCamelCase				=  original_esm_model.emb_layer_norm_before.weight
    _UpperCamelCase				=  original_esm_model.emb_layer_norm_before.bias

  _UpperCamelCase				=  original_esm_model.emb_layer_norm_after.weight
  _UpperCamelCase				=  original_esm_model.emb_layer_norm_after.bias

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

    # self attention
    _UpperCamelCase				=  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				=  esm_layer.self_attn.q_proj.weight
    _UpperCamelCase				=  esm_layer.self_attn.q_proj.bias
    _UpperCamelCase				=  esm_layer.self_attn.k_proj.weight
    _UpperCamelCase				=  esm_layer.self_attn.k_proj.bias
    _UpperCamelCase				=  esm_layer.self_attn.v_proj.weight
    _UpperCamelCase				=  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				=  esm_layer.self_attn.rot_emb.inv_freq

    # LayerNorm changes for pre-activation
    _UpperCamelCase				=  esm_layer.self_attn_layer_norm.weight
    _UpperCamelCase				=  esm_layer.self_attn_layer_norm.bias
    _UpperCamelCase				=  esm_layer.final_layer_norm.weight
    _UpperCamelCase				=  esm_layer.final_layer_norm.bias

    # self-attention output
    _UpperCamelCase				=  layer.attention.output
    assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
    _UpperCamelCase				=  esm_layer.self_attn.out_proj.weight
    _UpperCamelCase				=  esm_layer.self_attn.out_proj.bias

    # intermediate
    _UpperCamelCase				=  layer.intermediate
    assert intermediate.dense.weight.shape == esm_layer.fca.weight.shape
    _UpperCamelCase				=  esm_layer.fca.weight
    _UpperCamelCase				=  esm_layer.fca.bias

    # output
    _UpperCamelCase				=  layer.output
    assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
    _UpperCamelCase				=  esm_layer.fca.weight
    _UpperCamelCase				=  esm_layer.fca.bias
    # end of layer

  if is_folding_model:
    _UpperCamelCase				=  esm.esm_s_combine.data
    _UpperCamelCase				=  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				=  esm.esm.classification_heads["mnli"].dense.weight
    _UpperCamelCase				=  esm.classification_heads["mnli"].dense.bias
    _UpperCamelCase				=  esm.classification_heads["mnli"].out_proj.weight
    _UpperCamelCase				=  esm.classification_heads["mnli"].out_proj.bias
  else:
    # LM Head
    _UpperCamelCase				=  esm.lm_head.dense.weight
    _UpperCamelCase				=  esm.lm_head.dense.bias
    _UpperCamelCase				=  esm.lm_head.layer_norm.weight
    _UpperCamelCase				=  esm.lm_head.layer_norm.bias
    _UpperCamelCase				=  esm.lm_head.weight
    _UpperCamelCase				=  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				=  SAMPLE_DATA[:2]
  else:
    _UpperCamelCase				=  SAMPLE_DATA

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

    _UpperCamelCase				=  hf_tokenizer([row[1] for row in sample_data]  ,			return_tensors="pt"  ,			padding=a__      )
    _UpperCamelCase				=  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				=  esm.cuda().infer([row[1] for row in sample_data]      )
      _UpperCamelCase				=  model.cuda()(
          input_ids=hf_tokens["input_ids"].cuda()  ,			attention_mask=hf_tokens["attention_mask"].cuda()      )
    else:
      _UpperCamelCase				=  model(**a__  ,			output_hidden_states=a__      )
      _UpperCamelCase				=  our_output["logits"]
      if classification_head:
        _UpperCamelCase				=  esm.model.classification_heads["mnli"](esm.extract_features(a__      )      )
      else:
        _UpperCamelCase				=  esm(hf_tokens["input_ids"]  ,			repr_layers=list(range(999      )      )      )
        _UpperCamelCase				=  their_output["logits"]

    if is_folding_model:
      _UpperCamelCase				=  torch.max(torch.abs(our_output["positions"] - their_output["positions"]      )      ).item()
      _UpperCamelCase				=  torch.allclose(our_output["positions"]  ,			their_output["positions"]  ,			atol=1e-5      )
    else:
      _UpperCamelCase				=  torch.max(torch.abs(our_output - their_output      )      ).item()
      _UpperCamelCase				=  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				=  model.predict_contacts(hf_tokens["input_ids"]  ,			hf_tokens["attention_mask"]      )
      _UpperCamelCase				=  esm.predict_contacts(hf_tokens["input_ids"]      )
      _UpperCamelCase				=  torch.max(torch.abs(our_output - their_output      )      ).item()
      _UpperCamelCase				=  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__":
   lowerCamelCase__   = 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.''')
   lowerCamelCase__   = 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

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available, is_torch_available


lowerCamelCase__   = {
    '''configuration_ernie_m''': ['''ERNIE_M_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ErnieMConfig'''],
}

try:
   if not is_sentencepiece_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   lowerCamelCase__   = ['''ErnieMTokenizer''']

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   lowerCamelCase__   = [
       '''ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST''',
       '''ErnieMForMultipleChoice''',
       '''ErnieMForQuestionAnswering''',
       '''ErnieMForSequenceClassification''',
       '''ErnieMForTokenClassification''',
       '''ErnieMModel''',
       '''ErnieMPreTrainedModel''',
       '''ErnieMForInformationExtraction''',
   ]


if TYPE_CHECKING:
   from .configuration_ernie_m import ERNIE_M_PRETRAINED_CONFIG_ARCHIVE_MAP, ErnieMConfig

   try:
      if not is_sentencepiece_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_ernie_m import ErnieMTokenizer

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_ernie_m import (
          ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST,
          ErnieMForInformationExtraction,
          ErnieMForMultipleChoice,
          ErnieMForQuestionAnswering,
          ErnieMForSequenceClassification,
          ErnieMForTokenClassification,
          ErnieMModel,
          ErnieMPreTrainedModel,
      )


else:
   import sys

   lowerCamelCase__   = _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)




