




'''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       ):






       def __init__(							self,	lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase						):



              """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   lowercase__  (							self,	lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase=True,	lowerCAmelCase=0						):



              """simple docstring"""







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






       def __init__(							self,	lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase						):



              """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   lowercase__  (							self,	lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase=True,	lowerCAmelCase=0						):



              """simple docstring"""







              lowerCamelCase_							=self.mega(lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase						)
              return self.mlm_head(self.dropout(lowerCAmelCase						)						)





def   a_							(   __snake_case							:    str     , __snake_case							:    List[str]     , __snake_case							:    Any     )							->		Optional[int]:



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

       # load the original encoder
       lowerCamelCase_							=OriginalMegaForMaskedLM(**__snake_case     ).eval()

       # load its weights
       print(
           '''Original Mega encoder:'''     , original_mlm.mega.load_state_dict(
               torch.load(os.path.join(__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(__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=__snake_case     , norm_affine=__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=__snake_case     , add_lm_hidden_dense_layer=__snake_case     , )

       lowerCamelCase_							=MegaForMaskedLM(__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(__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(__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(__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(__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(__snake_case     )
       # mask a few tokens to make sure masking is applied appropriately :)
       lowerCamelCase_							=0

       # run forward passes
       lowerCamelCase_							=original_mlm(__snake_case     , __snake_case     , batch_first=__snake_case     , ignore_mask_value=0     )
       lowerCamelCase_							=hf_mlm(__snake_case     , __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(__snake_case     , __snake_case     , atol=1e-3     )

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


if __name__ == "__main__":
     a_   :	Optional[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""",
     )

     a_   :	Union[str, Any]							   = parser.parse_args()

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










'''simple docstring'''




from .hash_table import HashTable


class 						__UpperCamelCase (				lowerCamelCase__       ):






       def __init__(							self,	*lowerCAmelCase,	**lowerCAmelCase						):



              """simple docstring"""







              super().__init__(*lowerCAmelCase,	**lowerCAmelCase						)






       def   lowercase__  (							self,	lowerCAmelCase,	lowerCAmelCase=None						):



              """simple docstring"""







              lowerCamelCase_							=1
              lowerCamelCase_							=self.hash_function(key + i * i						)

              while self.values[new_key] is not None and self.values[new_key] != key:
                     i += 1
                     lowerCamelCase_							=(
                         self.hash_function(key + i * i						)
                         if not self.balanced_factor() >= self.lim_charge
                         else None
                     )

                     if new_key is None:
                            break

              return new_key




