code: 
  
  
  
  
  
  
  """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):
  
  
  
  
  
  
  
      """simple docstring"""
  
  
  
  
  
  
  
      def __init__(		self   :  int							,						UpperCamelCase_   :  Dict							,						UpperCamelCase_   :  Any							,						UpperCamelCase_   :  Dict   ):
  
  
  
  
  
             '''simple docstring'''
  
             super().__init__()
             __magic_name__										=						mega_args
             __magic_name__										=						nn.Embedding(UpperCamelCase_							,						self.mega_args.encoder_embed_dim   )
             __magic_name__										=						nn.ModuleList([MegaEncoderLayer(self.mega_args   ) for _ in range(UpperCamelCase_   )]   )
             __magic_name__										=						depth
  
  
  
  
  
  
  
      def 							a__ (		self   :  Any							,						UpperCamelCase_   :  str							,						UpperCamelCase_   :  Optional[Any]							,						UpperCamelCase_   :  Any=True							,						UpperCamelCase_   :  List[str]=0   ):
  
  
  
  
  
             '''simple docstring'''
  
             if batch_first:
                    __magic_name__										=						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:
                    __magic_name__										=						1 - attention_mask
  
             # get token embeddings from IDs
             __magic_name__										=						self.embedding_layer(UpperCamelCase_   )
  
             # pass through the Mega layers
             # input is (time, batch, encoder dim) and output is the same
             for encoder in self.encoders:
                    __magic_name__										=						encoder(UpperCamelCase_							,						UpperCamelCase_   )
  
             # return according to the shape specified
             if batch_first:
                    # (T, B, H) --> (B, T, H)
                    return torch.transpose(UpperCamelCase_							,						0							,						1   )
             else:
                    return embeds
  
  
  
  class 					UpperCamelCase__   (							nn.Module):
  
  
  
  
  
  
  
      """simple docstring"""
  
  
  
  
  
  
  
      def __init__(		self   :  Tuple							,						UpperCamelCase_   :  Optional[int]							,						UpperCamelCase_   :  Optional[int]							,						UpperCamelCase_   :  Optional[Any]   ):
  
  
  
  
  
             '''simple docstring'''
  
             super().__init__()
             __magic_name__										=						MegaLM(UpperCamelCase_							,						UpperCamelCase_							,						UpperCamelCase_   )
             __magic_name__										=						nn.Linear(mega_args.encoder_embed_dim							,						UpperCamelCase_   )
             __magic_name__										=						nn.Dropout(p=0.1   )
  
  
  
  
  
  
  
      def 							a__ (		self   :  str							,						UpperCamelCase_   :  List[str]							,						UpperCamelCase_   :  Optional[Any]							,						UpperCamelCase_   :  int=True							,						UpperCamelCase_   :  Tuple=0   ):
  
  
  
  
  
             '''simple docstring'''
  
             __magic_name__										=						self.mega(UpperCamelCase_							,						UpperCamelCase_							,						UpperCamelCase_							,						UpperCamelCase_   )
             return self.mlm_head(self.dropout(UpperCamelCase_   )   )
  
  
  def        A	(							__snake_case:   List[Any]   ,    __snake_case:   Any   ,    __snake_case:   int     )						->		int:
  
         """simple docstring"""
  
  
  
  
  
         with open(os.path.join(__snake_case   ,    'model_args.pkl'     )   ,    'rb'     ) as f:
                __magic_name__										=						pkl.load(__snake_case     )
  
         # load the original encoder
         __magic_name__										=						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
         __magic_name__										=						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   ,    )
  
         __magic_name__										=						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
         __magic_name__										=						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
         __magic_name__										=						original_mlm.mega.encoders.state_dict()
         __magic_name__										=						{}
         for module_name in original_state_dict.keys():
                __magic_name__										=						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:
                              __magic_name__										=						module_name.replace('move.beta'   ,    'ema_gate.ema_expansion_matrix'     )
                       elif "mega_layer.beta" in module_name:
                              __magic_name__										=						module_name.replace('beta'   ,    'qk_bias'     )
                       else:
                              __magic_name__										=						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:
                              __magic_name__										=						module_name.replace('move.gamma'   ,    'ema_gate.kernel_projection_matrix'     )
                       elif "mega_layer.gamma" in module_name:
                              __magic_name__										=						module_name.replace('gamma'   ,    'qk_weight'     )
                       else:
                              __magic_name__										=						module_name.replace('gamma'   ,    'g_param'     )
          # alpha is used in EMA and positional bias; renaming to improve readability
                elif "move.alpha" in module_name:
                       __magic_name__										=						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:
                       __magic_name__										=						module_name.replace('move.delta'   ,    'ema_gate.damping_factor'     )
                # omega is only used in EMA; renaming to improve readability
                elif "omega" in module_name:
                       __magic_name__										=						module_name.replace('move.omega'   ,    'ema_gate.residual_weight'     )
  
                if new_module_name:
                       __magic_name__										=						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():
                __magic_name__										=						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
         __magic_name__										=						torch.randint(0   ,    hf_config.vocab_size   ,    size=(4, 2_5_6)     )
         __magic_name__										=						torch.ones_like(__snake_case     )
         # mask a few tokens to make sure masking is applied appropriately :)
         __magic_name__										=						0
  
         # run forward passes
         __magic_name__										=						original_mlm(__snake_case   ,    __snake_case   ,    batch_first=__snake_case   ,    ignore_mask_value=0     )
         __magic_name__										=						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
         __magic_name__										=						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'     )
                __magic_name__										=						AutoTokenizer.from_pretrained(__snake_case     )
                tokenizer.save_pretrained(__snake_case     )
  
  
  if __name__ == "__main__":
   snake_case   :   List[Any]    		=		argparse.ArgumentParser()
  
   parser.add_argument(
       """--pretrained_checkpoint_path""",
       default=None,
       type=str,
       required=True,
       help="""Point to the directory containing your model weights using the official Mega repo""",
   )
  
   parser.add_argument(
       """--output_path""", default=None, type=str, required=True, help="""Location to save the Hugging Face version"""
   )
  
   parser.add_argument(
       """--includes_tokenizer""",
       action="""store_true""",
       help="""Use this flag if there is a Hugging Face tokenizer in the original checkpoint repo""",
   )
  
   snake_case   :   Union[str, Any]    		=		parser.parse_args()
  
   convert_checkpoint_to_huggingface(args.pretrained_checkpoint_path, args.output_path, args.includes_tokenizer)
  
  
style_context: 
  
  
  
  
  
  
  """simple docstring"""
  from typing import List, Optional, Union
  
  import numpy as np
  
  from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
  from ...feature_extraction_utils import BatchFeature
  from ...utils import PaddingStrategy, TensorType, logging
  
  
  snake_case   :   Tuple    		=		logging.get_logger(__name__)
  
  
  
  class 					UpperCamelCase__   (							a_):
  
  
  
  
  
  
  
      """simple docstring"""
      __UpperCAmelCase											=   ["""input_values""", """attention_mask"""]
  
  
  
  
  
  
  
      def __init__(		self   :  Union[str, Any]							,						UpperCamelCase_   :  int=1							,						UpperCamelCase_   :  Optional[int]=1_6_0_0_0							,						UpperCamelCase_   :  List[str]=0.0							,						UpperCamelCase_   :  Optional[int]=False							,						UpperCamelCase_   :  Tuple=True							,						**UpperCamelCase_   :  Tuple							,						):
  
  
  
  
  
             '''simple docstring'''
  
             super().__init__(feature_size=UpperCamelCase_							,						sampling_rate=UpperCamelCase_							,						padding_value=UpperCamelCase_							,						**UpperCamelCase_   )
             __magic_name__										=						return_attention_mask
             __magic_name__										=						do_normalize
  
  
  
  
  
  
  
      @staticmethod
      def 							a__ (		UpperCamelCase_   :  List[np.ndarray]							,						UpperCamelCase_   :  List[np.ndarray]							,						UpperCamelCase_   :  float = 0.0   ):
  
  
  
  
  
             '''simple docstring'''
  
             if attention_mask is not None:
                    __magic_name__										=						np.array(UpperCamelCase_							,						np.intaa   )
                    __magic_name__										=						[]
  
                    for vector, length in zip(UpperCamelCase_							,						attention_mask.sum(-1   )   ):
                           __magic_name__										=						(vector - vector[:length].mean()) / np.sqrt(vector[:length].var() + 1e-7   )
                           if length < normed_slice.shape[0]:
                                  __magic_name__										=						padding_value
  
                           normed_input_values.append(UpperCamelCase_   )
             else:
                    __magic_name__										=						[(x - x.mean()) / np.sqrt(x.var() + 1e-7   ) for x in input_values]
  
             return normed_input_values
  
  
  
  
  
  
  
      def __call__(		self   :  Any							,						UpperCamelCase_   :  Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]							,						UpperCamelCase_   :  Union[bool, str, PaddingStrategy] = False							,						UpperCamelCase_   :  Optional[int] = None							,						UpperCamelCase_   :  bool = False							,						UpperCamelCase_   :  Optional[int] = None							,						UpperCamelCase_   :  Optional[bool] = None							,						UpperCamelCase_   :  Optional[Union[str, TensorType]] = None							,						UpperCamelCase_   :  Optional[int] = None							,						**UpperCamelCase_   :  Any							,						):
  
  
  
  
  
             '''simple docstring'''
  
             if sampling_rate is not None:
                    if sampling_rate != self.sampling_rate:
                           raise ValueError(
                               f"""The model corresponding to this feature extractor: {self} was trained using a sampling rate of"""
                               f""" {self.sampling_rate}. Please make sure that the provided `raw_speech` input was sampled with"""
                               f""" {self.sampling_rate} and not {sampling_rate}."""   )
             else:
                    logger.warning(
                        'It is strongly recommended to pass the ``sampling_rate`` argument to this function. '
                        'Failing to do so can result in silent errors that might be hard to debug.'   )
  
             __magic_name__										=						isinstance(UpperCamelCase_							,						np.ndarray   ) and len(raw_speech.shape   ) > 1
             if is_batched_numpy and len(raw_speech.shape   ) > 2:
                    raise ValueError(f"""Only mono-channel audio is supported for input to {self}"""   )
             __magic_name__										=						is_batched_numpy or (
                 isinstance(UpperCamelCase_							,						(list, tuple)   ) and (isinstance(raw_speech[0]							,						(np.ndarray, tuple, list)   ))
             )
  
             # always return batch
             if not is_batched:
                    __magic_name__										=						[raw_speech]
  
             # convert into correct format for padding
             __magic_name__										=						BatchFeature({'input_values': raw_speech}   )
  
             __magic_name__										=						self.pad(
                 UpperCamelCase_							,						padding=UpperCamelCase_							,						max_length=UpperCamelCase_							,						truncation=UpperCamelCase_							,						pad_to_multiple_of=UpperCamelCase_							,						return_attention_mask=UpperCamelCase_							,						)
  
             # convert input values to correct format
             __magic_name__										=						padded_inputs['input_values']
             if not isinstance(input_values[0]							,						np.ndarray   ):
                    __magic_name__										=						[np.asarray(UpperCamelCase_							,						dtype=np.floataa   ) for array in input_values]
             elif (
                 not isinstance(UpperCamelCase_							,						np.ndarray   )
                 and isinstance(input_values[0]							,						np.ndarray   )
                 and input_values[0].dtype is np.dtype(np.floataa   )
             ):
                    __magic_name__										=						[array.astype(np.floataa   ) for array in input_values]
             elif isinstance(UpperCamelCase_							,						np.ndarray   ) and input_values.dtype is np.dtype(np.floataa   ):
                    __magic_name__										=						input_values.astype(np.floataa   )
  
             # convert attention_mask to correct format
             __magic_name__										=						padded_inputs.get('attention_mask'   )
             if attention_mask is not None:
                    __magic_name__										=						[np.asarray(UpperCamelCase_							,						dtype=np.intaa   ) for array in attention_mask]
  
             # zero-mean and unit-variance normalization
             if self.do_normalize:
                    __magic_name__										=						(
                        attention_mask
                        if self._get_padding_strategies(UpperCamelCase_							,						max_length=UpperCamelCase_   ) is not PaddingStrategy.DO_NOT_PAD
                        else None
                    )
                    __magic_name__										=						self.zero_mean_unit_var_norm(
                        padded_inputs['input_values']							,						attention_mask=UpperCamelCase_							,						padding_value=self.padding_value   )
  
             if return_tensors is not None:
                    __magic_name__										=						padded_inputs.convert_to_tensors(UpperCamelCase_   )
  
             return padded_inputs