


'''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
from accelerate.local_sgd import LocalSGD


########################################################################
# This is a fully working simple example to use Accelerate
# with LocalSGD, which is a method to synchronize model
# parameters every K batches. It is different, but complementary
# to 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_ :     Dict								=      1_6
UpperCAmelCase_ :     Dict								=      3_2


def 			_lowercase       (				UpperCamelCase__       :					Accelerator,				UpperCamelCase__       :					int = 16			):
 __A  :		Any       =			AutoTokenizer.from_pretrained('bert-base-cased'			)
 __A  :		Optional[int]       =			load_dataset('glue',				'mrpc'			)

 def tokenize_function(UpperCamelCase__       :					Any			):
  # max_length=None => use the model max length (it's actually the default)
  __A  :		Any       =			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():
  __A  :		List[Any]       =			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
 __A  :		Union[str, Any]       =			tokenized_datasets.rename_column('label',				'labels'			)

 def collate_fn(UpperCamelCase__       :					Any			):
  # On TPU it's best to pad everything to the same length or training will be very slow.
  __A  :		Tuple       =			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":
   __A  :		List[Any]       =			16
  elif accelerator.mixed_precision != "no":
   __A  :		Dict       =			8
  else:
   __A  :		Any       =			None

  return tokenizer.pad(
      UpperCamelCase__,				padding='longest',				max_length=UpperCamelCase__,				pad_to_multiple_of=UpperCamelCase__,				return_tensors='pt',				)

 # Instantiate dataloaders.
 __A  :		str       =			DataLoader(
     tokenized_datasets['train'],				shuffle=UpperCamelCase__,				collate_fn=UpperCamelCase__,				batch_size=UpperCamelCase__			)
 __A  :		int       =			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_ :     Optional[Any]								=      mocked_dataloaders  # noqa: F811


def 			_lowercase       (				UpperCamelCase__       :					Any,				UpperCamelCase__       :					Optional[int]			):
 # For testing only
 if os.environ.get('TESTING_MOCKED_DATALOADERS',				UpperCamelCase__			) == "1":
  __A  :		Union[str, Any]       =			2
 # New Code #
 __A  :		Union[str, Any]       =			int(args.gradient_accumulation_steps			)
 __A  :		str       =			int(args.local_sgd_steps			)
 # Initialize accelerator
 __A  :		Tuple       =			Accelerator(
     cpu=args.cpu,				mixed_precision=args.mixed_precision,				gradient_accumulation_steps=UpperCamelCase__			)
 if accelerator.distributed_type not in [DistributedType.NO, DistributedType.MULTI_CPU, DistributedType.MULTI_GPU]:
  raise NotImplementedError('LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)'			)
 # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
 __A  :		Dict       =			config['lr']
 __A  :		int       =			int(config['num_epochs']			)
 __A  :		List[Any]       =			int(config['seed']			)
 __A  :		str       =			int(config['batch_size']			)

 __A  :		Any       =			evaluate.load('glue',				'mrpc'			)

 set_seed(UpperCamelCase__			)
 __A  :		Tuple       =			get_dataloaders(UpperCamelCase__,				UpperCamelCase__			)
 # Instantiate the model (we build the model here so that the seed also control new weights initialization)
 __A  :		Optional[Any]       =			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).
 __A  :		str       =			model.to(accelerator.device			)

 # Instantiate optimizer
 __A  :		Union[str, Any]       =			AdamW(params=model.parameters(),				lr=UpperCamelCase__			)

 # Instantiate scheduler
 __A  :		List[Any]       =			get_linear_schedule_with_warmup(
     optimizer=UpperCamelCase__,				num_warmup_steps=100,				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.
 __A  :		Optional[int]       =			accelerator.prepare(
     UpperCamelCase__,				UpperCamelCase__,				UpperCamelCase__,				UpperCamelCase__,				UpperCamelCase__			)

 # Now we train the model
 for epoch in range(UpperCamelCase__			):
  model.train()
  with LocalSGD(
      accelerator=UpperCamelCase__,				model=UpperCamelCase__,				local_sgd_steps=UpperCamelCase__,				enabled=local_sgd_steps is not None			) as local_sgd:
   for step, batch in enumerate(UpperCamelCase__			):
    # 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(UpperCamelCase__			):
     __A  :		List[Any]       =			model(**UpperCamelCase__			)
     __A  :		Optional[Any]       =			output.loss
     accelerator.backward(UpperCamelCase__			)
     optimizer.step()
     lr_scheduler.step()
     optimizer.zero_grad()
     # LocalSGD-specific line
     local_sgd.step()

  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():
    __A  :		Union[str, Any]       =			model(**UpperCamelCase__			)
   __A  :		Any       =			outputs.logits.argmax(dim=-1			)
   __A  :		List[Any]       =			accelerator.gather_for_metrics((predictions, batch['labels'])			)
   metric.add_batch(
       predictions=UpperCamelCase__,				references=UpperCamelCase__,				)

  __A  :		Optional[int]       =			metric.compute()
  # Use accelerator.print to print only on the main process.
  accelerator.print(f"""epoch {epoch}:""",				UpperCamelCase__			)


def 			_lowercase       (				):
 __A  :		Dict       =			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.',				)
 # New Code #
 parser.add_argument(
     '--gradient_accumulation_steps',				type=UpperCamelCase__,				default=1,				help='The number of minibatches to be ran before gradients are accumulated.',				)
 parser.add_argument(
     '--local_sgd_steps',				type=UpperCamelCase__,				default=8,				help='Number of local SGD steps or None to disable local SGD'			)
 parser.add_argument('--cpu',				action='store_true',				help='If passed, will train on the CPU.'			)
 __A  :		List[Any]       =			parser.parse_args()
 __A  :		Optional[Any]       =			{'lr': 2E-5, 'num_epochs': 3, 'seed': 42, 'batch_size': 16}
 training_function(UpperCamelCase__,				UpperCamelCase__			)


if __name__ == "__main__":
  main()










'''simple docstring'''
import gc
import unittest

import numpy as np
import torch

from diffusers import (
    AudioDiffusionPipeline,
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    DiffusionPipeline,
    Mel,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()

class 						_lowerCamelCase					(    unittest.TestCase				):




 '''simple docstring'''


 def  snake_case__   (      self						):

  """simple docstring"""







  super().tearDown()
  gc.collect()
  torch.cuda.empty_cache()


 @property
 def  snake_case__   (      self						):

  """simple docstring"""







  torch.manual_seed(0						)
  __A  :		Optional[int]       =			UNetaDModel(
      sample_size=(32, 64)      ,      in_channels=1      ,      out_channels=1      ,      layers_per_block=2      ,      block_out_channels=(128, 128)      ,      down_block_types=('AttnDownBlock2D', 'DownBlock2D')      ,      up_block_types=('UpBlock2D', 'AttnUpBlock2D')      ,      )
  return model


 @property
 def  snake_case__   (      self						):

  """simple docstring"""







  torch.manual_seed(0						)
  __A  :		Dict       =			UNetaDConditionModel(
      sample_size=(64, 32)      ,      in_channels=1      ,      out_channels=1      ,      layers_per_block=2      ,      block_out_channels=(128, 128)      ,      down_block_types=('CrossAttnDownBlock2D', 'DownBlock2D')      ,      up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D')      ,      cross_attention_dim=10      ,      )
  return model


 @property
 def  snake_case__   (      self						):

  """simple docstring"""







  torch.manual_seed(0						)
  __A  :		Union[str, Any]       =			AutoencoderKL(
      sample_size=(128, 64)      ,      in_channels=1      ,      out_channels=1      ,      latent_channels=1      ,      layers_per_block=2      ,      block_out_channels=(128, 128)      ,      down_block_types=('DownEncoderBlock2D', 'DownEncoderBlock2D')      ,      up_block_types=('UpDecoderBlock2D', 'UpDecoderBlock2D')      ,      )
  __A  :		Union[str, Any]       =			UNetaDModel(
      sample_size=(64, 32)      ,      in_channels=1      ,      out_channels=1      ,      layers_per_block=2      ,      block_out_channels=(128, 128)      ,      down_block_types=('AttnDownBlock2D', 'DownBlock2D')      ,      up_block_types=('UpBlock2D', 'AttnUpBlock2D')      ,      )
  return vqvae, unet







 @slow
 def  snake_case__   (      self						):

  """simple docstring"""







  __A  :		str       =			'cpu'  # ensure determinism for the device-dependent torch.Generator
  __A  :		int       =			Mel(
      x_res=self.dummy_unet.config.sample_size[1]      ,      y_res=self.dummy_unet.config.sample_size[0]      ,      )

  __A  :		Tuple       =			DDPMScheduler()
  __A  :		Optional[int]       =			AudioDiffusionPipeline(vqvae=__lowercase      ,      unet=self.dummy_unet      ,      mel=__lowercase      ,      scheduler=__lowercase						)
  __A  :		Dict       =			pipe.to(__lowercase						)
  pipe.set_progress_bar_config(disable=__lowercase						)

  __A  :		Tuple       =			torch.Generator(device=__lowercase						).manual_seed(42						)
  __A  :		Tuple       =			pipe(generator=__lowercase      ,      steps=4						)
  __A  :		Union[str, Any]       =			output.audios[0]
  __A  :		Dict       =			output.images[0]

  __A  :		Optional[Any]       =			torch.Generator(device=__lowercase						).manual_seed(42						)
  __A  :		Optional[int]       =			pipe(generator=__lowercase      ,      steps=4      ,      return_dict=__lowercase						)
  __A  :		int       =			output[0][0]

  assert audio.shape == (1, (self.dummy_unet.config.sample_size[1] - 1) * mel.hop_length)
  assert (
      image.height == self.dummy_unet.config.sample_size[0]
      and image.width == self.dummy_unet.config.sample_size[1]
  )
  __A  :		Optional[int]       =			np.frombuffer(image.tobytes()      ,      dtype='uint8'						)[:10]
  __A  :		List[Any]       =			np.frombuffer(image_from_tuple.tobytes()      ,      dtype='uint8'						)[:10]
  __A  :		Optional[Any]       =			np.array([69, 255, 255, 255, 0, 0, 77, 181, 12, 127]						)

  assert np.abs(image_slice.flatten() - expected_slice						).max() == 0
  assert np.abs(image_from_tuple_slice.flatten() - expected_slice						).max() == 0

  __A  :		Optional[int]       =			Mel(
      x_res=self.dummy_vqvae_and_unet[0].config.sample_size[1]      ,      y_res=self.dummy_vqvae_and_unet[0].config.sample_size[0]      ,      )

  __A  :		Optional[Any]       =			DDIMScheduler()
  __A  :		List[Any]       =			self.dummy_vqvae_and_unet
  __A  :		int       =			AudioDiffusionPipeline(
      vqvae=self.dummy_vqvae_and_unet[0]      ,      unet=dummy_vqvae_and_unet[1]      ,      mel=__lowercase      ,      scheduler=__lowercase						)
  __A  :		Dict       =			pipe.to(__lowercase						)
  pipe.set_progress_bar_config(disable=__lowercase						)

  np.random.seed(0						)
  __A  :		Dict       =			np.random.uniform(-1      ,      1      ,      ((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)						)
  __A  :		Dict       =			torch.Generator(device=__lowercase						).manual_seed(42						)
  __A  :		List[Any]       =			pipe(raw_audio=__lowercase      ,      generator=__lowercase      ,      start_step=5      ,      steps=10						)
  __A  :		Tuple       =			output.images[0]

  assert (
      image.height == self.dummy_vqvae_and_unet[0].config.sample_size[0]
      and image.width == self.dummy_vqvae_and_unet[0].config.sample_size[1]
  )
  __A  :		int       =			np.frombuffer(image.tobytes()      ,      dtype='uint8'						)[:10]
  __A  :		List[str]       =			np.array([120, 117, 110, 109, 138, 167, 138, 148, 132, 121]						)

  assert np.abs(image_slice.flatten() - expected_slice						).max() == 0

  __A  :		Any       =			self.dummy_unet_condition
  __A  :		Optional[int]       =			AudioDiffusionPipeline(
      vqvae=self.dummy_vqvae_and_unet[0]      ,      unet=__lowercase      ,      mel=__lowercase      ,      scheduler=__lowercase						)
  __A  :		Union[str, Any]       =			pipe.to(__lowercase						)
  pipe.set_progress_bar_config(disable=__lowercase						)

  np.random.seed(0						)
  __A  :		Dict       =			torch.rand((1, 1, 10)						)
  __A  :		Any       =			pipe(generator=__lowercase      ,      encoding=__lowercase						)
  __A  :		Union[str, Any]       =			output.images[0]
  __A  :		List[Any]       =			np.frombuffer(image.tobytes()      ,      dtype='uint8'						)[:10]
  __A  :		int       =			np.array([107, 103, 120, 127, 142, 122, 113, 122, 97, 111]						)

  assert np.abs(image_slice.flatten() - expected_slice						).max() == 0


@slow
@require_torch_gpu
class 						_lowerCamelCase					(    unittest.TestCase				):




 '''simple docstring'''


 def  snake_case__   (      self						):

  """simple docstring"""







  super().tearDown()
  gc.collect()
  torch.cuda.empty_cache()







 def  snake_case__   (      self						):

  """simple docstring"""







  __A  :		Union[str, Any]       =			torch_device

  __A  :		Union[str, Any]       =			DiffusionPipeline.from_pretrained('teticio/audio-diffusion-ddim-256'						)
  __A  :		Dict       =			pipe.to(__lowercase						)
  pipe.set_progress_bar_config(disable=__lowercase						)

  __A  :		Any       =			torch.Generator(device=__lowercase						).manual_seed(42						)
  __A  :		Optional[Any]       =			pipe(generator=__lowercase						)
  __A  :		Union[str, Any]       =			output.audios[0]
  __A  :		List[Any]       =			output.images[0]

  assert audio.shape == (1, (pipe.unet.config.sample_size[1] - 1) * pipe.mel.hop_length)
  assert image.height == pipe.unet.config.sample_size[0] and image.width == pipe.unet.config.sample_size[1]
  __A  :		Union[str, Any]       =			np.frombuffer(image.tobytes()      ,      dtype='uint8'						)[:10]
  __A  :		Optional[Any]       =			np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26]						)

  assert np.abs(image_slice.flatten() - expected_slice						).max() == 0






