

'''simple docstring'''

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							(		):






     '''simple docstring'''





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

def 			__UpperCamelCase							(		lowercase__  :					str,   lowercase__  :					Optional[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							(		lowercase__  :					str,   lowercase__  :					str,   lowercase__  :					bool,   lowercase__  :					str,   lowercase__  :					str	):






     '''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(lowercase__,   lowercase__	) and key != "trunk":
                    setattr(lowercase__,   lowercase__,   lowercase__	)
          for key, val in esm.cfg.trunk.items():
               if hasattr(esmfold_config.trunk,   lowercase__	) and key != "structure_module":
                    setattr(esmfold_config.trunk,   lowercase__,   lowercase__	)
          for key, val in esm.cfg.trunk.structure_module.items():
               if hasattr(esmfold_config.trunk.structure_module,   lowercase__	):
                    setattr(esmfold_config.trunk.structure_module,   lowercase__,   lowercase__	)
     elif hasattr(lowercase__,   '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=lowercase__,   hidden_size=lowercase__,   num_hidden_layers=lowercase__,   num_attention_heads=lowercase__,   intermediate_size=lowercase__,   max_position_embeddings=10_26,   layer_norm_eps=1E-5,   attention_probs_dropout_prob=0.0,   hidden_dropout_prob=0.0,   pad_token_id=lowercase__,   emb_layer_norm_before=lowercase__,   token_dropout=lowercase__,   position_embedding_type=lowercase__,   is_folding_model=lowercase__,   esmfold_config=lowercase__,   vocab_list=lowercase__,   )
     if classification_head:
          __lowercase     						=esm.classification_heads['mnli'].out_proj.weight.shape[0]
     print('Our ESM config:',   lowercase__	)

     if model.startswith('esmfold'	):
          __lowercase     						=EsmForProteinFolding
     elif classification_head:
          __lowercase     						=EsmForSequenceClassification
     else:
          __lowercase     						=EsmForMaskedLM
     __lowercase     						=model_class(lowercase__	)
     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',   lowercase__	) 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=lowercase__,   add_special_tokens=lowercase__	)
          __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(lowercase__	)
          # Prepare tokenizer and make sure it matches
          with TemporaryDirectory() as tempdir:
               __lowercase     						='\n'.join(alphabet.all_toks	)
               __lowercase     						=Path(lowercase__	) / 'vocab.txt'
               vocab_file.write_text(lowercase__	)
               __lowercase     						=EsmTokenizer(vocab_file=str(lowercase__	)	)

          __lowercase     						=hf_tokenizer([row[1] for row in sample_data],   return_tensors='pt',   padding=lowercase__	)
          __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(**lowercase__,   output_hidden_states=lowercase__	)
               __lowercase     						=our_output['logits']
               if classification_head:
                    __lowercase     						=esm.model.classification_heads['mnli'](esm.extract_features(lowercase__	)	)
               else:
                    __lowercase     						=esm(hf_tokens['input_ids'],   repr_layers=list(range(9_99	)	)	)
                    __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(lowercase__,   lowercase__,   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(lowercase__,   lowercase__,   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(lowercase__	).mkdir(parents=lowercase__,   exist_ok=lowercase__	)
          print(F'''Saving model to {pytorch_dump_folder_path}'''	)
          model.save_pretrained(lowercase__	)

          del esm  # Free up some memory before continuing

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

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


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
       )





'''simple docstring'''

from __future__ import annotations

def 			__UpperCamelCase							(		lowercase__  :					list[int]	):






     '''simple docstring'''





     if not list_of_ints:
          return []

     if min(lowercase__	) < 0:
          raise ValueError('All numbers must be positive'	)

     __lowercase     						=max(len(bin(lowercase__	)[2:]	) for x in list_of_ints	)
     return _msd_radix_sort(lowercase__,   lowercase__	)

def 			__UpperCamelCase							(		lowercase__  :					list[int],   lowercase__  :					int	):






     '''simple docstring'''





     if bit_position == 0 or len(lowercase__	) in [0, 1]:
          return list_of_ints

     __lowercase     						=[]
     __lowercase     						=[]
     # Split numbers based on bit at bit_position from the right
     for number in list_of_ints:
          if (number >> (bit_position - 1)) & 1:
               # number has a one at bit bit_position
               ones.append(lowercase__	)
          else:
               # number has a zero at bit bit_position
               zeros.append(lowercase__	)

    # recursively split both lists further
     __lowercase     						=_msd_radix_sort(lowercase__,   bit_position - 1	)
     __lowercase     						=_msd_radix_sort(lowercase__,   bit_position - 1	)

     # recombine lists
     __lowercase     						=zeros
     res.extend(lowercase__	)

     return res

def 			__UpperCamelCase							(		lowercase__  :					list[int]	):






     '''simple docstring'''





     __lowercase     						=len(lowercase__	)
     if not list_of_ints or length == 1:
          return

     if min(lowercase__	) < 0:
          raise ValueError('All numbers must be positive'	)

     __lowercase     						=max(len(bin(lowercase__	)[2:]	) for x in list_of_ints	)
     _msd_radix_sort_inplace(lowercase__,   lowercase__,   0,   lowercase__	)

def 			__UpperCamelCase							(		lowercase__  :					list[int],   lowercase__  :					int,   lowercase__  :					int,   lowercase__  :					int	):






     '''simple docstring'''





     if bit_position == 0 or end_index - begin_index <= 1:
          return

     bit_position -= 1

     __lowercase     						=begin_index
     __lowercase     						=end_index - 1
     while i <= j:
          __lowercase     						=False
          if not (list_of_ints[i] >> bit_position) & 1:
               # found zero at the beginning
               i += 1
               __lowercase     						=True
          if (list_of_ints[j] >> bit_position) & 1:
               # found one at the end
               j -= 1
               __lowercase     						=True

          if changed:
               continue

          __lowercase						,       __lowercase     						=list_of_ints[j], list_of_ints[i]
          j -= 1
          if j != i:
               i += 1

     _msd_radix_sort_inplace(lowercase__,   lowercase__,   lowercase__,   lowercase__	)
     _msd_radix_sort_inplace(lowercase__,   lowercase__,   lowercase__,   lowercase__	)


if __name__ == "__main__":
       import doctest

       doctest.testmod()


