





'''simple docstring'''




import argparse
import re

import torch
from CLAP import create_model

from transformers import AutoFeatureExtractor, ClapConfig, ClapModel


lowercase__   :				Tuple      =   {
    '''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__   :				List[str]      =   AutoFeatureExtractor.from_pretrained('''laion/clap-htsat-unfused''', truncation='''rand_trunc''')






def 	_lowerCAmelCase      (  __snake_case  :			Optional[int]			,							__snake_case  :			Dict=False      )					->					int:
					__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  :			int      )					->					Tuple:
					__A						:			List[Any]								=	{}

					__A						:			Union[str, Any]								=	r'.*sequential.(\d+).*'
					__A						:			Optional[int]								=	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						:			int								=	key.replace(__snake_case			,							__snake_case      )

										if re.match(__snake_case			,							__snake_case      ):
															# replace sequential layers with list
															__A						:			Any								=	re.match(__snake_case			,							__snake_case      ).group(1      )

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

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

															__A						:			Optional[Any]								=	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						:			Optional[int]								=	value
															__A						:			Any								=	mixed_qkv.size(0      ) // 3

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

															__A						:			List[str]								=	query_layer
															__A						:			Optional[int]								=	key_layer
															__A						:			List[str]								=	value_layer
										else:
															__A						:			Any								=	value

					return model_state_dict






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

					clap_model.eval()
					__A						:			Optional[Any]								=	clap_model.state_dict()
					__A						:			Any								=	rename_state_dict(__snake_case      )

					__A						:			Optional[int]								=	ClapConfig()
					__A						:			List[str]								=	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[str]      =   parser.parse_args()

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






'''simple docstring'''




from tempfile import TemporaryDirectory
from unittest import TestCase
from unittest.mock import MagicMock, patch

from transformers import AutoModel, TFAutoModel
from transformers.onnx import FeaturesManager
from transformers.testing_utils import SMALL_MODEL_IDENTIFIER, require_tf, require_torch


@require_torch
@require_tf
class 			SCREAMING_SNAKE_CASE   (a__							):


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			Optional[int]								=	SMALL_MODEL_IDENTIFIER
          __A						:			Any								=	'pt'
          __A						:			str								=	'tf'


     def 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase):

          '''simple docstring'''






          __A						:			Any								=	AutoModel.from_pretrained(self.test_model)
          model_pt.save_pretrained(_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase):

          '''simple docstring'''






          __A						:			Union[str, Any]								=	TFAutoModel.from_pretrained(self.test_model						,     from_pt=_UpperCAmelCase)
          model_tf.save_pretrained(_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			Union[str, Any]								=	'mock_framework'

          # Framework provided - return whatever the user provides
          __A						:			Union[str, Any]								=	FeaturesManager.determine_framework(self.test_model						,     _UpperCAmelCase)
          self.assertEqual(_UpperCAmelCase						,     _UpperCAmelCase)

          # Local checkpoint and framework provided - return provided framework
          # PyTorch checkpoint
          with TemporaryDirectory() as local_pt_ckpt:
               self._setup_pt_ckpt(_UpperCAmelCase)
               __A						:			List[str]								=	FeaturesManager.determine_framework(_UpperCAmelCase						,     _UpperCAmelCase)
               self.assertEqual(_UpperCAmelCase						,     _UpperCAmelCase)

          # TensorFlow checkpoint
          with TemporaryDirectory() as local_tf_ckpt:
               self._setup_tf_ckpt(_UpperCAmelCase)
               __A						:			Tuple								=	FeaturesManager.determine_framework(_UpperCAmelCase						,     _UpperCAmelCase)
               self.assertEqual(_UpperCAmelCase						,     _UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          with TemporaryDirectory() as local_pt_ckpt:
               self._setup_pt_ckpt(_UpperCAmelCase)
               __A						:			Optional[int]								=	FeaturesManager.determine_framework(_UpperCAmelCase)
               self.assertEqual(_UpperCAmelCase						,     self.framework_pt)

          # TensorFlow checkpoint
          with TemporaryDirectory() as local_tf_ckpt:
               self._setup_tf_ckpt(_UpperCAmelCase)
               __A						:			List[str]								=	FeaturesManager.determine_framework(_UpperCAmelCase)
               self.assertEqual(_UpperCAmelCase						,     self.framework_tf)

          # Invalid local checkpoint
          with TemporaryDirectory() as local_invalid_ckpt:
               with self.assertRaises(_UpperCAmelCase):
                    __A						:			Tuple								=	FeaturesManager.determine_framework(_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			Tuple								=	MagicMock(return_value=_UpperCAmelCase)
          with patch('transformers.onnx.features.is_tf_available'						,     _UpperCAmelCase):
               __A						:			Union[str, Any]								=	FeaturesManager.determine_framework(self.test_model)
               self.assertEqual(_UpperCAmelCase						,     self.framework_pt)

          # PyTorch not in environment -> use TensorFlow
          __A						:			List[str]								=	MagicMock(return_value=_UpperCAmelCase)
          with patch('transformers.onnx.features.is_torch_available'						,     _UpperCAmelCase):
               __A						:			List[Any]								=	FeaturesManager.determine_framework(self.test_model)
               self.assertEqual(_UpperCAmelCase						,     self.framework_tf)

          # Both in environment -> use PyTorch
          __A						:			Any								=	MagicMock(return_value=_UpperCAmelCase)
          __A						:			Dict								=	MagicMock(return_value=_UpperCAmelCase)
          with patch('transformers.onnx.features.is_tf_available'						,     _UpperCAmelCase), patch(
              'transformers.onnx.features.is_torch_available'						,     _UpperCAmelCase):
               __A						:			int								=	FeaturesManager.determine_framework(self.test_model)
               self.assertEqual(_UpperCAmelCase						,     self.framework_pt)

          # Both not in environment -> raise error
          __A						:			List[str]								=	MagicMock(return_value=_UpperCAmelCase)
          __A						:			Tuple								=	MagicMock(return_value=_UpperCAmelCase)
          with patch('transformers.onnx.features.is_tf_available'						,     _UpperCAmelCase), patch(
              'transformers.onnx.features.is_torch_available'						,     _UpperCAmelCase):
               with self.assertRaises(_UpperCAmelCase):
                    __A						:			int								=	FeaturesManager.determine_framework(self.test_model)