





import argparse
import os

# New Code #
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
from accelerate.utils import find_executable_batch_size


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to ensure out-of-memory errors never
# interrupt training, 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)
#
# 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__						:			int									=16
UpperCAmelCase__						:			List[str]									=32





def      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase = 16				)						->  Union[str, Any]:
				lowerCamelCase	 =AutoTokenizer.from_pretrained("""bert-base-cased"""				)
				lowerCamelCase	 =load_dataset("""glue"""			,      """mrpc"""				)

				def tokenize_function(_UpperCAmelCase				):
								# max_length=None => use the model max length (it's actually the default)
								lowerCamelCase	 =tokenizer(examples["""sentence1"""]			,      examples["""sentence2"""]			,      truncation=_UpperCAmelCase			,      max_length=_UpperCAmelCase				)
								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():
								lowerCamelCase	 =datasets.map(
								    _UpperCAmelCase			,      batched=_UpperCAmelCase			,      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
				lowerCamelCase	 =tokenized_datasets.rename_column("""label"""			,      """labels"""				)

				def collate_fn(_UpperCAmelCase				):
								# On TPU it's best to pad everything to the same length or training will be very slow.
								lowerCamelCase	 =1_28 if accelerator.distributed_type == DistributedType.TPU else None
								# When using mixed precision we want round multiples of 8/16
								if accelerator.mixed_precision == "fp8":
												lowerCamelCase	 =16
								elif accelerator.mixed_precision != "no":
												lowerCamelCase	 =8
								else:
												lowerCamelCase	 =None

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

				# Instantiate dataloaders.
				lowerCamelCase	 =DataLoader(
				    tokenized_datasets["""train"""]			,      shuffle=_UpperCAmelCase			,      collate_fn=_UpperCAmelCase			,      batch_size=_UpperCAmelCase				)
				lowerCamelCase	 =DataLoader(
				    tokenized_datasets["""validation"""]			,      shuffle=_UpperCAmelCase			,      collate_fn=_UpperCAmelCase			,      batch_size=_UpperCAmelCase				)

				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__						:			Tuple									=mocked_dataloaders  # noqa: F811





def      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase				)						->  Union[str, Any]:
				# For testing only
				if os.environ.get("""TESTING_MOCKED_DATALOADERS"""			,      _UpperCAmelCase				) == "1":
								lowerCamelCase	 =2
				# Initialize accelerator
				lowerCamelCase	 =Accelerator(cpu=args.cpu			,      mixed_precision=args.mixed_precision				)
				# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
				lowerCamelCase	 =config["""lr"""]
				lowerCamelCase	 =int(config["""num_epochs"""]				)
				lowerCamelCase	 =int(config["""seed"""]				)
				lowerCamelCase	 =int(config["""batch_size"""]				)

				lowerCamelCase	 =evaluate.load("""glue"""			,      """mrpc"""				)

				# New Code #
				# We now can define an inner training loop function. It should take a batch size as the only parameter,
				# and build the dataloaders in there.
				# It also gets our decorator
				@find_executable_batch_size(starting_batch_size=_UpperCAmelCase				)
				def inner_training_loop(_UpperCAmelCase				):
								# And now just move everything below under this function
								# We need to bring in the Accelerator object from earlier
								nonlocal accelerator
								# And reset all of its attributes that could hold onto any memory:
								accelerator.free_memory()

								# Then we can declare the model, optimizer, and everything else:
								set_seed(_UpperCAmelCase				)

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

								# 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).
								lowerCamelCase	 =model.to(accelerator.device				)

								# Instantiate optimizer
								lowerCamelCase	 =AdamW(params=model.parameters()			,      lr=_UpperCAmelCase				)
								lowerCamelCase ,  lowerCamelCase	 =get_dataloaders(_UpperCAmelCase			,      _UpperCAmelCase				)

								# Instantiate scheduler
								lowerCamelCase	 =get_linear_schedule_with_warmup(
								    optimizer=_UpperCAmelCase			,      num_warmup_steps=1_00			,      num_training_steps=(len(_UpperCAmelCase				) * num_epochs)			,      )

								# 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.
								lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase	 =accelerator.prepare(
								    _UpperCAmelCase			,      _UpperCAmelCase			,      _UpperCAmelCase			,      _UpperCAmelCase			,      _UpperCAmelCase				)

								# Now we train the model
								for epoch in range(_UpperCAmelCase				):
												model.train()
												for step, batch in enumerate(_UpperCAmelCase				):
																# We could avoid this line since we set the accelerator with `device_placement=True`.
																batch.to(accelerator.device				)
																lowerCamelCase	 =model(**_UpperCAmelCase				)
																lowerCamelCase	 =outputs.loss
																accelerator.backward(_UpperCAmelCase				)
																optimizer.step()
																lr_scheduler.step()
																optimizer.zero_grad()

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

												lowerCamelCase	 =metric.compute()
												# Use accelerator.print to print only on the main process.
												accelerator.print(F"""epoch {epoch}:"""			,      _UpperCAmelCase				)

    # New Code #
    # And call it at the end with no arguments
    # Note: You could also refactor this outside of your training loop function
				inner_training_loop()





def      _lowercase					(						)						->  List[str]:
				lowerCamelCase	 =argparse.ArgumentParser(description="""Simple example of training script."""				)
				parser.add_argument(
				    """--mixed_precision"""			,      type=_UpperCAmelCase			,      default=_UpperCAmelCase			,      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."""				)
				lowerCamelCase	 =parser.parse_args()
				lowerCamelCase	 ={"""lr""": 2e-5, """num_epochs""": 3, """seed""": 42, """batch_size""": 16}
				training_function(_UpperCAmelCase			,      _UpperCAmelCase				)


if __name__ == "__main__":
							main()












import argparse
from pathlib import Path
from typing import Dict, OrderedDict, Tuple

import torch
from audiocraft.models import MusicGen

from transformers import (
    AutoFeatureExtractor,
    AutoTokenizer,
    EncodecModel,
    MusicgenDecoderConfig,
    MusicgenForConditionalGeneration,
    MusicgenProcessor,
    TaEncoderModel,
)
from transformers.models.musicgen.modeling_musicgen import MusicgenForCausalLM
from transformers.utils import logging


logging.set_verbosity_info()
UpperCAmelCase__						:			Tuple									=logging.get_logger(__name__)


UpperCAmelCase__						:			str									=['''model.decoder.embed_positions.weights''']





def      _lowercase					(						_UpperCAmelCase				)						->  List[Any]:
				if "emb" in name:
								lowerCamelCase	 =name.replace("""emb"""			,      """model.decoder.embed_tokens"""				)
				if "transformer" in name:
								lowerCamelCase	 =name.replace("""transformer"""			,      """model.decoder"""				)
				if "cross_attention" in name:
								lowerCamelCase	 =name.replace("""cross_attention"""			,      """encoder_attn"""				)
				if "linear1" in name:
								lowerCamelCase	 =name.replace("""linear1"""			,      """fc1"""				)
				if "linear2" in name:
								lowerCamelCase	 =name.replace("""linear2"""			,      """fc2"""				)
				if "norm1" in name:
								lowerCamelCase	 =name.replace("""norm1"""			,      """self_attn_layer_norm"""				)
				if "norm_cross" in name:
								lowerCamelCase	 =name.replace("""norm_cross"""			,      """encoder_attn_layer_norm"""				)
				if "norm2" in name:
								lowerCamelCase	 =name.replace("""norm2"""			,      """final_layer_norm"""				)
				if "out_norm" in name:
								lowerCamelCase	 =name.replace("""out_norm"""			,      """model.decoder.layer_norm"""				)
				if "linears" in name:
								lowerCamelCase	 =name.replace("""linears"""			,      """lm_heads"""				)
				if "condition_provider.conditioners.description.output_proj" in name:
								lowerCamelCase	 =name.replace("""condition_provider.conditioners.description.output_proj"""			,      """enc_to_dec_proj"""				)
				return name





def      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase				)						->  Tuple[Dict, Dict]:
				lowerCamelCase	 =list(state_dict.keys()				)
				lowerCamelCase	 ={}
				for key in keys:
								lowerCamelCase	 =state_dict.pop(_UpperCAmelCase				)
								lowerCamelCase	 =rename_keys(_UpperCAmelCase				)
								if "in_proj_weight" in key:
												# split fused qkv proj
												lowerCamelCase	 =val[:hidden_size, :]
												lowerCamelCase	 =val[hidden_size : 2 * hidden_size, :]
												lowerCamelCase	 =val[-hidden_size:, :]
								elif "enc_to_dec_proj" in key:
												lowerCamelCase	 =val
								else:
												lowerCamelCase	 =val
				return state_dict, enc_dec_proj_state_dict





def      _lowercase					(						_UpperCAmelCase				)						->  MusicgenDecoderConfig:
				if checkpoint == "small":
								# default config values
								lowerCamelCase	 =10_24
								lowerCamelCase	 =24
								lowerCamelCase	 =16
				elif checkpoint == "medium":
								lowerCamelCase	 =15_36
								lowerCamelCase	 =48
								lowerCamelCase	 =24
				elif checkpoint == "large":
								lowerCamelCase	 =20_48
								lowerCamelCase	 =48
								lowerCamelCase	 =32
				else:
								raise ValueError(F"""Checkpoint should be one of `['small', 'medium', 'large']`, got {checkpoint}."""				)
				lowerCamelCase	 =MusicgenDecoderConfig(
				    hidden_size=_UpperCAmelCase			,      ffn_dim=hidden_size * 4			,      num_hidden_layers=_UpperCAmelCase			,      num_attention_heads=_UpperCAmelCase			,      )
				return config





@torch.no_grad()
def      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase=None			,      _UpperCAmelCase=None			,      _UpperCAmelCase="cpu"				)						->  Dict:
				lowerCamelCase	 =MusicGen.get_pretrained(_UpperCAmelCase			,      device=_UpperCAmelCase				)
				lowerCamelCase	 =decoder_config_from_checkpoint(_UpperCAmelCase				)

				lowerCamelCase	 =fairseq_model.lm.state_dict()
				lowerCamelCase ,  lowerCamelCase	 =rename_state_dict(
				    _UpperCAmelCase			,      hidden_size=decoder_config.hidden_size				)

				lowerCamelCase	 =TaEncoderModel.from_pretrained("""t5-base"""				)
				lowerCamelCase	 =EncodecModel.from_pretrained("""facebook/encodec_32khz"""				)
				lowerCamelCase	 =MusicgenForCausalLM(_UpperCAmelCase				).eval()

				# load all decoder weights - expect that we'll be missing embeddings and enc-dec projection
				lowerCamelCase ,  lowerCamelCase	 =decoder.load_state_dict(_UpperCAmelCase			,      strict=_UpperCAmelCase				)

				for key in missing_keys.copy():
								if key.startswith(("""text_encoder""", """audio_encoder""")				) or key in EXPECTED_MISSING_KEYS:
												missing_keys.remove(_UpperCAmelCase				)

				if len(_UpperCAmelCase				) > 0:
								raise ValueError(F"""Missing key(s) in state_dict: {missing_keys}"""				)

				if len(_UpperCAmelCase				) > 0:
								raise ValueError(F"""Unexpected key(s) in state_dict: {unexpected_keys}"""				)

				# init the composite model
				lowerCamelCase	 =MusicgenForConditionalGeneration(text_encoder=_UpperCAmelCase			,      audio_encoder=_UpperCAmelCase			,      decoder=_UpperCAmelCase				)

				# load the pre-trained enc-dec projection (from the decoder state dict)
				model.enc_to_dec_proj.load_state_dict(_UpperCAmelCase				)

				# check we can do a forward pass
				lowerCamelCase	 =torch.arange(0			,      8			,      dtype=torch.long				).reshape(2			,      -1				)
				lowerCamelCase	 =input_ids.reshape(2 * 4			,      -1				)

				with torch.no_grad():
								lowerCamelCase	 =model(input_ids=_UpperCAmelCase			,      decoder_input_ids=_UpperCAmelCase				).logits

				if logits.shape != (8, 1, 20_48):
								raise ValueError("""Incorrect shape for logits"""				)

				# now construct the processor
				lowerCamelCase	 =AutoTokenizer.from_pretrained("""t5-base"""				)
				lowerCamelCase	 =AutoFeatureExtractor.from_pretrained("""facebook/encodec_32khz"""			,      padding_side="""left"""				)

				lowerCamelCase	 =MusicgenProcessor(feature_extractor=_UpperCAmelCase			,      tokenizer=_UpperCAmelCase				)

				# set the appropriate bos/pad token ids
				lowerCamelCase	 =20_48
				lowerCamelCase	 =20_48

				# set other default generation config params
				lowerCamelCase	 =int(30 * audio_encoder.config.frame_rate				)
				lowerCamelCase	 =True
				lowerCamelCase	 =3.0

				if pytorch_dump_folder is not None:
								Path(_UpperCAmelCase				).mkdir(exist_ok=_UpperCAmelCase				)
								logger.info(F"""Saving model {checkpoint} to {pytorch_dump_folder}"""				)
								model.save_pretrained(_UpperCAmelCase				)
								processor.save_pretrained(_UpperCAmelCase				)

				if repo_id:
								logger.info(F"""Pushing model {checkpoint} to {repo_id}"""				)
								model.push_to_hub(_UpperCAmelCase				)
								processor.push_to_hub(_UpperCAmelCase				)


if __name__ == "__main__":
							UpperCAmelCase__						:			List[Any]									=argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '''--checkpoint''',
							    default='''small''',
							    type=str,
							    help='''Checkpoint size of the MusicGen model you\'d like to convert. Can be one of: `[\'small\', \'medium\', \'large\']`.''',
							)
							parser.add_argument(
							    '''--pytorch_dump_folder''',
							    required=True,
							    default=None,
							    type=str,
							    help='''Path to the output PyTorch model directory.''',
							)
							parser.add_argument(
							    '''--push_to_hub''', default=None, type=str, help='''Where to upload the converted model on the 🤗 hub.'''
							)
							parser.add_argument(
							    '''--device''', default='''cpu''', type=str, help='''Torch device to run the conversion, either cpu or cuda.'''
							)

							UpperCAmelCase__						:			Optional[Any]									=parser.parse_args()
							convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)





