


'''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,
# specifically showcasing how to properly calculate the metrics on the
# validation dataset when in a distributed system, 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
#
########################################################################


_lowercase    :    Tuple			     =  16
_lowercase    :    Tuple			     =  32


def    lowerCamelCase__   ( A     :    Accelerator   ,     A     :    int = 16  ):




						'''simple docstring'''






						UpperCAmelCase         =			AutoTokenizer.from_pretrained('''bert-base-cased'''  )
						UpperCAmelCase         =			load_dataset('''glue'''   ,     '''mrpc'''  )

						def tokenize_function(A     :    Any  ):
												# 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     :    Optional[Any]  ):
												# On TPU it's best to pad everything to the same length or training will be very slow.
												UpperCAmelCase         =			1_28 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

	_lowercase    :    Union[str, Any]			     =  mocked_dataloaders  # noqa: F811


def    lowerCamelCase__   ( A     :    Dict   ,     A     :    int  ):




						'''simple docstring'''






						if os.environ.get('''TESTING_MOCKED_DATALOADERS'''   ,     A  ) == "1":
												UpperCAmelCase         =			2
						# Initialize accelerator
						UpperCAmelCase         =			Accelerator(cpu=args.cpu   ,     mixed_precision=args.mixed_precision  )
						# 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'''  )

						# If the batch size is too big we use gradient accumulation
						UpperCAmelCase         =			1
						if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
												UpperCAmelCase         =			batch_size // MAX_GPU_BATCH_SIZE
												UpperCAmelCase         =			MAX_GPU_BATCH_SIZE

						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=1_00   ,     num_training_steps=(len(A  ) * 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.
						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  )
																		UpperCAmelCase         =			model(**A  )
																		UpperCAmelCase         =			outputs.loss
																		UpperCAmelCase         =			loss / gradient_accumulation_steps
																		accelerator.backward(A  )
																		if step % gradient_accumulation_steps == 0:
																								optimizer.step()
																								lr_scheduler.step()
																								optimizer.zero_grad()

												model.eval()
												UpperCAmelCase         =			0
												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((predictions, batch['''labels'''])  )
																		# New Code #
																		# First we check if it's a distributed system
																		if accelerator.use_distributed:
																								# Then see if we're on the last batch of our eval dataloader
																								if step == len(A  ) - 1:
																														# Last batch needs to be truncated on distributed systems as it contains additional samples
																														UpperCAmelCase         =			predictions[: len(eval_dataloader.dataset  ) - samples_seen]
																														UpperCAmelCase         =			references[: len(eval_dataloader.dataset  ) - samples_seen]
																								else:
																														# Otherwise we add the number of samples seen
																														samples_seen += references.shape[0]
            # All of this can be avoided if you use `Accelerator.gather_for_metrics` instead of `Accelerator.gather`:
            # 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    lowerCamelCase__   ( ):




						'''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.'''   ,     )
						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()









'''simple docstring'''



import time
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, torch_device

from ..test_modeling_common import ids_tensor


if is_torch_available():
	import torch

	from transformers.generation import (
	    MaxLengthCriteria,
	    MaxNewTokensCriteria,
	    MaxTimeCriteria,
	    StoppingCriteriaList,
	    validate_stopping_criteria,
	)






@require_torch
class  UpperCamelCase__(     unittest.TestCase ):




				def   a__(	self    :				List[Any]   ,					lowerCAmelCase    :				Optional[Any]   )->  Optional[int]:

										"""simple docstring"""




										UpperCAmelCase         =			3
										UpperCAmelCase         =			250

										UpperCAmelCase         =			ids_tensor((batch_size, length)   ,					lowerCAmelCase   )
										UpperCAmelCase         =			torch.ones((batch_size, length)   ,					device=lowerCAmelCase   ,					dtype=torch.float   ) / length
										return input_ids, scores




				def   a__(	self    :				Dict   )->  Optional[int]:

										"""simple docstring"""




										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(5   )

										UpperCAmelCase         =			StoppingCriteriaList(
										    [
										        MaxLengthCriteria(max_length=10   ),
										        MaxTimeCriteria(max_time=0.1   ),
										    ]   )

										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(9   )
										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(10   )
										self.assertTrue(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )




				def   a__(	self    :				Union[str, Any]   )->  Union[str, Any]:

										"""simple docstring"""




										UpperCAmelCase         =			MaxLengthCriteria(max_length=10   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(5   )
										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(9   )
										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(10   )
										self.assertTrue(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )




				def   a__(	self    :				Optional[Any]   )->  List[Any]:

										"""simple docstring"""




										UpperCAmelCase         =			MaxNewTokensCriteria(start_length=5   ,					max_new_tokens=5   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(5   )
										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(9   )
										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(10   )
										self.assertTrue(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase         =			StoppingCriteriaList([criteria]   )
										self.assertEqual(criteria_list.max_length   ,					10   )




				def   a__(	self    :				Tuple   )->  Optional[Any]:

										"""simple docstring"""




										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(5   )

										UpperCAmelCase         =			MaxTimeCriteria(max_time=0.1   )
										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase         =			MaxTimeCriteria(max_time=0.1   ,					initial_timestamp=time.time() - 0.2   )
										self.assertTrue(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )





				def   a__(	self    :				int   )->  Any:

										"""simple docstring"""




										validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10   )]   )   ,					10   )

										with self.assertWarns(lowerCAmelCase   ):
																validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10   )]   )   ,					11   )

										UpperCAmelCase         =			validate_stopping_criteria(StoppingCriteriaList()   ,					11   )

										self.assertEqual(len(lowerCAmelCase   )   ,					1   )





