code: 
  
  
  import argparse
  
  import torch
  
  from transformers import MraConfig, MraForMaskedLM
  
  
  
  
  def 			__lowercase	(						_UpperCamelCase     )    ->List[Any]:
  
  
  
  
  
  
  
        """simple docstring"""
  
        if "model" in orig_key:
              lowercase			:   Any					   =		orig_key.replace('''model.''',    ''''''     )
        if "norm1" in orig_key:
              lowercase			:   List[str]					   =		orig_key.replace('''norm1''',    '''attention.output.LayerNorm'''     )
        if "norm2" in orig_key:
              lowercase			:   str					   =		orig_key.replace('''norm2''',    '''output.LayerNorm'''     )
        if "norm" in orig_key:
              lowercase			:   Any					   =		orig_key.replace('''norm''',    '''LayerNorm'''     )
        if "transformer" in orig_key:
              lowercase			:   Any					   =		orig_key.split('''.'''     )[0].split('''_'''     )[-1]
              lowercase			:   str					   =		orig_key.replace(f"""transformer_{layer_num}""",    f"""encoder.layer.{layer_num}"""     )
        if "mha.attn" in orig_key:
              lowercase			:   List[Any]					   =		orig_key.replace('''mha.attn''',    '''attention.self'''     )
        if "mha" in orig_key:
              lowercase			:   int					   =		orig_key.replace('''mha''',    '''attention'''     )
        if "W_q" in orig_key:
              lowercase			:   Union[str, Any]					   =		orig_key.replace('''W_q''',    '''self.query'''     )
        if "W_k" in orig_key:
              lowercase			:   str					   =		orig_key.replace('''W_k''',    '''self.key'''     )
        if "W_v" in orig_key:
              lowercase			:   Dict					   =		orig_key.replace('''W_v''',    '''self.value'''     )
        if "ff.0" in orig_key:
              lowercase			:   Tuple					   =		orig_key.replace('''ff.0''',    '''intermediate.dense'''     )
        if "ff.2" in orig_key:
              lowercase			:   Optional[int]					   =		orig_key.replace('''ff.2''',    '''output.dense'''     )
        if "ff" in orig_key:
              lowercase			:   str					   =		orig_key.replace('''ff''',    '''output.dense'''     )
        if "mlm_class" in orig_key:
              lowercase			:   Optional[Any]					   =		orig_key.replace('''mlm.mlm_class''',    '''cls.predictions.decoder'''     )
        if "mlm" in orig_key:
              lowercase			:   Tuple					   =		orig_key.replace('''mlm''',    '''cls.predictions.transform'''     )
        if "backbone.backbone.encoders" in orig_key:
              lowercase			:   str					   =		orig_key.replace('''backbone.backbone.encoders''',    '''encoder.layer'''     )
        if "cls" not in orig_key:
              lowercase			:   List[Any]					   =		'''mra.''' + orig_key
  
        return orig_key
  
  
  
  
  def 			__lowercase	(						_UpperCamelCase,    _UpperCamelCase     )    ->List[Any]:
  
  
  
  
  
  
  
        """simple docstring"""
  
        for key in orig_state_dict.copy().keys():
              lowercase			:   Dict					   =		orig_state_dict.pop(_UpperCamelCase     )
  
              if ("pooler" in key) or ("sen_class" in key):
                    continue
              else:
                    lowercase			:   List[Any]					   =		val
  
        lowercase			:   Union[str, Any]					   =		orig_state_dict['''cls.predictions.decoder.bias''']
        lowercase			:   List[Any]					   =		torch.arange(_UpperCamelCase     ).expand((1, -1)     ) + 2
  
        return orig_state_dict
  
  
  
  
  def 			__lowercase	(						_UpperCamelCase,    _UpperCamelCase,    _UpperCamelCase     )    ->int:
  
  
  
  
  
  
  
        """simple docstring"""
  
        lowercase			:   Optional[int]					   =		torch.load(_UpperCamelCase,    map_location='''cpu'''     )['''model_state_dict''']
        lowercase			:   Union[str, Any]					   =		MraConfig.from_json_file(_UpperCamelCase     )
        lowercase			:   List[str]					   =		MraForMaskedLM(_UpperCamelCase     )
  
        lowercase			:   List[str]					   =		convert_checkpoint_helper(config.max_position_embeddings,    _UpperCamelCase     )
  
        print(model.load_state_dict(_UpperCamelCase     )     )
        model.eval()
        model.save_pretrained(_UpperCamelCase     )
  
        print(f"""Checkpoint successfuly converted. Model saved at {pytorch_dump_path}"""     )
  
  
  if __name__ == "__main__":
    __a   				=      argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        '''--pytorch_model_path''', default=None, type=str, required=True, help='''Path to Mra pytorch checkpoint.'''
    )
    parser.add_argument(
        '''--config_file''',
        default=None,
        type=str,
        required=True,
        help='''The json file for Mra model config.''',
    )
    parser.add_argument(
        '''--pytorch_dump_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
    )
    __a   				=      parser.parse_args()
    convert_mra_checkpoint(args.pytorch_model_path, args.config_file, args.pytorch_dump_path)
  
  
style_context: 
  
  
  from ...utils import (
      OptionalDependencyNotAvailable,
      is_torch_available,
      is_transformers_available,
      is_transformers_version,
  )
  
  
  try:
    if not (is_transformers_available() and is_torch_available() and is_transformers_version('''>=''', '''4.27.0''')):
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    from ...utils.dummy_torch_and_transformers_objects import (
        AudioLDMPipeline,
    )
  else:
    from .pipeline_audioldm import AudioLDMPipeline
  
  