






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 the checkpointing capability,
# 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
#
########################################################################

lowerCamelCase_						     =  1_6
lowerCamelCase_						     =  3_2
def      lowerCamelCase					(						a_ ,						a_ = 16	)	->			Optional[Any]:
							lowerCAmelCase_										=     AutoTokenizer.from_pretrained('bert-base-cased'	)
							lowerCAmelCase_										=     load_dataset('glue' ,						'mrpc'	)

							def tokenize_function(a_	):
														# max_length=None => use the model max length (it's actually the default)
														lowerCAmelCase_										=     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():
														lowerCAmelCase_										=     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
							lowerCAmelCase_										=     tokenized_datasets.rename_column('label' ,						'labels'	)

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

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

							# Instantiate dataloaders.
							lowerCAmelCase_										=     DataLoader(
							    tokenized_datasets['train'] ,						shuffle=a_ ,						collate_fn=a_ ,						batch_size=a_	)
							lowerCAmelCase_										=     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

		lowerCamelCase_						     =  mocked_dataloaders  # noqa: F811
def      lowerCamelCase					(						a_ ,						a_	)	->			List[str]:
							# For testing only
							if os.environ.get('TESTING_MOCKED_DATALOADERS' ,						a_	) == "1":
														lowerCAmelCase_										=     2
							# Initialize accelerator
							lowerCAmelCase_										=     Accelerator(cpu=args.cpu ,						mixed_precision=args.mixed_precision	)
							# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
							lowerCAmelCase_										=     config['lr']
							lowerCAmelCase_										=     int(config['num_epochs']	)
							lowerCAmelCase_										=     int(config['seed']	)
							lowerCAmelCase_										=     int(config['batch_size']	)

							# New Code #
							# Parse out whether we are saving every epoch or after a certain number of batches
							if hasattr(args.checkpointing_steps ,						'isdigit'	):
														if args.checkpointing_steps == "epoch":
																					lowerCAmelCase_										=     args.checkpointing_steps
														elif args.checkpointing_steps.isdigit():
																					lowerCAmelCase_										=     int(args.checkpointing_steps	)
														else:
																					raise ValueError(
																					    F'''Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed.'''	)
							else:
														lowerCAmelCase_										=     None

							set_seed(a_	)

							lowerCAmelCase_						,   lowerCAmelCase_										=     get_dataloaders(a_ ,						a_	)
							lowerCAmelCase_										=     evaluate.load('glue' ,						'mrpc'	)

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

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

							# Instantiate optimizer
							lowerCAmelCase_										=     AdamW(params=model.parameters() ,						lr=a_	)

							# Instantiate scheduler
							lowerCAmelCase_										=     get_linear_schedule_with_warmup(
							    optimizer=a_ ,						num_warmup_steps=100 ,						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.
							lowerCAmelCase_						,   lowerCAmelCase_						,   lowerCAmelCase_						,   lowerCAmelCase_						,   lowerCAmelCase_										=     accelerator.prepare(
							    a_ ,						a_ ,						a_ ,						a_ ,						a_	)

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

							# We need to load the checkpoint back in before training here with `load_state`
							# The total number of epochs is adjusted based on where the state is being loaded from,
							# as we assume continuation of the same training script
							if args.resume_from_checkpoint:
														if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
																					accelerator.print(F'''Resumed from checkpoint: {args.resume_from_checkpoint}'''	)
																					accelerator.load_state(args.resume_from_checkpoint	)
																					lowerCAmelCase_										=     os.path.basename(args.resume_from_checkpoint	)
														else:
																					# Get the most recent checkpoint
																					lowerCAmelCase_										=     [f.name for f in os.scandir(os.getcwd()	) if f.is_dir()]
																					dirs.sort(key=os.path.getctime	)
																					lowerCAmelCase_										=     dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
														# Extract `epoch_{i}` or `step_{i}`
														lowerCAmelCase_										=     os.path.splitext(a_	)[0]

														if "epoch" in training_difference:
																					lowerCAmelCase_										=     int(training_difference.replace('epoch_' ,						''	)	) + 1
																					lowerCAmelCase_										=     None
														else:
																					lowerCAmelCase_										=     int(training_difference.replace('step_' ,						''	)	)
																					lowerCAmelCase_										=     resume_step // len(a_	)
																					resume_step -= starting_epoch * len(a_	)

    # Now we train the model
							for epoch in range(a_ ,						a_	):
														model.train()
														# New Code #
														if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:
																					# We need to skip steps until we reach the resumed step
																					lowerCAmelCase_										=     accelerator.skip_first_batches(a_ ,						a_	)
																					overall_step += resume_step
														else:
																					# After the first iteration though, we need to go back to the original dataloader
																					lowerCAmelCase_										=     train_dataloader
														for step, batch in enumerate(a_	):
																					# We could avoid this line since we set the accelerator with `device_placement=True`.
																					batch.to(accelerator.device	)
																					lowerCAmelCase_										=     model(**a_	)
																					lowerCAmelCase_										=     outputs.loss
																					lowerCAmelCase_										=     loss / gradient_accumulation_steps
																					accelerator.backward(a_	)
																					if step % gradient_accumulation_steps == 0:
																												optimizer.step()
																												lr_scheduler.step()
																												optimizer.zero_grad()
																					# New Code #
																					overall_step += 1

																					# New Code #
																					# We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`
																					# These are saved to folders named `step_{overall_step}`
																					# Will contain files: "pytorch_model.bin", "optimizer.bin", "scheduler.bin", and "random_states.pkl"
																					# If mixed precision was used, will also save a "scalar.bin" file
																					if isinstance(a_ ,						a_	):
																												lowerCAmelCase_										=     F'''step_{overall_step}'''
																												if overall_step % checkpointing_steps == 0:
																																			if args.output_dir is not None:
																																										lowerCAmelCase_										=     os.path.join(args.output_dir ,						a_	)
																																			accelerator.save_state(a_	)

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

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

														# New Code #
														# We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`
														# These are saved to folders named `epoch_{epoch}`
														# Will contain files: "pytorch_model.bin", "optimizer.bin", "scheduler.bin", and "random_states.pkl"
														# If mixed precision was used, will also save a "scalar.bin" file
														if checkpointing_steps == "epoch":
																					lowerCAmelCase_										=     F'''epoch_{epoch}'''
																					if args.output_dir is not None:
																												lowerCAmelCase_										=     os.path.join(args.output_dir ,						a_	)
																					accelerator.save_state(a_	)
def      lowerCamelCase					(						)	->			Union[str, Any]:
							lowerCAmelCase_										=     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.'	)
							parser.add_argument(
							    '--checkpointing_steps' ,						type=a_ ,						default=a_ ,						help='Whether the various states should be saved at the end of every n steps, or \'epoch\' for each epoch.' ,						)
							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(
							    '--resume_from_checkpoint' ,						type=a_ ,						default=a_ ,						help='If the training should continue from a checkpoint folder.' ,						)
							lowerCAmelCase_										=     parser.parse_args()
							lowerCAmelCase_										=     {'lr': 2e-5, 'num_epochs': 3, 'seed': 42, 'batch_size': 16}
							training_function(a_ ,						a_	)


if __name__ == "__main__":
		main()











from dataclasses import dataclass
from typing import Callable, Optional

import datasets




@dataclass
class 							a_ (						datasets.BuilderConfig		):


							'''simple docstring'''






							__a:	Optional[Callable]     =    None
							__a:	Optional[dict]     =    None
							__a:	Optional[datasets.Features]     =    None

							def      _lowercase       (		self   )					->					Optional[int]:
														'''simple docstring'''

														assert self.generator is not None, "generator must be specified"

														if self.gen_kwargs is None:
																					lowerCAmelCase_										=     {}




class 							a_ (						datasets.GeneratorBasedBuilder		):


							'''simple docstring'''






							__a:	Optional[Any]     =    GeneratorConfig

							def      _lowercase       (		self   )					->					Union[str, Any]:
														'''simple docstring'''

														return datasets.DatasetInfo(features=self.config.features   )

							def      _lowercase       (		self			,			lowercase_   )					->					Dict:
														'''simple docstring'''

														return [datasets.SplitGenerator(name=datasets.Split.TRAIN			,			gen_kwargs=self.config.gen_kwargs   )]






							def      _lowercase       (		self			,			**lowercase_   )					->					Union[str, Any]:
														'''simple docstring'''

														for idx, ex in enumerate(self.config.generator(**lowercase_   )   ):
																					yield idx, ex



