code: 
  
  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
  			)
  
  
  
  
  
  
style_context: 
  
  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__)
  
  
  
  
  
  