


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

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

UpperCAmelCase__        =				{
    '''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('''ARNDCQEGHILKMFPSTWYV''')

UpperCAmelCase__        =				restypes + ['''X''']
UpperCAmelCase__        =				restypes_with_x + ['''<pad>''', '''<mask>''', '''<cls>''', '''<sep>''', '''<eos>''']
def 							UpperCAmelCase_   (				)	->     Any:







 """simple docstring"""


 with TemporaryDirectory() as tempdir:
  _lowercase				 ='''\n'''.join(__snake_case		)
  _lowercase				 =Path(__snake_case		) / '''vocab.txt'''
  vocab_file.write_text(__snake_case		)
  _lowercase				 =EsmTokenizer(vocab_file=str(__snake_case		)		)
 _lowercase				 =0  # Overlaps with 'A' but that seems to be what they want
 return hf_tokenizer
def 							UpperCAmelCase_   (				__snake_case	,		__snake_case		)	->     int:







 """simple docstring"""


 _lowercase				 =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 							UpperCAmelCase_   (				__snake_case	,		__snake_case	,		__snake_case	,		__snake_case	,		__snake_case		)	->     List[Any]:







 """simple docstring"""


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

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

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

 if is_folding_model:
  _lowercase				 =esm.esm
 else:
  _lowercase				 =esm

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

 if model.startswith('''esmfold'''		):
  _lowercase				 =EsmForProteinFolding
 elif classification_head:
  _lowercase				 =EsmForSequenceClassification
 else:
  _lowercase				 =EsmForMaskedLM
 _lowercase				 =model_class(__snake_case		)
 model.eval()

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

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

 _lowercase				 =original_esm_model.emb_layer_norm_after.weight
 _lowercase				 =original_esm_model.emb_layer_norm_after.bias

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

  # self attention
  _lowercase				 =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)		)
  )

  _lowercase				 =esm_layer.self_attn.q_proj.weight
  _lowercase				 =esm_layer.self_attn.q_proj.bias
  _lowercase				 =esm_layer.self_attn.k_proj.weight
  _lowercase				 =esm_layer.self_attn.k_proj.bias
  _lowercase				 =esm_layer.self_attn.v_proj.weight
  _lowercase				 =esm_layer.self_attn.v_proj.bias

  if getattr(esm_layer.self_attn	,		'''rot_emb'''	,		__snake_case		) 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.
   _lowercase				 =esm_layer.self_attn.rot_emb.inv_freq

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

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

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

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

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

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

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

  if is_folding_model:
   _lowercase				 =torch.max(torch.abs(our_output['''positions'''] - their_output['''positions''']		)		).item()
   _lowercase				 =torch.allclose(our_output['''positions''']	,		their_output['''positions''']	,		atol=1e-5		)
  else:
   _lowercase				 =torch.max(torch.abs(our_output - their_output		)		).item()
   _lowercase				 =torch.allclose(__snake_case	,		__snake_case	,		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
   _lowercase				 =model.predict_contacts(hf_tokens['''input_ids''']	,		hf_tokens['''attention_mask''']		)
   _lowercase				 =esm.predict_contacts(hf_tokens['''input_ids''']		)
   _lowercase				 =torch.max(torch.abs(our_output - their_output		)		).item()
   _lowercase				 =torch.allclose(__snake_case	,		__snake_case	,		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(__snake_case		).mkdir(parents=__snake_case	,		exist_ok=__snake_case		)
  print(F"Saving model to {pytorch_dump_folder_path}"		)
  model.save_pretrained(__snake_case		)

  del esm  # Free up some memory before continuing

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

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


if __name__ == "__main__":
  UpperCAmelCase__        =				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__        =				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
  )





import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils import AddedToken
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_mpnet import MPNetTokenizer


UpperCAmelCase__        =				logging.get_logger(__name__)

UpperCAmelCase__        =				{'''vocab_file''': '''vocab.txt''', '''tokenizer_file''': '''tokenizer.json'''}

UpperCAmelCase__        =				{
    '''vocab_file''': {
        '''microsoft/mpnet-base''': '''https://huggingface.co/microsoft/mpnet-base/resolve/main/vocab.txt''',
    },
    '''tokenizer_file''': {
        '''microsoft/mpnet-base''': '''https://huggingface.co/microsoft/mpnet-base/resolve/main/tokenizer.json''',
    },
}

UpperCAmelCase__        =				{
    '''microsoft/mpnet-base''': 512,
}

UpperCAmelCase__        =				{
    '''microsoft/mpnet-base''': {'''do_lower_case''': True},
}







class  lowerCamelCase__  (			lowerCAmelCase):
     SCREAMING_SNAKE_CASE__      =     VOCAB_FILES_NAMES
     SCREAMING_SNAKE_CASE__      =     PRETRAINED_VOCAB_FILES_MAP
     SCREAMING_SNAKE_CASE__      =     PRETRAINED_INIT_CONFIGURATION
     SCREAMING_SNAKE_CASE__      =     PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     SCREAMING_SNAKE_CASE__      =     MPNetTokenizer
     SCREAMING_SNAKE_CASE__      =     ['''input_ids''', '''attention_mask''']





     def __init__(self						,	UpperCAmelCase=None						,	UpperCAmelCase=None						,	UpperCAmelCase=True						,	UpperCAmelCase="<s>"						,	UpperCAmelCase="</s>"						,	UpperCAmelCase="</s>"						,	UpperCAmelCase="<s>"						,	UpperCAmelCase="[UNK]"						,	UpperCAmelCase="<pad>"						,	UpperCAmelCase="<mask>"						,	UpperCAmelCase=True						,	UpperCAmelCase=None						,	**UpperCAmelCase						,	)	->  Optional[int]:
      super().__init__(
          UpperCAmelCase						,	tokenizer_file=UpperCAmelCase						,	do_lower_case=UpperCAmelCase						,	bos_token=UpperCAmelCase						,	eos_token=UpperCAmelCase						,	sep_token=UpperCAmelCase						,	cls_token=UpperCAmelCase						,	unk_token=UpperCAmelCase						,	pad_token=UpperCAmelCase						,	mask_token=UpperCAmelCase						,	tokenize_chinese_chars=UpperCAmelCase						,	strip_accents=UpperCAmelCase						,	**UpperCAmelCase						,	)

      _lowercase				 =json.loads(self.backend_tokenizer.normalizer.__getstate__()			)
      if (
          pre_tok_state.get('''lowercase'''						,	UpperCAmelCase			) != do_lower_case
          or pre_tok_state.get('''strip_accents'''						,	UpperCAmelCase			) != strip_accents
      ):
       _lowercase				 =getattr(UpperCAmelCase						,	pre_tok_state.pop('''type'''			)			)
       _lowercase				 =do_lower_case
       _lowercase				 =strip_accents
       _lowercase				 =pre_tok_class(**UpperCAmelCase			)

      _lowercase				 =do_lower_case





     @property
     def 	__A   (self			)	->  str:
      if self._mask_token is None:
       if self.verbose:
        logger.error('''Using mask_token, but it is not set yet.'''			)
       return None
      return str(self._mask_token			)





     @mask_token.setter
     def 	__A   (self						,	UpperCAmelCase			)	->  Optional[int]:
      _lowercase				 =AddedToken(UpperCAmelCase						,	lstrip=UpperCAmelCase						,	rstrip=UpperCAmelCase			) if isinstance(UpperCAmelCase						,	UpperCAmelCase			) else value
      _lowercase				 =value





     def 	__A   (self						,	UpperCAmelCase						,	UpperCAmelCase=None			)	->  List[str]:
      _lowercase				 =[self.bos_token_id] + token_ids_a + [self.eos_token_id]
      if token_ids_a is None:
       return output

      return output + [self.eos_token_id] + token_ids_a + [self.eos_token_id]





     def 	__A   (self						,	UpperCAmelCase						,	UpperCAmelCase = None			)	->  List[int]:
      _lowercase				 =[self.sep_token_id]
      _lowercase				 =[self.cls_token_id]

      if token_ids_a is None:
       return len(cls + token_ids_a + sep			) * [0]
      return len(cls + token_ids_a + sep + sep + token_ids_a + sep			) * [0]





     def 	__A   (self						,	UpperCAmelCase						,	UpperCAmelCase = None			)	->  Tuple[str]:
      _lowercase				 =self._tokenizer.model.save(UpperCAmelCase						,	name=UpperCAmelCase			)
      return tuple(UpperCAmelCase			)

