




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							(  ):
							with TemporaryDirectory() as tempdir:
														lowerCAmelCase =		'''\n'''.join(_snake_case				)
														lowerCAmelCase =		Path(_snake_case				) / '''vocab.txt'''
														vocab_file.write_text(_snake_case				)
														lowerCAmelCase =		EsmTokenizer(vocab_file=str(_snake_case				)				)
							lowerCAmelCase =		0  # Overlaps with 'A' but that seems to be what they want
							return hf_tokenizer



def    UpperCAmelCase							(  _snake_case    ,			_snake_case				):
							lowerCAmelCase =		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							(  _snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case				):
							if model.startswith('''esmfold'''				):
														lowerCAmelCase =		MODEL_MAPPING[model]()
							else:
														lowerCAmelCase ,  lowerCAmelCase =		MODEL_MAPPING[model]()
							esm.eval()  # disable dropout

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

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

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

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

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

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

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

							lowerCAmelCase =		original_esm_model.emb_layer_norm_after.weight
							lowerCAmelCase =		original_esm_model.emb_layer_norm_after.bias

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

														# self attention
														lowerCAmelCase =		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 =		esm_layer.self_attn.q_proj.weight
														lowerCAmelCase =		esm_layer.self_attn.q_proj.bias
														lowerCAmelCase =		esm_layer.self_attn.k_proj.weight
														lowerCAmelCase =		esm_layer.self_attn.k_proj.bias
														lowerCAmelCase =		esm_layer.self_attn.v_proj.weight
														lowerCAmelCase =		esm_layer.self_attn.v_proj.bias

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

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

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

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

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

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

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

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

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

														del esm  # Free up some memory before continuing

							print(F"""Saving tokenizer to {pytorch_dump_folder_path}"""				)
							hf_tokenizer.save_pretrained(_snake_case				)

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


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
		)









import argparse

from transformers import CanineConfig, CanineModel, CanineTokenizer, load_tf_weights_in_canine
from transformers.utils import logging


logging.set_verbosity_info()



def    UpperCAmelCase							(  _snake_case    ,			_snake_case				):
							# Initialize PyTorch model
							lowerCAmelCase =		CanineConfig()
							lowerCAmelCase =		CanineModel(_snake_case				)
							model.eval()

							print(F"""Building PyTorch model from configuration: {config}"""				)

							# Load weights from tf checkpoint
							load_tf_weights_in_canine(_snake_case    ,			_snake_case    ,			_snake_case				)

							# Save pytorch-model (weights and configuration)
							print(F"""Save PyTorch model to {pytorch_dump_path}"""				)
							model.save_pretrained(_snake_case				)

							# Save tokenizer files
							lowerCAmelCase =		CanineTokenizer()
							print(F"""Save tokenizer files to {pytorch_dump_path}"""				)
							tokenizer.save_pretrained(_snake_case				)


if __name__ == "__main__":
		UpperCAmelCase_    =argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    """--tf_checkpoint_path""",
		    default=None,
		    type=str,
		    required=True,
		    help="""Path to the TensorFlow checkpoint. Should end with model.ckpt""",
		)
		parser.add_argument(
		    """--pytorch_dump_path""",
		    default=None,
		    type=str,
		    required=True,
		    help="""Path to a folder where the PyTorch model will be placed.""",
		)
		UpperCAmelCase_    =parser.parse_args()
		convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.pytorch_dump_path)



