




import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import OPTConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, slow

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


if is_flax_available():
				import os

				# The slow tests are often failing with OOM error on GPU
				# This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
				# but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
				lowercase_		=							'''platform'''

				import jax
				import jax.numpy as jnp

				from transformers import FlaxOPTForCausalLM, FlaxOPTModel, GPTaTokenizer







def  lowerCAmelCase  (	UpperCAmelCase,	UpperCAmelCase,	UpperCAmelCase=None,	UpperCAmelCase=None       )  ->List[str]:






					"""simple docstring"""






					if attention_mask is None:
										__magic_name__					:							List[str]		   =							np.where(input_ids != config.pad_token_id,	1,	0       )
					return {
					    "input_ids": input_ids,
					    "attention_mask": attention_mask,
					}





@require_flax
class 		A__				:






	def __init__(	self				,     lowerCamelCase				,     lowerCamelCase=13				,     lowerCamelCase=7				,     lowerCamelCase=True				,     lowerCamelCase=False				,     lowerCamelCase=99				,     lowerCamelCase=16				,     lowerCamelCase=2				,     lowerCamelCase=4				,     lowerCamelCase=4				,     lowerCamelCase="gelu"				,     lowerCamelCase=0.1				,     lowerCamelCase=0.1				,     lowerCamelCase=20				,     lowerCamelCase=2				,     lowerCamelCase=1				,     lowerCamelCase=0				,     lowerCamelCase=16				,     lowerCamelCase=16				,     lowerCamelCase=0.0_2				,     )      ->   Tuple:




						"""simple docstring"""






						__magic_name__					:							Any		   =							parent
						__magic_name__					:							Tuple		   =							batch_size
						__magic_name__					:							Optional[int]		   =							seq_length
						__magic_name__					:							List[str]		   =							is_training
						__magic_name__					:							Optional[int]		   =							use_labels
						__magic_name__					:							List[str]		   =							vocab_size
						__magic_name__					:							Any		   =							hidden_size
						__magic_name__					:							Dict		   =							num_hidden_layers
						__magic_name__					:							str		   =							num_attention_heads
						__magic_name__					:							Tuple		   =							intermediate_size
						__magic_name__					:							Optional[int]		   =							hidden_act
						__magic_name__					:							str		   =							hidden_dropout_prob
						__magic_name__					:							int		   =							attention_probs_dropout_prob
						__magic_name__					:							Dict		   =							max_position_embeddings
						__magic_name__					:							Optional[int]		   =							eos_token_id
						__magic_name__					:							Any		   =							pad_token_id
						__magic_name__					:							Tuple		   =							bos_token_id
						__magic_name__					:							str		   =							embed_dim
						__magic_name__					:							Optional[int]		   =							word_embed_proj_dim
						__magic_name__					:							Any		   =							initializer_range
						__magic_name__					:							Optional[int]		   =							False






	def  lowercase      (	self						)      ->   List[str]:




						"""simple docstring"""






						__magic_name__					:							Union[str, Any]		   =							np.clip(ids_tensor([self.batch_size, self.seq_length - 1]				,     self.vocab_size						)				,     3				,     self.vocab_size						)
						__magic_name__					:							Dict		   =							np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)				,     dtype=np.intaa						))				,     -1						)

						__magic_name__					:							Tuple		   =							OPTConfig(
						    vocab_size=self.vocab_size				,     hidden_size=self.hidden_size				,     num_hidden_layers=self.num_hidden_layers				,     num_attention_heads=self.num_attention_heads				,     ffn_dim=self.intermediate_size				,     dropout=self.hidden_dropout_prob				,     attention_dropout=self.attention_probs_dropout_prob				,     max_position_embeddings=self.max_position_embeddings				,     eos_token_id=self.eos_token_id				,     bos_token_id=self.bos_token_id				,     pad_token_id=self.pad_token_id				,     embed_dim=self.embed_dim				,     is_encoder_decoder=lowerCamelCase				,     word_embed_proj_dim=self.word_embed_proj_dim				,     initializer_range=self.initializer_range				,     use_cache=lowerCamelCase				,     )
						__magic_name__					:							Tuple		   =							prepare_opt_inputs_dict(lowerCamelCase				,     lowerCamelCase						)
						return config, inputs_dict






	def  lowercase      (	self						)      ->   Any:




						"""simple docstring"""






						__magic_name__   , __magic_name__					:							Optional[int]		   =							self.prepare_config_and_inputs()
						return config, inputs_dict






	def  lowercase      (	self				,     lowerCamelCase				,     lowerCamelCase				,     lowerCamelCase						)      ->   List[Any]:




						"""simple docstring"""






						__magic_name__					:							int		   =							20
						__magic_name__					:							str		   =							model_class_name(lowerCamelCase						)

						__magic_name__					:							Union[str, Any]		   =							inputs_dict['''input_ids''']
						__magic_name__					:							Tuple		   =							inputs_dict['''attention_mask''']

						__magic_name__					:							Optional[Any]		   =							model.init_cache(input_ids.shape[0]				,     lowerCamelCase						)
						__magic_name__					:							str		   =							jnp.ones((input_ids.shape[0], max_length)				,     dtype='''i4'''						)

						__magic_name__					:							str		   =							jnp.broadcast_to(
						    jnp.arange(input_ids.shape[-1] - 1						)[None, :]				,     (input_ids.shape[0], input_ids.shape[-1] - 1)				,     )
						__magic_name__					:							str		   =							model(
						    input_ids[:, :-1]				,     attention_mask=lowerCamelCase				,     past_key_values=lowerCamelCase				,     position_ids=lowerCamelCase				,     )

						__magic_name__					:							Any		   =							jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]				,     dtype='''i4'''						)
						__magic_name__					:							Tuple		   =							model(
						    input_ids[:, -1:]				,     attention_mask=lowerCamelCase				,     past_key_values=outputs_cache.past_key_values				,     position_ids=lowerCamelCase				,     )

						__magic_name__					:							Union[str, Any]		   =							model(lowerCamelCase						)

						__magic_name__					:							str		   =							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  lowercase      (	self				,     lowerCamelCase				,     lowerCamelCase				,     lowerCamelCase						)      ->   Dict:




						"""simple docstring"""






						__magic_name__					:							List[Any]		   =							20
						__magic_name__					:							Dict		   =							model_class_name(lowerCamelCase						)

						__magic_name__   , __magic_name__					:							int		   =							(
						    inputs_dict['''input_ids'''],
						    inputs_dict['''attention_mask'''],
						)

						__magic_name__					:							List[str]		   =							jnp.concatenate(
						    [
						        attention_mask,
						        jnp.zeros((attention_mask.shape[0], max_length - attention_mask.shape[1])						),
						    ]				,     axis=-1				,     )

						__magic_name__					:							Optional[Any]		   =							model.init_cache(input_ids.shape[0]				,     lowerCamelCase						)
						__magic_name__					:							List[Any]		   =							jnp.broadcast_to(
						    jnp.arange(input_ids.shape[-1] - 1						)[None, :]				,     (input_ids.shape[0], input_ids.shape[-1] - 1)				,     )

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

						__magic_name__					:							List[Any]		   =							model(lowerCamelCase				,     attention_mask=lowerCamelCase						)

						__magic_name__					:							int		   =							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 		A__				(   __SCREAMING_SNAKE_CASE		,     unittest.TestCase		,     __SCREAMING_SNAKE_CASE ):
	lowerCamelCase__		: Dict				    =(FlaxOPTModel, FlaxOPTForCausalLM) if is_flax_available() else ()
	lowerCamelCase__		: Optional[int]				    =() if is_flax_available() else ()






	def  lowercase      (	self						)      ->   Dict:




						"""simple docstring"""






						__magic_name__					:							Union[str, Any]		   =							FlaxOPTModelTester(self						)






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




						"""simple docstring"""






						__magic_name__   , __magic_name__					:							int		   =							self.model_tester.prepare_config_and_inputs()
						for model_class in self.all_model_classes:
											self.model_tester.check_use_cache_forward(lowerCamelCase				,     lowerCamelCase				,     lowerCamelCase						)






	def  lowercase      (	self						)      ->   int:




						"""simple docstring"""






						__magic_name__   , __magic_name__					:							Optional[int]		   =							self.model_tester.prepare_config_and_inputs()
						for model_class in self.all_model_classes:
											self.model_tester.check_use_cache_forward_with_attn_mask(lowerCamelCase				,     lowerCamelCase				,     lowerCamelCase						)






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




						"""simple docstring"""






						for model_class_name in self.all_model_classes:
											__magic_name__					:							Dict		   =							model_class_name.from_pretrained('''facebook/opt-125m'''						)
											__magic_name__					:							str		   =							np.ones((1, 1)						) * model.config.eos_token_id
											__magic_name__					:							List[str]		   =							model(lowerCamelCase						)
											self.assertIsNotNone(lowerCamelCase						)





@require_sentencepiece
@require_flax
class 		A__				(   unittest.TestCase ):






	@slow
	def  lowercase      (	self						)      ->   Optional[Any]:




						"""simple docstring"""






						__magic_name__					:							str		   =							FlaxOPTModel.from_pretrained('''facebook/opt-350m'''						)
						__magic_name__					:							Tuple		   =							jnp.array([[0, 31414, 232, 328, 740, 1140, 12695, 69, 46078, 1588, 2]]						)
						__magic_name__					:							Optional[Any]		   =							model(input_ids=lowerCamelCase						).last_hidden_state
						__magic_name__					:							Optional[Any]		   =							(1, 11, 512)
						self.assertEqual(output.shape				,     lowerCamelCase						)
						__magic_name__					:							Tuple		   =							jnp.array(
						    [[-0.2_8_6_7, -1.9_2_5_6, -0.3_0_6_2], [-1.2_7_1_1, -0.1_3_3_7, -0.1_8_9_7], [0.4_1_0_9, 0.1_1_8_7, -1.3_1_4_2]]						)
						self.assertTrue(jnp.allclose(output[:, :3, :3]				,     lowerCamelCase				,     atol=4e-2						)						)





@require_flax
@slow
class 		A__				(   unittest.TestCase ):






	def  lowercase      (	self						)      ->   Dict:




						"""simple docstring"""






						super().setUp()
						__magic_name__					:							str		   =							'''facebook/opt-350m'''






	def  lowercase      (	self						)      ->   Tuple:




						"""simple docstring"""






						__magic_name__					:							List[Any]		   =							FlaxOPTForCausalLM.from_pretrained(self.path_model						)
						__magic_name__					:							int		   =							GPTaTokenizer.from_pretrained(self.path_model						)

						__magic_name__					:							int		   =							[
						    '''Today is a beautiful day and I want to''',
						    '''In the city of''',
						    '''Paris is the capital of France and''',
						    '''Computers and mobile phones have taken''',
						]
						# verify that prompt without BOS token is identical to Metaseq -> add_special_tokens=False
						__magic_name__					:							Dict		   =							tokenizer(lowerCamelCase				,     return_tensors='''jax'''				,     padding=lowerCamelCase				,     add_special_tokens=lowerCamelCase						)
						__magic_name__					:							List[str]		   =							model(inputs.input_ids				,     attention_mask=inputs.attention_mask						)[0].mean(axis=-1						)
						__magic_name__					:							Optional[int]		   =							jnp.array(
						    [
						        [1.3_8_5_1, -1_3.8_9_2_3, -1_0.5_2_2_9, -1_0.7_5_3_3, -0.2_3_0_9, -1_0.2_3_8_4, -0.5_3_6_5, -9.0_9_4_7, -5.1_6_7_0],
						        [-4.7_0_7_3, -1_0.6_2_7_6, -3.9_4_1_5, -2_1.5_2_4_2, -0.2_8_2_2, -0.2_8_2_2, -0.2_8_2_2, -0.2_8_2_2, -0.2_8_2_2],
						        [0.6_2_4_7, -3.4_2_2_9, -8.9_1_7_9, -1.4_2_9_7, -1_4.1_6_5_0, 1.4_1_4_6, -9.0_2_1_8, -0.2_7_0_3, -0.2_7_0_3],
						        [6.4_7_8_3, -1.9_9_1_3, -1_0.7_9_2_6, -2.3_3_3_6, 1.5_0_9_2, -0.9_9_7_4, -6.8_2_1_3, 1.3_4_7_7, 1.3_4_7_7],
						    ]						)
						self.assertTrue(jnp.allclose(lowerCamelCase				,     lowerCamelCase				,     atol=4e-2						)						)

						__magic_name__					:							List[str]		   =							jax.jit(lowerCamelCase						)
						__magic_name__					:							int		   =							model(inputs.input_ids				,     attention_mask=inputs.attention_mask						)[0].mean(axis=-1						)
						self.assertTrue(jnp.allclose(lowerCamelCase				,     lowerCamelCase				,     atol=4e-2						)						)





@require_flax
@slow
class 		A__				(   unittest.TestCase ):






	@property
	def  lowercase      (	self						)      ->   Any:




						"""simple docstring"""






						return [
						    "Today is a beautiful day and I want",
						    "In the city of",
						    "Paris is the capital of France and",
						    "Computers and mobile phones have taken",
						]






	def  lowercase      (	self						)      ->   Any:




						"""simple docstring"""






						__magic_name__					:							List[str]		   =							'''facebook/opt-125m'''

						__magic_name__					:							Union[str, Any]		   =							[
						    '''Today is a beautiful day and I want to''',
						    '''In the city of New York, the city''',
						    '''Paris is the capital of France and the capital''',
						    '''Computers and mobile phones have taken over the''',
						]

						__magic_name__					:							List[str]		   =							[]

						__magic_name__					:							Optional[Any]		   =							FlaxOPTForCausalLM.from_pretrained(lowerCamelCase						)
						__magic_name__					:							str		   =							GPTaTokenizer.from_pretrained(lowerCamelCase						)

						for prompt in self.prompts:
											__magic_name__					:							Optional[int]		   =							tokenizer(lowerCamelCase				,     return_tensors='''jax'''						).input_ids

											__magic_name__					:							List[str]		   =							model.generate(lowerCamelCase				,     max_length=10						)
											__magic_name__					:							Any		   =							generated_ids[0]

											__magic_name__					:							Tuple		   =							tokenizer.batch_decode(lowerCamelCase				,     skip_special_tokens=lowerCamelCase						)
											predicted_outputs += generated_string

						self.assertListEqual(lowerCamelCase				,     lowerCamelCase						)






	def  lowercase      (	self						)      ->   Any:




						"""simple docstring"""






						__magic_name__					:							str		   =							'''facebook/opt-350m'''

						__magic_name__					:							Tuple		   =							[
						    '''Today is a beautiful day and I want to''',
						    '''In the city of San Francisco, the city''',
						    '''Paris is the capital of France and the capital''',
						    '''Computers and mobile phones have taken over the''',
						]

						__magic_name__					:							Union[str, Any]		   =							[]
						__magic_name__					:							Any		   =							FlaxOPTForCausalLM.from_pretrained(lowerCamelCase						)
						__magic_name__					:							Dict		   =							GPTaTokenizer.from_pretrained(lowerCamelCase						)

						for prompt in self.prompts:
											__magic_name__					:							List[str]		   =							tokenizer(lowerCamelCase				,     return_tensors='''jax'''						).input_ids

											__magic_name__					:							Any		   =							model.generate(lowerCamelCase				,     max_length=10						)
											__magic_name__					:							List[Any]		   =							generated_ids[0]

											__magic_name__					:							Optional[int]		   =							tokenizer.batch_decode(lowerCamelCase				,     skip_special_tokens=lowerCamelCase						)
											predicted_outputs += generated_string

						self.assertListEqual(lowerCamelCase				,     lowerCamelCase						)






	def  lowercase      (	self						)      ->   Dict:




						"""simple docstring"""






						__magic_name__					:							Any		   =							'''facebook/opt-125m'''
						__magic_name__					:							str		   =							[
						    '''Today is a beautiful day and I want to thank''',
						    '''In the city of Rome Canaver Canaver Canaver Canaver''',
						]
						__magic_name__					:							Dict		   =							FlaxOPTForCausalLM.from_pretrained(lowerCamelCase						)
						__magic_name__					:							Any		   =							GPTaTokenizer.from_pretrained(lowerCamelCase						)
						__magic_name__					:							Dict		   =							tokenizer(
						    [
						        '''Today is a beautiful day and I want to''',
						        '''In the city of''',
						    ]				,     return_tensors='''jax'''				,     padding=lowerCamelCase				,     )

						__magic_name__					:							Dict		   =							jax.jit(model.generate						)

						__magic_name__					:							Dict		   =							jit_generate(inputs['''input_ids''']				,     attention_mask=inputs['''attention_mask''']						).sequences

						__magic_name__					:							int		   =							tokenizer.batch_decode(lowerCamelCase				,     skip_special_tokens=lowerCamelCase						)

						self.assertIsNotNone(lowerCamelCase				,     lowerCamelCase						)






	def  lowercase      (	self						)      ->   str:




						"""simple docstring"""






						__magic_name__					:							Dict		   =							'''facebook/opt-350m'''

						__magic_name__					:							Optional[int]		   =							GPTaTokenizer.from_pretrained(lowerCamelCase						)
						__magic_name__					:							Optional[int]		   =							FlaxOPTForCausalLM.from_pretrained(lowerCamelCase						)

						__magic_name__					:							Tuple		   =							'''left'''

						# use different length sentences to test batching
						__magic_name__					:							str		   =							[
						    '''Hello, my dog is a little''',
						    '''Today, I''',
						]

						__magic_name__					:							Optional[Any]		   =							tokenizer(lowerCamelCase				,     return_tensors='''jax'''				,     padding=lowerCamelCase						)
						__magic_name__					:							List[Any]		   =							inputs['''input_ids''']

						__magic_name__					:							str		   =							model.generate(input_ids=lowerCamelCase				,     attention_mask=inputs['''attention_mask''']				,     trace=lowerCamelCase						)

						__magic_name__					:							Optional[int]		   =							tokenizer(sentences[0]				,     return_tensors='''jax'''						).input_ids
						__magic_name__					:							Union[str, Any]		   =							model.generate(input_ids=lowerCamelCase						)

						__magic_name__					:							List[Any]		   =							inputs_non_padded.shape[-1] - inputs['''attention_mask'''][-1].sum()
						__magic_name__					:							List[str]		   =							tokenizer(sentences[1]				,     return_tensors='''jax'''						).input_ids
						__magic_name__					:							Dict		   =							model.generate(input_ids=lowerCamelCase				,     max_length=model.config.max_length - num_paddings						)

						__magic_name__					:							str		   =							tokenizer.batch_decode(outputs[0]				,     skip_special_tokens=lowerCamelCase						)
						__magic_name__					:							Dict		   =							tokenizer.decode(output_non_padded[0][0]				,     skip_special_tokens=lowerCamelCase						)
						__magic_name__					:							Optional[int]		   =							tokenizer.decode(output_padded[0][0]				,     skip_special_tokens=lowerCamelCase						)

						__magic_name__					:							Dict		   =							[
						    '''Hello, my dog is a little bit of a dork.\nI\'m a little bit''',
						    '''Today, I was in the middle of a conversation with a friend about the''',
						]
						self.assertListEqual(lowerCamelCase				,     lowerCamelCase						)
						self.assertListEqual(lowerCamelCase				,     [non_padded_sentence, padded_sentence]						)








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


						'''simple docstring'''



						def __init__(					self	,       A_	,       A_	,       A_      )	->	Dict:
							super().__init__()
							lowerCAmelCase    =	mega_args
							lowerCAmelCase    =	nn.Embedding(A_	,       self.mega_args.encoder_embed_dim      )
							lowerCAmelCase    =	nn.ModuleList([MegaEncoderLayer(self.mega_args      ) for _ in range(A_      )]      )
							lowerCAmelCase    =	depth




						def 				__snake_case    (					self	,       A_	,       A_	,       A_=True	,       A_=0      )	->	List[Any]:

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

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

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







class __snake_case(      nn.Module					):


						'''simple docstring'''



						def __init__(					self	,       A_	,       A_	,       A_      )	->	Optional[Any]:
							super().__init__()
							lowerCAmelCase    =	MegaLM(A_	,       A_	,       A_      )
							lowerCAmelCase    =	nn.Linear(mega_args.encoder_embed_dim	,       A_      )
							lowerCAmelCase    =	nn.Dropout(p=0.1      )




						def 				__snake_case    (					self	,       A_	,       A_	,       A_=True	,       A_=0      )	->	List[Any]:
							lowerCAmelCase    =	self.mega(A_	,       A_	,       A_	,       A_      )
							return self.mlm_head(self.dropout(A_      )      )







def 		_snake_case		(					_SCREAMING_SNAKE_CASE     :						Optional[int]							,		_SCREAMING_SNAKE_CASE     :						List[str]							,		_SCREAMING_SNAKE_CASE     :						Union[str, Any]							)			->						List[Any]:







	"""simple docstring"""
	with open(os.path.join(_SCREAMING_SNAKE_CASE							,		"""model_args.pkl"""							)							,		"""rb"""							) as f:
		lowerCAmelCase    =	pkl.load(_SCREAMING_SNAKE_CASE							)

	# load the original encoder
	lowerCAmelCase    =	OriginalMegaForMaskedLM(**_SCREAMING_SNAKE_CASE							).eval()

	# load its weights
	print(
	    """Original Mega encoder:"""							,		original_mlm.mega.load_state_dict(
	        torch.load(os.path.join(_SCREAMING_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(_SCREAMING_SNAKE_CASE							,		"""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=_SCREAMING_SNAKE_CASE							,		norm_affine=_SCREAMING_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=_SCREAMING_SNAKE_CASE							,		add_lm_hidden_dense_layer=_SCREAMING_SNAKE_CASE							,		)

	lowerCAmelCase    =	MegaForMaskedLM(_SCREAMING_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
	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(_SCREAMING_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():
		lowerCAmelCase    =	original_state_dict.pop(_SCREAMING_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(_SCREAMING_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(_SCREAMING_SNAKE_CASE							,		"""mlm_head_weights.pt"""							)							,		map_location="""cpu"""							)							)							,		)

	# test on a randomly generated input sequence
	lowerCAmelCase    =	torch.randint(0							,		hf_config.vocab_size							,		size=(4, 256)							)
	lowerCAmelCase    =	torch.ones_like(_SCREAMING_SNAKE_CASE							)
	# mask a few tokens to make sure masking is applied appropriately :)
	lowerCAmelCase    =	0

	# run forward passes
	lowerCAmelCase    =	original_mlm(_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		batch_first=_SCREAMING_SNAKE_CASE							,		ignore_mask_value=0							)
	lowerCAmelCase    =	hf_mlm(_SCREAMING_SNAKE_CASE							,		_SCREAMING_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
	lowerCAmelCase    =	torch.allclose(_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		atol=1E-3							)

	if success:
		print("""Yay!"""							)
		hf_mlm.save_pretrained(_SCREAMING_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"""							)
		lowerCAmelCase    =	AutoTokenizer.from_pretrained(_SCREAMING_SNAKE_CASE							)
		tokenizer.save_pretrained(_SCREAMING_SNAKE_CASE							)


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

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