






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

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

__lowerCAmelCase							={
    "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,
}

__lowerCAmelCase							=list("ARNDCQEGHILKMFPSTWYV")

__lowerCAmelCase							=restypes + ["X"]
__lowerCAmelCase							=restypes_with_x + ["<pad>", "<mask>", "<cls>", "<sep>", "<eos>"]

def 							a	(   )				->				Dict:






					"""simple docstring"""







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

def 							a	(   _UpperCAmelCase					,  _UpperCAmelCase )				->				str:






					"""simple docstring"""







					a_									=   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	(   _UpperCAmelCase					,  _UpperCAmelCase					,  _UpperCAmelCase					,  _UpperCAmelCase					,  _UpperCAmelCase )				->				int:






					"""simple docstring"""







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

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

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

					if is_folding_model:
										a_									=   esm.esm
					else:
										a_									=   esm

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

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

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

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

					a_									=   original_esm_model.emb_layer_norm_after.weight
					a_									=   original_esm_model.emb_layer_norm_after.bias

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

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

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

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

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

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

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

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

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

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

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

										del esm  # Free up some memory before continuing

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

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


if __name__ == "__main__":
						__lowerCAmelCase							=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.")
						__lowerCAmelCase							=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 							a	(   _UpperCAmelCase					,  _UpperCAmelCase					,  _UpperCAmelCase )				->				list[int]:






					"""simple docstring"""







					if hue < 0 or hue > 3_6_0:
										raise Exception('hue should be between 0 and 360' )

					if saturation < 0 or saturation > 1:
										raise Exception('saturation should be between 0 and 1' )

					if value < 0 or value > 1:
										raise Exception('value should be between 0 and 1' )

					a_									=   value * saturation
					a_									=   hue / 6_0
					a_									=   chroma * (1 - abs(hue_section % 2 - 1 ))
					a_									=   value - chroma

					if hue_section >= 0 and hue_section <= 1:
										a_									=   round(2_5_5 * (chroma + match_value) )
										a_									=   round(2_5_5 * (second_largest_component + match_value) )
										a_									=   round(2_5_5 * (match_value) )
					elif hue_section > 1 and hue_section <= 2:
										a_									=   round(2_5_5 * (second_largest_component + match_value) )
										a_									=   round(2_5_5 * (chroma + match_value) )
										a_									=   round(2_5_5 * (match_value) )
					elif hue_section > 2 and hue_section <= 3:
										a_									=   round(2_5_5 * (match_value) )
										a_									=   round(2_5_5 * (chroma + match_value) )
										a_									=   round(2_5_5 * (second_largest_component + match_value) )
					elif hue_section > 3 and hue_section <= 4:
										a_									=   round(2_5_5 * (match_value) )
										a_									=   round(2_5_5 * (second_largest_component + match_value) )
										a_									=   round(2_5_5 * (chroma + match_value) )
					elif hue_section > 4 and hue_section <= 5:
										a_									=   round(2_5_5 * (second_largest_component + match_value) )
										a_									=   round(2_5_5 * (match_value) )
										a_									=   round(2_5_5 * (chroma + match_value) )
					else:
										a_									=   round(2_5_5 * (chroma + match_value) )
										a_									=   round(2_5_5 * (match_value) )
										a_									=   round(2_5_5 * (second_largest_component + match_value) )

					return [red, green, blue]

def 							a	(   _UpperCAmelCase					,  _UpperCAmelCase					,  _UpperCAmelCase )				->				list[float]:






					"""simple docstring"""







					if red < 0 or red > 2_5_5:
										raise Exception('red should be between 0 and 255' )

					if green < 0 or green > 2_5_5:
										raise Exception('green should be between 0 and 255' )

					if blue < 0 or blue > 2_5_5:
										raise Exception('blue should be between 0 and 255' )

					a_									=   red / 2_5_5
					a_									=   green / 2_5_5
					a_									=   blue / 2_5_5
					a_									=   max(_UpperCAmelCase					,  _UpperCAmelCase					,  _UpperCAmelCase )
					a_									=   value - min(_UpperCAmelCase					,  _UpperCAmelCase					,  _UpperCAmelCase )
					a_									=   0 if value == 0 else chroma / value

					if chroma == 0:
										a_									=   0.0
					elif value == float_red:
										a_									=   6_0 * (0 + (float_green - float_blue) / chroma)
					elif value == float_green:
										a_									=   6_0 * (2 + (float_blue - float_red) / chroma)
					else:
										a_									=   6_0 * (4 + (float_red - float_green) / chroma)

					a_									=   (hue + 3_6_0) % 3_6_0

					return [hue, saturation, value]





def 							a	(   _UpperCAmelCase					,  _UpperCAmelCase )				->				bool:






					"""simple docstring"""







					a_									=   abs(hsv_a[0] - hsv_a[0] ) < 0.2
					a_									=   abs(hsv_a[1] - hsv_a[1] ) < 0.0_0_2
					a_									=   abs(hsv_a[2] - hsv_a[2] ) < 0.0_0_2

					return check_hue and check_saturation and check_value

