



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


UpperCamelCase     =							16
UpperCamelCase     =							32

def       _lowerCamelCase				(							UpperCAmelCase_ :       Accelerator,	UpperCAmelCase_ :       int = 16     )	-> List[str]:






	"""simple docstring"""





	A__											=					AutoTokenizer.from_pretrained("bert-base-cased"     )
	A__											=					load_dataset("glue",	"mrpc"     )

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

	def collate_fn(UpperCAmelCase_ :       str     ):
		# On TPU it's best to pad everything to the same length or training will be very slow.
		A__											=					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__											=					16
		elif accelerator.mixed_precision != "no":
			A__											=					8
		else:
			A__											=					None

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

	# Instantiate dataloaders.
	A__											=					DataLoader(
	    tokenized_datasets["train"],	shuffle=UpperCAmelCase_,	collate_fn=UpperCAmelCase_,	batch_size=UpperCAmelCase_     )
	A__											=					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     =							mocked_dataloaders  # noqa: F811

def       _lowerCamelCase				(							UpperCAmelCase_ :       List[Any],	UpperCAmelCase_ :       Tuple     )	-> Dict:






	"""simple docstring"""





	if os.environ.get("TESTING_MOCKED_DATALOADERS",	UpperCAmelCase_     ) == "1":
		A__											=					2
	# New Code #
	A__											=					int(args.gradient_accumulation_steps     )
	# Initialize accelerator
	A__											=					Accelerator(
	    cpu=args.cpu,	mixed_precision=args.mixed_precision,	gradient_accumulation_steps=UpperCAmelCase_     )
	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
	A__											=					config["lr"]
	A__											=					int(config["num_epochs"]     )
	A__											=					int(config["seed"]     )
	A__											=					int(config["batch_size"]     )

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

	set_seed(UpperCAmelCase_     )
	A__							,							A__											=					get_dataloaders(UpperCAmelCase_,	UpperCAmelCase_     )
	# Instantiate the model (we build the model here so that the seed also control new weights initialization)
	A__											=					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__											=					model.to(accelerator.device     )

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

	# Instantiate scheduler
	A__											=					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__							,							A__							,							A__							,							A__							,							A__											=					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_     ):
			# 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__											=					model(**UpperCAmelCase_     )
				A__											=					output.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():
				A__											=					model(**UpperCAmelCase_     )
			A__											=					outputs.logits.argmax(dim=-1     )
			A__							,							A__											=					accelerator.gather_for_metrics((predictions, batch["labels"])     )
			metric.add_batch(
			    predictions=UpperCAmelCase_,	references=UpperCAmelCase_,	)

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

def       _lowerCamelCase				(							)	-> int:






	"""simple docstring"""





	A__											=					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("--cpu",	action="store_true",	help="If passed, will train on the CPU."     )
	A__											=					parser.parse_args()
	A__											=					{"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
	training_function(UpperCAmelCase_,	UpperCAmelCase_     )


if __name__ == "__main__":
			main()











"""simple docstring"""





import argparse
from typing import List

import evaluate
import numpy as np
import torch
from datasets import DatasetDict, load_dataset

# New Code #
# We'll be using StratifiedKFold for this example
from sklearn.model_selection import StratifiedKFold
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,
# specifically showcasing how to perform Cross Validation,
# and builds off the `nlp_example.py` script.
#
# 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 help focus on the differences in the code, building `DataLoaders`
# was refactored into its own function.
# 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
#
########################################################################


_UpperCamelCase					      =   16
_UpperCamelCase					      =   32



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				Accelerator     ,    SCREAMING_SNAKE_CASE							:				DatasetDict     ,    SCREAMING_SNAKE_CASE							:				List[int]     ,    SCREAMING_SNAKE_CASE							:				List[int]     ,    SCREAMING_SNAKE_CASE							:				int = 16 ):


							'''simple docstring'''







							__lowerCamelCase   :   Union[str, Any]       =AutoTokenizer.from_pretrained('''bert-base-cased''' )
							__lowerCamelCase   :   List[str]       =DatasetDict(
							    {
							        '''train''': dataset['''train'''].select(SCREAMING_SNAKE_CASE ),
							        '''validation''': dataset['''train'''].select(SCREAMING_SNAKE_CASE ),
							        '''test''': dataset['''validation'''],
							    } )

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

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

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

							# Instantiate dataloaders.
							__lowerCamelCase   :   int       =DataLoader(
							    tokenized_datasets['''train''']     ,    shuffle=SCREAMING_SNAKE_CASE     ,    collate_fn=SCREAMING_SNAKE_CASE     ,    batch_size=SCREAMING_SNAKE_CASE )
							__lowerCamelCase   :   Optional[Any]       =DataLoader(
							    tokenized_datasets['''validation''']     ,    shuffle=SCREAMING_SNAKE_CASE     ,    collate_fn=SCREAMING_SNAKE_CASE     ,    batch_size=SCREAMING_SNAKE_CASE )

							__lowerCamelCase   :   Tuple       =DataLoader(
							    tokenized_datasets['''test''']     ,    shuffle=SCREAMING_SNAKE_CASE     ,    collate_fn=SCREAMING_SNAKE_CASE     ,    batch_size=SCREAMING_SNAKE_CASE )

							return train_dataloader, eval_dataloader, test_dataloader



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				Any     ,    SCREAMING_SNAKE_CASE							:				List[Any] ):


							'''simple docstring'''







							__lowerCamelCase   :   Optional[Any]       =[]
							# Download the dataset
							__lowerCamelCase   :   Any       =load_dataset('''glue'''     ,    '''mrpc''' )
							# Create our splits
							__lowerCamelCase   :   Optional[Any]       =StratifiedKFold(n_splits=int(args.num_folds ) )
							# Initialize accelerator
							__lowerCamelCase   :   str       =Accelerator(cpu=args.cpu     ,    mixed_precision=args.mixed_precision )
							# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
							__lowerCamelCase   :   Any       =config['''lr''']
							__lowerCamelCase   :   int       =int(config['''num_epochs'''] )
							__lowerCamelCase   :   Optional[int]       =int(config['''seed'''] )
							__lowerCamelCase   :   Any       =int(config['''batch_size'''] )

							__lowerCamelCase   :   List[str]       =evaluate.load('''glue'''     ,    '''mrpc''' )

							# If the batch size is too big we use gradient accumulation
							__lowerCamelCase   :   List[Any]       =1
							if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
														__lowerCamelCase   :   Optional[Any]       =batch_size // MAX_GPU_BATCH_SIZE
														__lowerCamelCase   :   List[Any]       =MAX_GPU_BATCH_SIZE

							set_seed(SCREAMING_SNAKE_CASE )

							# New Code #
							# Create our folds:
							__lowerCamelCase   :   str       =kfold.split(np.zeros(datasets['''train'''].num_rows )     ,    datasets['''train''']['''label'''] )
							__lowerCamelCase   :   Tuple       =[]
							# Iterate over them
							for i, (train_idxs, valid_idxs) in enumerate(SCREAMING_SNAKE_CASE ):
														__lowerCamelCase     ,    __lowerCamelCase     ,    __lowerCamelCase   :   Any       =get_fold_dataloaders(
														    SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE     ,    )
														# Instantiate the model (we build the model here so that the seed also control new weights initialization)
														__lowerCamelCase   :   Optional[int]       =AutoModelForSequenceClassification.from_pretrained('''bert-base-cased'''     ,    return_dict=SCREAMING_SNAKE_CASE )

														# 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).
														__lowerCamelCase   :   Optional[Any]       =model.to(accelerator.device )

														# Instantiate optimizer
														__lowerCamelCase   :   Any       =AdamW(params=model.parameters()     ,    lr=SCREAMING_SNAKE_CASE )

														# Instantiate scheduler
														__lowerCamelCase   :   Dict       =get_linear_schedule_with_warmup(
														    optimizer=SCREAMING_SNAKE_CASE     ,    num_warmup_steps=100     ,    num_training_steps=(len(SCREAMING_SNAKE_CASE ) * num_epochs) // gradient_accumulation_steps     ,    )

														# 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   :   Any       =accelerator.prepare(
														    SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE )

														# Now we train the model
														for epoch in range(SCREAMING_SNAKE_CASE ):
																					model.train()
																					for step, batch in enumerate(SCREAMING_SNAKE_CASE ):
																												# We could avoid this line since we set the accelerator with `device_placement=True`.
																												batch.to(accelerator.device )
																												__lowerCamelCase   :   Dict       =model(**SCREAMING_SNAKE_CASE )
																												__lowerCamelCase   :   Optional[Any]       =outputs.loss
																												__lowerCamelCase   :   int       =loss / gradient_accumulation_steps
																												accelerator.backward(SCREAMING_SNAKE_CASE )
																												if step % gradient_accumulation_steps == 0:
																																			optimizer.step()
																																			lr_scheduler.step()
																																			optimizer.zero_grad()

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

																					__lowerCamelCase   :   str       =metric.compute()
																					# Use accelerator.print to print only on the main process.
																					accelerator.print(F'epoch {epoch}:'     ,    SCREAMING_SNAKE_CASE )

														# New Code #
														# We also run predictions on the test set at the very end
														__lowerCamelCase   :   Tuple       =[]
														for step, batch in enumerate(SCREAMING_SNAKE_CASE ):
																					# We could avoid this line since we set the accelerator with `device_placement=True`.
																					batch.to(accelerator.device )
																					with torch.no_grad():
																												__lowerCamelCase   :   List[Any]       =model(**SCREAMING_SNAKE_CASE )
																					__lowerCamelCase   :   int       =outputs.logits
																					__lowerCamelCase     ,    __lowerCamelCase   :   Union[str, Any]       =accelerator.gather_for_metrics((predictions, batch['''labels''']) )
																					fold_predictions.append(predictions.cpu() )
																					if i == 0:
																												# We need all of the test predictions
																												test_references.append(references.cpu() )
        # Use accelerator.print to print only on the main process.
														test_predictions.append(torch.cat(SCREAMING_SNAKE_CASE     ,    dim=0 ) )
														# We now need to release all our memory and get rid of the current model, optimizer, etc
														accelerator.free_memory()
							# New Code #
							# Finally we check the accuracy of our folded results:
							__lowerCamelCase   :   Optional[Any]       =torch.cat(SCREAMING_SNAKE_CASE     ,    dim=0 )
							__lowerCamelCase   :   int       =torch.stack(SCREAMING_SNAKE_CASE     ,    dim=0 ).sum(dim=0 ).div(int(args.num_folds ) ).argmax(dim=-1 )
							__lowerCamelCase   :   Tuple       =metric.compute(predictions=SCREAMING_SNAKE_CASE     ,    references=SCREAMING_SNAKE_CASE )
							accelerator.print('''Average test metrics from all folds:'''     ,    SCREAMING_SNAKE_CASE )



def 					lowerCAmelCase_      (   ):


							'''simple docstring'''







							__lowerCamelCase   :   int       =argparse.ArgumentParser(description='''Simple example of training script.''' )
							parser.add_argument(
							    '''--mixed_precision'''     ,    type=SCREAMING_SNAKE_CASE     ,    default=SCREAMING_SNAKE_CASE     ,    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.''' )
							# New Code #
							parser.add_argument('''--num_folds'''     ,    type=SCREAMING_SNAKE_CASE     ,    default=3     ,    help='''The number of splits to perform across the dataset''' )
							__lowerCamelCase   :   Any       =parser.parse_args()
							__lowerCamelCase   :   Optional[int]       ={'''lr''': 2E-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
							training_function(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE )


if __name__ == "__main__":
							main()




