


'''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
# 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
#
########################################################################


a__  :					Optional[Any]        =       1_6
a__  :					str        =       3_2







def _UpperCamelCase       (     __A      ,			__A = 16    )							->		List[Any]:


      '''simple docstring'''



      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

  a__  :					Any        =       mocked_dataloaders  # noqa: F811







def _UpperCamelCase       (     __A      ,			__A    )							->		Tuple:


      '''simple docstring'''



      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 _UpperCamelCase       (     )							->		List[str]:


      '''simple docstring'''



      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()












import itertools
import random
import unittest

import numpy as np

from transformers import is_speech_available
from transformers.testing_utils import require_torch, require_torchaudio

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
			from transformers import SpeechaTextFeatureExtractor

_a 						=      random.Random()





def lowerCAmelCase__(__snake_case       ,__snake_case=1.0       ,__snake_case=None       ,__snake_case=None							)				->      List[Any]:







	'''simple docstring'''







	if rng is None:
		lowerCamelCase__								=		global_rng

	lowerCamelCase__								=		[]
	for batch_idx in range(shape[0]							):
		values.append([]							)
		for _ in range(shape[1]							):
			values[-1].append(rng.random() * scale							)

	return values






@require_torch
@require_torchaudio
class 					__A    (    unittest.TestCase ):

			'''simple docstring'''







			def __init__(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase=7       ,  __lowerCAmelCase=4_0_0       ,  __lowerCAmelCase=2_0_0_0       ,  __lowerCAmelCase=2_4       ,  __lowerCAmelCase=2_4       ,  __lowerCAmelCase=0.0       ,  __lowerCAmelCase=1_6_0_0_0       ,  __lowerCAmelCase=True       ,  __lowerCAmelCase=True       ,  ):



				'''simple docstring'''







				lowerCamelCase__								=		parent
				lowerCamelCase__								=		batch_size
				lowerCamelCase__								=		min_seq_length
				lowerCamelCase__								=		max_seq_length
				lowerCamelCase__								=		(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
				lowerCamelCase__								=		feature_size
				lowerCamelCase__								=		num_mel_bins
				lowerCamelCase__								=		padding_value
				lowerCamelCase__								=		sampling_rate
				lowerCamelCase__								=		return_attention_mask
				lowerCamelCase__								=		do_normalize





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				return {
				    "feature_size": self.feature_size,
				    "num_mel_bins": self.num_mel_bins,
				    "padding_value": self.padding_value,
				    "sampling_rate": self.sampling_rate,
				    "return_attention_mask": self.return_attention_mask,
				    "do_normalize": self.do_normalize,
				}





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase=False       ,  __lowerCAmelCase=False							):



				'''simple docstring'''







				def _flatten(__lowerCAmelCase							):
					return list(itertools.chain(*__lowerCAmelCase							)							)

				if equal_length:
					lowerCamelCase__								=		[floats_list((self.max_seq_length, self.feature_size)							) for _ in range(self.batch_size							)]
				else:
					# make sure that inputs increase in size
					lowerCamelCase__								=		[
					    floats_list((x, self.feature_size)							)
					    for x in range(self.min_seq_length       ,  self.max_seq_length       ,  self.seq_length_diff							)
					]
				if numpify:
					lowerCamelCase__								=		[np.asarray(__lowerCAmelCase							) for x in speech_inputs]
				return speech_inputs






@require_torch
@require_torchaudio
class 					__A    (    lowerCAmelCase			,				unittest.TestCase ):

			'''simple docstring'''







			lowerCAmelCase_							     =						SpeechaTextFeatureExtractor if is_speech_available() else None





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		SpeechaTextFeatureExtractionTester(self							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase							):



				'''simple docstring'''







				self.assertTrue(np.all(np.mean(__lowerCAmelCase       ,  axis=0							) < 1E-3							)							)
				self.assertTrue(np.all(np.abs(np.var(__lowerCAmelCase       ,  axis=0							) - 1							) < 1E-3							)							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
				# create three inputs of length 800, 1000, and 1200
				lowerCamelCase__								=		[floats_list((1, x)							)[0] for x in range(8_0_0       ,  1_4_0_0       ,  2_0_0							)]
				lowerCamelCase__								=		[np.asarray(__lowerCAmelCase							) for speech_input in speech_inputs]

				# Test feature size
				lowerCamelCase__								=		feature_extractor(__lowerCAmelCase       ,  padding=__lowerCAmelCase       ,  return_tensors='''np'''							).input_features
				self.assertTrue(input_features.ndim == 3							)
				self.assertTrue(input_features.shape[-1] == feature_extractor.feature_size							)

				# Test not batched input
				lowerCamelCase__								=		feature_extractor(speech_inputs[0]       ,  return_tensors='''np'''							).input_features
				lowerCamelCase__								=		feature_extractor(np_speech_inputs[0]       ,  return_tensors='''np'''							).input_features
				self.assertTrue(np.allclose(__lowerCAmelCase       ,  __lowerCAmelCase       ,  atol=1E-3							)							)

				# Test batched
				lowerCamelCase__								=		feature_extractor(__lowerCAmelCase       ,  return_tensors='''np'''							).input_features
				lowerCamelCase__								=		feature_extractor(__lowerCAmelCase       ,  return_tensors='''np'''							).input_features
				for enc_seq_a, enc_seq_a in zip(__lowerCAmelCase       ,  __lowerCAmelCase							):
					self.assertTrue(np.allclose(__lowerCAmelCase       ,  __lowerCAmelCase       ,  atol=1E-3							)							)

				# Test 2-D numpy arrays are batched.
				lowerCamelCase__								=		[floats_list((1, x)							)[0] for x in (8_0_0, 8_0_0, 8_0_0)]
				lowerCamelCase__								=		np.asarray(__lowerCAmelCase							)
				lowerCamelCase__								=		feature_extractor(__lowerCAmelCase       ,  return_tensors='''np'''							).input_features
				lowerCamelCase__								=		feature_extractor(__lowerCAmelCase       ,  return_tensors='''np'''							).input_features
				for enc_seq_a, enc_seq_a in zip(__lowerCAmelCase       ,  __lowerCAmelCase							):
					self.assertTrue(np.allclose(__lowerCAmelCase       ,  __lowerCAmelCase       ,  atol=1E-3							)							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
				lowerCamelCase__								=		[floats_list((1, x)							)[0] for x in range(8_0_0       ,  1_4_0_0       ,  2_0_0							)]

				lowerCamelCase__								=		['''longest''', '''max_length''', '''do_not_pad''']
				lowerCamelCase__								=		[None, 1_6, None]
				for max_length, padding in zip(__lowerCAmelCase       ,  __lowerCAmelCase							):
					lowerCamelCase__								=		feature_extractor(
					    __lowerCAmelCase       ,  padding=__lowerCAmelCase       ,  max_length=__lowerCAmelCase       ,  return_attention_mask=__lowerCAmelCase							)
					lowerCamelCase__								=		inputs.input_features
					lowerCamelCase__								=		inputs.attention_mask
					lowerCamelCase__								=		[np.sum(__lowerCAmelCase							) for x in attention_mask]

					self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]							)
					self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]							)
					self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
				lowerCamelCase__								=		[floats_list((1, x)							)[0] for x in range(8_0_0       ,  1_4_0_0       ,  2_0_0							)]

				lowerCamelCase__								=		['''longest''', '''max_length''', '''do_not_pad''']
				lowerCamelCase__								=		[None, 1_6, None]
				for max_length, padding in zip(__lowerCAmelCase       ,  __lowerCAmelCase							):
					lowerCamelCase__								=		feature_extractor(
					    __lowerCAmelCase       ,  max_length=__lowerCAmelCase       ,  padding=__lowerCAmelCase       ,  return_tensors='''np'''       ,  return_attention_mask=__lowerCAmelCase							)
					lowerCamelCase__								=		inputs.input_features
					lowerCamelCase__								=		inputs.attention_mask
					lowerCamelCase__								=		[np.sum(__lowerCAmelCase							) for x in attention_mask]

					self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]							)
					self.assertTrue(input_features[0][fbank_feat_lengths[0] :].sum() < 1E-6							)
					self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]							)
					self.assertTrue(input_features[0][fbank_feat_lengths[1] :].sum() < 1E-6							)
					self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
				lowerCamelCase__								=		[floats_list((1, x)							)[0] for x in range(8_0_0       ,  1_4_0_0       ,  2_0_0							)]
				lowerCamelCase__								=		feature_extractor(
				    __lowerCAmelCase       ,  padding='''max_length'''       ,  max_length=4       ,  truncation=__lowerCAmelCase       ,  return_tensors='''np'''       ,  return_attention_mask=__lowerCAmelCase       ,  )
				lowerCamelCase__								=		inputs.input_features
				lowerCamelCase__								=		inputs.attention_mask
				lowerCamelCase__								=		np.sum(attention_mask == 1       ,  axis=1							)

				self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]							)
				self._check_zero_mean_unit_variance(input_features[1]							)
				self._check_zero_mean_unit_variance(input_features[2]							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
				lowerCamelCase__								=		[floats_list((1, x)							)[0] for x in range(8_0_0       ,  1_4_0_0       ,  2_0_0							)]
				lowerCamelCase__								=		feature_extractor(
				    __lowerCAmelCase       ,  padding='''longest'''       ,  max_length=4       ,  truncation=__lowerCAmelCase       ,  return_tensors='''np'''       ,  return_attention_mask=__lowerCAmelCase       ,  )
				lowerCamelCase__								=		inputs.input_features
				lowerCamelCase__								=		inputs.attention_mask
				lowerCamelCase__								=		np.sum(attention_mask == 1       ,  axis=1							)

				self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]							)
				self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]							)
				self._check_zero_mean_unit_variance(input_features[2]							)

				# make sure that if max_length < longest -> then pad to max_length
				self.assertEqual(input_features.shape       ,  (3, 4, 2_4)							)

				lowerCamelCase__								=		[floats_list((1, x)							)[0] for x in range(8_0_0       ,  1_4_0_0       ,  2_0_0							)]
				lowerCamelCase__								=		feature_extractor(
				    __lowerCAmelCase       ,  padding='''longest'''       ,  max_length=1_6       ,  truncation=__lowerCAmelCase       ,  return_tensors='''np'''       ,  return_attention_mask=__lowerCAmelCase       ,  )
				lowerCamelCase__								=		inputs.input_features
				lowerCamelCase__								=		inputs.attention_mask
				lowerCamelCase__								=		np.sum(attention_mask == 1       ,  axis=1							)

				self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]							)
				self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]							)
				self._check_zero_mean_unit_variance(input_features[2]							)

				# make sure that if max_length < longest -> then pad to max_length
				self.assertEqual(input_features.shape       ,  (3, 6, 2_4)							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				import torch

				lowerCamelCase__								=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
				lowerCamelCase__								=		np.random.rand(1_0_0       ,  3_2							).astype(np.floataa							)
				lowerCamelCase__								=		np_speech_inputs.tolist()

				for inputs in [py_speech_inputs, np_speech_inputs]:
					lowerCamelCase__								=		feature_extractor.pad([{'''input_features''': inputs}]       ,  return_tensors='''np'''							)
					self.assertTrue(np_processed.input_features.dtype == np.floataa							)
					lowerCamelCase__								=		feature_extractor.pad([{'''input_features''': inputs}]       ,  return_tensors='''pt'''							)
					self.assertTrue(pt_processed.input_features.dtype == torch.floataa							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase							):



				'''simple docstring'''







				from datasets import load_dataset

				lowerCamelCase__								=		load_dataset('''hf-internal-testing/librispeech_asr_dummy'''       ,  '''clean'''       ,  split='''validation'''							)
				# automatic decoding with librispeech
				lowerCamelCase__								=		ds.sort('''id'''							).select(range(__lowerCAmelCase							)							)[:num_samples]['''audio''']

				return [x["array"] for x in speech_samples]





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		np.array([
				    -1.5745, -1.7713, -1.7020, -1.6069, -1.2250, -1.1105, -0.9072, -0.8241,
				    -1.2310, -0.8098, -0.3320, -0.4101, -0.7985, -0.4996, -0.8213, -0.9128,
				    -1.0420, -1.1286, -1.0440, -0.7999, -0.8405, -1.2275, -1.5443, -1.4625,
				]							)
				# fmt: on

				lowerCamelCase__								=		self._load_datasamples(1							)
				lowerCamelCase__								=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
				lowerCamelCase__								=		feature_extractor(__lowerCAmelCase       ,  return_tensors='''pt'''							).input_features
				self.assertEquals(input_features.shape       ,  (1, 5_8_4, 2_4)							)
				self.assertTrue(np.allclose(input_features[0, 0, :3_0]       ,  __lowerCAmelCase       ,  atol=1E-4							)							)



