

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

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

UpperCAmelCase_       :				int		   =      {
    '''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_       :				Optional[Any]		   =      list('''ARNDCQEGHILKMFPSTWYV''')

UpperCAmelCase_       :				List[Any]		   =      restypes + ['''X''']
UpperCAmelCase_       :				List[Any]		   =      restypes_with_x + ['''<pad>''', '''<mask>''', '''<cls>''', '''<sep>''', '''<eos>''']
def        __SCREAMING_SNAKE_CASE    (      )				->							Tuple:
	with TemporaryDirectory() as tempdir:
		__A   : Union[str, Any]								= """\n""".join(a__					)
		__A   : Tuple								= Path(a__					) / """vocab.txt"""
		vocab_file.write_text(a__					)
		__A   : str								= EsmTokenizer(vocab_file=str(a__					)					)
	__A   : Optional[int]								= 0  # Overlaps with 'A' but that seems to be what they want
	return hf_tokenizer
def        __SCREAMING_SNAKE_CASE    (      a__						:      Tuple ,a__						:      List[str]					)				->							Any:
	__A   : 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        __SCREAMING_SNAKE_CASE    (      a__						:      str ,a__						:      str ,a__						:      bool ,a__						:      str ,a__						:      str					)				->							str:
	if model.startswith("""esmfold"""					):
		__A   : Tuple								= MODEL_MAPPING[model]()
	else:
		__A						,							__A   : Any								= MODEL_MAPPING[model]()
	esm.eval()  # disable dropout

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

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

	if is_folding_model:
		__A   : Any								= esm.esm
	else:
		__A   : int								= esm

	__A   : str								= 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   : Dict								= esm.classification_heads["""mnli"""].out_proj.weight.shape[0]
	print("""Our ESM config:""" ,a__					)

	if model.startswith("""esmfold"""					):
		__A   : str								= EsmForProteinFolding
	elif classification_head:
		__A   : Dict								= EsmForSequenceClassification
	else:
		__A   : Any								= EsmForMaskedLM
	__A   : Dict								= model_class(a__					)
	model.eval()

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

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

	__A   : Any								= original_esm_model.emb_layer_norm_after.weight
	__A   : int								= original_esm_model.emb_layer_norm_after.bias

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

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

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

		# LayerNorm changes for pre-activation
		__A   : Optional[int]								= esm_layer.self_attn_layer_norm.weight
		__A   : List[str]								= esm_layer.self_attn_layer_norm.bias
		__A   : Union[str, Any]								= esm_layer.final_layer_norm.weight
		__A   : List[Any]								= esm_layer.final_layer_norm.bias

		# self-attention output
		__A   : EsmSelfOutput								= layer.attention.output
		assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
		__A   : int								= esm_layer.self_attn.out_proj.weight
		__A   : Optional[Any]								= esm_layer.self_attn.out_proj.bias

		# intermediate
		__A   : EsmIntermediate								= layer.intermediate
		assert intermediate.dense.weight.shape == esm_layer.fca.weight.shape
		__A   : Optional[int]								= esm_layer.fca.weight
		__A   : List[Any]								= esm_layer.fca.bias

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

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

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

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

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

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


UpperCAmelCase_       :				Optional[Any]		   =      {
    '''configuration_prophetnet''': ['''PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ProphetNetConfig'''],
    '''tokenization_prophetnet''': ['''ProphetNetTokenizer'''],
}

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			UpperCAmelCase_       :				List[str]		   =      [
			    '''PROPHETNET_PRETRAINED_MODEL_ARCHIVE_LIST''',
			    '''ProphetNetDecoder''',
			    '''ProphetNetEncoder''',
			    '''ProphetNetForCausalLM''',
			    '''ProphetNetForConditionalGeneration''',
			    '''ProphetNetModel''',
			    '''ProphetNetPreTrainedModel''',
			]


if TYPE_CHECKING:
			from .configuration_prophetnet import PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP, ProphetNetConfig
			from .tokenization_prophetnet import ProphetNetTokenizer

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_prophetnet import (
						    PROPHETNET_PRETRAINED_MODEL_ARCHIVE_LIST,
						    ProphetNetDecoder,
						    ProphetNetEncoder,
						    ProphetNetForCausalLM,
						    ProphetNetForConditionalGeneration,
						    ProphetNetModel,
						    ProphetNetPreTrainedModel,
						)

else:
			import sys

			UpperCAmelCase_       :				Any		   =      _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)





