



"""simple docstring"""






import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTNeoConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
		import jax
		import jax.numpy as jnp

		from transformers.modeling_flax_pytorch_utils import (
		    convert_pytorch_state_dict_to_flax,
		    load_flax_weights_in_pytorch_model,
		)
		from transformers.models.gpt_neo.modeling_flax_gpt_neo import FlaxGPTNeoForCausalLM, FlaxGPTNeoModel

if is_torch_available():
		import torch

class 			__SCREAMING_SNAKE_CASE      :


					'''simple docstring'''





					def __init__(			self				:    Optional[Any],    lowerCamelCase				:    Any,    lowerCamelCase				:    Any=14,    lowerCamelCase				:    Tuple=7,    lowerCamelCase				:    List[Any]=True,    lowerCamelCase				:    Optional[int]=True,    lowerCamelCase				:    List[str]=False,    lowerCamelCase				:    List[Any]=True,    lowerCamelCase				:    List[str]=99,    lowerCamelCase				:    Dict=32,    lowerCamelCase				:    int=4,    lowerCamelCase				:    Optional[int]=4,    lowerCamelCase				:    int=[[["global", "local"], 2]],    lowerCamelCase				:    List[Any]=37,    lowerCamelCase				:    Dict="gelu",    lowerCamelCase				:    Union[str, Any]=0.1,    lowerCamelCase				:    Tuple=0.1,    lowerCamelCase				:    List[str]=512,    lowerCamelCase				:    Optional[int]=7,    lowerCamelCase				:    Tuple=0.02,    )->   int:
								lowerCamelCase__						:						Dict						      =parent
								lowerCamelCase__						:						Tuple						      =batch_size
								lowerCamelCase__						:						Dict						      =seq_length
								lowerCamelCase__						:						Dict						      =is_training
								lowerCamelCase__						:						Union[str, Any]						      =use_input_mask
								lowerCamelCase__						:						Union[str, Any]						      =use_token_type_ids
								lowerCamelCase__						:						str						      =use_labels
								lowerCamelCase__						:						int						      =vocab_size
								lowerCamelCase__						:						List[str]						      =hidden_size
								lowerCamelCase__						:						Optional[Any]						      =num_hidden_layers
								lowerCamelCase__						:						Dict						      =num_attention_heads
								lowerCamelCase__						:						Tuple						      =attention_types
								lowerCamelCase__						:						str						      =intermediate_size
								lowerCamelCase__						:						str						      =hidden_act
								lowerCamelCase__						:						Optional[int]						      =hidden_dropout_prob
								lowerCamelCase__						:						str						      =attention_probs_dropout_prob
								lowerCamelCase__						:						Optional[Any]						      =max_position_embeddings
								lowerCamelCase__						:						Union[str, Any]						      =window_size
								lowerCamelCase__						:						Optional[int]						      =initializer_range
								lowerCamelCase__						:						List[Any]						      =None
								lowerCamelCase__						:						Tuple						      =vocab_size - 1
								lowerCamelCase__						:						Union[str, Any]						      =vocab_size - 1
								lowerCamelCase__						:						Dict						      =vocab_size - 1



					def 						snake_case			(			self				:    Optional[int]	)->   str:
								lowerCamelCase__						:						Optional[Any]						      =ids_tensor([self.batch_size, self.seq_length],    self.vocab_size	)

								lowerCamelCase__						:						Dict						      =None
								if self.use_input_mask:
											lowerCamelCase__						:						List[str]						      =random_attention_mask([self.batch_size, self.seq_length]	)

								lowerCamelCase__						:						Union[str, Any]						      =GPTNeoConfig(
								    vocab_size=self.vocab_size,    hidden_size=self.hidden_size,    num_layers=self.num_hidden_layers,    num_heads=self.num_attention_heads,    max_position_embeddings=self.max_position_embeddings,    use_cache=__snake_case,    bos_token_id=self.bos_token_id,    eos_token_id=self.eos_token_id,    pad_token_id=self.pad_token_id,    window_size=self.window_size,    attention_types=self.attention_types,    )

								return (config, input_ids, input_mask)



					def 						snake_case			(			self				:    Dict	)->   List[str]:
								lowerCamelCase__						:						str						      =self.prepare_config_and_inputs()
								lowerCamelCase__						:						List[str]						      =config_and_inputs
								lowerCamelCase__						:						str						      ={"input_ids": input_ids, "attention_mask": attention_mask}
								return config, inputs_dict



					def 						snake_case			(			self				:    Dict,    lowerCamelCase				:    int,    lowerCamelCase				:    int,    lowerCamelCase				:    List[str],    lowerCamelCase				:    int	)->   Union[str, Any]:
								lowerCamelCase__						:						Any						      =20
								lowerCamelCase__						:						Optional[int]						      =model_class_name(__snake_case	)

								lowerCamelCase__						:						Optional[int]						      =model.init_cache(input_ids.shape[0],    __snake_case	)
								lowerCamelCase__						:						Union[str, Any]						      =jnp.ones((input_ids.shape[0], max_decoder_length),    dtype='''i4'''	)

								lowerCamelCase__						:						Any						      =jnp.broadcast_to(
								    jnp.arange(input_ids.shape[-1] - 1	)[None, :],    (input_ids.shape[0], input_ids.shape[-1] - 1)	)
								lowerCamelCase__						:						Dict						      =model(
								    input_ids[:, :-1],    attention_mask=__snake_case,    past_key_values=__snake_case,    position_ids=__snake_case,    )

								lowerCamelCase__						:						Optional[Any]						      =jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]],    dtype='''i4'''	)
								lowerCamelCase__						:						Union[str, Any]						      =model(
								    input_ids[:, -1:],    attention_mask=__snake_case,    past_key_values=outputs_cache.past_key_values,    position_ids=__snake_case,    )

								lowerCamelCase__						:						int						      =model(__snake_case	)

								lowerCamelCase__						:						Dict						      =np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])	)	)
								self.parent.assertTrue(diff < 1E-3,    msg=F'''Max diff is {diff}'''	)




					def 						snake_case			(			self				:    Optional[int],    lowerCamelCase				:    Union[str, Any],    lowerCamelCase				:    Union[str, Any],    lowerCamelCase				:    Optional[int],    lowerCamelCase				:    Any	)->   Union[str, Any]:
								lowerCamelCase__						:						Tuple						      =20
								lowerCamelCase__						:						int						      =model_class_name(__snake_case	)

								lowerCamelCase__						:						Tuple						      =jnp.concatenate(
								    [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])	)],    axis=-1,    )

								lowerCamelCase__						:						Optional[int]						      =model.init_cache(input_ids.shape[0],    __snake_case	)
								lowerCamelCase__						:						Optional[Any]						      =jnp.broadcast_to(
								    jnp.arange(input_ids.shape[-1] - 1	)[None, :],    (input_ids.shape[0], input_ids.shape[-1] - 1)	)

								lowerCamelCase__						:						Union[str, Any]						      =model(
								    input_ids[:, :-1],    attention_mask=__snake_case,    past_key_values=__snake_case,    position_ids=__snake_case,    )
								lowerCamelCase__						:						Dict						      =jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]],    dtype='''i4'''	)
								lowerCamelCase__						:						List[str]						      =model(
								    input_ids[:, -1:],    past_key_values=outputs_cache.past_key_values,    attention_mask=__snake_case,    position_ids=__snake_case,    )

								lowerCamelCase__						:						Optional[int]						      =model(__snake_case,    attention_mask=__snake_case	)

								lowerCamelCase__						:						Tuple						      =np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])	)	)
								self.parent.assertTrue(diff < 1E-3,    msg=F'''Max diff is {diff}'''	)





@require_flax
class 			__SCREAMING_SNAKE_CASE      ( lowerCAmelCase_						,  lowerCAmelCase_						,  unittest.TestCase  ):


					'''simple docstring'''





					_a       =				(FlaxGPTNeoModel, FlaxGPTNeoForCausalLM) if is_flax_available() else ()
					_a       =				(FlaxGPTNeoForCausalLM,) if is_flax_available() else ()



					def 						snake_case			(			self				:    Tuple	)->   List[Any]:
								lowerCamelCase__						:						Union[str, Any]						      =FlaxGPTNeoModelTester(self	)



					def 						snake_case			(			self				:    int	)->   Optional[int]:



								for model_class_name in self.all_model_classes:
											lowerCamelCase__						:						Union[str, Any]						      =self.model_tester.prepare_config_and_inputs()
											self.model_tester.check_use_cache_forward(__snake_case,    __snake_case,    __snake_case,    __snake_case	)



					def 						snake_case			(			self				:    Optional[Any]	)->   Union[str, Any]:



								for model_class_name in self.all_model_classes:
											lowerCamelCase__						:						List[str]						      =self.model_tester.prepare_config_and_inputs()
											self.model_tester.check_use_cache_forward_with_attn_mask(
											    __snake_case,    __snake_case,    __snake_case,    __snake_case	)



					@slow
					def 						snake_case			(			self				:    Union[str, Any]	)->   Dict:
								lowerCamelCase__						:						List[Any]						      =GPTaTokenizer.from_pretrained('''gpt2''',    pad_token='''<|endoftext|>''',    padding_side='''left'''	)
								lowerCamelCase__						:						List[str]						      =tokenizer(['''Hello this is a long string''', '''Hey'''],    return_tensors='''np''',    padding=__snake_case,    truncation=__snake_case	)

								lowerCamelCase__						:						Optional[Any]						      =FlaxGPTNeoForCausalLM.from_pretrained('''EleutherAI/gpt-neo-125M'''	)
								lowerCamelCase__						:						Tuple						      =False
								lowerCamelCase__						:						int						      =model.config.eos_token_id

								lowerCamelCase__						:						Union[str, Any]						      =jax.jit(model.generate	)

								lowerCamelCase__						:						Optional[int]						      =jit_generate(
								    inputs['''input_ids'''],    attention_mask=inputs['''attention_mask'''],    pad_token_id=tokenizer.pad_token_id	).sequences

								lowerCamelCase__						:						Dict						      =tokenizer.batch_decode(__snake_case,    skip_special_tokens=__snake_case	)

								lowerCamelCase__						:						List[str]						      =[
								    "Hello this is a long string of text.\n\nI'm trying to get the text of the",
								    "Hey, I'm a little late to the party. I'm going to",
								]

								self.assertListEqual(__snake_case,    __snake_case	)



					@is_pt_flax_cross_test
					def 						snake_case			(			self				:    Optional[Any]	)->   List[str]:
								lowerCamelCase__						:						Tuple						      =self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
											with self.subTest(model_class.__name__	):
														# prepare inputs
														lowerCamelCase__						:						str						      =self._prepare_for_class(__snake_case,    __snake_case	)
														lowerCamelCase__						:						Dict						      ={k: torch.tensor(v.tolist()	) for k, v in prepared_inputs_dict.items()}

														# load corresponding PyTorch class
														lowerCamelCase__						:						List[str]						      =model_class.__name__[4:]  # Skip the "Flax" at the beginning
														lowerCamelCase__						:						List[Any]						      =getattr(__snake_case,    __snake_case	)

														lowerCamelCase__						:						int						      =pt_inputs["input_ids"].shape
														lowerCamelCase__						:						str						      =np.random.randint(0,    seq_length - 1,    size=(batch_size,)	)
														for batch_idx, start_index in enumerate(__snake_case	):
																	lowerCamelCase__						:						Optional[int]						      =0
																	lowerCamelCase__						:						List[Any]						      =1
																	lowerCamelCase__						:						int						      =0
																	lowerCamelCase__						:						Optional[int]						      =1
														lowerCamelCase__						:						List[Any]						      =pt_model_class(__snake_case	).eval()
														lowerCamelCase__						:						Any						      =model_class(__snake_case,    dtype=jnp.floataa	)

														lowerCamelCase__						:						Tuple						      =convert_pytorch_state_dict_to_flax(pt_model.state_dict(),    __snake_case	)
														lowerCamelCase__						:						Tuple						      =fx_state

														with torch.no_grad():
																	lowerCamelCase__						:						List[str]						      =pt_model(**__snake_case	).to_tuple()

														lowerCamelCase__						:						str						      =fx_model(**__snake_case	).to_tuple()
														self.assertEqual(len(__snake_case	),    len(__snake_case	),    '''Output lengths differ between Flax and PyTorch'''	)
														for fx_output, pt_output in zip(__snake_case,    __snake_case	):
																	self.assert_almost_equals(fx_output[:, -1],    pt_output[:, -1].numpy(),    4E-2	)

														with tempfile.TemporaryDirectory() as tmpdirname:
																	pt_model.save_pretrained(__snake_case	)
																	lowerCamelCase__						:						List[str]						      =model_class.from_pretrained(__snake_case,    from_pt=__snake_case	)

														lowerCamelCase__						:						str						      =fx_model_loaded(**__snake_case	).to_tuple()
														self.assertEqual(
														    len(__snake_case	),    len(__snake_case	),    '''Output lengths differ between Flax and PyTorch'''	)
														for fx_output_loaded, pt_output in zip(__snake_case,    __snake_case	):
																	self.assert_almost_equals(fx_output_loaded[:, -1],    pt_output[:, -1].numpy(),    4E-2	)



					@is_pt_flax_cross_test
					def 						snake_case			(			self				:    str	)->   Optional[int]:
								lowerCamelCase__						:						Optional[Any]						      =self.model_tester.prepare_config_and_inputs_for_common()
								for model_class in self.all_model_classes:
											with self.subTest(model_class.__name__	):
														# prepare inputs
														lowerCamelCase__						:						Optional[int]						      =self._prepare_for_class(__snake_case,    __snake_case	)
														lowerCamelCase__						:						Any						      ={k: torch.tensor(v.tolist()	) for k, v in prepared_inputs_dict.items()}

														# load corresponding PyTorch class
														lowerCamelCase__						:						Optional[Any]						      =model_class.__name__[4:]  # Skip the "Flax" at the beginning
														lowerCamelCase__						:						Any						      =getattr(__snake_case,    __snake_case	)

														lowerCamelCase__						:						Optional[Any]						      =pt_model_class(__snake_case	).eval()
														lowerCamelCase__						:						Tuple						      =model_class(__snake_case,    dtype=jnp.floataa	)

														lowerCamelCase__						:						Tuple						      =load_flax_weights_in_pytorch_model(__snake_case,    fx_model.params	)
														lowerCamelCase__						:						Optional[Any]						      =pt_inputs["input_ids"].shape
														lowerCamelCase__						:						str						      =np.random.randint(0,    seq_length - 1,    size=(batch_size,)	)
														for batch_idx, start_index in enumerate(__snake_case	):
																	lowerCamelCase__						:						str						      =0
																	lowerCamelCase__						:						Optional[int]						      =1
																	lowerCamelCase__						:						Optional[Any]						      =0
																	lowerCamelCase__						:						Optional[int]						      =1

														# make sure weights are tied in PyTorch
														pt_model.tie_weights()

														with torch.no_grad():
																	lowerCamelCase__						:						List[Any]						      =pt_model(**__snake_case	).to_tuple()

														lowerCamelCase__						:						List[Any]						      =fx_model(**__snake_case	).to_tuple()
														self.assertEqual(len(__snake_case	),    len(__snake_case	),    '''Output lengths differ between Flax and PyTorch'''	)
														for fx_output, pt_output in zip(__snake_case,    __snake_case	):
																	self.assert_almost_equals(fx_output[:, -1],    pt_output[:, -1].numpy(),    4E-2	)

														with tempfile.TemporaryDirectory() as tmpdirname:
																	fx_model.save_pretrained(__snake_case	)
																	lowerCamelCase__						:						Any						      =pt_model_class.from_pretrained(__snake_case,    from_flax=__snake_case	)

														with torch.no_grad():
																	lowerCamelCase__						:						Optional[Any]						      =pt_model_loaded(**__snake_case	).to_tuple()

														self.assertEqual(
														    len(__snake_case	),    len(__snake_case	),    '''Output lengths differ between Flax and PyTorch'''	)
														for fx_output, pt_output in zip(__snake_case,    __snake_case	):
																	self.assert_almost_equals(fx_output[:, -1],    pt_output[:, -1].numpy(),    4E-2	)




					@slow
					def 						snake_case			(			self				:    List[str]	)->   int:



								for model_class_name in self.all_model_classes:
											lowerCamelCase__						:						Union[str, Any]						      =model_class_name.from_pretrained('''EleutherAI/gpt-neo-125M'''	)
											lowerCamelCase__						:						Optional[int]						      =model(np.ones((1, 1)	)	)
											self.assertIsNotNone(__snake_case	)






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	,				__snake_case	,				__snake_case	,				__snake_case							)						->     Dict:






								"""simple docstring"""



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




							def     A__ (     self	,				__snake_case	,				__snake_case	,				__snake_case=True	,				__snake_case=0							)						->     List[str]:






								"""simple docstring"""



								if batch_first:
									UpperCAmelCase:					int 						=     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:
									UpperCAmelCase:					Union[str, Any] 						=     1 - attention_mask

								# get token embeddings from IDs
								UpperCAmelCase:					Dict 						=     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:
									UpperCAmelCase:					Any 						=     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 		__lowerCamelCase     (   nn.Module		):



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






								"""simple docstring"""



								super().__init__()
								UpperCAmelCase:					int 						=     MegaLM(__snake_case	,				__snake_case	,				__snake_case							)
								UpperCAmelCase:					Tuple 						=     nn.Linear(mega_args.encoder_embed_dim	,				__snake_case							)
								UpperCAmelCase:					List[str] 						=     nn.Dropout(p=0.1							)




							def     A__ (     self	,				__snake_case	,				__snake_case	,				__snake_case=True	,				__snake_case=0							)						->     Optional[Any]:






								"""simple docstring"""



								UpperCAmelCase:					str 						=     self.mega(__snake_case	,				__snake_case	,				__snake_case	,				__snake_case							)
								return self.mlm_head(self.dropout(__snake_case							)							)



def      __UpperCAmelCase    (  snake_case_							:		Optional[Any]				,					snake_case_							:		Dict				,					snake_case_							:		Tuple				):



	'''simple docstring'''



	with open(os.path.join(snake_case_				,					"model_args.pkl"				)				,					"rb"				) as f:
		UpperCAmelCase:					Union[str, Any] 						=     pkl.load(snake_case_				)

	# load the original encoder
	UpperCAmelCase:					Tuple 						=     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
	UpperCAmelCase:					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=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_				,					)

	UpperCAmelCase:					Union[str, Any] 						=     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
	UpperCAmelCase:					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
	UpperCAmelCase:					List[str] 						=     original_mlm.mega.encoders.state_dict()
	UpperCAmelCase:					Optional[Any] 						=     {}
	for module_name in original_state_dict.keys():
		UpperCAmelCase:					Optional[Any] 						=     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:
				UpperCAmelCase:					Any 						=     module_name.replace("move.beta"				,					"ema_gate.ema_expansion_matrix"				)
			elif "mega_layer.beta" in module_name:
				UpperCAmelCase:					Optional[int] 						=     module_name.replace("beta"				,					"qk_bias"				)
			else:
				UpperCAmelCase:					List[str] 						=     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:
				UpperCAmelCase:					Union[str, Any] 						=     module_name.replace("move.gamma"				,					"ema_gate.kernel_projection_matrix"				)
			elif "mega_layer.gamma" in module_name:
				UpperCAmelCase:					Any 						=     module_name.replace("gamma"				,					"qk_weight"				)
			else:
				UpperCAmelCase:					List[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:
			UpperCAmelCase:					Tuple 						=     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:
			UpperCAmelCase:					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:
			UpperCAmelCase:					Optional[int] 						=     module_name.replace("move.omega"				,					"ema_gate.residual_weight"				)

		if new_module_name:
			UpperCAmelCase:					List[Any] 						=     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():
		UpperCAmelCase:					str 						=     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
	UpperCAmelCase:					Optional[int] 						=     torch.randint(0				,					hf_config.vocab_size				,					size=(4, 2_5_6)				)
	UpperCAmelCase:					int 						=     torch.ones_like(snake_case_				)
	# mask a few tokens to make sure masking is applied appropriately :)
	UpperCAmelCase:					int 						=     0

	# run forward passes
	UpperCAmelCase:					Optional[int] 						=     original_mlm(snake_case_				,					snake_case_				,					batch_first=snake_case_				,					ignore_mask_value=0				)
	UpperCAmelCase:					Optional[int] 						=     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
	UpperCAmelCase:					Tuple 						=     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"				)
		UpperCAmelCase:					List[str] 						=     AutoTokenizer.from_pretrained(snake_case_				)
		tokenizer.save_pretrained(snake_case_				)


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

	snake_case_      :					str				      =     parser.parse_args()

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