



'''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					:			Any							,					lowerCAmelCase__					:			Tuple							,					lowerCAmelCase__					:			Dict							,					lowerCAmelCase__					:			str						):





					"""simple docstring"""






					super().__init__()
					__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		mega_args
					__SCREAMING_SNAKE_CASE	:		int   	=		nn.Embedding(lowerCAmelCase__							,					self.mega_args.encoder_embed_dim						)
					__SCREAMING_SNAKE_CASE	:		Any   	=		nn.ModuleList([MegaEncoderLayer(self.mega_args						) for _ in range(lowerCAmelCase__						)]						)
					__SCREAMING_SNAKE_CASE	:		str   	=		depth



			def UpperCamelCase__ (							self					:			Optional[Any]							,					lowerCAmelCase__					:			Union[str, Any]							,					lowerCAmelCase__					:			Union[str, Any]							,					lowerCAmelCase__					:			Union[str, Any]=True							,					lowerCAmelCase__					:			int=0						):





					"""simple docstring"""






					if batch_first:
							__SCREAMING_SNAKE_CASE	:		List[str]   	=		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:
							__SCREAMING_SNAKE_CASE	:		int   	=		1 - attention_mask

					# get token embeddings from IDs
					__SCREAMING_SNAKE_CASE	:		List[Any]   	=		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:
							__SCREAMING_SNAKE_CASE	:		List[Any]   	=		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   ):


			'''simple docstring'''

			def __init__(							self					:			Tuple							,					lowerCAmelCase__					:			Optional[Any]							,					lowerCAmelCase__					:			Tuple							,					lowerCAmelCase__					:			int						):





					"""simple docstring"""






					super().__init__()
					__SCREAMING_SNAKE_CASE	:		Any   	=		MegaLM(lowerCAmelCase__							,					lowerCAmelCase__							,					lowerCAmelCase__						)
					__SCREAMING_SNAKE_CASE	:		List[Any]   	=		nn.Linear(mega_args.encoder_embed_dim							,					lowerCAmelCase__						)
					__SCREAMING_SNAKE_CASE	:		List[str]   	=		nn.Dropout(p=0.1						)



			def UpperCamelCase__ (							self					:			Tuple							,					lowerCAmelCase__					:			List[str]							,					lowerCAmelCase__					:			List[str]							,					lowerCAmelCase__					:			int=True							,					lowerCAmelCase__					:			str=0						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		Any   	=		self.mega(lowerCAmelCase__							,					lowerCAmelCase__							,					lowerCAmelCase__							,					lowerCAmelCase__						)
					return self.mlm_head(self.dropout(lowerCAmelCase__						)						)


def     lowerCAmelCase_       (     _lowerCamelCase:       Union[str, Any]	,				_lowerCamelCase:       int	,				_lowerCamelCase:       List[str] ):
		with open(os.path.join(_lowerCamelCase	,				"""model_args.pkl""" )	,				"""rb""" ) as f:
				__SCREAMING_SNAKE_CASE	:		int   	=		pkl.load(_lowerCamelCase )

		# load the original encoder
		__SCREAMING_SNAKE_CASE	:		str   	=		OriginalMegaForMaskedLM(**_lowerCamelCase ).eval()

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

		# create a new config from the old one
		__SCREAMING_SNAKE_CASE	:		int   	=		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=_lowerCamelCase	,				norm_affine=_lowerCamelCase	,				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=_lowerCamelCase	,				add_lm_hidden_dense_layer=_lowerCamelCase	,				)

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

				if new_module_name:
						__SCREAMING_SNAKE_CASE	:		str   	=		new_module_name

		if len(_lowerCamelCase ) != 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():
				__SCREAMING_SNAKE_CASE	:		List[str]   	=		original_state_dict.pop(_lowerCamelCase )

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

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

		# test on a randomly generated input sequence
		__SCREAMING_SNAKE_CASE	:		Dict   	=		torch.randint(0	,				hf_config.vocab_size	,				size=(4, 2_56) )
		__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		torch.ones_like(_lowerCamelCase )
		# mask a few tokens to make sure masking is applied appropriately :)
		__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		0

		# run forward passes
		__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		original_mlm(_lowerCamelCase	,				_lowerCamelCase	,				batch_first=_lowerCamelCase	,				ignore_mask_value=0 )
		__SCREAMING_SNAKE_CASE	:		List[Any]   	=		hf_mlm(_lowerCamelCase	,				_lowerCamelCase )[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
		__SCREAMING_SNAKE_CASE	:		int   	=		torch.allclose(_lowerCamelCase	,				_lowerCamelCase	,				atol=1E-3 )

		if success:
				print("""Yay!""" )
				hf_mlm.save_pretrained(_lowerCamelCase )
		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""" )
				__SCREAMING_SNAKE_CASE	:		List[Any]   	=		AutoTokenizer.from_pretrained(_lowerCamelCase )
				tokenizer.save_pretrained(_lowerCamelCase )


if __name__ == "__main__":
						UpperCamelCase__	:     int    =			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__	:     Optional[int]    =			parser.parse_args()

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




'''simple docstring'''




UpperCamelCase__	:     List[Any]    =			10_00


def     lowerCAmelCase_       (     _lowerCamelCase:       float ):
		return x * x


def     lowerCAmelCase_       (     _lowerCamelCase:       Any	,				_lowerCamelCase:       float	,				_lowerCamelCase:       float	,				_lowerCamelCase:       int = 4 ):
		assert callable(
		    _lowerCamelCase ), F"the function(object) passed should be callable your input : {function}"
		assert isinstance(_lowerCamelCase	,				(float, int) ), F"a should be float or integer your input : {a}"
		assert isinstance(function(_lowerCamelCase )	,				(float, int) ), (
		    "the function should return integer or float return type of your function, "
		    F"{type(_lowerCamelCase )}"
		)
		assert isinstance(_lowerCamelCase	,				(float, int) ), F"b should be float or integer your input : {b}"
		assert (
		    isinstance(_lowerCamelCase	,				_lowerCamelCase ) and precision > 0
		), F"precision should be positive integer your input : {precision}"

		# just applying the formula of simpson for approximate integration written in
		# mentioned article in first comment of this file and above this function

		__SCREAMING_SNAKE_CASE	:		List[Any]   	=		(b - a) / N_STEPS
		__SCREAMING_SNAKE_CASE	:		List[Any]   	=		function(_lowerCamelCase ) + function(_lowerCamelCase )

		for i in range(1	,				_lowerCamelCase ):
				__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		a + h * i
				result += function(_lowerCamelCase ) * (4 if i % 2 else 2)

		result *= h / 3
		return round(_lowerCamelCase	,				_lowerCamelCase )


if __name__ == "__main__":
						import doctest

						doctest.testmod()