





'''simple docstring'''




import argparse
import re

import torch
from CLAP import create_model

from transformers import AutoFeatureExtractor, ClapConfig, ClapModel


lowercase__    :      Dict						=    {
    '''text_branch''': '''text_model''',
    '''audio_branch''': '''audio_model.audio_encoder''',
    '''attn''': '''attention.self''',
    '''self.proj''': '''output.dense''',
    '''attention.self_mask''': '''attn_mask''',
    '''mlp.fc1''': '''intermediate.dense''',
    '''mlp.fc2''': '''output.dense''',
    '''norm1''': '''layernorm_before''',
    '''norm2''': '''layernorm_after''',
    '''bn0''': '''batch_norm''',
}

lowercase__    :      int						=    AutoFeatureExtractor.from_pretrained('''laion/clap-htsat-unfused''', truncation='''rand_trunc''')






def        _lowerCAmelCase			(  __snake_case  : str ,							__snake_case  : Union[str, Any]=False			)     ->     Optional[int]:
     __A ,__A						: List[Any]      =       create_model(
         'HTSAT-tiny' ,							'roberta' ,							__snake_case ,							precision='fp32' ,							device='cuda:0' if torch.cuda.is_available() else 'cpu' ,							enable_fusion=__snake_case ,							fusion_type='aff_2d' if enable_fusion else None ,							)
     return model, model_cfg






def        _lowerCAmelCase			(  __snake_case  : Optional[int]			)     ->     Optional[Any]:
     __A						: Tuple      =       {}

     __A						: str      =       r'.*sequential.(\d+).*'
     __A						: List[str]      =       r'.*_projection.(\d+).*'

     for key, value in state_dict.items():
          # check if any key needs to be modified
          for key_to_modify, new_key in KEYS_TO_MODIFY_MAPPING.items():
               if key_to_modify in key:
                    __A						: str      =       key.replace(__snake_case ,							__snake_case			)

          if re.match(__snake_case ,							__snake_case			):
               # replace sequential layers with list
               __A						: Optional[int]      =       re.match(__snake_case ,							__snake_case			).group(1			)

               __A						: str      =       key.replace(f'sequential.{sequential_layer}.' ,							f'layers.{int(__snake_case			)//3}.linear.'			)
          elif re.match(__snake_case ,							__snake_case			):
               __A						: Any      =       int(re.match(__snake_case ,							__snake_case			).group(1			)			)

               # Because in CLAP they use `nn.Sequential`...
               __A						: Any      =       1 if projecton_layer == 0 else 2

               __A						: Dict      =       key.replace(f'_projection.{projecton_layer}.' ,							f'_projection.linear{transformers_projection_layer}.'			)

          if "audio" and "qkv" in key:
               # split qkv into query key and value
               __A						: int      =       value
               __A						: Union[str, Any]      =       mixed_qkv.size(0			) // 3

               __A						: List[Any]      =       mixed_qkv[:qkv_dim]
               __A						: List[str]      =       mixed_qkv[qkv_dim : qkv_dim * 2]
               __A						: Optional[int]      =       mixed_qkv[qkv_dim * 2 :]

               __A						: str      =       query_layer
               __A						: Union[str, Any]      =       key_layer
               __A						: Union[str, Any]      =       value_layer
          else:
               __A						: int      =       value

     return model_state_dict






def        _lowerCAmelCase			(  __snake_case  : List[str] ,							__snake_case  : Optional[int] ,							__snake_case  : Tuple ,							__snake_case  : List[Any]=False			)     ->     Dict:
     __A ,__A						: Optional[Any]      =       init_clap(__snake_case ,							enable_fusion=__snake_case			)

     clap_model.eval()
     __A						: Tuple      =       clap_model.state_dict()
     __A						: Any      =       rename_state_dict(__snake_case			)

     __A						: List[str]      =       ClapConfig()
     __A						: Dict      =       enable_fusion
     __A						: Tuple      =       ClapModel(__snake_case			)

     # ignore the spectrogram embedding layer
     model.load_state_dict(__snake_case ,							strict=__snake_case			)

     model.save_pretrained(__snake_case			)
     transformers_config.save_pretrained(__snake_case			)


if __name__ == "__main__":
  lowercase__    :      str						=    argparse.ArgumentParser()
  parser.add_argument('''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model.''')
  parser.add_argument('''--checkpoint_path''', default=None, type=str, help='''Path to fairseq checkpoint''')
  parser.add_argument('''--config_path''', default=None, type=str, help='''Path to hf config.json of model to convert''')
  parser.add_argument('''--enable_fusion''', action='''store_true''', help='''Whether to enable fusion or not''')
  lowercase__    :      List[Any]						=    parser.parse_args()

  convert_clap_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.enable_fusion)






'''simple docstring'''




lowercase__    :      Dict						=    {
    '''Pillow''': '''Pillow''',
    '''accelerate''': '''accelerate>=0.11.0''',
    '''compel''': '''compel==0.1.8''',
    '''black''': '''black~=23.1''',
    '''datasets''': '''datasets''',
    '''filelock''': '''filelock''',
    '''flax''': '''flax>=0.4.1''',
    '''hf-doc-builder''': '''hf-doc-builder>=0.3.0''',
    '''huggingface-hub''': '''huggingface-hub>=0.13.2''',
    '''requests-mock''': '''requests-mock==1.10.0''',
    '''importlib_metadata''': '''importlib_metadata''',
    '''invisible-watermark''': '''invisible-watermark''',
    '''isort''': '''isort>=5.5.4''',
    '''jax''': '''jax>=0.2.8,!=0.3.2''',
    '''jaxlib''': '''jaxlib>=0.1.65''',
    '''Jinja2''': '''Jinja2''',
    '''k-diffusion''': '''k-diffusion>=0.0.12''',
    '''torchsde''': '''torchsde''',
    '''note_seq''': '''note_seq''',
    '''librosa''': '''librosa''',
    '''numpy''': '''numpy''',
    '''omegaconf''': '''omegaconf''',
    '''parameterized''': '''parameterized''',
    '''protobuf''': '''protobuf>=3.20.3,<4''',
    '''pytest''': '''pytest''',
    '''pytest-timeout''': '''pytest-timeout''',
    '''pytest-xdist''': '''pytest-xdist''',
    '''ruff''': '''ruff>=0.0.241''',
    '''safetensors''': '''safetensors''',
    '''sentencepiece''': '''sentencepiece>=0.1.91,!=0.1.92''',
    '''scipy''': '''scipy''',
    '''onnx''': '''onnx''',
    '''regex''': '''regex!=2019.12.17''',
    '''requests''': '''requests''',
    '''tensorboard''': '''tensorboard''',
    '''torch''': '''torch>=1.4''',
    '''torchvision''': '''torchvision''',
    '''transformers''': '''transformers>=4.25.1''',
    '''urllib3''': '''urllib3<=2.0.0''',
}