



"""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()
__snake_case	:				List[Any]  					=			logging.get_logger(__name__)

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

__snake_case	:				Optional[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,
}

__snake_case	:				Optional[int]  					=			list('ARNDCQEGHILKMFPSTWYV')

__snake_case	:				Union[str, Any]  					=			restypes + ['X']
__snake_case	:				List[Any]  					=			restypes_with_x + ['<pad>', '<mask>', '<cls>', '<sep>', '<eos>']



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



def   a_	(			__a		,							__a      ):
     A__				=   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'''      ):
          A__				=   MODEL_MAPPING[model]()
     else:
          A__	,     A__				=   MODEL_MAPPING[model]()
     esm.eval()  # disable dropout

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

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

     if is_folding_model:
          A__				=   esm.esm
     else:
          A__				=   esm

     A__				=   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=1026		,							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:
          A__				=   esm.classification_heads['''mnli'''].out_proj.weight.shape[0]
     print('''Our ESM config:'''		,							__a      )

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

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

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

     A__				=   original_esm_model.emb_layer_norm_after.weight
     A__				=   original_esm_model.emb_layer_norm_after.bias

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

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

          A__				=   esm_layer.self_attn.q_proj.weight
          A__				=   esm_layer.self_attn.q_proj.bias
          A__				=   esm_layer.self_attn.k_proj.weight
          A__				=   esm_layer.self_attn.k_proj.bias
          A__				=   esm_layer.self_attn.v_proj.weight
          A__				=   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.
               A__				=   esm_layer.self_attn.rot_emb.inv_freq

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

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

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

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

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

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

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

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







import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


__snake_case	:				Optional[Any]  					=			500_000

__snake_case					,  __snake_case	:				Any  					=			os.path.split(__file__)
__snake_case	:				Optional[Any]  					=			os.path.join(RESULTS_BASEPATH, 'results', RESULTS_FILENAME.replace('.py', '.json'))



@get_duration
def   a_	(			__a      ):
     A__				=   dataset.select(range(0		,							len(__a      )		,							2      )      )



@get_duration
def   a_	(			__a      ):
     A__				=   dataset.sort('''numbers'''      )



@get_duration
def   a_	(			__a      ):
     A__				=   dataset.shuffle()



@get_duration
def   a_	(			__a      ):
     A__				=   dataset.train_test_split(0.1      )



@get_duration
def   a_	(			__a		,							__a=10      ):
     for shard_id in range(__a      ):
          A__				=   dataset.shard(__a		,							__a      )



def   a_	(			):
     A__				=   {'''num examples''': SPEED_TEST_N_EXAMPLES}
     A__				=   (select, sort, shuffle, train_test_split, shard)
     with tempfile.TemporaryDirectory() as tmp_dir:
          print('''generating dataset'''      )
          A__				=   datasets.Features({'''text''': datasets.Value('''string'''      ), '''numbers''': datasets.Value('''float32'''      )}      )
          A__				=   generate_example_dataset(
              os.path.join(__a		,							'''dataset.arrow'''      )		,							__a		,							num_examples=__a      )
          print('''Functions'''      )
          for func in functions:
               print(func.__name__      )
               A__				=   func(__a      )

     with open(__a		,							'''wb'''      ) as f:
          f.write(json.dumps(__a      ).encode('''utf-8'''      )      )


if __name__ == "__main__":  # useful to run the profiler
 benchmark_indices_mapping()



