



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




 def __init__(  self		,			__UpperCAmelCase		,			__UpperCAmelCase		,			__UpperCAmelCase			):



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




 def     lowerCamelCase (  self		,			__UpperCAmelCase		,			__UpperCAmelCase		,			__UpperCAmelCase=True		,			__UpperCAmelCase=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(lowercase_			)

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

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



class 	__lowerCAmelCase       (		nn.Module				):




 def __init__(  self		,			__UpperCAmelCase		,			__UpperCAmelCase		,			__UpperCAmelCase			):



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




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



    '''simple docstring'''
    __lowerCamelCase  							=      self.mega(lowercase_		,			lowercase_		,			lowercase_		,			lowercase_			)
    return self.mlm_head(self.dropout(lowercase_			)			)





def 						a__     (						_UpperCamelCase     :	Any	,_UpperCamelCase     :	str	,_UpperCamelCase     :	List[str]							):




   with open(os.path.join(_UpperCamelCase	,'''model_args.pkl'''							)	,'''rb'''							) as f:
      __lowerCamelCase  							=      pkl.load(_UpperCamelCase							)

   # load the original encoder
   __lowerCamelCase  							=      OriginalMegaForMaskedLM(**_UpperCamelCase							).eval()

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

   __lowerCamelCase  							=      MegaForMaskedLM(_UpperCamelCase							).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(_UpperCamelCase							) != 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(_UpperCamelCase							)

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

   # load the MLM head weights directly
   print(
       '''HF Mega MLM layer:'''	,hf_mlm.mlm_head.load_state_dict(
           torch.load(os.path.join(_UpperCamelCase	,'''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(_UpperCamelCase							)
   # mask a few tokens to make sure masking is applied appropriately :)
   __lowerCamelCase  							=      0

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

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


if __name__ == "__main__":
  a_								=   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_								=   parser.parse_args()

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











import json
from typing import List, Optional, Tuple

from tokenizers import pre_tokenizers, processors

from ...tokenization_utils_base import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_longformer import LongformerTokenizer


a_								=   logging.get_logger(__name__)

a_								=   {"""vocab_file""": """vocab.json""", """merges_file""": """merges.txt""", """tokenizer_file""": """tokenizer.json"""}

a_								=   {
    """vocab_file""": {
        """allenai/longformer-base-4096""": """https://huggingface.co/allenai/longformer-base-4096/resolve/main/vocab.json""",
        """allenai/longformer-large-4096""": (
            """https://huggingface.co/allenai/longformer-large-4096/resolve/main/vocab.json"""
        ),
        """allenai/longformer-large-4096-finetuned-triviaqa""": (
            """https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/vocab.json"""
        ),
        """allenai/longformer-base-4096-extra.pos.embd.only""": (
            """https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/vocab.json"""
        ),
        """allenai/longformer-large-4096-extra.pos.embd.only""": (
            """https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/vocab.json"""
        ),
    },
    """merges_file""": {
        """allenai/longformer-base-4096""": """https://huggingface.co/allenai/longformer-base-4096/resolve/main/merges.txt""",
        """allenai/longformer-large-4096""": (
            """https://huggingface.co/allenai/longformer-large-4096/resolve/main/merges.txt"""
        ),
        """allenai/longformer-large-4096-finetuned-triviaqa""": (
            """https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/merges.txt"""
        ),
        """allenai/longformer-base-4096-extra.pos.embd.only""": (
            """https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/merges.txt"""
        ),
        """allenai/longformer-large-4096-extra.pos.embd.only""": (
            """https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/merges.txt"""
        ),
    },
    """tokenizer_file""": {
        """allenai/longformer-base-4096""": (
            """https://huggingface.co/allenai/longformer-base-4096/resolve/main/tokenizer.json"""
        ),
        """allenai/longformer-large-4096""": (
            """https://huggingface.co/allenai/longformer-large-4096/resolve/main/tokenizer.json"""
        ),
        """allenai/longformer-large-4096-finetuned-triviaqa""": (
            """https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/tokenizer.json"""
        ),
        """allenai/longformer-base-4096-extra.pos.embd.only""": (
            """https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/tokenizer.json"""
        ),
        """allenai/longformer-large-4096-extra.pos.embd.only""": (
            """https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/tokenizer.json"""
        ),
    },
}

a_								=   {
    """allenai/longformer-base-4096""": 4_096,
    """allenai/longformer-large-4096""": 4_096,
    """allenai/longformer-large-4096-finetuned-triviaqa""": 4_096,
    """allenai/longformer-base-4096-extra.pos.embd.only""": 4_096,
    """allenai/longformer-large-4096-extra.pos.embd.only""": 4_096,
}



class 	__lowerCAmelCase       (		lowerCAmelCase__				):
 lowerCAmelCase__       =					VOCAB_FILES_NAMES
 lowerCAmelCase__       =					PRETRAINED_VOCAB_FILES_MAP
 lowerCAmelCase__       =					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 lowerCAmelCase__       =					["""input_ids""", """attention_mask"""]
 lowerCAmelCase__       =					LongformerTokenizer




 def __init__(  self		,			__UpperCAmelCase=None		,			__UpperCAmelCase=None		,			__UpperCAmelCase=None		,			__UpperCAmelCase="replace"		,			__UpperCAmelCase="<s>"		,			__UpperCAmelCase="</s>"		,			__UpperCAmelCase="</s>"		,			__UpperCAmelCase="<s>"		,			__UpperCAmelCase="<unk>"		,			__UpperCAmelCase="<pad>"		,			__UpperCAmelCase="<mask>"		,			__UpperCAmelCase=False		,			__UpperCAmelCase=True		,			**__UpperCAmelCase		,			):



    '''simple docstring'''
    super().__init__(
        __UpperCAmelCase		,			__UpperCAmelCase		,			tokenizer_file=__UpperCAmelCase		,			errors=__UpperCAmelCase		,			bos_token=__UpperCAmelCase		,			eos_token=__UpperCAmelCase		,			sep_token=__UpperCAmelCase		,			cls_token=__UpperCAmelCase		,			unk_token=__UpperCAmelCase		,			pad_token=__UpperCAmelCase		,			mask_token=__UpperCAmelCase		,			add_prefix_space=__UpperCAmelCase		,			trim_offsets=__UpperCAmelCase		,			**__UpperCAmelCase		,			)

    __lowerCamelCase  							=      json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()			)
    if pre_tok_state.get('''add_prefix_space'''		,			__UpperCAmelCase			) != add_prefix_space:
       __lowerCamelCase  							=      getattr(__UpperCAmelCase		,			pre_tok_state.pop('''type'''			)			)
       __lowerCamelCase  							=      add_prefix_space
       __lowerCamelCase  							=      pre_tok_class(**__UpperCAmelCase			)

    __lowerCamelCase  							=      add_prefix_space

    __lowerCamelCase  							=      '''post_processor'''
    __lowerCamelCase  							=      getattr(self.backend_tokenizer		,			__UpperCAmelCase		,			__UpperCAmelCase			)
    if tokenizer_component_instance:
       __lowerCamelCase  							=      json.loads(tokenizer_component_instance.__getstate__()			)

       # The lists 'sep' and 'cls' must be cased in tuples for the object `post_processor_class`
       if "sep" in state:
          __lowerCamelCase  							=      tuple(state['''sep''']			)
       if "cls" in state:
          __lowerCamelCase  							=      tuple(state['''cls''']			)

       __lowerCamelCase  							=      False

       if state.get('''add_prefix_space'''		,			__UpperCAmelCase			) != add_prefix_space:
          __lowerCamelCase  							=      add_prefix_space
          __lowerCamelCase  							=      True

       if state.get('''trim_offsets'''		,			__UpperCAmelCase			) != trim_offsets:
          __lowerCamelCase  							=      trim_offsets
          __lowerCamelCase  							=      True

       if changes_to_apply:
          __lowerCamelCase  							=      getattr(__UpperCAmelCase		,			state.pop('''type'''			)			)
          __lowerCamelCase  							=      component_class(**__UpperCAmelCase			)
          setattr(self.backend_tokenizer		,			__UpperCAmelCase		,			__UpperCAmelCase			)




 @property
 def     lowerCamelCase (  self			):



    '''simple docstring'''
    if self._mask_token is None:
       if self.verbose:
          logger.error('''Using mask_token, but it is not set yet.'''			)
       return None
    return str(self._mask_token			)




 @mask_token.setter
 def     lowerCamelCase (  self		,			__UpperCAmelCase			):



    '''simple docstring'''
    __lowerCamelCase  							=      AddedToken(__UpperCAmelCase		,			lstrip=__UpperCAmelCase		,			rstrip=__UpperCAmelCase			) if isinstance(__UpperCAmelCase		,			__UpperCAmelCase			) else value
    __lowerCamelCase  							=      value




 def     lowerCamelCase (  self		,			*__UpperCAmelCase		,			**__UpperCAmelCase			):



    '''simple docstring'''
    __lowerCamelCase  							=      kwargs.get('''is_split_into_words'''		,			__UpperCAmelCase			)
    assert self.add_prefix_space or not is_split_into_words, (
        F"""You need to instantiate {self.__class__.__name__} with add_prefix_space=True """
        "to use it with pretokenized inputs."
    )

    return super()._batch_encode_plus(*__UpperCAmelCase		,			**__UpperCAmelCase			)




 def     lowerCamelCase (  self		,			*__UpperCAmelCase		,			**__UpperCAmelCase			):



    '''simple docstring'''
    __lowerCamelCase  							=      kwargs.get('''is_split_into_words'''		,			__UpperCAmelCase			)

    assert self.add_prefix_space or not is_split_into_words, (
        F"""You need to instantiate {self.__class__.__name__} with add_prefix_space=True """
        "to use it with pretokenized inputs."
    )

    return super()._encode_plus(*__UpperCAmelCase		,			**__UpperCAmelCase			)




 def     lowerCamelCase (  self		,			__UpperCAmelCase		,			__UpperCAmelCase = None			):



    '''simple docstring'''
    __lowerCamelCase  							=      self._tokenizer.model.save(__UpperCAmelCase		,			name=__UpperCAmelCase			)
    return tuple(__UpperCAmelCase			)




 def     lowerCamelCase (  self		,			__UpperCAmelCase		,			__UpperCAmelCase=None			):



    '''simple docstring'''
    __lowerCamelCase  							=      [self.bos_token_id] + token_ids_a + [self.eos_token_id]
    if token_ids_a is None:
       return output

    return output + [self.eos_token_id] + token_ids_a + [self.eos_token_id]




 def     lowerCamelCase (  self		,			__UpperCAmelCase		,			__UpperCAmelCase = None			):



    '''simple docstring'''
    __lowerCamelCase  							=      [self.sep_token_id]
    __lowerCamelCase  							=      [self.cls_token_id]

    if token_ids_a is None:
       return len(cls + token_ids_a + sep			) * [0]
    return len(cls + token_ids_a + sep + sep + token_ids_a + sep			) * [0]






