



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

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

_lowercase					:					Union[str, 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,
}

_lowercase					:					Union[str, Any] 		=		list("ARNDCQEGHILKMFPSTWYV")

_lowercase					:					Any 		=		restypes + ["X"]
_lowercase					:					Union[str, Any] 		=		restypes_with_x + ["<pad>", "<mask>", "<cls>", "<sep>", "<eos>"]






def 	snake_case__  (						):
			"""simple docstring"""
			with TemporaryDirectory() as tempdir:
						lowerCamelCase__						:						Optional[int]						      ='''\n'''.join(__lowerCamelCase      )
						lowerCamelCase__						:						Any						      =Path(__lowerCamelCase      ) / '''vocab.txt'''
						vocab_file.write_text(__lowerCamelCase      )
						lowerCamelCase__						:						Union[str, Any]						      =EsmTokenizer(vocab_file=str(__lowerCamelCase      )      )
			lowerCamelCase__						:						Optional[Any]						      =0  # Overlaps with 'A' but that seems to be what they want
			return hf_tokenizer






def 	snake_case__  (						__lowerCamelCase							:       str				,						__lowerCamelCase							:       Tuple      ):
			"""simple docstring"""
			lowerCamelCase__						:						int						      =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 	snake_case__  (						__lowerCamelCase							:       str				,						__lowerCamelCase							:       str				,						__lowerCamelCase							:       bool				,						__lowerCamelCase							:       str				,						__lowerCamelCase							:       str      ):
			"""simple docstring"""
			if model.startswith('''esmfold'''      ):
						lowerCamelCase__						:						str						      =MODEL_MAPPING[model]()
			else:
						lowerCamelCase__      ,  lowerCamelCase__						:						Dict						      =MODEL_MAPPING[model]()
			esm.eval()  # disable dropout

			if model.startswith('''esmfold'''      ):
						lowerCamelCase__						:						Dict						      =esm.esm.embed_dim
						lowerCamelCase__						:						Optional[int]						      =esm.esm.num_layers
						lowerCamelCase__						:						Any						      =esm.esm.attention_heads
						lowerCamelCase__						:						Union[str, Any]						      =4 * embed_dim
						lowerCamelCase__						:						List[str]						      =esm.esm.token_dropout
						lowerCamelCase__						:						Optional[Any]						      =False  # This code path does not exist in ESM-2
						lowerCamelCase__						:						Tuple						      ='''rotary'''
						lowerCamelCase__						:						List[str]						      =True
						lowerCamelCase__						:						Optional[int]						      =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
						lowerCamelCase__						:						str						      =esm.args.embed_dim
						lowerCamelCase__						:						Any						      =esm.args.layers
						lowerCamelCase__						:						Dict						      =esm.args.attention_heads
						lowerCamelCase__						:						List[str]						      =esm.args.ffn_embed_dim
						lowerCamelCase__						:						Optional[Any]						      =esm.args.token_dropout
						lowerCamelCase__						:						int						      =True if esm.emb_layer_norm_before else False
						lowerCamelCase__						:						Optional[Any]						      ='''absolute'''
						lowerCamelCase__						:						Optional[int]						      =False
						lowerCamelCase__						:						str						      =None
			else:
						# Indicates an ESM-2 model
						lowerCamelCase__						:						Dict						      =esm.embed_dim
						lowerCamelCase__						:						Optional[int]						      =esm.num_layers
						lowerCamelCase__						:						List[Any]						      =esm.attention_heads
						lowerCamelCase__						:						Optional[Any]						      =4 * embed_dim  # This is hardcoded in ESM-2
						lowerCamelCase__						:						List[Any]						      =esm.token_dropout
						lowerCamelCase__						:						Optional[int]						      =False  # This code path does not exist in ESM-2
						lowerCamelCase__						:						List[Any]						      ='''rotary'''
						lowerCamelCase__						:						str						      =False
						lowerCamelCase__						:						Any						      =None

			if is_folding_model:
						lowerCamelCase__						:						List[str]						      =esm.esm.alphabet
			lowerCamelCase__						:						Union[str, Any]						      =tuple(alphabet.all_toks      )
			lowerCamelCase__						:						List[str]						      =alphabet.mask_idx
			lowerCamelCase__						:						Union[str, Any]						      =alphabet.padding_idx

			if is_folding_model:
						lowerCamelCase__						:						List[str]						      =esm.esm
			else:
						lowerCamelCase__						:						List[str]						      =esm

			lowerCamelCase__						:						Dict						      =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=1026				,						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:
						lowerCamelCase__						:						Union[str, Any]						      =esm.classification_heads['''mnli'''].out_proj.weight.shape[0]
			print('''Our ESM config:'''				,						__lowerCamelCase      )

			if model.startswith('''esmfold'''      ):
						lowerCamelCase__						:						str						      =EsmForProteinFolding
			elif classification_head:
						lowerCamelCase__						:						Any						      =EsmForSequenceClassification
			else:
						lowerCamelCase__						:						Optional[int]						      =EsmForMaskedLM
			lowerCamelCase__						:						Optional[Any]						      =model_class(__lowerCamelCase      )
			model.eval()

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

			if config.emb_layer_norm_before:
						lowerCamelCase__						:						Optional[Any]						      =original_esm_model.emb_layer_norm_before.weight
						lowerCamelCase__						:						Optional[Any]						      =original_esm_model.emb_layer_norm_before.bias

			lowerCamelCase__						:						Optional[Any]						      =original_esm_model.emb_layer_norm_after.weight
			lowerCamelCase__						:						Any						      =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__						:						Dict						      =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__						:						Union[str, Any]						      =esm_layer.self_attn.q_proj.weight
						lowerCamelCase__						:						Dict						      =esm_layer.self_attn.q_proj.bias
						lowerCamelCase__						:						Optional[Any]						      =esm_layer.self_attn.k_proj.weight
						lowerCamelCase__						:						str						      =esm_layer.self_attn.k_proj.bias
						lowerCamelCase__						:						Union[str, Any]						      =esm_layer.self_attn.v_proj.weight
						lowerCamelCase__						:						Tuple						      =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.
									lowerCamelCase__						:						str						      =esm_layer.self_attn.rot_emb.inv_freq

						# LayerNorm changes for pre-activation
						lowerCamelCase__						:						Optional[int]						      =esm_layer.self_attn_layer_norm.weight
						lowerCamelCase__						:						Optional[Any]						      =esm_layer.self_attn_layer_norm.bias
						lowerCamelCase__						:						Union[str, Any]						      =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__						:						Any						      =esm_layer.self_attn.out_proj.weight
						lowerCamelCase__						:						Optional[int]						      =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__						:						Optional[Any]						      =esm_layer.fca.bias

						# output
						lowerCamelCase__						:						EsmOutput						      =layer.output
						assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
						lowerCamelCase__						:						Optional[Any]						      =esm_layer.fca.weight
						lowerCamelCase__						:						Optional[int]						      =esm_layer.fca.bias
						# end of layer

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

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

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

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






def 	snake_case__  (						__lowerCamelCase							:       str				,						__lowerCamelCase							:       int = 13      ):
			"""simple docstring"""
			lowerCamelCase__						:						List[str]						      =''''''
			for c in s:
						if "A" <= c <= "Z":
									out += chr(ord('''A'''      ) + (ord(__lowerCamelCase      ) - ord('''A'''      ) + n) % 26      )
						elif "a" <= c <= "z":
									out += chr(ord('''a'''      ) + (ord(__lowerCamelCase      ) - ord('''a'''      ) + n) % 26      )
						else:
									out += c
			return out






def 	snake_case__  (						):
			"""simple docstring"""
			lowerCamelCase__						:						List[str]						      =input('''Enter message: '''      )

			lowerCamelCase__						:						str						      =dencrypt(__lowerCamelCase				,						13      )
			print('''Encryption:'''				,						__lowerCamelCase      )

			lowerCamelCase__						:						Any						      =dencrypt(__lowerCamelCase				,						13      )
			print('''Decryption: '''				,						__lowerCamelCase      )


if __name__ == "__main__":
		import doctest

		doctest.testmod()
		main()




