






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


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to ensure out-of-memory errors never
# interrupt training, 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)
#
# 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       		=      16
__lowerCAmelCase       		=      32





def 		_UpperCAmelCase   (							__A							:   Accelerator						,		__A							:   int = 16      ):
				a_		:  int											=      AutoTokenizer.from_pretrained('''bert-base-cased'''      )
				a_		:  List[Any]											=      load_dataset('''glue'''						,		'''mrpc'''      )

				def tokenize_function(__A							:   Optional[Any]      ):
								# max_length=None => use the model max length (it's actually the default)
								a_		:  Tuple											=      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_		:  List[str]											=      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_		:  Any											=      tokenized_datasets.rename_column('''label'''						,		'''labels'''      )

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

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

				# Instantiate dataloaders.
				a_		:  List[str]											=      DataLoader(
				    tokenized_datasets['''train''']						,		shuffle=UpperCamelCase__						,		collate_fn=UpperCamelCase__						,		batch_size=UpperCamelCase__      )
				a_		:  Union[str, Any]											=      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 		_UpperCAmelCase   (							__A							:   str						,		__A							:   str      ):
				# For testing only
				if os.environ.get('''TESTING_MOCKED_DATALOADERS'''						,		UpperCamelCase__      ) == "1":
								a_		:  Tuple											=      2
				# Initialize accelerator
				a_		:  Optional[int]											=      Accelerator(cpu=args.cpu						,		mixed_precision=args.mixed_precision      )
				# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
				a_		:  Any											=      config['''lr''']
				a_		:  Dict											=      int(config['''num_epochs''']      )
				a_		:  Union[str, Any]											=      int(config['''seed''']      )
				a_		:  List[str]											=      int(config['''batch_size''']      )

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

				# New Code #
				# We now can define an inner training loop function. It should take a batch size as the only parameter,
				# and build the dataloaders in there.
				# It also gets our decorator
				@find_executable_batch_size(starting_batch_size=UpperCamelCase__      )
				def inner_training_loop(__A							:   Any      ):
								# And now just move everything below under this function
								# We need to bring in the Accelerator object from earlier
								nonlocal accelerator
								# And reset all of its attributes that could hold onto any memory:
								accelerator.free_memory()

								# Then we can declare the model, optimizer, and everything else:
								set_seed(UpperCamelCase__      )

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

								# Instantiate optimizer
								a_		:  List[str]											=      AdamW(params=model.parameters()						,		lr=UpperCamelCase__      )
								a_				,			a_		:  List[Any]											=      get_dataloaders(UpperCamelCase__						,		UpperCamelCase__      )

								# Instantiate scheduler
								a_		:  Optional[Any]											=      get_linear_schedule_with_warmup(
								    optimizer=UpperCamelCase__						,		num_warmup_steps=1_00						,		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_		:  Union[str, Any]											=      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      )
																a_		:  Tuple											=      model(**UpperCamelCase__      )
																a_		:  List[Any]											=      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():
																				a_		:  Union[str, Any]											=      model(**UpperCamelCase__      )
																a_		:  Optional[Any]											=      outputs.logits.argmax(dim=-1      )
																a_				,			a_		:  List[str]											=      accelerator.gather_for_metrics((predictions, batch['''labels'''])      )
																metric.add_batch(
																    predictions=UpperCamelCase__						,		references=UpperCamelCase__						,		)

												a_		:  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 		_UpperCAmelCase   (							):
				a_		:  str											=      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.'''      )
				a_		:  Optional[Any]											=      parser.parse_args()
				a_		:  Union[str, Any]											=      {'''lr''': 2E-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
				training_function(UpperCamelCase__						,		UpperCamelCase__      )


if __name__ == "__main__":
		main()












'''simple docstring'''







import re
import string

import numpy as np

import datasets


__lowerCAmelCase       		=      '\nReturns the rate at which the input predicted strings exactly match their references, ignoring any strings input as part of the regexes_to_ignore list.\n'

__lowerCAmelCase       		=      '\nArgs:\n    predictions: List of predicted texts.\n    references: List of reference texts.\n    regexes_to_ignore: List, defaults to None. Regex expressions of characters to\n        ignore when calculating the exact matches. Note: these regexes are removed\n        from the input data before the changes based on the options below (e.g. ignore_case,\n        ignore_punctuation, ignore_numbers) are applied.\n    ignore_case: Boolean, defaults to False. If true, turns everything\n        to lowercase so that capitalization differences are ignored.\n    ignore_punctuation: Boolean, defaults to False. If true, removes all punctuation before\n        comparing predictions and references.\n    ignore_numbers: Boolean, defaults to False. If true, removes all punctuation before\n        comparing predictions and references.\nReturns:\n    exact_match: Dictionary containing exact_match rate. Possible values are between 0.0 and 100.0, inclusive.\nExamples:\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds)\n    >>> print(round(results["exact_match"], 1))\n    25.0\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=["the ", "yell"], ignore_case=True, ignore_punctuation=True)\n    >>> print(round(results["exact_match"], 1))\n    50.0\n\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=["the ", "yell", "YELL"], ignore_case=True, ignore_punctuation=True)\n    >>> print(round(results["exact_match"], 1))\n    75.0\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=["the ", "yell", "YELL"], ignore_case=True, ignore_punctuation=True, ignore_numbers=True)\n    >>> print(round(results["exact_match"], 1))\n    100.0\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["The cat sat on the mat.", "Theaters are great.", "It\'s like comparing oranges and apples."]\n    >>> preds = ["The cat sat on the mat?", "Theaters are great.", "It\'s like comparing apples and oranges."]\n    >>> results = exact_match.compute(references=refs, predictions=preds)\n    >>> print(round(results["exact_match"], 1))\n    33.3\n\n'

__lowerCAmelCase       		=      '\n'





@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION	,   _KWARGS_DESCRIPTION       )
class 	SCREAMING_SNAKE_CASE					(					datasets.Metric       ):






			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[Any]					)			->			int:
							return datasets.MetricInfo(
							    description=_DESCRIPTION  ,						citation=_CITATION  ,						inputs_description=_KWARGS_DESCRIPTION  ,						features=datasets.Features(
							        {
							            '''predictions''': datasets.Value('''string'''  ,						id='''sequence'''					),
							            '''references''': datasets.Value('''string'''  ,						id='''sequence'''					),
							        }					)  ,						reference_urls=[]  ,						)







			def 			SCREAMING_SNAKE_CASE      ( self      :			str  ,						__SCREAMING_SNAKE_CASE      :			List[Any]  ,						__SCREAMING_SNAKE_CASE      :			Optional[int]  ,						__SCREAMING_SNAKE_CASE      :			Dict=None  ,						__SCREAMING_SNAKE_CASE      :			int=False  ,						__SCREAMING_SNAKE_CASE      :			Optional[int]=False  ,						__SCREAMING_SNAKE_CASE      :			Dict=False  ,						)			->			str:
							if regexes_to_ignore is not None:
											for s in regexes_to_ignore:
															a_		:  Optional[Any]											=      np.array([re.sub(__SCREAMING_SNAKE_CASE  ,						''''''  ,						__SCREAMING_SNAKE_CASE					) for x in predictions]					)
															a_		:  int											=      np.array([re.sub(__SCREAMING_SNAKE_CASE  ,						''''''  ,						__SCREAMING_SNAKE_CASE					) for x in references]					)
							else:
											a_		:  List[str]											=      np.asarray(__SCREAMING_SNAKE_CASE					)
											a_		:  Any											=      np.asarray(__SCREAMING_SNAKE_CASE					)

							if ignore_case:
											a_		:  List[str]											=      np.char.lower(__SCREAMING_SNAKE_CASE					)
											a_		:  List[Any]											=      np.char.lower(__SCREAMING_SNAKE_CASE					)

							if ignore_punctuation:
											a_		:  Any											=      string.punctuation.maketrans(''''''  ,						''''''  ,						string.punctuation					)
											a_		:  Union[str, Any]											=      np.char.translate(__SCREAMING_SNAKE_CASE  ,						table=__SCREAMING_SNAKE_CASE					)
											a_		:  int											=      np.char.translate(__SCREAMING_SNAKE_CASE  ,						table=__SCREAMING_SNAKE_CASE					)

							if ignore_numbers:
											a_		:  int											=      string.digits.maketrans(''''''  ,						''''''  ,						string.digits					)
											a_		:  Optional[int]											=      np.char.translate(__SCREAMING_SNAKE_CASE  ,						table=__SCREAMING_SNAKE_CASE					)
											a_		:  Dict											=      np.char.translate(__SCREAMING_SNAKE_CASE  ,						table=__SCREAMING_SNAKE_CASE					)

							a_		:  Optional[Any]											=      predictions == references

							return {"exact_match": np.mean(__SCREAMING_SNAKE_CASE					) * 100}




