


import argparse
import gc
import json
import os
import shutil
import warnings

import torch

from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer


try:
			from transformers import LlamaTokenizerFast
except ImportError as e:
			warnings.warn(e)
			warnings.warn(
			    "The converted tokenizer will be the `slow` tokenizer. To use the fast, update your `tokenizers` library and re-run the tokenizer conversion"
			)
			_snake_case  :	Tuple			      =       None
_snake_case  :	Dict			      =       {
    "7B": 11_008,
    "13B": 13_824,
    "30B": 17_920,
    "65B": 22_016,
    "70B": 28_672,
}
_snake_case  :	Any			      =       {
    "7B": 1,
    "7Bf": 1,
    "13B": 2,
    "13Bf": 2,
    "30B": 4,
    "65B": 8,
    "70B": 8,
    "70Bf": 8,
}







def        lowerCAmelCase_  (      __lowerCamelCase							,			__lowerCamelCase=1							,			__lowerCamelCase=2_5_6			):
							return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3			)			) + multiple_of - 1) // multiple_of)







def        lowerCAmelCase_  (      __lowerCamelCase			):
							with open(__lowerCamelCase							,			"r"			) as f:
														return json.load(__lowerCamelCase			)







def        lowerCAmelCase_  (      __lowerCamelCase							,			__lowerCamelCase			):
							with open(__lowerCamelCase							,			"w"			) as f:
														json.dump(__lowerCamelCase							,			__lowerCamelCase			)







def        lowerCAmelCase_  (      __lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase=True			):
							os.makedirs(__lowerCamelCase							,			exist_ok=__lowerCamelCase			)
							__snake_case	:    Union[str, Any]											=     os.path.join(__lowerCamelCase							,			"tmp"			)
							os.makedirs(__lowerCamelCase							,			exist_ok=__lowerCamelCase			)

							__snake_case	:    Union[str, Any]											=     read_json(os.path.join(__lowerCamelCase							,			"params.json"			)			)
							__snake_case	:    Any											=     NUM_SHARDS[model_size]
							__snake_case	:    List[Any]											=     params["n_layers"]
							__snake_case	:    Any											=     params["n_heads"]
							__snake_case	:    Optional[Any]											=     n_heads // num_shards
							__snake_case	:    Dict											=     params["dim"]
							__snake_case	:    Dict											=     dim // n_heads
							__snake_case	:    Any											=     1_0_0_0_0.0
							__snake_case	:    List[Any]											=     1.0 / (base ** (torch.arange(0							,			__lowerCamelCase							,			2			).float() / dims_per_head))

							if "n_kv_heads" in params:
														__snake_case	:    Tuple											=     params["n_kv_heads"]  # for GQA / MQA
														__snake_case	:    Optional[int]											=     n_heads_per_shard // num_key_value_heads
														__snake_case	:    List[Any]											=     dim // num_key_value_heads
							else:  # compatibility with other checkpoints
														__snake_case	:    Optional[Any]											=     n_heads
														__snake_case	:    Tuple											=     n_heads_per_shard
														__snake_case	:    Optional[Any]											=     dim

							# permute for sliced rotary
							def permute(__lowerCamelCase							,			__lowerCamelCase=n_heads							,			__lowerCamelCase=dim							,			__lowerCamelCase=dim			):
														return w.view(__lowerCamelCase							,			dima // n_heads // 2							,			2							,			__lowerCamelCase			).transpose(1							,			2			).reshape(__lowerCamelCase							,			__lowerCamelCase			)

							print(F'Fetching all parameters from the checkpoint at {input_base_path}.'			)
							# Load weights
							if model_size == "7B":
														# Not sharded
														# (The sharded implementation would also work, but this is simpler.)
														__snake_case	:    Optional[Any]											=     torch.load(os.path.join(__lowerCamelCase							,			"consolidated.00.pth"			)							,			map_location="cpu"			)
							else:
														# Sharded
														__snake_case	:    int											=     [
														    torch.load(os.path.join(__lowerCamelCase							,			F'consolidated.{i:02d}.pth'			)							,			map_location="cpu"			)
														    for i in range(__lowerCamelCase			)
														]
							__snake_case	:    Dict											=     0
							__snake_case	:    Optional[int]											=     {"weight_map": {}}
							for layer_i in range(__lowerCamelCase			):
														__snake_case	:    str											=     F'pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin'
														if model_size == "7B":
																					# Unsharded
																					__snake_case	:    List[Any]											=     {
																					    F'model.layers.{layer_i}.self_attn.q_proj.weight': permute(
																					        loaded[F'layers.{layer_i}.attention.wq.weight']			),
																					    F'model.layers.{layer_i}.self_attn.k_proj.weight': permute(
																					        loaded[F'layers.{layer_i}.attention.wk.weight']			),
																					    F'model.layers.{layer_i}.self_attn.v_proj.weight': loaded[F'layers.{layer_i}.attention.wv.weight'],
																					    F'model.layers.{layer_i}.self_attn.o_proj.weight': loaded[F'layers.{layer_i}.attention.wo.weight'],
																					    F'model.layers.{layer_i}.mlp.gate_proj.weight': loaded[F'layers.{layer_i}.feed_forward.w1.weight'],
																					    F'model.layers.{layer_i}.mlp.down_proj.weight': loaded[F'layers.{layer_i}.feed_forward.w2.weight'],
																					    F'model.layers.{layer_i}.mlp.up_proj.weight': loaded[F'layers.{layer_i}.feed_forward.w3.weight'],
																					    F'model.layers.{layer_i}.input_layernorm.weight': loaded[F'layers.{layer_i}.attention_norm.weight'],
																					    F'model.layers.{layer_i}.post_attention_layernorm.weight': loaded[F'layers.{layer_i}.ffn_norm.weight'],
																					}
														else:
																					# Sharded
																					# Note that attention.w{q,k,v,o}, feed_fordward.w[1,2,3], attention_norm.weight and ffn_norm.weight share
																					# the same storage object, saving attention_norm and ffn_norm will save other weights too, which is
																					# redundant as other weights will be stitched from multiple shards. To avoid that, they are cloned.

																					__snake_case	:    Tuple											=     {
																					    F'model.layers.{layer_i}.input_layernorm.weight': loaded[0][
																					        F'layers.{layer_i}.attention_norm.weight'
																					    ].clone(),
																					    F'model.layers.{layer_i}.post_attention_layernorm.weight': loaded[0][
																					        F'layers.{layer_i}.ffn_norm.weight'
																					    ].clone(),
																					}
																					__snake_case	:    Dict											=     permute(
																					    torch.cat(
																					        [
																					            loaded[i][F'layers.{layer_i}.attention.wq.weight'].view(__lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase			)
																					            for i in range(__lowerCamelCase			)
																					        ]							,			dim=0							,			).reshape(__lowerCamelCase							,			__lowerCamelCase			)			)
																					__snake_case	:    List[Any]											=     permute(
																					    torch.cat(
																					        [
																					            loaded[i][F'layers.{layer_i}.attention.wk.weight'].view(
																					                __lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase			)
																					            for i in range(__lowerCamelCase			)
																					        ]							,			dim=0							,			).reshape(__lowerCamelCase							,			__lowerCamelCase			)							,			__lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase							,			)
																					__snake_case	:    List[Any]											=     torch.cat(
																					    [
																					        loaded[i][F'layers.{layer_i}.attention.wv.weight'].view(
																					            __lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase			)
																					        for i in range(__lowerCamelCase			)
																					    ]							,			dim=0							,			).reshape(__lowerCamelCase							,			__lowerCamelCase			)

																					__snake_case	:    Dict											=     torch.cat(
																					    [loaded[i][F'layers.{layer_i}.attention.wo.weight'] for i in range(__lowerCamelCase			)]							,			dim=1			)
																					__snake_case	:    List[Any]											=     torch.cat(
																					    [loaded[i][F'layers.{layer_i}.feed_forward.w1.weight'] for i in range(__lowerCamelCase			)]							,			dim=0			)
																					__snake_case	:    Tuple											=     torch.cat(
																					    [loaded[i][F'layers.{layer_i}.feed_forward.w2.weight'] for i in range(__lowerCamelCase			)]							,			dim=1			)
																					__snake_case	:    List[str]											=     torch.cat(
																					    [loaded[i][F'layers.{layer_i}.feed_forward.w3.weight'] for i in range(__lowerCamelCase			)]							,			dim=0			)

														__snake_case	:    Tuple											=     inv_freq
														for k, v in state_dict.items():
																					__snake_case	:    Tuple											=     filename
																					param_count += v.numel()
														torch.save(__lowerCamelCase							,			os.path.join(__lowerCamelCase							,			__lowerCamelCase			)			)

							__snake_case	:    Optional[int]											=     F'pytorch_model-{n_layers + 1}-of-{n_layers + 1}.bin'
							if model_size == "7B":
														# Unsharded
														__snake_case	:    Optional[Any]											=     {
														    "model.embed_tokens.weight": loaded["tok_embeddings.weight"],
														    "model.norm.weight": loaded["norm.weight"],
														    "lm_head.weight": loaded["output.weight"],
														}
							else:
														__snake_case	:    Optional[int]											=     {
														    "model.norm.weight": loaded[0]["norm.weight"],
														    "model.embed_tokens.weight": torch.cat(
														        [loaded[i]["tok_embeddings.weight"] for i in range(__lowerCamelCase			)]							,			dim=1			),
														    "lm_head.weight": torch.cat([loaded[i]["output.weight"] for i in range(__lowerCamelCase			)]							,			dim=0			),
														}

							for k, v in state_dict.items():
														__snake_case	:    str											=     filename
														param_count += v.numel()
							torch.save(__lowerCamelCase							,			os.path.join(__lowerCamelCase							,			__lowerCamelCase			)			)

							# Write configs
							__snake_case	:    str											=     {"total_size": param_count * 2}
							write_json(__lowerCamelCase							,			os.path.join(__lowerCamelCase							,			"pytorch_model.bin.index.json"			)			)
							__snake_case	:    str											=     params["ffn_dim_multiplier"] if "ffn_dim_multiplier" in params else 1
							__snake_case	:    Dict											=     params["multiple_of"] if "multiple_of" in params else 2_5_6
							__snake_case	:    List[str]											=     LlamaConfig(
							    hidden_size=__lowerCamelCase							,			intermediate_size=compute_intermediate_size(__lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase			)							,			num_attention_heads=params["n_heads"]							,			num_hidden_layers=params["n_layers"]							,			rms_norm_eps=params["norm_eps"]							,			num_key_value_heads=__lowerCamelCase							,			)
							config.save_pretrained(__lowerCamelCase			)

							# Make space so we can load the model properly now.
							del state_dict
							del loaded
							gc.collect()

							print("Loading the checkpoint in a Llama model."			)
							__snake_case	:    Union[str, Any]											=     LlamaForCausalLM.from_pretrained(__lowerCamelCase							,			torch_dtype=torch.floataa							,			low_cpu_mem_usage=__lowerCamelCase			)
							# Avoid saving this as part of the config.
							del model.config._name_or_path

							print("Saving in the Transformers format."			)
							model.save_pretrained(__lowerCamelCase							,			safe_serialization=__lowerCamelCase			)
							shutil.rmtree(__lowerCamelCase			)







def        lowerCAmelCase_  (      __lowerCamelCase							,			__lowerCamelCase			):
							# Initialize the tokenizer based on the `spm` model
							__snake_case	:    int											=     LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
							print(F'Saving a {tokenizer_class.__name__} to {tokenizer_path}.'			)
							__snake_case	:    List[Any]											=     tokenizer_class(__lowerCamelCase			)
							tokenizer.save_pretrained(__lowerCamelCase			)







def        lowerCAmelCase_  (      ):
							__snake_case	:    List[Any]											=     argparse.ArgumentParser()
							parser.add_argument(
							    "--input_dir"							,			help="Location of LLaMA weights, which contains tokenizer.model and model folders"							,			)
							parser.add_argument(
							    "--model_size"							,			choices=["7B", "7Bf", "13B", "13Bf", "30B", "65B", "70B", "70Bf", "tokenizer_only"]							,			)
							parser.add_argument(
							    "--output_dir"							,			help="Location to write HF model and tokenizer"							,			)
							parser.add_argument("--safe_serialization"							,			type=__lowerCamelCase							,			help="Whether or not to save using `safetensors`."			)
							__snake_case	:    int											=     parser.parse_args()
							if args.model_size != "tokenizer_only":
														write_model(
														    model_path=args.output_dir							,			input_base_path=os.path.join(args.input_dir							,			args.model_size			)							,			model_size=args.model_size							,			safe_serialization=args.safe_serialization							,			)
							__snake_case	:    Union[str, Any]											=     os.path.join(args.input_dir							,			"tokenizer.model"			)
							write_tokenizer(args.output_dir							,			__lowerCamelCase			)


if __name__ == "__main__":
			main()





import argparse
import re

import torch
from CLAP import create_model

from transformers import AutoFeatureExtractor, ClapConfig, ClapModel


_snake_case  :	int			      =       {
    "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",
}

_snake_case  :	Dict			      =       AutoFeatureExtractor.from_pretrained("laion/clap-htsat-unfused", truncation="rand_trunc")







def        lowerCAmelCase_  (      __lowerCamelCase							,			__lowerCamelCase=False			):
							__snake_case				,    __snake_case	:    Any											=     create_model(
							    "HTSAT-tiny"							,			"roberta"							,			__lowerCamelCase							,			precision="fp32"							,			device="cuda:0" if torch.cuda.is_available() else "cpu"							,			enable_fusion=__lowerCamelCase							,			fusion_type="aff_2d" if enable_fusion else None							,			)
							return model, model_cfg







def        lowerCAmelCase_  (      __lowerCamelCase			):
							__snake_case	:    int											=     {}

							__snake_case	:    List[Any]											=     R".*sequential.(\d+).*"
							__snake_case	:    Any											=     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:
																												__snake_case	:    Optional[int]											=     key.replace(__lowerCamelCase							,			__lowerCamelCase			)

														if re.match(__lowerCamelCase							,			__lowerCamelCase			):
																					# replace sequential layers with list
																					__snake_case	:    List[Any]											=     re.match(__lowerCamelCase							,			__lowerCamelCase			).group(1			)

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

																					# Because in CLAP they use `nn.Sequential`...
																					__snake_case	:    List[str]											=     1 if projecton_layer == 0 else 2

																					__snake_case	:    Tuple											=     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
																					__snake_case	:    List[str]											=     value
																					__snake_case	:    Optional[int]											=     mixed_qkv.size(0			) // 3

																					__snake_case	:    List[str]											=     mixed_qkv[:qkv_dim]
																					__snake_case	:    int											=     mixed_qkv[qkv_dim : qkv_dim * 2]
																					__snake_case	:    str											=     mixed_qkv[qkv_dim * 2 :]

																					__snake_case	:    int											=     query_layer
																					__snake_case	:    Tuple											=     key_layer
																					__snake_case	:    List[str]											=     value_layer
														else:
																					__snake_case	:    Tuple											=     value

							return model_state_dict







def        lowerCAmelCase_  (      __lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase=False			):
							__snake_case				,    __snake_case	:    List[str]											=     init_clap(__lowerCamelCase							,			enable_fusion=__lowerCamelCase			)

							clap_model.eval()
							__snake_case	:    Union[str, Any]											=     clap_model.state_dict()
							__snake_case	:    Any											=     rename_state_dict(__lowerCamelCase			)

							__snake_case	:    Dict											=     ClapConfig()
							__snake_case	:    Dict											=     enable_fusion
							__snake_case	:    Optional[Any]											=     ClapModel(__lowerCamelCase			)

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

							model.save_pretrained(__lowerCamelCase			)
							transformers_config.save_pretrained(__lowerCamelCase			)


if __name__ == "__main__":
			_snake_case  :	Union[str, Any]			      =       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")
			_snake_case  :	Tuple			      =       parser.parse_args()

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

