

'''simple docstring'''


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


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to combine both the gradient accumulation
# and automatic batch size finder utilities of Accelerate to perfrom
# automatic 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)
#
# 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
#
########################################################################

lowerCAmelCase__      				=			32







def     SCREAMING_SNAKE_CASE(	UpperCamelCase ,UpperCamelCase = 1_6	)     ->     Optional[Any]:
	UpperCAmelCase_		:   str													=       AutoTokenizer.from_pretrained('bert-base-cased'	)
	UpperCAmelCase_		:   Optional[Any]													=       load_dataset('glue' ,'mrpc'	)

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

	def collate_fn(UpperCamelCase	):
		# When using mixed precision we want round multiples of 8/16
		if accelerator.mixed_precision == "fp8":
			UpperCAmelCase_		:   Tuple													=       1_6
		elif accelerator.mixed_precision != "no":
			UpperCAmelCase_		:   str													=       8
		else:
			UpperCAmelCase_		:   List[Any]													=       None

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

	# Instantiate dataloaders.
	UpperCAmelCase_		:   Optional[Any]													=       DataLoader(
	    tokenized_datasets['train'] ,shuffle=UpperCamelCase ,collate_fn=UpperCamelCase ,batch_size=UpperCamelCase	)
	UpperCAmelCase_		:   Optional[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

						lowerCAmelCase__      				=			mocked_dataloaders  # noqa: F811







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

	UpperCAmelCase_		:   List[Any]													=       evaluate.load('glue' ,'mrpc'	)

	# New Code #
	# We use the `find_executable_batch_size` decorator, passing in the desired observed batch size
	# to train on. If a CUDA OOM error occurs, it will retry this loop cutting the batch size in
	# half each time. From this, we can calculate the number of gradient accumulation steps needed
	# and modify the Accelerator object as a result
	@find_executable_batch_size(starting_batch_size=int(UpperCamelCase	)	)
	def inner_training_loop(UpperCamelCase	):
		# Since we need to modify the outside accelerator object, we need to bring it
		# to the local scope
		nonlocal accelerator

		# We can calculate the number of gradient accumulation steps based on the current
		# batch size vs the starting batch size
		UpperCAmelCase_		:   str													=       observed_batch_size // batch_size

		# And then set it in the Accelerator directly:
		UpperCAmelCase_		:   List[Any]													=       num_gradient_accumulation_steps

		# Next we need to free all of the stored model references in the Accelerator each time
		accelerator.free_memory()

		# And set the seed so our results are reproducable each reset
		set_seed(UpperCamelCase	)

		# Instantiate the model (we build the model here so that the seed also control new weights initialization)
		UpperCAmelCase_		:   List[str]													=       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).
		UpperCAmelCase_		:   Optional[int]													=       model.to(accelerator.device	)

		# Instantiate optimizer
		UpperCAmelCase_		:   Union[str, Any]													=       AdamW(params=model.parameters() ,lr=UpperCamelCase	)
		UpperCAmelCase_						,     UpperCAmelCase_		:   str													=       get_dataloaders(UpperCamelCase ,UpperCamelCase	)

		# Instantiate scheduler
		UpperCAmelCase_		:   Optional[int]													=       get_linear_schedule_with_warmup(
		    optimizer=UpperCamelCase ,num_warmup_steps=1_0_0 ,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.
		UpperCAmelCase_						,     UpperCAmelCase_						,     UpperCAmelCase_						,     UpperCAmelCase_						,     UpperCAmelCase_		:   Tuple													=       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	):
				# And perform gradient accumulation
				with accelerator.accumulate(UpperCamelCase	):
					# We could avoid this line since we set the accelerator with `device_placement=True`.
					batch.to(accelerator.device	)
					UpperCAmelCase_		:   Optional[int]													=       model(**UpperCamelCase	)
					UpperCAmelCase_		:   int													=       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():
					UpperCAmelCase_		:   List[str]													=       model(**UpperCamelCase	)
				UpperCAmelCase_		:   Optional[int]													=       outputs.logits.argmax(dim=-1	)
				UpperCAmelCase_						,     UpperCAmelCase_		:   Union[str, Any]													=       accelerator.gather_for_metrics((predictions, batch['labels'])	)
				metric.add_batch(
				    predictions=UpperCamelCase ,references=UpperCamelCase ,)

			UpperCAmelCase_		:   List[Any]													=       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     SCREAMING_SNAKE_CASE(	)     ->     Union[str, Any]:
	UpperCAmelCase_		:   int													=       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.'	)
	UpperCAmelCase_		:   List[str]													=       parser.parse_args()
	# New Code #
	# We modify the starting batch size to be an observed batch size of 256, to guarentee an initial CUDA OOM
	UpperCAmelCase_		:   str													=       {'lr': 2e-5, 'num_epochs': 3, 'seed': 4_2, 'batch_size': 2_5_6}
	training_function(UpperCamelCase ,UpperCamelCase	)


if __name__ == "__main__":
						main()









'''simple docstring'''


import unittest

from transformers import is_torch_available, load_tool

from .test_tools_common import ToolTesterMixin


if is_torch_available():
						import torch



class        lowercase    (	unittest.TestCase,       a_			):


							def     _snake_case				(      self)				->					Optional[int]:
								UpperCAmelCase_		:   int													=       load_tool('speech-to-text')
								self.tool.setup()


							def     _snake_case				(      self)				->					Any:
								UpperCAmelCase_		:   str													=       self.tool(torch.ones(3000))
								self.assertEqual(_snake_case     ,      ' you')





							def     _snake_case				(      self)				->					str:
								UpperCAmelCase_		:   List[Any]													=       self.tool(audio=torch.ones(3000))
								self.assertEqual(_snake_case     ,      ' you')






