




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 						SCREAMING_SNAKE_CASE__		(nn.Module ):
     def __init__(	self  :		Dict					,	__lowerCamelCase  :		Tuple					,	__lowerCamelCase  :		Optional[int]					,	__lowerCamelCase  :		Dict	):



       """simple docstring"""
       super().__init__()
       lowerCAmelCase__								=			mega_args
       lowerCAmelCase__								=			nn.Embedding(__lowerCamelCase					,	self.mega_args.encoder_embed_dim	)
       lowerCAmelCase__								=			nn.ModuleList([MegaEncoderLayer(self.mega_args	) for _ in range(__lowerCamelCase	)]	)
       lowerCAmelCase__								=			depth



     def 					A__			(	self  :		List[Any]					,	__lowerCamelCase  :		Any					,	__lowerCamelCase  :		Dict					,	__lowerCamelCase  :		Optional[int]=True					,	__lowerCamelCase  :		str=0	):



       """simple docstring"""

       # 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(__lowerCamelCase	)

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

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





class 						SCREAMING_SNAKE_CASE__		(nn.Module ):
     def __init__(	self  :		str					,	__lowerCamelCase  :		Union[str, Any]					,	__lowerCamelCase  :		Any					,	__lowerCamelCase  :		str	):



       """simple docstring"""
       super().__init__()
       lowerCAmelCase__								=			MegaLM(__lowerCamelCase					,	__lowerCamelCase					,	__lowerCamelCase	)
       lowerCAmelCase__								=			nn.Linear(mega_args.encoder_embed_dim					,	__lowerCamelCase	)
       lowerCAmelCase__								=			nn.Dropout(p=0.1	)



     def 					A__			(	self  :		Optional[int]					,	__lowerCamelCase  :		Any					,	__lowerCamelCase  :		Tuple					,	__lowerCamelCase  :		Any=True					,	__lowerCamelCase  :		Optional[int]=0	):



       """simple docstring"""
       lowerCAmelCase__								=			self.mega(__lowerCamelCase					,	__lowerCamelCase					,	__lowerCamelCase					,	__lowerCamelCase	)
       return self.mlm_head(self.dropout(__lowerCamelCase	)	)



def 		a_       (   __lowerCAmelCase      ,   __lowerCAmelCase      ,   __lowerCAmelCase							):
  with open(os.path.join(__lowerCAmelCase      ,   '''model_args.pkl'''							)      ,   '''rb'''							) as f:
    lowerCAmelCase__								=			pkl.load(__lowerCAmelCase							)

  # load the original encoder
  lowerCAmelCase__								=			OriginalMegaForMaskedLM(**__lowerCAmelCase							).eval()

  # load its weights
  print(
      '''Original Mega encoder:'''      ,   original_mlm.mega.load_state_dict(
          torch.load(os.path.join(__lowerCAmelCase      ,   '''encoder_weights.pt'''							)      ,   map_location='''cpu'''							)							)      ,   )
  print(
      '''Original Mega MLM layer:'''      ,   original_mlm.mlm_head.load_state_dict(
          torch.load(os.path.join(__lowerCAmelCase      ,   '''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=__lowerCAmelCase      ,   norm_affine=__lowerCAmelCase      ,   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=__lowerCAmelCase      ,   add_lm_hidden_dense_layer=__lowerCAmelCase      ,   )

  lowerCAmelCase__								=			MegaForMaskedLM(__lowerCAmelCase							).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(__lowerCAmelCase							) != 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(__lowerCAmelCase							)

  # 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(__lowerCAmelCase							)							)

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

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

  # run forward passes
  lowerCAmelCase__								=			original_mlm(__lowerCAmelCase      ,   __lowerCAmelCase      ,   batch_first=__lowerCAmelCase      ,   ignore_mask_value=0							)
  lowerCAmelCase__								=			hf_mlm(__lowerCAmelCase      ,   __lowerCAmelCase							)[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(__lowerCAmelCase      ,   __lowerCAmelCase      ,   atol=1E-3							)

  if success:
    print('''Yay!'''							)
    hf_mlm.save_pretrained(__lowerCAmelCase							)
  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(__lowerCAmelCase							)
    tokenizer.save_pretrained(__lowerCAmelCase							)


if __name__ == "__main__":
      __magic_name__     :     List[Any]  							=					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""",
      )

      __magic_name__     :     int  							=					parser.parse_args()

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







def 		a_       (   __lowerCAmelCase      ,   __lowerCAmelCase							):
  lowerCAmelCase__								=			0
  for ele in arr:
    total_sum += product_sum(__lowerCAmelCase      ,   depth + 1							) if isinstance(__lowerCAmelCase      ,   __lowerCAmelCase							) else ele
  return total_sum * depth
def 		a_       (   __lowerCAmelCase							):
  return product_sum(__lowerCAmelCase      ,   1							)


if __name__ == "__main__":
      import doctest

      doctest.testmod()

