





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





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





						def    _snake_case      (	self				,						lowercase				,						lowercase				,						lowercase=True				,						lowercase=0			)		->					Dict:






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

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

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







class 		lowercase ( nn.Module ):





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





						def    _snake_case      (	self				,						lowercase				,						lowercase				,						lowercase=True				,						lowercase=0			)		->					Optional[int]:
												lowerCAmelCase  						=	self.mega(lowercase				,						lowercase				,						lowercase				,						lowercase			)
												return self.mlm_head(self.dropout(lowercase			)			)







def       UpperCAmelCase__							(  SCREAMING_SNAKE_CASE					:						Tuple							,						SCREAMING_SNAKE_CASE					:						List[str]							,						SCREAMING_SNAKE_CASE					:						str      ):







						'''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, 2_56)      )
						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__":
							SCREAMING_SNAKE_CASE__ =       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",
							)

							SCREAMING_SNAKE_CASE__ =       parser.parse_args()

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













"""simple docstring"""





import unittest
from dataclasses import dataclass
from typing import Optional

from transformers.utils import ModelOutput

@dataclass
class 		lowercase ( _UpperCAmelCase ):
      _SCREAMING_SNAKE_CASE         =				42
      _SCREAMING_SNAKE_CASE         =				None
      _SCREAMING_SNAKE_CASE         =				None







class 		lowercase ( unittest.TestCase ):





      def    _snake_case      (	self			)		->					Optional[Any]:
            lowerCAmelCase  						=	ModelOutputTest(a=30			)
            self.assertEqual(x.a				,						30			)
            self.assertIsNone(x.b			)
            self.assertIsNone(x.c			)
            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	x.d





      def    _snake_case      (	self			)		->					Optional[Any]:
            lowerCAmelCase  						=	ModelOutputTest(a=30				,						b=10			)
            self.assertEqual(x[0]				,						30			)
            self.assertEqual(x[1]				,						10			)
            self.assertEqual(x[:2]				,						(30, 10)			)
            self.assertEqual(x[:]				,						(30, 10)			)

            lowerCAmelCase  						=	ModelOutputTest(a=30				,						c=10			)
            self.assertEqual(x[0]				,						30			)
            self.assertEqual(x[1]				,						10			)
            self.assertEqual(x[:2]				,						(30, 10)			)
            self.assertEqual(x[:]				,						(30, 10)			)





      def    _snake_case      (	self			)		->					Tuple:
            lowerCAmelCase  						=	ModelOutputTest(a=30				,						b=10			)
            self.assertEqual(x["""a"""]				,						30			)
            self.assertEqual(x["""b"""]				,						10			)
            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	x["""c"""]

            lowerCAmelCase  						=	ModelOutputTest(a=30				,						c=10			)
            self.assertEqual(x["""a"""]				,						30			)
            self.assertEqual(x["""c"""]				,						10			)
            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	x["""b"""]





      def    _snake_case      (	self			)		->					Union[str, Any]:
            lowerCAmelCase  						=	ModelOutputTest(a=30			)
            self.assertEqual(list(x.keys()			)				,						["""a"""]			)
            self.assertEqual(list(x.values()			)				,						[30]			)
            self.assertEqual(list(x.items()			)				,						[("""a""", 30)]			)
            self.assertEqual(list(lowercase			)				,						["""a"""]			)

            lowerCAmelCase  						=	ModelOutputTest(a=30				,						b=10			)
            self.assertEqual(list(x.keys()			)				,						["""a""", """b"""]			)
            self.assertEqual(list(x.values()			)				,						[30, 10]			)
            self.assertEqual(list(x.items()			)				,						[("""a""", 30), ("""b""", 10)]			)
            self.assertEqual(list(lowercase			)				,						["""a""", """b"""]			)

            lowerCAmelCase  						=	ModelOutputTest(a=30				,						c=10			)
            self.assertEqual(list(x.keys()			)				,						["""a""", """c"""]			)
            self.assertEqual(list(x.values()			)				,						[30, 10]			)
            self.assertEqual(list(x.items()			)				,						[("""a""", 30), ("""c""", 10)]			)
            self.assertEqual(list(lowercase			)				,						["""a""", """c"""]			)

            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	x.update({"""d""": 20}			)
            with self.assertRaises(lowercase			):
                  del x["a"]
            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	x.pop("""a"""			)
            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	x.setdefault("""d"""				,						32			)





      def    _snake_case      (	self			)		->					Dict:
            lowerCAmelCase  						=	ModelOutputTest(a=30			)
            lowerCAmelCase  						=	10
            self.assertEqual(x.a				,						10			)
            self.assertEqual(x["""a"""]				,						10			)





      def    _snake_case      (	self			)		->					List[str]:
            lowerCAmelCase  						=	ModelOutputTest(a=30			)
            lowerCAmelCase  						=	10
            self.assertEqual(x.a				,						10			)
            self.assertEqual(x["""a"""]				,						10			)





      def    _snake_case      (	self			)		->					str:
            lowerCAmelCase  						=	ModelOutputTest({"""a""": 30, """b""": 10}			)
            self.assertEqual(list(x.keys()			)				,						["""a""", """b"""]			)
            self.assertEqual(x.a				,						30			)
            self.assertEqual(x.b				,						10			)





      def    _snake_case      (	self			)		->					Dict:
            lowerCAmelCase  						=	ModelOutputTest([("""a""", 30), ("""b""", 10)]			)
            self.assertEqual(list(x.keys()			)				,						["""a""", """b"""]			)
            self.assertEqual(x.a				,						30			)
            self.assertEqual(x.b				,						10			)

            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	ModelOutputTest([("""a""", 30), (10, 10)]			)

            lowerCAmelCase  						=	ModelOutputTest(a=(30, 30)			)
            self.assertEqual(list(x.keys()			)				,						["""a"""]			)
            self.assertEqual(x.a				,						(30, 30)			)






