





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

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

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

_snake_case				= list("""ARNDCQEGHILKMFPSTWYV""")

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


def _A   (    ):
   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   (    __magic_name__  ,					__magic_name__       ):
   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   (    __magic_name__  ,					__magic_name__  ,					__magic_name__  ,					__magic_name__  ,					__magic_name__       ):
   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=1026  ,					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__":
  _snake_case				= 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.""")
  _snake_case				= 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'''

import argparse

# utilities to import the model weights and config file
import os
import pickle as pkl

# PyTorch + new model classes
import torch
from torch import nn

from transformers import AutoTokenizer, MegaConfig, MegaForMaskedLM


# import the EncoderLayer class used to pretrain
# !! NOTE !! this requires the version of fairseq that is built when you install the Mega source
try:
				from fairseq.modules.mega_layer import MegaEncoderLayer
except ImportError:
				raise ImportError('You need to install the version of fairseq from the Mega repo!')

class __snake_case(      nn.Module					):


						'''simple docstring'''



						def __init__(					self	,       A_	,       A_	,       A_      )	->	Dict:
							super().__init__()
							lowerCAmelCase    =	mega_args
							lowerCAmelCase    =	nn.Embedding(A_	,       self.mega_args.encoder_embed_dim      )
							lowerCAmelCase    =	nn.ModuleList([MegaEncoderLayer(self.mega_args      ) for _ in range(A_      )]      )
							lowerCAmelCase    =	depth




						def 				__snake_case    (					self	,       A_	,       A_	,       A_=True	,       A_=0      )	->	List[Any]:

							# Mega expects embeddings to be (time, batch, embedding size), but
							# Hugging Face returns tokens as (batch, time)
							if batch_first:
								lowerCAmelCase    =	input_ids.T

							# to make things more confusing, Mega expects the attention mask to
							# be (batch, time), but with values of 0 (normal token) and 1 (ignore token)
							# which is the opposite of what HF returns
							if ignore_mask_value == 0:
								lowerCAmelCase    =	1 - attention_mask

							# get token embeddings from IDs
							lowerCAmelCase    =	self.embedding_layer(A_      )

							# pass through the Mega layers
							# input is (time, batch, encoder dim) and output is the same
							for encoder in self.encoders:
								lowerCAmelCase    =	encoder(A_	,       A_      )

							# return according to the shape specified
							if batch_first:
								# (T, B, H) --> (B, T, H)
								return torch.transpose(A_	,       0	,       1      )
							else:
								return embeds







class __snake_case(      nn.Module					):


						'''simple docstring'''



						def __init__(					self	,       A_	,       A_	,       A_      )	->	Optional[Any]:
							super().__init__()
							lowerCAmelCase    =	MegaLM(A_	,       A_	,       A_      )
							lowerCAmelCase    =	nn.Linear(mega_args.encoder_embed_dim	,       A_      )
							lowerCAmelCase    =	nn.Dropout(p=0.1      )




						def 				__snake_case    (					self	,       A_	,       A_	,       A_=True	,       A_=0      )	->	List[Any]:
							lowerCAmelCase    =	self.mega(A_	,       A_	,       A_	,       A_      )
							return self.mlm_head(self.dropout(A_      )      )







def 		_snake_case		(					_SCREAMING_SNAKE_CASE     :						Optional[int]							,		_SCREAMING_SNAKE_CASE     :						List[str]							,		_SCREAMING_SNAKE_CASE     :						Union[str, Any]							)			->						List[Any]:







	"""simple docstring"""
	with open(os.path.join(_SCREAMING_SNAKE_CASE							,		"""model_args.pkl"""							)							,		"""rb"""							) as f:
		lowerCAmelCase    =	pkl.load(_SCREAMING_SNAKE_CASE							)

	# load the original encoder
	lowerCAmelCase    =	OriginalMegaForMaskedLM(**_SCREAMING_SNAKE_CASE							).eval()

	# load its weights
	print(
	    """Original Mega encoder:"""							,		original_mlm.mega.load_state_dict(
	        torch.load(os.path.join(_SCREAMING_SNAKE_CASE							,		"""encoder_weights.pt"""							)							,		map_location="""cpu"""							)							)							,		)
	print(
	    """Original Mega MLM layer:"""							,		original_mlm.mlm_head.load_state_dict(
	        torch.load(os.path.join(_SCREAMING_SNAKE_CASE							,		"""mlm_head_weights.pt"""							)							,		map_location="""cpu"""							)							)							,		)

	# create a new config from the old one
	lowerCAmelCase    =	MegaConfig(
	    num_hidden_layers=mega_original_args["""depth"""]							,		vocab_size=mega_original_args["""vocab_size"""]							,		hidden_size=mega_original_args["""mega_args"""].encoder_embed_dim							,		shared_representation_size=mega_original_args["""mega_args"""].encoder_z_dim							,		intermediate_size=mega_original_args["""mega_args"""].encoder_hidden_dim							,		ema_projection_size=mega_original_args["""mega_args"""].encoder_n_dim							,		dropout_prob=mega_original_args["""mega_args"""].dropout							,		attention_probs_dropout_prob=mega_original_args["""mega_args"""].attention_dropout							,		hidden_dropout_prob=mega_original_args["""mega_args"""].hidden_dropout							,		activation=mega_original_args["""mega_args"""].activation_fn							,		attention_activation=mega_original_args["""mega_args"""].attention_activation_fn							,		bidirectional=mega_original_args["""mega_args"""].bidirectional							,		use_chunking=mega_original_args["""mega_args"""].encoder_chunk_size > 0							,		chunk_size=mega_original_args["""mega_args"""].encoder_chunk_size							,		truncation=mega_original_args["""mega_args"""].truncation_length							,		normalization_type=mega_original_args["""mega_args"""].normalization_type							,		normalize_before_mega=_SCREAMING_SNAKE_CASE							,		norm_affine=_SCREAMING_SNAKE_CASE							,		use_feature_dropout=mega_original_args["""mega_args"""].feature_dropout							,		relative_positional_bias=mega_original_args["""mega_args"""].rel_pos_bias							,		max_positions=mega_original_args["""mega_args"""].max_source_positions							,		nffn_hidden_size=mega_original_args["""mega_args"""].encoder_ffn_embed_dim							,		normalize_before_ffn=mega_original_args["""mega_args"""].normalize_before							,		nffn_activation_dropout_prob=0.0							,		add_token_type_embeddings=_SCREAMING_SNAKE_CASE							,		add_lm_hidden_dense_layer=_SCREAMING_SNAKE_CASE							,		)

	lowerCAmelCase    =	MegaForMaskedLM(_SCREAMING_SNAKE_CASE							).eval()

	# the originl checkpoint just uses nn.Embedding for the word embeddings
	# we use a wrapper module for embeddings to add support for positional embeddings
	lowerCAmelCase    =	original_mlm.mega.embedding_layer.weight

	# modify the state dictionary of the original checkpoint to account for naming issues in the Hugging Face
	# ecosystem -- any names containing "beta" or "gamma" aren't safe to use and are renamed upon _load_pretrained,
	# also renaming previously confusing parameter names
	lowerCAmelCase    =	original_mlm.mega.encoders.state_dict()
	lowerCAmelCase    =	{}
	for module_name in original_state_dict.keys():
		lowerCAmelCase    =	None
		# have to handle gamma, beta, and alpha differently due to their use
		# in multiple modules within the original repository;
		# beta is used in EMA, MovingAverageGatedAttention, and RotaryRelativePositionalBias, and must be renamed due to flax/tf weights
		# the EMA sublayer was renamed from "move" to "ema_gate" for readability, so that is also done here
		if "beta" in module_name:
			# EMA sub-layers were always called "move" in the original repo
			if "move.beta" in module_name:
				lowerCAmelCase    =	module_name.replace("""move.beta"""							,		"""ema_gate.ema_expansion_matrix"""							)
			elif "mega_layer.beta" in module_name:
				lowerCAmelCase    =	module_name.replace("""beta"""							,		"""qk_bias"""							)
			else:
				lowerCAmelCase    =	module_name.replace("""beta"""							,		"""b_param"""							)
        # beta is used in EMA and MovingAverageGatedAttention, and must be renamed due to flax/tf weights
		elif "gamma" in module_name:
			if "move.gamma" in module_name:
				lowerCAmelCase    =	module_name.replace("""move.gamma"""							,		"""ema_gate.kernel_projection_matrix"""							)
			elif "mega_layer.gamma" in module_name:
				lowerCAmelCase    =	module_name.replace("""gamma"""							,		"""qk_weight"""							)
			else:
				lowerCAmelCase    =	module_name.replace("""gamma"""							,		"""g_param"""							)
        # alpha is used in EMA and positional bias; renaming to improve readability
		elif "move.alpha" in module_name:
			lowerCAmelCase    =	module_name.replace("""move.alpha"""							,		"""ema_gate.decay_factor"""							)
		# delta is only used in EMA; renaming to improve readability
		elif "move.delta" in module_name:
			lowerCAmelCase    =	module_name.replace("""move.delta"""							,		"""ema_gate.damping_factor"""							)
		# omega is only used in EMA; renaming to improve readability
		elif "omega" in module_name:
			lowerCAmelCase    =	module_name.replace("""move.omega"""							,		"""ema_gate.residual_weight"""							)

		if new_module_name:
			lowerCAmelCase    =	new_module_name

	if len(_SCREAMING_SNAKE_CASE							) != 0:
		print(f'Renaming these keys: {updated_keys.keys()}'							)
	else:
		print("""No need to rename state dict entries"""							)
	for old, new in updated_keys.items():
		lowerCAmelCase    =	original_state_dict.pop(_SCREAMING_SNAKE_CASE							)

	# now attempt to load the state dictionary with updated names
	# note that we now call it `mega.layers` instead of `mega.encoders` due to hugging face style
	print("""HF Mega encoder:"""							,		hf_mlm.mega.layers.load_state_dict(_SCREAMING_SNAKE_CASE							)							)

	# load the MLM head weights directly
	print(
	    """HF Mega MLM layer:"""							,		hf_mlm.mlm_head.load_state_dict(
	        torch.load(os.path.join(_SCREAMING_SNAKE_CASE							,		"""mlm_head_weights.pt"""							)							,		map_location="""cpu"""							)							)							,		)

	# test on a randomly generated input sequence
	lowerCAmelCase    =	torch.randint(0							,		hf_config.vocab_size							,		size=(4, 256)							)
	lowerCAmelCase    =	torch.ones_like(_SCREAMING_SNAKE_CASE							)
	# mask a few tokens to make sure masking is applied appropriately :)
	lowerCAmelCase    =	0

	# run forward passes
	lowerCAmelCase    =	original_mlm(_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		batch_first=_SCREAMING_SNAKE_CASE							,		ignore_mask_value=0							)
	lowerCAmelCase    =	hf_mlm(_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							)[0]

	# print shapes and diff
	print(f'original output {original_output.shape}'							)
	print(f'hf output {hf_output.shape}'							)
	print(f'max diff: {(original_output - hf_output).max()}'							)  # 0.0
	lowerCAmelCase    =	torch.allclose(_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		atol=1E-3							)

	if success:
		print("""Yay!"""							)
		hf_mlm.save_pretrained(_SCREAMING_SNAKE_CASE							)
	else:
		raise RuntimeError(f'Something\'s broken :(\nOriginal:\n{original_output}\n\nHF\n{hf_output}\n{hf_mlm}'							)

	if includes_tokenizer:
		print("""Transferring tokenizer"""							)
		lowerCAmelCase    =	AutoTokenizer.from_pretrained(_SCREAMING_SNAKE_CASE							)
		tokenizer.save_pretrained(_SCREAMING_SNAKE_CASE							)


if __name__ == "__main__":
				UpperCAmelCase			=    argparse.ArgumentParser()

				parser.add_argument(
				    '--pretrained_checkpoint_path',
				    default=None,
				    type=str,
				    required=True,
				    help='Point to the directory containing your model weights using the official Mega repo',
				)

				parser.add_argument(
				    '--output_path', default=None, type=str, required=True, help='Location to save the Hugging Face version'
				)

				parser.add_argument(
				    '--includes_tokenizer',
				    action='store_true',
				    help='Use this flag if there is a Hugging Face tokenizer in the original checkpoint repo',
				)

				UpperCAmelCase			=    parser.parse_args()

				convert_checkpoint_to_huggingface(args.pretrained_checkpoint_path, args.output_path, args.includes_tokenizer)