




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



  def __init__(			self		,     __lowerCAmelCase		,     __lowerCAmelCase		,     __lowerCAmelCase ):







      '''simple docstring'''


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



  def     a_							(			self		,     __lowerCAmelCase		,     __lowerCAmelCase		,     __lowerCAmelCase=True		,     __lowerCAmelCase=0 ):







      '''simple docstring'''


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

      # get token embeddings from IDs
      _a				:							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:
          _a				:							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      UpperCamelCase       (	nn.Module       ):



  def __init__(			self		,     __lowerCAmelCase		,     __lowerCAmelCase		,     __lowerCAmelCase ):







      '''simple docstring'''


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



  def     a_							(			self		,     __lowerCAmelCase		,     __lowerCAmelCase		,     __lowerCAmelCase=True		,     __lowerCAmelCase=0 ):







      '''simple docstring'''


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






def UpperCAmelCase    (       UpperCAmelCase_   ,	UpperCAmelCase_   ,	UpperCAmelCase_						):



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

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

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

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

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


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

  a		:	str       =      parser.parse_args()

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




import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, slow, torch_device

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin





class 		__lowerCamelCase     (   lowercase							, lowercase							, unittest.TestCase		):
							lowerCamelCase__:				int      =							UNetaDModel
							lowerCamelCase__:				Any      =							'''sample'''



							@property
							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								UpperCAmelCase:					Union[str, Any] 						=     4
								UpperCAmelCase:					Optional[Any] 						=     1_4
								UpperCAmelCase:					str 						=     1_6

								UpperCAmelCase:					Optional[int] 						=     floats_tensor((batch_size, num_features, seq_len)							).to(__snake_case							)
								UpperCAmelCase:					Any 						=     torch.tensor([1_0] * batch_size							).to(__snake_case							)

								return {"sample": noise, "timestep": time_step}



							@property
							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								return (4, 1_4, 1_6)



							@property
							def     A__ (     self							)						->     Union[str, Any]:






								"""simple docstring"""



								return (4, 1_4, 1_6)



							def     A__ (     self							)						->     int:






								"""simple docstring"""



								pass



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






								"""simple docstring"""



								pass



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     Optional[Any]:






								"""simple docstring"""



								super().test_determinism()



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								super().test_outputs_equivalence()



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     Dict:






								"""simple docstring"""



								super().test_from_save_pretrained()



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     str:






								"""simple docstring"""



								super().test_from_save_pretrained_variant()



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     List[Any]:






								"""simple docstring"""



								super().test_model_from_pretrained()



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     int:






								"""simple docstring"""



								super().test_output()



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






								"""simple docstring"""



								UpperCAmelCase:					Tuple 						=     {
								    "block_out_channels": (3_2, 6_4, 1_2_8, 2_5_6),
								    "in_channels": 1_4,
								    "out_channels": 1_4,
								    "time_embedding_type": "positional",
								    "use_timestep_embedding": True,
								    "flip_sin_to_cos": False,
								    "freq_shift": 1.0,
								    "out_block_type": "OutConv1DBlock",
								    "mid_block_type": "MidResTemporalBlock1D",
								    "down_block_types": ("DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D"),
								    "up_block_types": ("UpResnetBlock1D", "UpResnetBlock1D", "UpResnetBlock1D"),
								    "act_fn": "mish",
								}
								UpperCAmelCase:					Tuple 						=     self.dummy_input
								return init_dict, inputs_dict



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								UpperCAmelCase						,   UpperCAmelCase:					Any 						=     UNetaDModel.from_pretrained(
								    "bglick13/hopper-medium-v2-value-function-hor32"	,				output_loading_info=__snake_case	,				subfolder="unet"							)
								self.assertIsNotNone(__snake_case							)
								self.assertEqual(len(loading_info["missing_keys"]							)	,				0							)

								model.to(__snake_case							)
								UpperCAmelCase:					Tuple 						=     model(**self.dummy_input							)

								assert image is not None, "Make sure output is not None"



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     List[str]:






								"""simple docstring"""



								UpperCAmelCase:					Union[str, Any] 						=     UNetaDModel.from_pretrained("bglick13/hopper-medium-v2-value-function-hor32"	,				subfolder="unet"							)
								torch.manual_seed(0							)
								if torch.cuda.is_available():
									torch.cuda.manual_seed_all(0							)

								UpperCAmelCase:					int 						=     model.config.in_channels
								UpperCAmelCase:					List[str] 						=     1_6
								UpperCAmelCase:					Union[str, Any] 						=     torch.randn((1, seq_len, num_features)							).permute(
								    0	,				2	,				1							)  # match original, we can update values and remove
								UpperCAmelCase:					List[str] 						=     torch.full((num_features,)	,				0							)

								with torch.no_grad():
									UpperCAmelCase:					Union[str, Any] 						=     model(__snake_case	,				__snake_case							).sample.permute(0	,				2	,				1							)

								UpperCAmelCase:					str 						=     output[0, -3:, -3:].flatten()
								# fmt: off
								UpperCAmelCase:					Dict 						=     torch.tensor([-2.13_71_72, 1.1_42_60_16, 0.3_68_86_87, -0.76_69_22, 0.7_30_31_46, 0.11_03_88_64, -0.4_76_06_33, 0.13_27_01_72, 0.02_59_13_48]							)
								# fmt: on
								self.assertTrue(torch.allclose(__snake_case	,				__snake_case	,				rtol=1e-3							)							)



							def     A__ (     self							)						->     Any:






								"""simple docstring"""



								pass




							@slow
							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								UpperCAmelCase:					List[str] 						=     "harmonai/maestro-150k"
								UpperCAmelCase:					Optional[Any] 						=     UNetaDModel.from_pretrained(__snake_case	,				subfolder="unet"							)
								model.to(__snake_case							)

								UpperCAmelCase:					Any 						=     6_5_5_3_6
								UpperCAmelCase:					str 						=     torch.sin(torch.arange(__snake_case							)[None, None, :].repeat(1	,				2	,				1							)							).to(__snake_case							)
								UpperCAmelCase:					Dict 						=     torch.tensor([1]							).to(__snake_case							)

								with torch.no_grad():
									UpperCAmelCase:					Optional[Any] 						=     model(__snake_case	,				__snake_case							).sample

								UpperCAmelCase:					List[Any] 						=     output.abs().sum()
								UpperCAmelCase:					Dict 						=     output.abs().max()

								assert (output_sum - 2_24.08_96).abs() < 0.5
								assert (output_max - 0.06_07).abs() < 4e-4





class 		__lowerCamelCase     (   lowercase							, lowercase							, unittest.TestCase		):
							lowerCamelCase__:				Dict      =							UNetaDModel
							lowerCamelCase__:				int      =							'''sample'''



							@property
							def     A__ (     self							)						->     int:






								"""simple docstring"""



								UpperCAmelCase:					List[Any] 						=     4
								UpperCAmelCase:					int 						=     1_4
								UpperCAmelCase:					Optional[int] 						=     1_6

								UpperCAmelCase:					List[str] 						=     floats_tensor((batch_size, num_features, seq_len)							).to(__snake_case							)
								UpperCAmelCase:					Any 						=     torch.tensor([1_0] * batch_size							).to(__snake_case							)

								return {"sample": noise, "timestep": time_step}



							@property
							def     A__ (     self							)						->     List[Any]:






								"""simple docstring"""



								return (4, 1_4, 1_6)



							@property
							def     A__ (     self							)						->     List[Any]:






								"""simple docstring"""



								return (4, 1_4, 1)



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     Dict:






								"""simple docstring"""



								super().test_determinism()



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     Optional[int]:






								"""simple docstring"""



								super().test_outputs_equivalence()



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     List[str]:






								"""simple docstring"""



								super().test_from_save_pretrained()



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     int:






								"""simple docstring"""



								super().test_from_save_pretrained_variant()



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     Optional[int]:






								"""simple docstring"""



								super().test_model_from_pretrained()



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     str:






								"""simple docstring"""



								UpperCAmelCase						,   UpperCAmelCase:					Optional[int] 						=     self.prepare_init_args_and_inputs_for_common()
								UpperCAmelCase:					Tuple 						=     self.model_class(**__snake_case							)
								model.to(__snake_case							)
								model.eval()

								with torch.no_grad():
									UpperCAmelCase:					Union[str, Any] 						=     model(**__snake_case							)

									if isinstance(__snake_case	,				__snake_case							):
										UpperCAmelCase:					List[str] 						=     output.sample

								self.assertIsNotNone(__snake_case							)
								UpperCAmelCase:					Optional[Any] 						=     torch.Size((inputs_dict["sample"].shape[0], 1)							)
								self.assertEqual(output.shape	,				__snake_case	,				"Input and output shapes do not match"							)



							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								pass



							def     A__ (     self							)						->     Any:






								"""simple docstring"""



								pass



							def     A__ (     self							)						->     int:






								"""simple docstring"""



								UpperCAmelCase:					Tuple 						=     {
								    "in_channels": 1_4,
								    "out_channels": 1_4,
								    "down_block_types": ["DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D"],
								    "up_block_types": [],
								    "out_block_type": "ValueFunction",
								    "mid_block_type": "ValueFunctionMidBlock1D",
								    "block_out_channels": [3_2, 6_4, 1_2_8, 2_5_6],
								    "layers_per_block": 1,
								    "downsample_each_block": True,
								    "use_timestep_embedding": True,
								    "freq_shift": 1.0,
								    "flip_sin_to_cos": False,
								    "time_embedding_type": "positional",
								    "act_fn": "mish",
								}
								UpperCAmelCase:					Union[str, Any] 						=     self.dummy_input
								return init_dict, inputs_dict



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								UpperCAmelCase						,   UpperCAmelCase:					Optional[Any] 						=     UNetaDModel.from_pretrained(
								    "bglick13/hopper-medium-v2-value-function-hor32"	,				output_loading_info=__snake_case	,				subfolder="value_function"							)
								self.assertIsNotNone(__snake_case							)
								self.assertEqual(len(vf_loading_info["missing_keys"]							)	,				0							)

								value_function.to(__snake_case							)
								UpperCAmelCase:					Union[str, Any] 						=     value_function(**self.dummy_input							)

								assert image is not None, "Make sure output is not None"



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     Any:






								"""simple docstring"""



								UpperCAmelCase						,   UpperCAmelCase:					Dict 						=     UNetaDModel.from_pretrained(
								    "bglick13/hopper-medium-v2-value-function-hor32"	,				output_loading_info=__snake_case	,				subfolder="value_function"							)
								torch.manual_seed(0							)
								if torch.cuda.is_available():
									torch.cuda.manual_seed_all(0							)

								UpperCAmelCase:					str 						=     value_function.config.in_channels
								UpperCAmelCase:					int 						=     1_4
								UpperCAmelCase:					int 						=     torch.randn((1, seq_len, num_features)							).permute(
								    0	,				2	,				1							)  # match original, we can update values and remove
								UpperCAmelCase:					Union[str, Any] 						=     torch.full((num_features,)	,				0							)

								with torch.no_grad():
									UpperCAmelCase:					Dict 						=     value_function(__snake_case	,				__snake_case							).sample

								# fmt: off
								UpperCAmelCase:					Tuple 						=     torch.tensor([1_65.25] * seq_len							)
								# fmt: on
								self.assertTrue(torch.allclose(__snake_case	,				__snake_case	,				rtol=1e-3							)							)




							def     A__ (     self							)						->     Optional[Any]:






								"""simple docstring"""



								pass
