





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





     """simple docstring"""



     def __init__(							self       :Any			,      snake_case       :int			,      snake_case       :Tuple			,      snake_case       :Dict  ):





        '''simple docstring'''




        super().__init__()
        A_					:				Optional[Any]          =			mega_args
        A_					:				List[str]          =			nn.Embedding(snake_case			,      self.mega_args.encoder_embed_dim  )
        A_					:				str          =			nn.ModuleList([MegaEncoderLayer(self.mega_args  ) for _ in range(snake_case  )]  )
        A_					:				Optional[Any]          =			depth





     def 						SCREAMING_SNAKE_CASE		(							self       :str			,      snake_case       :int			,      snake_case       :str			,      snake_case       :int=True			,      snake_case       :List[Any]=0  ):





        '''simple docstring'''




        if batch_first:
           A_					:				Optional[Any]          =			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_					:				Optional[int]          =			1 - attention_mask

        # get token embeddings from IDs
        A_					:				int          =			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 					__magic_name__  (							nn.Module							):





     """simple docstring"""



     def __init__(							self       :str			,      snake_case       :str			,      snake_case       :int			,      snake_case       :Dict  ):





        '''simple docstring'''




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





     def 						SCREAMING_SNAKE_CASE		(							self       :str			,      snake_case       :List[Any]			,      snake_case       :Optional[int]			,      snake_case       :Union[str, Any]=True			,      snake_case       :List[Any]=0  ):





        '''simple docstring'''




        A_					:				List[Any]          =			self.mega(snake_case			,      snake_case			,      snake_case			,      snake_case  )
        return self.mlm_head(self.dropout(snake_case  )  )




def      __snake_case						(					_lowerCAmelCase : Union[str, Any]				,				_lowerCAmelCase : Any				,				_lowerCAmelCase : List[str]	) ->						str:
   with open(os.path.join(_lowerCAmelCase				,				"model_args.pkl"	)				,				"rb"	) as f:
      A_					:				Optional[int]          =			pkl.load(_lowerCAmelCase	)

   # load the original encoder
   A_					:				Tuple          =			OriginalMegaForMaskedLM(**_lowerCAmelCase	).eval()

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

   # create a new config from the old one
   A_					:				Any          =			MegaConfig(
       num_hidden_layers=mega_original_args["depth"]				,				vocab_size=mega_original_args["vocab_size"]				,				hidden_size=mega_original_args["mega_args"].encoder_embed_dim				,				shared_representation_size=mega_original_args["mega_args"].encoder_z_dim				,				intermediate_size=mega_original_args["mega_args"].encoder_hidden_dim				,				ema_projection_size=mega_original_args["mega_args"].encoder_n_dim				,				dropout_prob=mega_original_args["mega_args"].dropout				,				attention_probs_dropout_prob=mega_original_args["mega_args"].attention_dropout				,				hidden_dropout_prob=mega_original_args["mega_args"].hidden_dropout				,				activation=mega_original_args["mega_args"].activation_fn				,				attention_activation=mega_original_args["mega_args"].attention_activation_fn				,				bidirectional=mega_original_args["mega_args"].bidirectional				,				use_chunking=mega_original_args["mega_args"].encoder_chunk_size > 0				,				chunk_size=mega_original_args["mega_args"].encoder_chunk_size				,				truncation=mega_original_args["mega_args"].truncation_length				,				normalization_type=mega_original_args["mega_args"].normalization_type				,				normalize_before_mega=_lowerCAmelCase				,				norm_affine=_lowerCAmelCase				,				use_feature_dropout=mega_original_args["mega_args"].feature_dropout				,				relative_positional_bias=mega_original_args["mega_args"].rel_pos_bias				,				max_positions=mega_original_args["mega_args"].max_source_positions				,				nffn_hidden_size=mega_original_args["mega_args"].encoder_ffn_embed_dim				,				normalize_before_ffn=mega_original_args["mega_args"].normalize_before				,				nffn_activation_dropout_prob=0.0				,				add_token_type_embeddings=_lowerCAmelCase				,				add_lm_hidden_dense_layer=_lowerCAmelCase				,				)

   A_					:				List[Any]          =			MegaForMaskedLM(_lowerCAmelCase	).eval()

   # the originl checkpoint just uses nn.Embedding for the word embeddings
   # we use a wrapper module for embeddings to add support for positional embeddings
   A_					:				Optional[int]          =			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_					:				Optional[int]          =			original_mlm.mega.encoders.state_dict()
   A_					:				Any          =			{}
   for module_name in original_state_dict.keys():
      A_					:				int          =			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_					:				List[str]          =			module_name.replace("move.beta"				,				"ema_gate.ema_expansion_matrix"	)
         elif "mega_layer.beta" in module_name:
            A_					:				int          =			module_name.replace("beta"				,				"qk_bias"	)
         else:
            A_					:				Optional[int]          =			module_name.replace("beta"				,				"b_param"	)
        # beta is used in EMA and MovingAverageGatedAttention, and must be renamed due to flax/tf weights
      elif "gamma" in module_name:
         if "move.gamma" in module_name:
            A_					:				Any          =			module_name.replace("move.gamma"				,				"ema_gate.kernel_projection_matrix"	)
         elif "mega_layer.gamma" in module_name:
            A_					:				Optional[Any]          =			module_name.replace("gamma"				,				"qk_weight"	)
         else:
            A_					:				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_					:				Optional[int]          =			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_					:				int          =			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_					:				List[Any]          =			module_name.replace("move.omega"				,				"ema_gate.residual_weight"	)

      if new_module_name:
         A_					:				List[Any]          =			new_module_name

   if len(_lowerCAmelCase	) != 0:
      print(f"Renaming these keys: {updated_keys.keys()}"	)
   else:
      print("No need to rename state dict entries"	)
   for old, new in updated_keys.items():
      A_					:				Optional[Any]          =			original_state_dict.pop(_lowerCAmelCase	)

   # now attempt to load the state dictionary with updated names
   # note that we now call it `mega.layers` instead of `mega.encoders` due to hugging face style
   print("HF Mega encoder:"				,				hf_mlm.mega.layers.load_state_dict(_lowerCAmelCase	)	)

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

   # test on a randomly generated input sequence
   A_					:				Union[str, Any]          =			torch.randint(0				,				hf_config.vocab_size				,				size=(4, 256)	)
   A_					:				List[str]          =			torch.ones_like(_lowerCAmelCase	)
   # mask a few tokens to make sure masking is applied appropriately :)
   A_					:				Tuple          =			0

   # run forward passes
   A_					:				Union[str, Any]          =			original_mlm(_lowerCAmelCase				,				_lowerCAmelCase				,				batch_first=_lowerCAmelCase				,				ignore_mask_value=0	)
   A_					:				Dict          =			hf_mlm(_lowerCAmelCase				,				_lowerCAmelCase	)[0]

   # print shapes and diff
   print(f"original output {original_output.shape}"	)
   print(f"hf output {hf_output.shape}"	)
   print(f"max diff: {(original_output - hf_output).max()}"	)  # 0.0
   A_					:				Tuple          =			torch.allclose(_lowerCAmelCase				,				_lowerCAmelCase				,				atol=1e-3	)

   if success:
      print("Yay!"	)
      hf_mlm.save_pretrained(_lowerCAmelCase	)
   else:
      raise RuntimeError(f"Something's broken :(\nOriginal:\n{original_output}\n\nHF\n{hf_output}\n{hf_mlm}"	)

   if includes_tokenizer:
      print("Transferring tokenizer"	)
      A_					:				Tuple          =			AutoTokenizer.from_pretrained(_lowerCAmelCase	)
      tokenizer.save_pretrained(_lowerCAmelCase	)


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

 _lowerCAmelCase							:      Optional[Any]						=	parser.parse_args()

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







import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
 import torch

if is_vision_available():
 from PIL import Image

 from transformers import ViltImageProcessor


class 					__magic_name__  (							unittest.TestCase							):





     """simple docstring"""



     def __init__(							self       :Dict			,      snake_case       :str			,      snake_case       :Tuple=7			,      snake_case       :Union[str, Any]=3			,      snake_case       :Any=18			,      snake_case       :Tuple=30			,      snake_case       :List[str]=400			,      snake_case       :Optional[Any]=True			,      snake_case       :Optional[int]=None			,      snake_case       :str=2			,      snake_case       :Tuple=True			,      snake_case       :List[str]=[0.5, 0.5, 0.5]			,      snake_case       :int=[0.5, 0.5, 0.5]			,      ):





        '''simple docstring'''




        A_					:				Optional[Any]          =			size if size is not None else {"shortest_edge": 30}
        A_					:				Tuple          =			parent
        A_					:				Tuple          =			batch_size
        A_					:				Any          =			num_channels
        A_					:				Dict          =			image_size
        A_					:				int          =			min_resolution
        A_					:				Optional[Any]          =			max_resolution
        A_					:				Union[str, Any]          =			do_resize
        A_					:				List[Any]          =			size
        A_					:				List[str]          =			size_divisor
        A_					:				str          =			do_normalize
        A_					:				int          =			image_mean
        A_					:				str          =			image_std



     def 						SCREAMING_SNAKE_CASE		(							self       :List[Any]  ):





        '''simple docstring'''




        return {
            "image_mean": self.image_mean,
            "image_std": self.image_std,
            "do_normalize": self.do_normalize,
            "do_resize": self.do_resize,
            "size": self.size,
            "size_divisor": self.size_divisor,
        }





     def 						SCREAMING_SNAKE_CASE		(							self       :Optional[Any]			,      snake_case       :Tuple			,      snake_case       :List[str]=False  ):





        '''simple docstring'''




        if not batched:
           A_					:				List[Any]          =			self.size["shortest_edge"]
           A_					:				Tuple          =			image_inputs[0]
           if isinstance(snake_case			,      Image.Image  ):
              A_      ,						A_					:				str          =			image.size
           else:
              A_      ,						A_					:				Optional[int]          =			image.shape[1], image.shape[2]
           A_					:				int          =			size / min(snake_case			,      snake_case  )
           if h < w:
              A_      ,						A_					:				Dict          =			size, scale * w
           else:
              A_      ,						A_					:				Tuple          =			scale * h, size

           A_					:				List[Any]          =			int((1_333 / 800) * size  )
           if max(snake_case			,      snake_case  ) > max_size:
              A_					:				Union[str, Any]          =			max_size / max(snake_case			,      snake_case  )
              A_					:				Tuple          =			newh * scale
              A_					:				List[str]          =			neww * scale

           A_      ,						A_					:				Optional[int]          =			int(newh + 0.5  ), int(neww + 0.5  )
           A_      ,						A_					:				Tuple          =			(
               newh // self.size_divisor * self.size_divisor,
               neww // self.size_divisor * self.size_divisor,
           )

        else:
           A_					:				List[Any]          =			[]
           for image in image_inputs:
              A_      ,						A_					:				str          =			self.get_expected_values([image]  )
              expected_values.append((expected_height, expected_width)  )
           A_					:				Optional[Any]          =			max(snake_case			,      key=lambda snake_case  : item[0]  )[0]
           A_					:				Optional[int]          =			max(snake_case			,      key=lambda snake_case  : item[1]  )[1]

        return expected_height, expected_width




@require_torch
@require_vision
class 					__magic_name__  (							lowerCamelCase__    ,				unittest.TestCase							):





     """simple docstring"""



     __UpperCamelCase					    =     ViltImageProcessor if is_vision_available() else None



     def 						SCREAMING_SNAKE_CASE		(							self       :Union[str, Any]  ):





        '''simple docstring'''




        A_					:				Dict          =			ViltImageProcessingTester(self  )



     @property
     def 						SCREAMING_SNAKE_CASE		(							self       :Union[str, Any]  ):





        '''simple docstring'''




        return self.image_processor_tester.prepare_image_processor_dict()



     def 						SCREAMING_SNAKE_CASE		(							self       :List[Any]  ):





        '''simple docstring'''




        A_					:				Optional[int]          =			self.image_processing_class(**self.image_processor_dict  )
        self.assertTrue(hasattr(snake_case			,      "image_mean"  )  )
        self.assertTrue(hasattr(snake_case			,      "image_std"  )  )
        self.assertTrue(hasattr(snake_case			,      "do_normalize"  )  )
        self.assertTrue(hasattr(snake_case			,      "do_resize"  )  )
        self.assertTrue(hasattr(snake_case			,      "size"  )  )
        self.assertTrue(hasattr(snake_case			,      "size_divisor"  )  )



     def 						SCREAMING_SNAKE_CASE		(							self       :Tuple  ):





        '''simple docstring'''




        A_					:				Tuple          =			self.image_processing_class.from_dict(self.image_processor_dict  )
        self.assertEqual(image_processor.size			,      {"shortest_edge": 30}  )

        A_					:				Any          =			self.image_processing_class.from_dict(self.image_processor_dict			,      size=42  )
        self.assertEqual(image_processor.size			,      {"shortest_edge": 42}  )



     def 						SCREAMING_SNAKE_CASE		(							self       :List[str]  ):





        '''simple docstring'''




        pass



     def 						SCREAMING_SNAKE_CASE		(							self       :Optional[int]  ):





        '''simple docstring'''




        A_					:				Dict          =			self.image_processing_class(**self.image_processor_dict  )
        # create random PIL images
        A_					:				str          =			prepare_image_inputs(self.image_processor_tester			,      equal_resolution=snake_case  )
        for image in image_inputs:
           self.assertIsInstance(snake_case			,      Image.Image  )

        # Test not batched input
        A_					:				Optional[Any]          =			image_processing(image_inputs[0]			,      return_tensors="pt"  ).pixel_values

        A_      ,						A_					:				str          =			self.image_processor_tester.get_expected_values(snake_case  )
        self.assertEqual(
            encoded_images.shape			,      (1, self.image_processor_tester.num_channels, expected_height, expected_width)			,      )

        # Test batched
        A_					:				int          =			image_processing(snake_case			,      return_tensors="pt"  ).pixel_values

        A_      ,						A_					:				Dict          =			self.image_processor_tester.get_expected_values(snake_case			,      batched=snake_case  )
        self.assertEqual(
            encoded_images.shape			,      (
                self.image_processor_tester.batch_size,
                self.image_processor_tester.num_channels,
                expected_height,
                expected_width,
            )			,      )



     def 						SCREAMING_SNAKE_CASE		(							self       :Optional[Any]  ):





        '''simple docstring'''




        A_					:				Any          =			self.image_processing_class(**self.image_processor_dict  )
        # create random numpy tensors
        A_					:				List[str]          =			prepare_image_inputs(self.image_processor_tester			,      equal_resolution=snake_case			,      numpify=snake_case  )
        for image in image_inputs:
           self.assertIsInstance(snake_case			,      np.ndarray  )

        # Test not batched input
        A_					:				List[Any]          =			image_processing(image_inputs[0]			,      return_tensors="pt"  ).pixel_values

        A_      ,						A_					:				int          =			self.image_processor_tester.get_expected_values(snake_case  )
        self.assertEqual(
            encoded_images.shape			,      (1, self.image_processor_tester.num_channels, expected_height, expected_width)			,      )

        # Test batched
        A_					:				Dict          =			image_processing(snake_case			,      return_tensors="pt"  ).pixel_values

        A_      ,						A_					:				str          =			self.image_processor_tester.get_expected_values(snake_case			,      batched=snake_case  )
        self.assertEqual(
            encoded_images.shape			,      (
                self.image_processor_tester.batch_size,
                self.image_processor_tester.num_channels,
                expected_height,
                expected_width,
            )			,      )





     def 						SCREAMING_SNAKE_CASE		(							self       :str  ):





        '''simple docstring'''




        A_					:				Any          =			self.image_processing_class(**self.image_processor_dict  )
        # create random PyTorch tensors
        A_					:				Optional[Any]          =			prepare_image_inputs(self.image_processor_tester			,      equal_resolution=snake_case			,      torchify=snake_case  )
        for image in image_inputs:
           self.assertIsInstance(snake_case			,      torch.Tensor  )

        # Test not batched input
        A_					:				int          =			image_processing(image_inputs[0]			,      return_tensors="pt"  ).pixel_values

        A_      ,						A_					:				Union[str, Any]          =			self.image_processor_tester.get_expected_values(snake_case  )
        self.assertEqual(
            encoded_images.shape			,      (1, self.image_processor_tester.num_channels, expected_height, expected_width)			,      )

        # Test batched
        A_					:				List[str]          =			image_processing(snake_case			,      return_tensors="pt"  ).pixel_values

        A_      ,						A_					:				Optional[int]          =			self.image_processor_tester.get_expected_values(snake_case			,      batched=snake_case  )
        self.assertEqual(
            encoded_images.shape			,      (
                self.image_processor_tester.batch_size,
                self.image_processor_tester.num_channels,
                expected_height,
                expected_width,
            )			,      )
