

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





			def __init__(	self      ,						__snake_case      ,						__snake_case      ,						__snake_case      )     ->				List[Any]:




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





			def 	__magic_name__					(	self      ,						__snake_case      ,						__snake_case      ,						__snake_case=True      ,						__snake_case=0      )     ->				Dict:




						'''simple docstring'''

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

						# get token embeddings from IDs
						__a						   =self.embedding_layer(__snake_case      )

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

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






class 			__magic_name__					(    nn.Module   ):





			def __init__(	self      ,						__snake_case      ,						__snake_case      ,						__snake_case      )     ->				Any:




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





			def 	__magic_name__					(	self      ,						__snake_case      ,						__snake_case      ,						__snake_case=True      ,						__snake_case=0      )     ->				Union[str, Any]:




						'''simple docstring'''
						__a						   =self.mega(__snake_case      ,						__snake_case      ,						__snake_case      ,						__snake_case      )
						return self.mlm_head(self.dropout(__snake_case      )      )






def  UpperCamelCase_(							_snake_case     :						List[str] ,  _snake_case     :						List[Any] ,  _snake_case     :						str			):




			"""simple docstring"""






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

			# load the original encoder
			__a						   =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
			__a						   =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 ,  )

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

						if new_module_name:
									__a						   =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():
						__a						   =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
			__a						   =torch.randint(0 ,  hf_config.vocab_size ,  size=(4, 256)			)
			__a						   =torch.ones_like(_snake_case			)
			# mask a few tokens to make sure masking is applied appropriately :)
			__a						   =0

			# run forward passes
			__a						   =original_mlm(_snake_case ,  _snake_case ,  batch_first=_snake_case ,  ignore_mask_value=0			)
			__a						   =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
			__a						   =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'			)
						__a						   =AutoTokenizer.from_pretrained(_snake_case			)
						tokenizer.save_pretrained(_snake_case			)


if __name__ == "__main__":
						_lowerCAmelCase      :    str						       =    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",
						)

						_lowerCAmelCase      :    List[Any]						       =    parser.parse_args()

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









from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


_lowerCAmelCase      :    str						       =    logging.get_logger(__name__)

_lowerCAmelCase      :    List[str]						       =    {
    "facebook/mbart-large-cc25": "https://huggingface.co/facebook/mbart-large-cc25/resolve/main/config.json",
    # See all MBART models at https://huggingface.co/models?filter=mbart
}






class 			__magic_name__					(    lowerCAmelCase_   ):
			SCREAMING_SNAKE_CASE							     =		'mbart'
			SCREAMING_SNAKE_CASE							     =		['past_key_values']
			SCREAMING_SNAKE_CASE							     =		{'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}





			def __init__(	self      ,						__snake_case=5_0265      ,						__snake_case=1024      ,						__snake_case=12      ,						__snake_case=4096      ,						__snake_case=16      ,						__snake_case=12      ,						__snake_case=4096      ,						__snake_case=16      ,						__snake_case=0.0      ,						__snake_case=0.0      ,						__snake_case=True      ,						__snake_case=True      ,						__snake_case="gelu"      ,						__snake_case=1024      ,						__snake_case=0.1      ,						__snake_case=0.0      ,						__snake_case=0.0      ,						__snake_case=0.02      ,						__snake_case=0.0      ,						__snake_case=False      ,						__snake_case=1      ,						__snake_case=0      ,						__snake_case=2      ,						__snake_case=2      ,						**__snake_case      ,						)     ->				Optional[Any]:




						'''simple docstring'''
						__a						   =vocab_size
						__a						   =max_position_embeddings
						__a						   =d_model
						__a						   =encoder_ffn_dim
						__a						   =encoder_layers
						__a						   =encoder_attention_heads
						__a						   =decoder_ffn_dim
						__a						   =decoder_layers
						__a						   =decoder_attention_heads
						__a						   =dropout
						__a						   =attention_dropout
						__a						   =activation_dropout
						__a						   =activation_function
						__a						   =init_std
						__a						   =encoder_layerdrop
						__a						   =decoder_layerdrop
						__a						   =classifier_dropout
						__a						   =use_cache
						__a						   =encoder_layers
						__a						   =scale_embedding  # scale factor will be sqrt(d_model) if True
						super().__init__(
						    pad_token_id=__snake_case      ,						bos_token_id=__snake_case      ,						eos_token_id=__snake_case      ,						is_encoder_decoder=__snake_case      ,						forced_eos_token_id=__snake_case      ,						**__snake_case      ,						)






class 			__magic_name__					(    lowerCAmelCase_   ):





			@property
			def 	__magic_name__					(	self      )     ->				Mapping[str, Mapping[int, str]]:




						'''simple docstring'''
						if self.task in ["default", "seq2seq-lm"]:
									__a						   =OrderedDict(
									    [
									        ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
									        ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
									    ]      )

									if self.use_past:
												__a						   ={0: 'batch'}
												__a						   ={0: 'batch', 1: 'past_decoder_sequence + sequence'}
									else:
												__a						   ={0: 'batch', 1: 'decoder_sequence'}
												__a						   ={0: 'batch', 1: 'decoder_sequence'}

									if self.use_past:
												self.fill_with_past_key_values_(__snake_case      ,						direction='inputs'      )
						elif self.task == "causal-lm":
									# TODO: figure this case out.
									__a						   =OrderedDict(
									    [
									        ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
									        ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
									    ]      )
									if self.use_past:
												__a   ,					__a						   =self.num_layers
												for i in range(__snake_case      ):
															__a						   ={0: 'batch', 2: 'past_sequence + sequence'}
															__a						   ={0: 'batch', 2: 'past_sequence + sequence'}
						else:
									__a						   =OrderedDict(
									    [
									        ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
									        ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
									        ('decoder_input_ids', {0: 'batch', 1: 'decoder_sequence'}),
									        ('decoder_attention_mask', {0: 'batch', 1: 'decoder_sequence'}),
									    ]      )

						return common_inputs





			@property
			def 	__magic_name__					(	self      )     ->				Mapping[str, Mapping[int, str]]:




						'''simple docstring'''
						if self.task in ["default", "seq2seq-lm"]:
									__a						   =super().outputs
						else:
									__a						   =super(__snake_case      ,						self      ).outputs
									if self.use_past:
												__a   ,					__a						   =self.num_layers
												for i in range(__snake_case      ):
															__a						   ={0: 'batch', 2: 'past_sequence + sequence'}
															__a						   ={0: 'batch', 2: 'past_sequence + sequence'}
						return common_outputs





			def 	__magic_name__					(	self      ,						__snake_case      ,						__snake_case = -1      ,						__snake_case = -1      ,						__snake_case = False      ,						__snake_case = None      ,						)     ->				Mapping[str, Any]:




						'''simple docstring'''
						__a						   =self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
						    __snake_case      ,						__snake_case      ,						__snake_case      ,						__snake_case      ,						__snake_case      )

						# Generate decoder inputs
						__a						   =seq_length if not self.use_past else 1
						__a						   =self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
						    __snake_case      ,						__snake_case      ,						__snake_case      ,						__snake_case      ,						__snake_case      )
						__a						   ={f'decoder_{name}': tensor for name, tensor in decoder_inputs.items()}
						__a						   =dict(**__snake_case      ,						**__snake_case      )

						if self.use_past:
									if not is_torch_available():
												raise ValueError('Cannot generate dummy past_keys inputs without PyTorch installed.'      )
									else:
												import torch
									__a   ,					__a						   =common_inputs['input_ids'].shape
									__a						   =common_inputs['decoder_input_ids'].shape[1]
									__a   ,					__a						   =self.num_attention_heads
									__a						   =(
									    batch,
									    num_encoder_attention_heads,
									    encoder_seq_length,
									    self._config.hidden_size // num_encoder_attention_heads,
									)
									__a						   =decoder_seq_length + 3
									__a						   =(
									    batch,
									    num_decoder_attention_heads,
									    decoder_past_length,
									    self._config.hidden_size // num_decoder_attention_heads,
									)

									__a						   =torch.cat(
									    [common_inputs['decoder_attention_mask'], torch.ones(__snake_case      ,						__snake_case      )]      ,						dim=1      )

									__a						   =[]
									# If the number of encoder and decoder layers are present in the model configuration, both are considered
									__a   ,					__a						   =self.num_layers
									__a						   =min(__snake_case      ,						__snake_case      )
									__a						   =max(__snake_case      ,						__snake_case      ) - min_num_layers
									__a						   ='encoder' if num_encoder_layers > num_decoder_layers else 'decoder'

									for _ in range(__snake_case      ):
												common_inputs["past_key_values"].append(
												    (
												        torch.zeros(__snake_case      ),
												        torch.zeros(__snake_case      ),
												        torch.zeros(__snake_case      ),
												        torch.zeros(__snake_case      ),
												    )      )
									# TODO: test this.
									__a						   =encoder_shape if remaining_side_name == 'encoder' else decoder_shape
									for _ in range(__snake_case      ,						__snake_case      ):
												common_inputs["past_key_values"].append((torch.zeros(__snake_case      ), torch.zeros(__snake_case      ))      )
						return common_inputs





			def 	__magic_name__					(	self      ,						__snake_case      ,						__snake_case = -1      ,						__snake_case = -1      ,						__snake_case = False      ,						__snake_case = None      ,						)     ->				Mapping[str, Any]:




						'''simple docstring'''
						__a						   =self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
						    __snake_case      ,						__snake_case      ,						__snake_case      ,						__snake_case      ,						__snake_case      )

						if self.use_past:
									if not is_torch_available():
												raise ValueError('Cannot generate dummy past_keys inputs without PyTorch installed.'      )
									else:
												import torch
									__a   ,					__a						   =common_inputs['input_ids'].shape
									# Not using the same length for past_key_values
									__a						   =seqlen + 2
									__a   ,					__a						   =self.num_layers
									__a   ,					__a						   =self.num_attention_heads
									__a						   =(
									    batch,
									    num_encoder_attention_heads,
									    past_key_values_length,
									    self._config.hidden_size // num_encoder_attention_heads,
									)

									__a						   =common_inputs['attention_mask'].dtype
									__a						   =torch.cat(
									    [common_inputs['attention_mask'], torch.ones(__snake_case      ,						__snake_case      ,						dtype=__snake_case      )]      ,						dim=1      )
									__a						   =[
									    (torch.zeros(__snake_case      ), torch.zeros(__snake_case      )) for _ in range(__snake_case      )
									]
						return common_inputs





			def 	__magic_name__					(	self      ,						__snake_case      ,						__snake_case = -1      ,						__snake_case = -1      ,						__snake_case = False      ,						__snake_case = None      ,						)     ->				Mapping[str, Any]:




						'''simple docstring'''
						# Copied from OnnxConfig.generate_dummy_inputs
						# Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
						# If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
						__a						   =compute_effective_axis_dimension(
						    __snake_case      ,						fixed_dimension=OnnxConfig.default_fixed_batch      ,						num_token_to_add=0      )

						# If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
						__a						   =tokenizer.num_special_tokens_to_add(__snake_case      )
						__a						   =compute_effective_axis_dimension(
						    __snake_case      ,						fixed_dimension=OnnxConfig.default_fixed_sequence      ,						num_token_to_add=__snake_case      )

						# Generate dummy inputs according to compute batch and sequence
						__a						   =[' '.join([tokenizer.unk_token]      ) * seq_length] * batch_size
						__a						   =dict(tokenizer(__snake_case      ,						return_tensors=__snake_case      )      )
						return common_inputs





			def 	__magic_name__					(	self      ,						__snake_case      ,						__snake_case = -1      ,						__snake_case = -1      ,						__snake_case = False      ,						__snake_case = None      ,						)     ->				Mapping[str, Any]:




						'''simple docstring'''
						if self.task in ["default", "seq2seq-lm"]:
									__a						   =self._generate_dummy_inputs_for_default_and_seqaseq_lm(
									    __snake_case      ,						batch_size=__snake_case      ,						seq_length=__snake_case      ,						is_pair=__snake_case      ,						framework=__snake_case      )

						elif self.task == "causal-lm":
									__a						   =self._generate_dummy_inputs_for_causal_lm(
									    __snake_case      ,						batch_size=__snake_case      ,						seq_length=__snake_case      ,						is_pair=__snake_case      ,						framework=__snake_case      )
						else:
									__a						   =self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
									    __snake_case      ,						batch_size=__snake_case      ,						seq_length=__snake_case      ,						is_pair=__snake_case      ,						framework=__snake_case      )

						return common_inputs





			def 	__magic_name__					(	self      ,						__snake_case      ,						__snake_case      ,						__snake_case      ,						__snake_case      )     ->				str:




						'''simple docstring'''
						if self.task in ["default", "seq2seq-lm"]:
									__a						   =super()._flatten_past_key_values_(__snake_case      ,						__snake_case      ,						__snake_case      ,						__snake_case      )
						else:
									__a						   =super(__snake_case      ,						self      )._flatten_past_key_values_(
									    __snake_case      ,						__snake_case      ,						__snake_case      ,						__snake_case      )






