
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()
__SCREAMING_SNAKE_CASE			:   Any         =     logging.get_logger(__name__)

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

__SCREAMING_SNAKE_CASE			:   Tuple         =     {
    '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,
}

__SCREAMING_SNAKE_CASE			:   str         =     list('ARNDCQEGHILKMFPSTWYV')

__SCREAMING_SNAKE_CASE			:   Dict         =     restypes + ['X']
__SCREAMING_SNAKE_CASE			:   List[str]         =     restypes_with_x + ['<pad>', '<mask>', '<cls>', '<sep>', '<eos>']
def     UpperCAmelCase__							(       ):


					'''simple docstring'''

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


					'''simple docstring'''

					lowerCAmelCase		:					Optional[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     UpperCAmelCase__							(       __magic_name__						:     str							,  __magic_name__						:     str							,  __magic_name__						:     bool							,  __magic_name__						:     str							,  __magic_name__						:     str  ):


					'''simple docstring'''

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

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

					if is_folding_model:
										lowerCAmelCase		:					Dict  	=     esm.esm.alphabet
					lowerCAmelCase		:					Optional[int]  	=     tuple(alphabet.all_toks  )
					lowerCAmelCase		:					Optional[int]  	=     alphabet.mask_idx
					lowerCAmelCase		:					str  	=     alphabet.padding_idx

					if is_folding_model:
										lowerCAmelCase		:					Any  	=     esm.esm
					else:
										lowerCAmelCase		:					Any  	=     esm

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

					if model.startswith('''esmfold'''  ):
										lowerCAmelCase		:					int  	=     EsmForProteinFolding
					elif classification_head:
										lowerCAmelCase		:					Optional[int]  	=     EsmForSequenceClassification
					else:
										lowerCAmelCase		:					Optional[Any]  	=     EsmForMaskedLM
					lowerCAmelCase		:					Tuple  	=     model_class(__magic_name__  )
					model.eval()

					# Now let's copy all the weights.
					# Embeddings
					lowerCAmelCase		:					int  	=     original_esm_model.embed_tokens.weight
					if position_embedding_type == "absolute":
										lowerCAmelCase		:					Optional[int]  	=     original_esm_model.embed_positions.weight

					if config.emb_layer_norm_before:
										lowerCAmelCase		:					Union[str, Any]  	=     original_esm_model.emb_layer_norm_before.weight
										lowerCAmelCase		:					Union[str, Any]  	=     original_esm_model.emb_layer_norm_before.bias

					lowerCAmelCase		:					Optional[int]  	=     original_esm_model.emb_layer_norm_after.weight
					lowerCAmelCase		:					str  	=     original_esm_model.emb_layer_norm_after.bias

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

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

										lowerCAmelCase		:					Tuple  	=     esm_layer.self_attn.q_proj.weight
										lowerCAmelCase		:					Union[str, Any]  	=     esm_layer.self_attn.q_proj.bias
										lowerCAmelCase		:					str  	=     esm_layer.self_attn.k_proj.weight
										lowerCAmelCase		:					List[Any]  	=     esm_layer.self_attn.k_proj.bias
										lowerCAmelCase		:					Tuple  	=     esm_layer.self_attn.v_proj.weight
										lowerCAmelCase		:					Tuple  	=     esm_layer.self_attn.v_proj.bias

										if getattr(esm_layer.self_attn							,  '''rot_emb'''							,  __magic_name__  ) 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.
															lowerCAmelCase		:					Dict  	=     esm_layer.self_attn.rot_emb.inv_freq

										# LayerNorm changes for pre-activation
										lowerCAmelCase		:					Optional[int]  	=     esm_layer.self_attn_layer_norm.weight
										lowerCAmelCase		:					str  	=     esm_layer.self_attn_layer_norm.bias
										lowerCAmelCase		:					List[str]  	=     esm_layer.final_layer_norm.weight
										lowerCAmelCase		:					Dict  	=     esm_layer.final_layer_norm.bias

										# self-attention output
										lowerCAmelCase		:					EsmSelfOutput  	=     layer.attention.output
										assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
										lowerCAmelCase		:					int  	=     esm_layer.self_attn.out_proj.weight
										lowerCAmelCase		:					Tuple  	=     esm_layer.self_attn.out_proj.bias

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

										# output
										lowerCAmelCase		:					EsmOutput  	=     layer.output
										assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
										lowerCAmelCase		:					Dict  	=     esm_layer.fca.weight
										lowerCAmelCase		:					Union[str, Any]  	=     esm_layer.fca.bias
										# end of layer

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

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

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

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

										del esm  # Free up some memory before continuing

					print(f'''Saving tokenizer to {pytorch_dump_folder_path}'''  )
					hf_tokenizer.save_pretrained(__magic_name__  )

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


if __name__ == "__main__":
		__SCREAMING_SNAKE_CASE			:   Union[str, Any]         =     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.')
		__SCREAMING_SNAKE_CASE			:   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
		)


from ..utils import requires_backends


__SCREAMING_SNAKE_CASE			:   Union[str, Any]         =     None


class 		__magic_name__   :





		def __init__(							self    :	Optional[int]						,  *lowerCamelCase__    :	List[str]						,  **lowerCamelCase__    :	Optional[Any]				):
							requires_backends(self						,  ['''detectron2''']				)





		@classmethod
		def       _A							(							cls    :	Dict						,  *lowerCamelCase__    :	List[str]						,  **lowerCamelCase__    :	Optional[Any]				):
							requires_backends(cls						,  ['''detectron2''']				)
