





'''simple docstring'''






import argparse
import gc
import json
import os

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.deepspeed import DummyOptim, DummyScheduler


lowerCAmelCase									=   16
lowerCAmelCase									=   32






def   __A    (				a_    :	List[str]						):
		return int(x / 2**2_0						)
class 		lowerCamelCase	:



	def __enter__( self						):
			gc.collect()
			torch.cuda.empty_cache()
			torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
			lowerCAmelCase  :     Dict       						= torch.cuda.memory_allocated()
			return self





	def __exit__( self		,   *a_						):
			gc.collect()
			torch.cuda.empty_cache()
			lowerCAmelCase  :     List[str]       						= torch.cuda.memory_allocated()
			lowerCAmelCase  :     str       						= torch.cuda.max_memory_allocated()
			lowerCAmelCase  :     Tuple       						= bamb(self.end - self.begin						)
			lowerCAmelCase  :     Dict       						= bamb(self.peak - self.begin						)
			# print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")






def   __A    (				a_    :	Accelerator   ,a_    :	int = 1_6   ,a_    :	str = "bert-base-cased"   ,a_    :	int = 3_2_0   ,a_    :	int = 1_6_0   ,):
		lowerCAmelCase  :     str       						= AutoTokenizer.from_pretrained(a_						)
		lowerCAmelCase  :     Any       						= load_dataset(
		    "glue"   ,"mrpc"   ,split={"train": f'''train[:{n_train}]''', "validation": f'''validation[:{n_val}]'''}						)

		def tokenize_function(a_    :	List[str]						):
				# max_length=None => use the model max length (it's actually the default)
				lowerCAmelCase  :     Tuple       						= 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
		lowerCAmelCase  :     Optional[int]       						= datasets.map(
		    a_   ,batched=a_   ,remove_columns=["idx", "sentence1", "sentence2"]   ,load_from_cache_file=a_						)

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

		def collate_fn(a_    :	int						):
				# On TPU it's best to pad everything to the same length or training will be very slow.
				if accelerator.distributed_type == DistributedType.TPU:
						return tokenizer.pad(a_   ,padding="max_length"   ,max_length=1_2_8   ,return_tensors="pt"						)
				return tokenizer.pad(a_   ,padding="longest"   ,return_tensors="pt"						)

		# Instantiate dataloaders.
		lowerCAmelCase  :     Any       						= DataLoader(
		    tokenized_datasets["train"]   ,shuffle=a_   ,collate_fn=a_   ,batch_size=a_						)
		lowerCAmelCase  :     Any       						= DataLoader(
		    tokenized_datasets["validation"]   ,shuffle=a_   ,collate_fn=a_   ,batch_size=a_						)

		return train_dataloader, eval_dataloader






def   __A    (				a_    :	Optional[int]   ,a_    :	int						):
		# Initialize accelerator
		lowerCAmelCase  :     Dict       						= Accelerator()

		# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
		lowerCAmelCase  :     Optional[int]       						= config["lr"]
		lowerCAmelCase  :     Dict       						= int(config["num_epochs"]						)
		lowerCAmelCase  :     Union[str, Any]       						= int(config["seed"]						)
		lowerCAmelCase  :     Any       						= int(config["batch_size"]						)
		lowerCAmelCase  :     Optional[Any]       						= args.model_name_or_path

		set_seed(a_						)
		lowerCAmelCase						,		lowerCAmelCase  :     List[Any]       						= get_dataloaders(a_   ,a_   ,a_   ,args.n_train   ,args.n_val						)

		# Instantiate the model (we build the model here so that the seed also control new weights initialization)
		lowerCAmelCase  :     Union[str, Any]       						= AutoModelForSequenceClassification.from_pretrained(a_   ,return_dict=a_						)

		# Instantiate optimizer
		lowerCAmelCase  :     str       						= (
		    AdamW
		    if accelerator.state.deepspeed_plugin is None
		    or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
		    else DummyOptim
		)
		lowerCAmelCase  :     str       						= optimizer_cls(params=model.parameters()   ,lr=a_						)

		if accelerator.state.deepspeed_plugin is not None:
				lowerCAmelCase  :     Dict       						= accelerator.state.deepspeed_plugin.deepspeed_config[
				    "gradient_accumulation_steps"
				]
		else:
				lowerCAmelCase  :     str       						= 1
		lowerCAmelCase  :     str       						= (len(a_						) * num_epochs) // gradient_accumulation_steps

		# Instantiate scheduler
		if (
		    accelerator.state.deepspeed_plugin is None
		    or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
		):
				lowerCAmelCase  :     Union[str, Any]       						= get_linear_schedule_with_warmup(
				    optimizer=a_   ,num_warmup_steps=0   ,num_training_steps=a_   ,)
		else:
				lowerCAmelCase  :     str       						= DummyScheduler(a_   ,total_num_steps=a_   ,warmup_num_steps=0						)

		# 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  :     List[str]       						= accelerator.prepare(
		    a_   ,a_   ,a_   ,a_   ,a_						)

		# We need to keep track of how many total steps we have iterated over
		lowerCAmelCase  :     Tuple       						= 0
		# We also need to keep track of the stating epoch so files are named properly
		lowerCAmelCase  :     Tuple       						= 0

		# Now we train the model
		lowerCAmelCase  :     List[str]       						= {}
		for epoch in range(a_   ,a_						):
				with TorchTracemalloc() as tracemalloc:
						model.train()
						for step, batch in enumerate(a_						):
								lowerCAmelCase  :     List[Any]       						= model(**a_						)
								lowerCAmelCase  :     Dict       						= outputs.loss
								lowerCAmelCase  :     Optional[int]       						= loss / gradient_accumulation_steps
								accelerator.backward(a_						)
								if step % gradient_accumulation_steps == 0:
										optimizer.step()
										lr_scheduler.step()
										optimizer.zero_grad()

								overall_step += 1

        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
				accelerator.print("Memory before entering the train : {}".format(bamb(tracemalloc.begin						)						)						)
				accelerator.print("Memory consumed at the end of the train (end-begin): {}".format(tracemalloc.used						)						)
				accelerator.print("Peak Memory consumed during the train (max-begin): {}".format(tracemalloc.peaked						)						)
				accelerator.print(
				    "Total Peak Memory consumed during the train (max): {}".format(
				        tracemalloc.peaked + bamb(tracemalloc.begin						)						)						)
				lowerCAmelCase  :     Tuple       						= tracemalloc.peaked + bamb(tracemalloc.begin						)
				if args.peak_memory_upper_bound is not None:
						assert (
						    train_total_peak_memory[f'''epoch-{epoch}'''] <= args.peak_memory_upper_bound
						), "Peak memory usage exceeded the upper bound"

		accelerator.wait_for_everyone()
		if accelerator.is_main_process:
				with open(os.path.join(args.output_dir   ,"peak_memory_utilization.json"						)   ,"w"						) as f:
						json.dump(a_   ,a_						)






def   __A    (				):
		lowerCAmelCase  :     Optional[int]       						= argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage."						)
		parser.add_argument(
		    "--model_name_or_path"   ,type=a_   ,default="bert-base-cased"   ,help="Path to pretrained model or model identifier from huggingface.co/models."   ,required=a_   ,)
		parser.add_argument(
		    "--output_dir"   ,type=a_   ,default="."   ,help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory."   ,)
		parser.add_argument(
		    "--peak_memory_upper_bound"   ,type=a_   ,default=a_   ,help="The upper bound of peak memory usage in MB. If set, the training will throw an error if the peak memory usage exceeds this value."   ,)
		parser.add_argument(
		    "--n_train"   ,type=a_   ,default=3_2_0   ,help="Number of training examples to use."   ,)
		parser.add_argument(
		    "--n_val"   ,type=a_   ,default=1_6_0   ,help="Number of validation examples to use."   ,)
		parser.add_argument(
		    "--num_epochs"   ,type=a_   ,default=1   ,help="Number of train epochs."   ,)
		lowerCAmelCase  :     Tuple       						= parser.parse_args()
		lowerCAmelCase  :     Dict       						= {"lr": 2e-5, "num_epochs": args.num_epochs, "seed": 4_2, "batch_size": 1_6}
		training_function(a_   ,a_						)


if __name__ == "__main__":
					main()










'''simple docstring'''






import itertools
import random
import unittest

import numpy as np

from transformers import WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST, WavaVecaConfig, WavaVecaFeatureExtractor
from transformers.testing_utils import require_torch, slow

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


lowerCAmelCase									=   random.Random()






def   __A    (				a_    :	Union[str, Any]   ,a_    :	Tuple=1.0   ,a_    :	Optional[int]=None   ,a_    :	Union[str, Any]=None						):
		if rng is None:
				lowerCAmelCase  :     str       						= global_rng

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

		return values
class 		lowerCamelCase	(   unittest.TestCase							):



	def __init__( self		,   a_		,   a_=7		,   a_=400		,   a_=2_000		,   a_=1		,   a_=0.0		,   a_=16_000		,   a_=True		,   a_=True		,   ):
			lowerCAmelCase  :     Dict       						= parent
			lowerCAmelCase  :     Optional[int]       						= batch_size
			lowerCAmelCase  :     Dict       						= min_seq_length
			lowerCAmelCase  :     Dict       						= max_seq_length
			lowerCAmelCase  :     Optional[Any]       						= (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
			lowerCAmelCase  :     Union[str, Any]       						= feature_size
			lowerCAmelCase  :     Tuple       						= padding_value
			lowerCAmelCase  :     Dict       						= sampling_rate
			lowerCAmelCase  :     int       						= return_attention_mask
			lowerCAmelCase  :     Optional[Any]       						= do_normalize



	def 	_lowerCamelCase     ( self						):
			return {
			    "feature_size": self.feature_size,
			    "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		,   a_=False		,   a_=False						):
			def _flatten(a_						):
					return list(itertools.chain(*a_						)						)

			if equal_length:
					lowerCAmelCase  :     List[Any]       						= floats_list((self.batch_size, self.max_seq_length)						)
			else:
					# make sure that inputs increase in size
					lowerCAmelCase  :     Any       						= [
					    _flatten(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  :     List[str]       						= [np.asarray(a_						) for x in speech_inputs]

			return speech_inputs






class 		lowerCamelCase	(   _A   , unittest.TestCase							):
	snake_case_   = WavaVecaFeatureExtractor



	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase  :     Tuple       						= WavaVecaFeatureExtractionTester(self						)



	def 	_lowerCamelCase     ( self		,   a_						):
			self.assertTrue(np.all(np.mean(a_		,   axis=0						) < 1e-3						)						)
			self.assertTrue(np.all(np.abs(np.var(a_		,   axis=0						) - 1						) < 1e-3						)						)



	def 	_lowerCamelCase     ( self						):
			# Tests that all call wrap to encode_plus and batch_encode_plus
			lowerCAmelCase  :     int       						= self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()						)
			# create three inputs of length 800, 1000, and 1200
			lowerCAmelCase  :     str       						= [floats_list((1, x)						)[0] for x in range(800		,   1_400		,   200						)]
			lowerCAmelCase  :     Dict       						= [np.asarray(a_						) for speech_input in speech_inputs]

			# Test not batched input
			lowerCAmelCase  :     Optional[int]       						= feat_extract(speech_inputs[0]		,   return_tensors="np"						).input_values
			lowerCAmelCase  :     Any       						= feat_extract(np_speech_inputs[0]		,   return_tensors="np"						).input_values
			self.assertTrue(np.allclose(a_		,   a_		,   atol=1e-3						)						)

			# Test batched
			lowerCAmelCase  :     Any       						= feat_extract(a_		,   return_tensors="np"						).input_values
			lowerCAmelCase  :     int       						= feat_extract(a_		,   return_tensors="np"						).input_values
			for enc_seq_a, enc_seq_a in zip(a_		,   a_						):
					self.assertTrue(np.allclose(a_		,   a_		,   atol=1e-3						)						)

			# Test 2-D numpy arrays are batched.
			lowerCAmelCase  :     Dict       						= [floats_list((1, x)						)[0] for x in (800, 800, 800)]
			lowerCAmelCase  :     Tuple       						= np.asarray(a_						)
			lowerCAmelCase  :     int       						= feat_extract(a_		,   return_tensors="np"						).input_values
			lowerCAmelCase  :     Union[str, Any]       						= feat_extract(a_		,   return_tensors="np"						).input_values
			for enc_seq_a, enc_seq_a in zip(a_		,   a_						):
					self.assertTrue(np.allclose(a_		,   a_		,   atol=1e-3						)						)



	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase  :     int       						= self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()						)
			lowerCAmelCase  :     Any       						= [floats_list((1, x)						)[0] for x in range(800		,   1_400		,   200						)]

			lowerCAmelCase  :     Union[str, Any]       						= ["longest", "max_length", "do_not_pad"]
			lowerCAmelCase  :     Optional[int]       						= [None, 1_600, None]
			for max_length, padding in zip(a_		,   a_						):
					lowerCAmelCase  :     List[Any]       						= feat_extract(a_		,   padding=a_		,   max_length=a_		,   return_tensors="np"						)
					lowerCAmelCase  :     Any       						= processed.input_values

					self._check_zero_mean_unit_variance(input_values[0][:800]						)
					self.assertTrue(input_values[0][800:].sum() < 1e-6						)
					self._check_zero_mean_unit_variance(input_values[1][:1_000]						)
					self.assertTrue(input_values[0][1_000:].sum() < 1e-6						)
					self._check_zero_mean_unit_variance(input_values[2][:1_200]						)



	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase  :     List[Any]       						= self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()						)
			lowerCAmelCase  :     Optional[int]       						= range(800		,   1_400		,   200						)
			lowerCAmelCase  :     Tuple       						= [floats_list((1, x)						)[0] for x in lengths]

			lowerCAmelCase  :     str       						= ["longest", "max_length", "do_not_pad"]
			lowerCAmelCase  :     Any       						= [None, 1_600, None]

			for max_length, padding in zip(a_		,   a_						):
					lowerCAmelCase  :     Dict       						= feat_extract(a_		,   max_length=a_		,   padding=a_						)
					lowerCAmelCase  :     int       						= processed.input_values

					self._check_zero_mean_unit_variance(input_values[0][:800]						)
					self._check_zero_mean_unit_variance(input_values[1][:1_000]						)
					self._check_zero_mean_unit_variance(input_values[2][:1_200]						)



	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase  :     List[Any]       						= self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()						)
			lowerCAmelCase  :     Tuple       						= [floats_list((1, x)						)[0] for x in range(800		,   1_400		,   200						)]
			lowerCAmelCase  :     Dict       						= feat_extract(
			    a_		,   truncation=a_		,   max_length=1_000		,   padding="max_length"		,   return_tensors="np"						)
			lowerCAmelCase  :     Union[str, Any]       						= processed.input_values

			self._check_zero_mean_unit_variance(input_values[0, :800]						)
			self._check_zero_mean_unit_variance(input_values[1]						)
			self._check_zero_mean_unit_variance(input_values[2]						)



	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase  :     Optional[int]       						= self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()						)
			lowerCAmelCase  :     Tuple       						= [floats_list((1, x)						)[0] for x in range(800		,   1_400		,   200						)]
			lowerCAmelCase  :     Union[str, Any]       						= feat_extract(
			    a_		,   truncation=a_		,   max_length=1_000		,   padding="longest"		,   return_tensors="np"						)
			lowerCAmelCase  :     Dict       						= processed.input_values

			self._check_zero_mean_unit_variance(input_values[0, :800]						)
			self._check_zero_mean_unit_variance(input_values[1, :1_000]						)
			self._check_zero_mean_unit_variance(input_values[2]						)

			# make sure that if max_length < longest -> then pad to max_length
			self.assertTrue(input_values.shape == (3, 1_000)						)

			lowerCAmelCase  :     List[str]       						= [floats_list((1, x)						)[0] for x in range(800		,   1_400		,   200						)]
			lowerCAmelCase  :     Union[str, Any]       						= feat_extract(
			    a_		,   truncation=a_		,   max_length=2_000		,   padding="longest"		,   return_tensors="np"						)
			lowerCAmelCase  :     Optional[Any]       						= processed.input_values

			self._check_zero_mean_unit_variance(input_values[0, :800]						)
			self._check_zero_mean_unit_variance(input_values[1, :1_000]						)
			self._check_zero_mean_unit_variance(input_values[2]						)

			# make sure that if max_length > longest -> then pad to longest
			self.assertTrue(input_values.shape == (3, 1_200)						)



	@require_torch
	def 	_lowerCamelCase     ( self						):
			import torch

			lowerCAmelCase  :     Optional[int]       						= self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()						)
			lowerCAmelCase  :     List[Any]       						= np.random.rand(100						).astype(np.floataa						)
			lowerCAmelCase  :     List[Any]       						= np_speech_inputs.tolist()

			for inputs in [py_speech_inputs, np_speech_inputs]:
					lowerCAmelCase  :     Union[str, Any]       						= feature_extractor.pad([{"input_values": inputs}]		,   return_tensors="np"						)
					self.assertTrue(np_processed.input_values.dtype == np.floataa						)
					lowerCAmelCase  :     Tuple       						= feature_extractor.pad([{"input_values": inputs}]		,   return_tensors="pt"						)
					self.assertTrue(pt_processed.input_values.dtype == torch.floataa						)





	@slow
	@require_torch
	def 	_lowerCamelCase     ( self						):
			# this test makes sure that models that are using
			# group norm don't have their feature extractor return the
			# attention_mask
			for model_id in WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST:
					lowerCAmelCase  :     Any       						= WavaVecaConfig.from_pretrained(a_						)
					lowerCAmelCase  :     List[str]       						= WavaVecaFeatureExtractor.from_pretrained(a_						)

					# only "layer" feature extraction norm should make use of
					# attention_mask
					self.assertEqual(feat_extract.return_attention_mask		,   config.feat_extract_norm == "layer"						)



