






'''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 					__UpperCAmelCase (		nn.Module	):

      '''simple docstring'''





      def __init__(					self  ,		_UpperCAmelCase  ,		_UpperCAmelCase  ,		_UpperCAmelCase						):
             super().__init__()
             UpperCAmelCase__	:   Union[str, Any]												=	mega_args
             UpperCAmelCase__	:   Optional[int]												=	nn.Embedding(_UpperCAmelCase  ,		self.mega_args.encoder_embed_dim						)
             UpperCAmelCase__	:   Dict												=	nn.ModuleList([MegaEncoderLayer(self.mega_args						) for _ in range(_UpperCAmelCase						)]						)
             UpperCAmelCase__	:   Optional[Any]												=	depth




      def 							lowerCamelCase     (					self  ,		_UpperCAmelCase  ,		_UpperCAmelCase  ,		_UpperCAmelCase=True  ,		_UpperCAmelCase=0						):

             # Mega expects embeddings to be (time, batch, embedding size), but
             # Hugging Face returns tokens as (batch, time)
             if batch_first:
                    UpperCAmelCase__	:   Tuple												=	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:
                    UpperCAmelCase__	:   Union[str, Any]												=	1 - attention_mask

             # get token embeddings from IDs
             UpperCAmelCase__	:   str												=	self.embedding_layer(_UpperCAmelCase						)

             # pass through the Mega layers
             # input is (time, batch, encoder dim) and output is the same
             for encoder in self.encoders:
                    UpperCAmelCase__	:   Optional[int]												=	encoder(_UpperCAmelCase  ,		_UpperCAmelCase						)

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




class 					__UpperCAmelCase (		nn.Module	):

      '''simple docstring'''





      def __init__(					self  ,		_UpperCAmelCase  ,		_UpperCAmelCase  ,		_UpperCAmelCase						):
             super().__init__()
             UpperCAmelCase__	:   Dict												=	MegaLM(_UpperCAmelCase  ,		_UpperCAmelCase  ,		_UpperCAmelCase						)
             UpperCAmelCase__	:   int												=	nn.Linear(mega_args.encoder_embed_dim  ,		_UpperCAmelCase						)
             UpperCAmelCase__	:   Any												=	nn.Dropout(p=0.1						)




      def 							lowerCamelCase     (					self  ,		_UpperCAmelCase  ,		_UpperCAmelCase  ,		_UpperCAmelCase=True  ,		_UpperCAmelCase=0						):
             UpperCAmelCase__	:   int												=	self.mega(_UpperCAmelCase  ,		_UpperCAmelCase  ,		_UpperCAmelCase  ,		_UpperCAmelCase						)
             return self.mlm_head(self.dropout(_UpperCAmelCase						)						)






def  lowerCAmelCase__		(					a_  :     Dict			,						a_  :     Optional[int]			,						a_  :     Tuple			)   -> Optional[Any]:
       with open(os.path.join(a_			,						'''model_args.pkl'''			)			,						'''rb'''			) as f:
              UpperCAmelCase__	:   List[Any]												=	pkl.load(a_			)

       # load the original encoder
       UpperCAmelCase__	:   Tuple												=	OriginalMegaForMaskedLM(**a_			).eval()

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

       # create a new config from the old one
       UpperCAmelCase__	:   str												=	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=a_			,						norm_affine=a_			,						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=a_			,						add_lm_hidden_dense_layer=a_			,						)

       UpperCAmelCase__	:   int												=	MegaForMaskedLM(a_			).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
       UpperCAmelCase__	:   str												=	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
       UpperCAmelCase__	:   Dict												=	original_mlm.mega.encoders.state_dict()
       UpperCAmelCase__	:   Union[str, Any]												=	{}
       for module_name in original_state_dict.keys():
              UpperCAmelCase__	:   Dict												=	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:
                            UpperCAmelCase__	:   Union[str, Any]												=	module_name.replace('''move.beta'''			,						'''ema_gate.ema_expansion_matrix'''			)
                     elif "mega_layer.beta" in module_name:
                            UpperCAmelCase__	:   Tuple												=	module_name.replace('''beta'''			,						'''qk_bias'''			)
                     else:
                            UpperCAmelCase__	:   List[str]												=	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:
                            UpperCAmelCase__	:   Optional[int]												=	module_name.replace('''move.gamma'''			,						'''ema_gate.kernel_projection_matrix'''			)
                     elif "mega_layer.gamma" in module_name:
                            UpperCAmelCase__	:   Tuple												=	module_name.replace('''gamma'''			,						'''qk_weight'''			)
                     else:
                            UpperCAmelCase__	:   str												=	module_name.replace('''gamma'''			,						'''g_param'''			)
        # alpha is used in EMA and positional bias; renaming to improve readability
              elif "move.alpha" in module_name:
                     UpperCAmelCase__	:   List[Any]												=	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:
                     UpperCAmelCase__	:   List[Any]												=	module_name.replace('''move.delta'''			,						'''ema_gate.damping_factor'''			)
              # omega is only used in EMA; renaming to improve readability
              elif "omega" in module_name:
                     UpperCAmelCase__	:   str												=	module_name.replace('''move.omega'''			,						'''ema_gate.residual_weight'''			)

              if new_module_name:
                     UpperCAmelCase__	:   Any												=	new_module_name

       if len(a_			) != 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():
              UpperCAmelCase__	:   List[Any]												=	original_state_dict.pop(a_			)

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

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

       # test on a randomly generated input sequence
       UpperCAmelCase__	:   Optional[int]												=	torch.randint(0			,						hf_config.vocab_size			,						size=(4, 2_5_6)			)
       UpperCAmelCase__	:   int												=	torch.ones_like(a_			)
       # mask a few tokens to make sure masking is applied appropriately :)
       UpperCAmelCase__	:   List[Any]												=	0

       # run forward passes
       UpperCAmelCase__	:   Union[str, Any]												=	original_mlm(a_			,						a_			,						batch_first=a_			,						ignore_mask_value=0			)
       UpperCAmelCase__	:   str												=	hf_mlm(a_			,						a_			)[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
       UpperCAmelCase__	:   str												=	torch.allclose(a_			,						a_			,						atol=1E-3			)

       if success:
              print('''Yay!'''			)
              hf_mlm.save_pretrained(a_			)
       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'''			)
              UpperCAmelCase__	:   int												=	AutoTokenizer.from_pretrained(a_			)
              tokenizer.save_pretrained(a_			)


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)





from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__A    :				List[str]    						=	logging.get_logger(__name__)

__A    :				List[Any]    						=	{
    """mnaylor/mega-base-wikitext""": """https://huggingface.co/mnaylor/mega-base-wikitext/resolve/main/config.json""",
}




class       UpperCAmelCase_   ( A							):




							'''simple docstring'''



							a__    	=  '''mega'''





							def __init__(       self							:						Any							,  a							:						List[str]=30_522							,  a							:						Dict=128							,  a							:						Dict=4							,  a							:						Union[str, Any]=256							,  a							:						Tuple=16							,  a							:						Any=True							,  a							:						Dict=64							,  a							:						Union[str, Any]=False							,  a							:						Tuple=-1							,  a							:						str=None							,  a							:						int=True							,  a							:						Optional[int]="scalenorm"							,  a							:						Optional[int]=True							,  a							:						Dict="silu"							,  a							:						Optional[Any]="softmax"							,  a							:						Any=0.1							,  a							:						List[Any]=0.1							,  a							:						Tuple=0.1							,  a							:						Optional[int]=False							,  a							:						List[str]=True							,  a							:						Optional[int]=256							,  a							:						List[str]=True							,  a							:						int=0.1							,  a							:						str=2_048							,  a							:						Tuple=False							,  a							:						Optional[int]=2							,  a							:						Optional[int]=0.02							,  a							:						Optional[int]=0.2							,  a							:						Union[str, Any]=0.02							,  a							:						str=1.0							,  a							:						Dict=1							,  a							:						Union[str, Any]=0							,  a							:						Dict=2							,  a							:						List[str]="rotary"							,  a							:						Union[str, Any]=None							,  a							:						List[Any]=True							,  a							:						Tuple=True							,  **a							:						Optional[int]							,  )      ->       Dict:
										super().__init__(pad_token_id=a							,  bos_token_id=a							,  eos_token_id=a							,  **a      )

										SCREAMING_SNAKE_CASE     =       vocab_size
										SCREAMING_SNAKE_CASE     =       hidden_size
										SCREAMING_SNAKE_CASE     =       num_hidden_layers
										SCREAMING_SNAKE_CASE     =       activation
										SCREAMING_SNAKE_CASE     =       attention_activation
										SCREAMING_SNAKE_CASE     =       intermediate_size
										SCREAMING_SNAKE_CASE     =       ema_projection_size
										SCREAMING_SNAKE_CASE     =       bidirectional
										SCREAMING_SNAKE_CASE     =       shared_representation_size
										SCREAMING_SNAKE_CASE     =       use_chunking
										SCREAMING_SNAKE_CASE     =       chunk_size
										SCREAMING_SNAKE_CASE     =       truncation
										SCREAMING_SNAKE_CASE     =       normalize_before_mega
										SCREAMING_SNAKE_CASE     =       normalization_type
										SCREAMING_SNAKE_CASE     =       norm_affine
										SCREAMING_SNAKE_CASE     =       dropout_prob
										SCREAMING_SNAKE_CASE     =       hidden_dropout_prob
										SCREAMING_SNAKE_CASE     =       attention_probs_dropout_prob
										SCREAMING_SNAKE_CASE     =       use_feature_dropout
										SCREAMING_SNAKE_CASE     =       use_normalized_ffn
										SCREAMING_SNAKE_CASE     =       nffn_hidden_size
										SCREAMING_SNAKE_CASE     =       normalize_before_ffn
										SCREAMING_SNAKE_CASE     =       nffn_activation_dropout_prob
										SCREAMING_SNAKE_CASE     =       max_positions
										SCREAMING_SNAKE_CASE     =       add_token_type_embeddings
										SCREAMING_SNAKE_CASE     =       type_vocab_size
										SCREAMING_SNAKE_CASE     =       initializer_range
										SCREAMING_SNAKE_CASE     =       ema_delta_alpha_range
										SCREAMING_SNAKE_CASE     =       ema_beta_range
										SCREAMING_SNAKE_CASE     =       ema_gamma_omega_range
										SCREAMING_SNAKE_CASE     =       relative_positional_bias
										SCREAMING_SNAKE_CASE     =       use_cache
										SCREAMING_SNAKE_CASE     =       classifier_dropout
										SCREAMING_SNAKE_CASE     =       add_lm_hidden_dense_layer
										SCREAMING_SNAKE_CASE     =       1  # not used but required by Hugging Face





class       UpperCAmelCase_   ( A							):




							'''simple docstring'''



							@property
							def  _UpperCAmelCase  (       self							:						Union[str, Any]      )      ->       Mapping[str, Mapping[int, str]]:
										if self.task == "multiple-choice":
													SCREAMING_SNAKE_CASE     =       {0: """batch""", 1: """choice""", 2: """sequence"""}
										else:
													SCREAMING_SNAKE_CASE     =       {0: """batch""", 1: """sequence"""}
										return OrderedDict(
										    [
										        ("""input_ids""", dynamic_axis),
										        ("""attention_mask""", dynamic_axis),
										    ]      )






