



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







	'''simple docstring'''




	def __init__(	self   ,       lowerCamelCase   ,       lowerCamelCase   ,       lowerCamelCase       ):



				'''simple docstring'''


				super().__init__()
				a__					=  mega_args
				a__					=  nn.Embedding(lowerCamelCase   ,       self.mega_args.encoder_embed_dim       )
				a__					=  nn.ModuleList([MegaEncoderLayer(self.mega_args       ) for _ in range(lowerCamelCase       )]       )
				a__					=  depth







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



				'''simple docstring'''


				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(lowerCamelCase       )

				# pass through the Mega layers
				# input is (time, batch, encoder dim) and output is the same
				for encoder in self.encoders:
							a__					=  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 lowerCAmelCase_       (   nn.Module			):







	'''simple docstring'''




	def __init__(	self   ,       lowerCamelCase   ,       lowerCamelCase   ,       lowerCamelCase       ):



				'''simple docstring'''


				super().__init__()
				a__					=  MegaLM(lowerCamelCase   ,       lowerCamelCase   ,       lowerCamelCase       )
				a__					=  nn.Linear(mega_args.encoder_embed_dim   ,       lowerCamelCase       )
				a__					=  nn.Dropout(p=0.1       )







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



				'''simple docstring'''


				a__					=  self.mega(lowerCamelCase   ,       lowerCamelCase   ,       lowerCamelCase   ,       lowerCamelCase       )
				return self.mlm_head(self.dropout(lowerCamelCase       )       )




def    UpperCAmelCase					(  lowercase__						:					Optional[int]				,					lowercase__						:					Dict				,					lowercase__						:					List[Any]	):



			'''simple docstring'''





			with open(os.path.join(lowercase__				,					"""model_args.pkl"""	)				,					"""rb"""	) as f:
						a__					=  pkl.load(lowercase__	)

			# load the original encoder
			a__					=  OriginalMegaForMaskedLM(**lowercase__	).eval()

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

			a__					=  MegaForMaskedLM(lowercase__	).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(lowercase__	) != 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(lowercase__	)

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

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

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

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


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

	_lowercase				:					int     					=parser.parse_args()

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






from typing import TYPE_CHECKING

from ...utils import _LazyModule


_lowercase				:					List[Any]     					={"""tokenization_bert_japanese""": ["""BertJapaneseTokenizer""", """CharacterTokenizer""", """MecabTokenizer"""]}


if TYPE_CHECKING:
	from .tokenization_bert_japanese import BertJapaneseTokenizer, CharacterTokenizer, MecabTokenizer

else:
	import sys

	_lowercase				:					Union[str, Any]     					=_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

