






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

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

A:   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,
}

A:   Optional[Any]						=					list("ARNDCQEGHILKMFPSTWYV")

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






def 						_UpperCAmelCase     (	)      ->      Optional[int]:





						"""simple docstring"""
						with TemporaryDirectory() as tempdir:
												lowercase_   : Any   	=      '\n'.join(snake_case_				)
												lowercase_   : List[str]   	=      Path(snake_case_				) / 'vocab.txt'
												vocab_file.write_text(snake_case_				)
												lowercase_   : Dict   	=      EsmTokenizer(vocab_file=str(snake_case_				)				)
						lowercase_   : List[Any]   	=      0  # Overlaps with 'A' but that seems to be what they want
						return hf_tokenizer






def 						_UpperCAmelCase     (	a		:						Tuple    ,      a		:						Any				)      ->      Union[str, Any]:





						"""simple docstring"""
						lowercase_   : List[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 						_UpperCAmelCase     (	a		:						str    ,      a		:						str    ,      a		:						bool    ,      a		:						str    ,      a		:						str				)      ->      int:





						"""simple docstring"""
						if model.startswith('esmfold'				):
												lowercase_   : int   	=      MODEL_MAPPING[model]()
						else:
												lowercase_   : List[Any]   	=      MODEL_MAPPING[model]()
						esm.eval()  # disable dropout

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

						if is_folding_model:
												lowercase_   : Optional[Any]   	=      esm.esm.alphabet
						lowercase_   : Dict   	=      tuple(alphabet.all_toks				)
						lowercase_   : List[str]   	=      alphabet.mask_idx
						lowercase_   : List[str]   	=      alphabet.padding_idx

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

						lowercase_   : str   	=      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=1_0_2_6    ,      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:
												lowercase_   : int   	=      esm.classification_heads['mnli'].out_proj.weight.shape[0]
						print('Our ESM config:'    ,      snake_case_				)

						if model.startswith('esmfold'				):
												lowercase_   : Optional[int]   	=      EsmForProteinFolding
						elif classification_head:
												lowercase_   : List[str]   	=      EsmForSequenceClassification
						else:
												lowercase_   : List[str]   	=      EsmForMaskedLM
						lowercase_   : Dict   	=      model_class(snake_case_				)
						model.eval()

						# Now let's copy all the weights.
						# Embeddings
						lowercase_   : Dict   	=      original_esm_model.embed_tokens.weight
						if position_embedding_type == "absolute":
												lowercase_   : Union[str, Any]   	=      original_esm_model.embed_positions.weight

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

						lowercase_   : Any   	=      original_esm_model.emb_layer_norm_after.weight
						lowercase_   : Union[str, Any]   	=      original_esm_model.emb_layer_norm_after.bias

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

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

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

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

												# self-attention output
												lowercase_   : EsmSelfOutput   	=      layer.attention.output
												assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
												lowercase_   : List[str]   	=      esm_layer.self_attn.out_proj.weight
												lowercase_   : Union[str, Any]   	=      esm_layer.self_attn.out_proj.bias

												# intermediate
												lowercase_   : EsmIntermediate   	=      layer.intermediate
												assert intermediate.dense.weight.shape == esm_layer.fca.weight.shape
												lowercase_   : List[str]   	=      esm_layer.fca.weight
												lowercase_   : List[str]   	=      esm_layer.fca.bias

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

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

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

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

												if is_folding_model:
																		lowercase_   : List[str]   	=      torch.max(torch.abs(our_output['positions'] - their_output['positions']				)				).item()
																		lowercase_   : Tuple   	=      torch.allclose(our_output['positions']    ,      their_output['positions']    ,      atol=1e-5				)
												else:
																		lowercase_   : Tuple   	=      torch.max(torch.abs(our_output - their_output				)				).item()
																		lowercase_   : str   	=      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
																		lowercase_   : Tuple   	=      model.predict_contacts(hf_tokens['input_ids']    ,      hf_tokens['attention_mask']				)
																		lowercase_   : List[Any]   	=      esm.predict_contacts(hf_tokens['input_ids']				)
																		lowercase_   : Union[str, Any]   	=      torch.max(torch.abs(our_output - their_output				)				).item()
																		lowercase_   : Optional[Any]   	=      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__":
						A:   Optional[Any]						=					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:   Dict						=					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        SCREAMING_SNAKE_CASE_  (	snake_case_				:	int							)       ->	int:
 if number != int(snake_case_							):
  raise ValueError('factorial() only accepts integral values'							)
 if number < 0:
  raise ValueError('factorial() not defined for negative values'							)
 SCREAMING_SNAKE_CASE  :  str								=		1
 for i in range(1   ,     number + 1							):
  value *= i
 return value



def        SCREAMING_SNAKE_CASE_  (	snake_case_				:	int							)       ->	int:
 if not isinstance(snake_case_   ,     snake_case_							):
  raise ValueError('factorial() only accepts integral values'							)
 if n < 0:
  raise ValueError('factorial() not defined for negative values'							)
 return 1 if n == 0 or n == 1 else n * factorial(n - 1							)


if __name__ == "__main__":
   import doctest

   doctest.testmod()

   __UpperCAmelCase												=   int(input('Enter a positive integer: ').strip() or 0)
   print(f"""factorial{n} is {factorial(n)}""")


