





"""simple docstring"""






import argparse
import os

import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing the experiment tracking capability,
# and builds off the `nlp_example.py` script.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To help focus on the differences in the code, building `DataLoaders`
# was refactored into its own function.
# New additions from the base script can be found quickly by
# looking for the # New Code # tags
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################

__UpperCAmelCase						=			16
__UpperCAmelCase						=			32






def      lowercase__ (					lowerCAmelCase__   :			Accelerator							,			lowerCAmelCase__   :			int = 1_6			)    ->			Optional[int]:



			'''simple docstring'''
			a__  :	List[Any]		 =			AutoTokenizer.from_pretrained("bert-base-cased"			)
			a__  :	Dict		 =			load_dataset("glue"							,			"mrpc"			)

			def tokenize_function(lowerCAmelCase__   :			Optional[int]			):
						# max_length=None => use the model max length (it's actually the default)
						a__  :	Any		 =			tokenizer(examples["sentence1"]							,			examples["sentence2"]							,			truncation=lowerCAmelCase__							,			max_length=lowerCAmelCase__			)
						return outputs

			# Apply the method we just defined to all the examples in all the splits of the dataset
			# starting with the main process first:
			with accelerator.main_process_first():
						a__  :	List[Any]		 =			datasets.map(
						    lowerCAmelCase__							,			batched=lowerCAmelCase__							,			remove_columns=["idx", "sentence1", "sentence2"]							,			)

			# We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
			# transformers library
			a__  :	Tuple		 =			tokenized_datasets.rename_column("label"							,			"labels"			)

			def collate_fn(lowerCAmelCase__   :			Tuple			):
						# On TPU it's best to pad everything to the same length or training will be very slow.
						a__  :	Dict		 =			1_2_8 if accelerator.distributed_type == DistributedType.TPU else None
						# When using mixed precision we want round multiples of 8/16
						if accelerator.mixed_precision == "fp8":
									a__  :	Tuple		 =			1_6
						elif accelerator.mixed_precision != "no":
									a__  :	Optional[int]		 =			8
						else:
									a__  :	List[Any]		 =			None

						return tokenizer.pad(
						    lowerCAmelCase__							,			padding="longest"							,			max_length=lowerCAmelCase__							,			pad_to_multiple_of=lowerCAmelCase__							,			return_tensors="pt"							,			)

			# Instantiate dataloaders.
			a__  :	Optional[Any]		 =			DataLoader(
			    tokenized_datasets["train"]							,			shuffle=lowerCAmelCase__							,			collate_fn=lowerCAmelCase__							,			batch_size=lowerCAmelCase__			)
			a__  :	List[Any]		 =			DataLoader(
			    tokenized_datasets["validation"]							,			shuffle=lowerCAmelCase__							,			collate_fn=lowerCAmelCase__							,			batch_size=lowerCAmelCase__			)

			return train_dataloader, eval_dataloader


# For testing only
if os.environ.get('''TESTING_MOCKED_DATALOADERS''', None) == "1":
		from accelerate.test_utils.training import mocked_dataloaders

		__UpperCAmelCase						=			mocked_dataloaders  # noqa: F811






def      lowercase__ (					lowerCAmelCase__   :			Tuple							,			lowerCAmelCase__   :			int			)    ->			List[str]:



			'''simple docstring'''
			# For testing only
			if os.environ.get("TESTING_MOCKED_DATALOADERS"							,			lowerCAmelCase__			) == "1":
						a__  :	str		 =			2
			# Initialize Accelerator

			# New Code #
			# We pass in "all" to `log_with` to grab all available trackers in the environment
			# Note: If using a custom `Tracker` class, should be passed in here such as:
			# >>> log_with = ["all", MyCustomTrackerClassInstance()]
			if args.with_tracking:
						a__  :	Optional[int]		 =			Accelerator(
						    cpu=args.cpu							,			mixed_precision=args.mixed_precision							,			log_with="all"							,			project_dir=args.project_dir			)
			else:
						a__  :	Any		 =			Accelerator(cpu=args.cpu							,			mixed_precision=args.mixed_precision			)
			# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
			a__  :	Optional[Any]		 =			config["lr"]
			a__  :	Dict		 =			int(config["num_epochs"]			)
			a__  :	Optional[int]		 =			int(config["seed"]			)
			a__  :	Union[str, Any]		 =			int(config["batch_size"]			)
			set_seed(lowerCAmelCase__			)

			a__ ,	a__  :	List[Any]		 =			get_dataloaders(lowerCAmelCase__							,			lowerCAmelCase__			)
			a__  :	List[str]		 =			evaluate.load("glue"							,			"mrpc"			)

			# If the batch size is too big we use gradient accumulation
			a__  :	Union[str, Any]		 =			1
			if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
						a__  :	List[str]		 =			batch_size // MAX_GPU_BATCH_SIZE
						a__  :	List[str]		 =			MAX_GPU_BATCH_SIZE

			# Instantiate the model (we build the model here so that the seed also control new weights initialization)
			a__  :	Dict		 =			AutoModelForSequenceClassification.from_pretrained("bert-base-cased"							,			return_dict=lowerCAmelCase__			)

			# We could avoid this line since the accelerator is set with `device_placement=True` (default value).
			# Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
			# creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
			a__  :	Union[str, Any]		 =			model.to(accelerator.device			)

			# Instantiate optimizer
			a__  :	Optional[Any]		 =			AdamW(params=model.parameters()							,			lr=lowerCAmelCase__			)

			# Instantiate scheduler
			a__  :	List[Any]		 =			get_linear_schedule_with_warmup(
			    optimizer=lowerCAmelCase__							,			num_warmup_steps=1_0_0							,			num_training_steps=(len(lowerCAmelCase__			) * num_epochs) // gradient_accumulation_steps							,			)

			# Prepare everything
			# There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
			# prepare method.
			a__ ,	a__ ,	a__ ,	a__ ,	a__  :	Union[str, Any]		 =			accelerator.prepare(
			    lowerCAmelCase__							,			lowerCAmelCase__							,			lowerCAmelCase__							,			lowerCAmelCase__							,			lowerCAmelCase__			)

			# New Code #
			# We need to initialize the trackers we use. Overall configurations can also be stored
			if args.with_tracking:
						a__  :	List[Any]		 =			os.path.split(lowerCAmelCase__			)[-1].split("."			)[0]
						accelerator.init_trackers(lowerCAmelCase__							,			lowerCAmelCase__			)

			# Now we train the model
			for epoch in range(lowerCAmelCase__			):
						model.train()
						# New Code #
						# For our tracking example, we will log the total loss of each epoch
						if args.with_tracking:
									a__  :	Optional[Any]		 =			0
						for step, batch in enumerate(lowerCAmelCase__			):
									# We could avoid this line since we set the accelerator with `device_placement=True`.
									batch.to(accelerator.device			)
									a__  :	List[str]		 =			model(**lowerCAmelCase__			)
									a__  :	int		 =			outputs.loss
									# New Code #
									if args.with_tracking:
												total_loss += loss.detach().float()
									a__  :	Tuple		 =			loss / gradient_accumulation_steps
									accelerator.backward(lowerCAmelCase__			)
									if step % gradient_accumulation_steps == 0:
												optimizer.step()
												lr_scheduler.step()
												optimizer.zero_grad()

						model.eval()
						for step, batch in enumerate(lowerCAmelCase__			):
									# We could avoid this line since we set the accelerator with `device_placement=True` (the default).
									batch.to(accelerator.device			)
									with torch.no_grad():
												a__  :	Union[str, Any]		 =			model(**lowerCAmelCase__			)
									a__  :	Union[str, Any]		 =			outputs.logits.argmax(dim=-1			)
									a__ ,	a__  :	Optional[Any]		 =			accelerator.gather_for_metrics((predictions, batch["labels"])			)
									metric.add_batch(
									    predictions=lowerCAmelCase__							,			references=lowerCAmelCase__							,			)

						a__  :	Union[str, Any]		 =			metric.compute()
						# Use accelerator.print to print only on the main process.
						accelerator.print(F"epoch {epoch}:"							,			lowerCAmelCase__			)

						# New Code #
						# To actually log, we call `Accelerator.log`
						# The values passed can be of `str`, `int`, `float` or `dict` of `str` to `float`/`int`
						if args.with_tracking:
									accelerator.log(
									    {
									        "accuracy": eval_metric["accuracy"],
									        "f1": eval_metric["f1"],
									        "train_loss": total_loss.item() / len(lowerCAmelCase__			),
									        "epoch": epoch,
									    }							,			step=lowerCAmelCase__							,			)

    # New Code #
    # When a run is finished, you should call `accelerator.end_training()`
    # to close all of the open trackers
			if args.with_tracking:
						accelerator.end_training()






def      lowercase__ (					)    ->			Tuple:



			'''simple docstring'''
			a__  :	Optional[Any]		 =			argparse.ArgumentParser(description="Simple example of training script."			)
			parser.add_argument(
			    "--mixed_precision"							,			type=lowerCAmelCase__							,			default=lowerCAmelCase__							,			choices=["no", "fp16", "bf16", "fp8"]							,			help="Whether to use mixed precision. Choose"
			    "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
			    "and an Nvidia Ampere GPU."							,			)
			parser.add_argument("--cpu"							,			action="store_true"							,			help="If passed, will train on the CPU."			)
			parser.add_argument(
			    "--with_tracking"							,			action="store_true"							,			help="Whether to load in all available experiment trackers from the environment and use them for logging."							,			)
			parser.add_argument(
			    "--project_dir"							,			type=lowerCAmelCase__							,			default="logs"							,			help="Location on where to store experiment tracking logs` and relevent project information"							,			)
			a__  :	List[Any]		 =			parser.parse_args()
			a__  :	Optional[Any]		 =			{"lr": 2E-5, "num_epochs": 3, "seed": 4_2, "batch_size": 1_6}
			training_function(lowerCAmelCase__							,			lowerCAmelCase__			)


if __name__ == "__main__":
		main()






"""simple docstring"""






import argparse
import re

import torch
from CLAP import create_model

from transformers import AutoFeatureExtractor, ClapConfig, ClapModel


__UpperCAmelCase						=			{
    '''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''',
}

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






def      lowercase__ (					lowerCAmelCase__   :			Optional[int]							,			lowerCAmelCase__   :			Optional[int]=False			)    ->			List[str]:



			'''simple docstring'''
			a__ ,	a__  :	Optional[int]		 =			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      lowercase__ (					lowerCAmelCase__   :			Optional[int]			)    ->			Dict:



			'''simple docstring'''
			a__  :	Any		 =			{}

			a__  :	Tuple		 =			R".*sequential.(\d+).*"
			a__  :	Tuple		 =			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(lowerCAmelCase__							,			lowerCAmelCase__			)

						if re.match(lowerCAmelCase__							,			lowerCAmelCase__			):
									# replace sequential layers with list
									a__  :	Optional[Any]		 =			re.match(lowerCAmelCase__							,			lowerCAmelCase__			).group(1			)

									a__  :	int		 =			key.replace(F"sequential.{sequential_layer}."							,			F"layers.{int(lowerCAmelCase__			)//3}.linear."			)
						elif re.match(lowerCAmelCase__							,			lowerCAmelCase__			):
									a__  :	Tuple		 =			int(re.match(lowerCAmelCase__							,			lowerCAmelCase__			).group(1			)			)

									# Because in CLAP they use `nn.Sequential`...
									a__  :	List[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__  :	List[Any]		 =			value
									a__  :	List[Any]		 =			mixed_qkv.size(0			) // 3

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

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

			return model_state_dict






def      lowercase__ (					lowerCAmelCase__   :			List[str]							,			lowerCAmelCase__   :			Any							,			lowerCAmelCase__   :			Optional[Any]							,			lowerCAmelCase__   :			Tuple=False			)    ->			Tuple:



			'''simple docstring'''
			a__ ,	a__  :	Tuple		 =			init_clap(lowerCAmelCase__							,			enable_fusion=lowerCAmelCase__			)

			clap_model.eval()
			a__  :	Optional[int]		 =			clap_model.state_dict()
			a__  :	Optional[Any]		 =			rename_state_dict(lowerCAmelCase__			)

			a__  :	Union[str, Any]		 =			ClapConfig()
			a__  :	Dict		 =			enable_fusion
			a__  :	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__":
		__UpperCAmelCase						=			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''')
		__UpperCAmelCase						=			parser.parse_args()

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