


import math
from typing import Any, Callable, List, Optional, Tuple, Union

import numpy as np
import torch

from ...models import TaFilmDecoder
from ...schedulers import DDPMScheduler
from ...utils import is_onnx_available, logging, randn_tensor


if is_onnx_available():
				from ..onnx_utils import OnnxRuntimeModel

from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline
from .continous_encoder import SpectrogramContEncoder
from .notes_encoder import SpectrogramNotesEncoder


UpperCamelCase__   : List[str]					  = logging.get_logger(__name__)  # pylint: disable=invalid-name

UpperCamelCase__   : Optional[Any]					  = 2_56







class   _UpperCamelCase  (       lowercase_  ):







					'''simple docstring'''

					lowerCamelCase     :     str			    =       ['''melgan''']





					def __init__(						self    :  int ,					__lowercase    :  SpectrogramNotesEncoder ,					__lowercase    :  SpectrogramContEncoder ,					__lowercase    :  TaFilmDecoder ,					__lowercase    :  DDPMScheduler ,					__lowercase    :  OnnxRuntimeModel if is_onnx_available() else Any ,					):






						'''simple docstring'''







						super().__init__()

						# From MELGAN
						UpperCAmelCase_								=     math.log(1e-5		)  # Matches MelGAN training.
						UpperCAmelCase_								=     4.0  # Largest value for most examples
						UpperCAmelCase_								=     1_28

						self.register_modules(
						    notes_encoder=UpperCamelCase__ ,					continuous_encoder=UpperCamelCase__ ,					decoder=UpperCamelCase__ ,					scheduler=UpperCamelCase__ ,					melgan=UpperCamelCase__ ,					)





					def        SCREAMING_SNAKE_CASE							(						self    :  Tuple ,					__lowercase    :  Any ,					__lowercase    :  int=(-1.0, 1.0) ,					__lowercase    :  Tuple=False		):






						'''simple docstring'''







						UpperCAmelCase_						,	UpperCAmelCase_								=     output_range
						if clip:
							UpperCAmelCase_								=     torch.clip(UpperCamelCase__ ,					self.min_value ,					self.max_value		)
						# Scale to [0, 1].
						UpperCAmelCase_								=     (features - self.min_value) / (self.max_value - self.min_value)
						# Scale to [min_out, max_out].
						return zero_one * (max_out - min_out) + min_out





					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[Any] ,					__lowercase    :  str ,					__lowercase    :  List[str]=(-1.0, 1.0) ,					__lowercase    :  Optional[Any]=False		):






						'''simple docstring'''







						UpperCAmelCase_						,	UpperCAmelCase_								=     input_range
						UpperCAmelCase_								=     torch.clip(UpperCamelCase__ ,					UpperCamelCase__ ,					UpperCamelCase__		) if clip else outputs
						# Scale to [0, 1].
						UpperCAmelCase_								=     (outputs - min_out) / (max_out - min_out)
						# Scale to [self.min_value, self.max_value].
						return zero_one * (self.max_value - self.min_value) + self.min_value





					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[int] ,					__lowercase    :  Union[str, Any] ,					__lowercase    :  Optional[int] ,					__lowercase    :  Any		):






						'''simple docstring'''







						UpperCAmelCase_								=     input_tokens > 0
						UpperCAmelCase_						,	UpperCAmelCase_								=     self.notes_encoder(
						    encoder_input_tokens=UpperCamelCase__ ,					encoder_inputs_mask=UpperCamelCase__		)

						UpperCAmelCase_						,	UpperCAmelCase_								=     self.continuous_encoder(
						    encoder_inputs=UpperCamelCase__ ,					encoder_inputs_mask=UpperCamelCase__		)

						return [(tokens_encoded, tokens_mask), (continuous_encoded, continuous_mask)]





					def        SCREAMING_SNAKE_CASE							(						self    :  List[Any] ,					__lowercase    :  Dict ,					__lowercase    :  Any ,					__lowercase    :  int		):






						'''simple docstring'''







						UpperCAmelCase_								=     noise_time
						if not torch.is_tensor(UpperCamelCase__		):
							UpperCAmelCase_								=     torch.tensor([timesteps] ,					dtype=torch.long ,					device=input_tokens.device		)
						elif torch.is_tensor(UpperCamelCase__		) and len(timesteps.shape		) == 0:
							UpperCAmelCase_								=     timesteps[None].to(input_tokens.device		)

						# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
						UpperCAmelCase_								=     timesteps * torch.ones(input_tokens.shape[0] ,					dtype=timesteps.dtype ,					device=timesteps.device		)

						UpperCAmelCase_								=     self.decoder(
						    encodings_and_masks=UpperCamelCase__ ,					decoder_input_tokens=UpperCamelCase__ ,					decoder_noise_time=UpperCamelCase__		)
						return logits





					@torch.no_grad()
					def __call__(						self    :  Dict ,					__lowercase    :  List[List[int]] ,					__lowercase    :  Optional[torch.Generator] = None ,					__lowercase    :  int = 1_00 ,					__lowercase    :  bool = True ,					__lowercase    :  str = "numpy" ,					__lowercase    :  Optional[Callable[[int, int, torch.FloatTensor], None]] = None ,					__lowercase    :  int = 1 ,					):






						'''simple docstring'''







						if (callback_steps is None) or (
						    callback_steps is not None and (not isinstance(UpperCamelCase__ ,					UpperCamelCase__		) or callback_steps <= 0)
						):
							raise ValueError(
							    F"""`callback_steps` has to be a positive integer but is {callback_steps} of type"""
							    F""" {type(UpperCamelCase__		)}."""		)

						UpperCAmelCase_								=     np.zeros([1, TARGET_FEATURE_LENGTH, self.n_dims] ,					dtype=np.floataa		)
						UpperCAmelCase_								=     np.zeros([1, 0, self.n_dims] ,					np.floataa		)
						UpperCAmelCase_								=     torch.ones((1, TARGET_FEATURE_LENGTH) ,					dtype=UpperCamelCase__ ,					device=self.device		)

						for i, encoder_input_tokens in enumerate(UpperCamelCase__		):
							if i == 0:
								UpperCAmelCase_								=     torch.from_numpy(pred_mel[:1].copy()		).to(
								    device=self.device ,					dtype=self.decoder.dtype		)
								# The first chunk has no previous context.
								UpperCAmelCase_								=     torch.zeros((1, TARGET_FEATURE_LENGTH) ,					dtype=UpperCamelCase__ ,					device=self.device		)
							else:
								# The full song pipeline does not feed in a context feature, so the mask
								# will be all 0s after the feature converter. Because we know we're
								# feeding in a full context chunk from the previous prediction, set it
								# to all 1s.
								UpperCAmelCase_								=     ones

							UpperCAmelCase_								=     self.scale_features(
							    UpperCamelCase__ ,					output_range=[-1.0, 1.0] ,					clip=UpperCamelCase__		)

							UpperCAmelCase_								=     self.encode(
							    input_tokens=torch.IntTensor([encoder_input_tokens]		).to(device=self.device		) ,					continuous_inputs=UpperCamelCase__ ,					continuous_mask=UpperCamelCase__ ,					)

							# Sample encoder_continuous_inputs shaped gaussian noise to begin loop
							UpperCAmelCase_								=     randn_tensor(
							    shape=encoder_continuous_inputs.shape ,					generator=UpperCamelCase__ ,					device=self.device ,					dtype=self.decoder.dtype ,					)

							# set step values
							self.scheduler.set_timesteps(UpperCamelCase__		)

							# Denoising diffusion loop
							for j, t in enumerate(self.progress_bar(self.scheduler.timesteps		)		):
								UpperCAmelCase_								=     self.decode(
								    encodings_and_masks=UpperCamelCase__ ,					input_tokens=UpperCamelCase__ ,					noise_time=t / self.scheduler.config.num_train_timesteps ,					)

								# Compute previous output: x_t -> x_t-1
								UpperCAmelCase_								=     self.scheduler.step(UpperCamelCase__ ,					UpperCamelCase__ ,					UpperCamelCase__ ,					generator=UpperCamelCase__		).prev_sample

							UpperCAmelCase_								=     self.scale_to_features(UpperCamelCase__ ,					input_range=[-1.0, 1.0]		)
							UpperCAmelCase_								=     mel[:1]
							UpperCAmelCase_								=     mel.cpu().float().numpy()

							UpperCAmelCase_								=     np.concatenate([full_pred_mel, pred_mel[:1]] ,					axis=1		)

							# call the callback, if provided
							if callback is not None and i % callback_steps == 0:
								callback(UpperCamelCase__ ,					UpperCamelCase__		)

							logger.info("""Generated segment""" ,					UpperCamelCase__		)

						if output_type == "numpy" and not is_onnx_available():
							raise ValueError(
							    """Cannot return output in 'np' format if ONNX is not available. Make sure to have ONNX installed or set 'output_type' to 'mel'."""		)
						elif output_type == "numpy" and self.melgan is None:
							raise ValueError(
							    """Cannot return output in 'np' format if melgan component is not defined. Make sure to define `self.melgan` or set 'output_type' to 'mel'."""		)

						if output_type == "numpy":
							UpperCAmelCase_								=     self.melgan(input_features=full_pred_mel.astype(np.floataa		)		)
						else:
							UpperCAmelCase_								=     full_pred_mel

						if not return_dict:
							return (output,)

						return AudioPipelineOutput(audios=UpperCamelCase__		)









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
# and perform gradient accumulation
#
# 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 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__   : Union[str, Any]					  = 16
UpperCamelCase__   : Tuple					  = 32




def    A_(			A				,      A = 16   ):
	UpperCAmelCase_								=     AutoTokenizer.from_pretrained("""bert-base-cased"""   )
	UpperCAmelCase_								=     load_dataset("""glue"""				,      """mrpc"""   )

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

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

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

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

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




def    A_(			A				,      A   ):
	# For testing only
	if os.environ.get("""TESTING_MOCKED_DATALOADERS"""				,      A   ) == "1":
		UpperCAmelCase_								=     2
	# New Code #
	UpperCAmelCase_								=     int(args.gradient_accumulation_steps   )
	# Initialize accelerator
	UpperCAmelCase_								=     Accelerator(
	    cpu=args.cpu				,      mixed_precision=args.mixed_precision				,      gradient_accumulation_steps=A   )
	if accelerator.distributed_type == DistributedType.TPU and gradient_accumulation_steps > 1:
		raise NotImplementedError(
		    """Gradient accumulation on TPUs is currently not supported. Pass `gradient_accumulation_steps=1`"""   )
	# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
	UpperCAmelCase_								=     config["""lr"""]
	UpperCAmelCase_								=     int(config["""num_epochs"""]   )
	UpperCAmelCase_								=     int(config["""seed"""]   )
	UpperCAmelCase_								=     int(config["""batch_size"""]   )

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

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

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

	# Instantiate optimizer
	UpperCAmelCase_								=     AdamW(params=model.parameters()				,      lr=A   )

	# Instantiate scheduler
	UpperCAmelCase_								=     get_linear_schedule_with_warmup(
	    optimizer=A				,      num_warmup_steps=100				,      num_training_steps=(len(A   ) * 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.
	UpperCAmelCase_						,	UpperCAmelCase_						,	UpperCAmelCase_						,	UpperCAmelCase_						,	UpperCAmelCase_								=     accelerator.prepare(
	    A				,      A				,      A				,      A				,      A   )

	# Now we train the model
	for epoch in range(A   ):
		model.train()
		for step, batch in enumerate(A   ):
			# We could avoid this line since we set the accelerator with `device_placement=True`.
			batch.to(accelerator.device   )
			# New code #
			# We use the new `accumulate` context manager to perform gradient accumulation
			# We also currently do not support TPUs nor advise it as bugs were found on the XLA side when running our tests.
			with accelerator.accumulate(A   ):
				UpperCAmelCase_								=     model(**A   )
				UpperCAmelCase_								=     output.loss
				accelerator.backward(A   )
				optimizer.step()
				lr_scheduler.step()
				optimizer.zero_grad()

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

		UpperCAmelCase_								=     metric.compute()
		# Use accelerator.print to print only on the main process.
		accelerator.print(f"""epoch {epoch}:"""				,      A   )




def    A_(			):
	UpperCAmelCase_								=     argparse.ArgumentParser(description="""Simple example of training script."""   )
	parser.add_argument(
	    """--mixed_precision"""				,      type=A				,      default=A				,      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."""				,      )
	# New Code #
	parser.add_argument(
	    """--gradient_accumulation_steps"""				,      type=A				,      default=1				,      help="""The number of minibatches to be ran before gradients are accumulated."""				,      )
	parser.add_argument("""--cpu"""				,      action="""store_true"""				,      help="""If passed, will train on the CPU."""   )
	UpperCAmelCase_								=     parser.parse_args()
	UpperCAmelCase_								=     {"""lr""": 2E-5, """num_epochs""": 3, """seed""": 42, """batch_size""": 16}
	training_function(A				,      A   )


if __name__ == "__main__":
				main()


