



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

UpperCAmelCase :				Dict							=[
    (
        """protein1""",
        """MNGTEGPNFYVPFSNATGVVRSPFEYPQYYLAEPWQFSMLAAYMFLLIVLGFPINFLTLYVTVQHKKLRTPLNYILLNLAVADLFMVLGGFTSTLYTSLHGYFVFGPTGCNLEGFFATLGGEIALWSLVVLAIERYVVVCKPMSNFRFGENHAIMGVAFTWVMALACAAPPLAGWSRYIPEGLQCSCGIDYYTLKPEVNNESFVIYMFVVHFTIPMIIIFFCYGQLVFTVKEAAAQQQESATTQKAEKEVTRMVIIMVIAFLICWVPYASVAFYIFTHQGSNFGPIFMTIPAFFAKSAAIYNPVIYIMMNKQFRNCMLTTICCGKNPLGDDEASATVSKTETSQVAPA""",
    ),
    ("""protein2""", """MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLA"""),
    ("""protein3""", """MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLAGG"""),
    ("""protein4""", """MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLA"""),
]

UpperCAmelCase :				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,
}

UpperCAmelCase :				Tuple							=list("""ARNDCQEGHILKMFPSTWYV""")

UpperCAmelCase :				Dict							=restypes + ["""X"""]
UpperCAmelCase :				int							=restypes_with_x + ["""<pad>""", """<mask>""", """<cls>""", """<sep>""", """<eos>"""]



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



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

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

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

										del esm  # Free up some memory before continuing

					print(f"""Saving tokenizer to {pytorch_dump_folder_path}""")
					hf_tokenizer.save_pretrained(_lowerCAmelCase)

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


if __name__ == "__main__":
	UpperCAmelCase :				Optional[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.""")
	UpperCAmelCase :				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'''







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


SCREAMING_SNAKE_CASE__     : Dict				=   {"""configuration_detr""": ["""DETR_PRETRAINED_CONFIG_ARCHIVE_MAP""", """DetrConfig""", """DetrOnnxConfig"""]}

try:
							if not is_vision_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							SCREAMING_SNAKE_CASE__     : List[Any]				=   ["""DetrFeatureExtractor"""]
							SCREAMING_SNAKE_CASE__     : str				=   ["""DetrImageProcessor"""]

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							SCREAMING_SNAKE_CASE__     : Optional[Any]				=   [
							    """DETR_PRETRAINED_MODEL_ARCHIVE_LIST""",
							    """DetrForObjectDetection""",
							    """DetrForSegmentation""",
							    """DetrModel""",
							    """DetrPreTrainedModel""",
							]


if TYPE_CHECKING:
							from .configuration_detr import DETR_PRETRAINED_CONFIG_ARCHIVE_MAP, DetrConfig, DetrOnnxConfig

							try:
														if not is_vision_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .feature_extraction_detr import DetrFeatureExtractor
														from .image_processing_detr import DetrImageProcessor

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_detr import (
														    DETR_PRETRAINED_MODEL_ARCHIVE_LIST,
														    DetrForObjectDetection,
														    DetrForSegmentation,
														    DetrModel,
														    DetrPreTrainedModel,
														)

else:
							import sys

							SCREAMING_SNAKE_CASE__     : Optional[Any]				=   _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
