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

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

__a        = {
    '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,
}

__a        = list('ARNDCQEGHILKMFPSTWYV')

__a        = restypes + ['X']
__a        = restypes_with_x + ['<pad>', '<mask>', '<cls>', '<sep>', '<eos>']




def     a						(							):


				'''simple docstring'''





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




def     a						(							snake_case__: Optional[Any]						,       snake_case__: List[Any] ):


				'''simple docstring'''
				lowercase_					  = 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     a						(							snake_case__: str						,       snake_case__: str						,       snake_case__: bool						,       snake_case__: str						,       snake_case__: str ):


				'''simple docstring'''





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

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

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

				if is_folding_model:
								lowercase_					  = esm.esm
				else:
								lowercase_					  = esm

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

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

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

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

				lowercase_					  = original_esm_model.emb_layer_norm_after.weight
				lowercase_					  = original_esm_model.emb_layer_norm_after.bias

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

								# self attention
								lowercase_					  = 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) )
								)

								lowercase_					  = esm_layer.self_attn.q_proj.weight
								lowercase_					  = esm_layer.self_attn.q_proj.bias
								lowercase_					  = esm_layer.self_attn.k_proj.weight
								lowercase_					  = esm_layer.self_attn.k_proj.bias
								lowercase_					  = esm_layer.self_attn.v_proj.weight
								lowercase_					  = esm_layer.self_attn.v_proj.bias

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

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

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

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

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

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

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

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

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

								del esm  # Free up some memory before continuing

				print(F'''Saving tokenizer to {pytorch_dump_folder_path}''' )
				hf_tokenizer.save_pretrained(SCREAMING_SNAKE_CASE_ )

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


if __name__ == "__main__":
					__a        = 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.')
					__a        = 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 __future__ import annotations

__a        = 'Alexander Joslin'




def     a						(							snake_case__: int ):


				'''simple docstring'''

				if number <= 0:
								lowercase_					  = F'''n must be greater than 0. Got n = {number}'''
								raise ValueError(snake_case__ )

				lowercase_					  = [number + 1] * (number + 1)

				# starting position
				lowercase_					  = 0
				for i in range(1						,       snake_case__ ):
								lowercase_					  = min(table[i + 1]						,       table[i] + 1 )
								# check if out of bounds
								if i * 2 <= number:
												lowercase_					  = min(table[i * 2]						,       table[i] + 1 )
								# check if out of bounds
								if i * 3 <= number:
												lowercase_					  = min(table[i * 3]						,       table[i] + 1 )
				return table[number]


if __name__ == "__main__":
						import doctest

						doctest.testmod()




