
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
from accelerate.local_sgd import LocalSGD


########################################################################
# This is a fully working simple example to use Accelerate
# with LocalSGD, which is a method to synchronize model
# parameters every K batches. It is different, but complementary
# to 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
#
########################################################################


__SCREAMING_SNAKE_CASE						       =						16
__SCREAMING_SNAKE_CASE						       =						32
def    SCREAMING_SNAKE_CASE__		( lowerCAmelCase_				:			Tuple       ,lowerCAmelCase_				:			List[str] = 16							)    ->   Dict:







							"""simple docstring"""






							SCREAMING_SNAKE_CASE_	:							Tuple											=AutoTokenizer.from_pretrained('bert-base-cased'							)
							SCREAMING_SNAKE_CASE_	:							List[str]											=load_dataset('glue'       ,'mrpc'							)

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

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

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

							# Instantiate dataloaders.
							SCREAMING_SNAKE_CASE_	:							Optional[int]											=DataLoader(
							    tokenized_datasets['train']       ,shuffle=lowerCAmelCase_       ,collate_fn=lowerCAmelCase_       ,batch_size=lowerCAmelCase_							)
							SCREAMING_SNAKE_CASE_	:							Tuple											=DataLoader(
							    tokenized_datasets['validation']       ,shuffle=lowerCAmelCase_       ,collate_fn=lowerCAmelCase_       ,batch_size=lowerCAmelCase_							)

							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

							__SCREAMING_SNAKE_CASE						       =						mocked_dataloaders  # noqa: F811
def    SCREAMING_SNAKE_CASE__		( lowerCAmelCase_				:			List[str]       ,lowerCAmelCase_				:			Tuple							)    ->   str:







							"""simple docstring"""






							if os.environ.get('TESTING_MOCKED_DATALOADERS'       ,lowerCAmelCase_							) == "1":
														SCREAMING_SNAKE_CASE_	:							int											=2
							# New Code #
							SCREAMING_SNAKE_CASE_	:							List[str]											=int(args.gradient_accumulation_steps							)
							SCREAMING_SNAKE_CASE_	:							int											=int(args.local_sgd_steps							)
							# Initialize accelerator
							SCREAMING_SNAKE_CASE_	:							int											=Accelerator(
							    cpu=args.cpu       ,mixed_precision=args.mixed_precision       ,gradient_accumulation_steps=lowerCAmelCase_							)
							if accelerator.distributed_type not in [DistributedType.NO, DistributedType.MULTI_CPU, DistributedType.MULTI_GPU]:
														raise NotImplementedError('LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)'							)
							# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
							SCREAMING_SNAKE_CASE_	:							Tuple											=config["lr"]
							SCREAMING_SNAKE_CASE_	:							Union[str, Any]											=int(config['num_epochs']							)
							SCREAMING_SNAKE_CASE_	:							Optional[int]											=int(config['seed']							)
							SCREAMING_SNAKE_CASE_	:							Optional[Any]											=int(config['batch_size']							)

							SCREAMING_SNAKE_CASE_	:							Optional[int]											=evaluate.load('glue'       ,'mrpc'							)

							set_seed(lowerCAmelCase_							)
							SCREAMING_SNAKE_CASE_	:							Tuple											=get_dataloaders(lowerCAmelCase_       ,lowerCAmelCase_							)
							# Instantiate the model (we build the model here so that the seed also control new weights initialization)
							SCREAMING_SNAKE_CASE_	:							List[str]											=AutoModelForSequenceClassification.from_pretrained('bert-base-cased'       ,return_dict=lowerCAmelCase_							)

							# 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).
							SCREAMING_SNAKE_CASE_	:							Union[str, Any]											=model.to(accelerator.device							)

							# Instantiate optimizer
							SCREAMING_SNAKE_CASE_	:							str											=AdamW(params=model.parameters()       ,lr=lowerCAmelCase_							)

							# Instantiate scheduler
							SCREAMING_SNAKE_CASE_	:							Optional[int]											=get_linear_schedule_with_warmup(
							    optimizer=lowerCAmelCase_       ,num_warmup_steps=100       ,num_training_steps=(len(lowerCAmelCase_							) * 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.
							SCREAMING_SNAKE_CASE_	:							List[str]											=accelerator.prepare(
							    lowerCAmelCase_       ,lowerCAmelCase_       ,lowerCAmelCase_       ,lowerCAmelCase_       ,lowerCAmelCase_							)

							# Now we train the model
							for epoch in range(lowerCAmelCase_							):
														model.train()
														with LocalSGD(
														    accelerator=lowerCAmelCase_       ,model=lowerCAmelCase_       ,local_sgd_steps=lowerCAmelCase_       ,enabled=local_sgd_steps is not None							) as local_sgd:
																					for step, batch in enumerate(lowerCAmelCase_							):
																												# 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(lowerCAmelCase_							):
																																			SCREAMING_SNAKE_CASE_	:							List[Any]											=model(**lowerCAmelCase_							)
																																			SCREAMING_SNAKE_CASE_	:							List[str]											=output.loss
																																			accelerator.backward(lowerCAmelCase_							)
																																			optimizer.step()
																																			lr_scheduler.step()
																																			optimizer.zero_grad()
																																			# LocalSGD-specific line
																																			local_sgd.step()

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

														SCREAMING_SNAKE_CASE_	:							List[Any]											=metric.compute()
														# Use accelerator.print to print only on the main process.
														accelerator.print(F"""epoch {epoch}:"""       ,lowerCAmelCase_							)
def    SCREAMING_SNAKE_CASE__		( )    ->   Dict:







							"""simple docstring"""






							SCREAMING_SNAKE_CASE_	:							Dict											=argparse.ArgumentParser(description='Simple example of training script.'							)
							parser.add_argument(
							    '--mixed_precision'       ,type=lowerCAmelCase_       ,default=lowerCAmelCase_       ,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=lowerCAmelCase_       ,default=1       ,help='The number of minibatches to be ran before gradients are accumulated.'       ,)
							parser.add_argument(
							    '--local_sgd_steps'       ,type=lowerCAmelCase_       ,default=8       ,help='Number of local SGD steps or None to disable local SGD'							)
							parser.add_argument('--cpu'       ,action='store_true'       ,help='If passed, will train on the CPU.'							)
							SCREAMING_SNAKE_CASE_	:							List[str]											=parser.parse_args()
							SCREAMING_SNAKE_CASE_	:							List[str]											={"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
							training_function(lowerCAmelCase_       ,lowerCAmelCase_							)


if __name__ == "__main__":
							main()









'''simple docstring'''


import math
import time

from transformers import Trainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput, speed_metrics


if is_torch_tpu_available(check_device=False):
      import torch_xla.core.xla_model as xm
      import torch_xla.debug.metrics as met







class 							__magic_name__	(       lowerCAmelCase ):

     """simple docstring"""






     def __init__(   self			,     *lowerCamelCase			,     lowerCamelCase=None			,     lowerCamelCase=None			,     **lowerCamelCase		):


          '''simple docstring'''





          super().__init__(*lowerCamelCase			,     **lowerCamelCase		)
          __A	:       Tuple      =							eval_examples
          __A	:       List[str]      =							post_process_function


     def lowerCAmelCase__  (   self			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase = "eval"		):


          '''simple docstring'''





          __A	:       Any      =							self.eval_dataset if eval_dataset is None else eval_dataset
          __A	:       Dict      =							self.get_eval_dataloader(lowerCamelCase		)
          __A	:       str      =							self.eval_examples if eval_examples is None else eval_examples

          # Temporarily disable metric computation, we will do it in the loop here.
          __A	:       Any      =							self.compute_metrics
          __A	:       Optional[int]      =							None
          __A	:       Tuple      =							self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
          __A	:       Dict      =							time.time()
          try:
               __A	:       List[str]      =							eval_loop(
                   lowerCamelCase			,     description="Evaluation"			,     prediction_loss_only=True if compute_metrics is None else None			,     ignore_keys=lowerCamelCase			,     metric_key_prefix=lowerCamelCase			,     )
          finally:
               __A	:       List[str]      =							compute_metrics
          __A	:       List[Any]      =							self.args.eval_batch_size * self.args.world_size
          if f"{metric_key_prefix}_jit_compilation_time" in output.metrics:
               start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"]
          output.metrics.update(
              speed_metrics(
                  lowerCamelCase			,     lowerCamelCase			,     num_samples=output.num_samples			,     num_steps=math.ceil(output.num_samples / total_batch_size		)			,     )		)
          if self.post_process_function is not None and self.compute_metrics is not None and self.args.should_save:
               # Only the main node write the results by default
               __A	:       List[Any]      =							self.post_process_function(lowerCamelCase			,     lowerCamelCase			,     output.predictions		)
               __A	:       List[str]      =							self.compute_metrics(lowerCamelCase		)

               # Prefix all keys with metric_key_prefix + '_'
               for key in list(metrics.keys()		):
                    if not key.startswith(f"{metric_key_prefix}_"		):
                         __A	:       Union[str, Any]      =							metrics.pop(lowerCamelCase		)
               metrics.update(output.metrics		)
          else:
               __A	:       Any      =							output.metrics

          if self.args.should_log:
               # Only the main node log the results by default
               self.log(lowerCamelCase		)

          if self.args.tpu_metrics_debug or self.args.debug:
               # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
               xm.master_print(met.metrics_report()		)

          __A	:       str      =							self.callback_handler.on_evaluate(self.args			,     self.state			,     self.control			,     lowerCamelCase		)
          return metrics


     def lowerCAmelCase__  (   self			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase=None			,     lowerCamelCase = "test"		):


          '''simple docstring'''





          __A	:       Tuple      =							self.get_test_dataloader(lowerCamelCase		)

          # Temporarily disable metric computation, we will do it in the loop here.
          __A	:       Union[str, Any]      =							self.compute_metrics
          __A	:       Union[str, Any]      =							None
          __A	:       Optional[Any]      =							self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
          __A	:       Tuple      =							time.time()
          try:
               __A	:       List[str]      =							eval_loop(
                   lowerCamelCase			,     description="Prediction"			,     prediction_loss_only=True if compute_metrics is None else None			,     ignore_keys=lowerCamelCase			,     metric_key_prefix=lowerCamelCase			,     )
          finally:
               __A	:       int      =							compute_metrics
          __A	:       List[Any]      =							self.args.eval_batch_size * self.args.world_size
          if f"{metric_key_prefix}_jit_compilation_time" in output.metrics:
               start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"]
          output.metrics.update(
              speed_metrics(
                  lowerCamelCase			,     lowerCamelCase			,     num_samples=output.num_samples			,     num_steps=math.ceil(output.num_samples / total_batch_size		)			,     )		)

          if self.post_process_function is None or self.compute_metrics is None:
               return output

          __A	:       Dict      =							self.post_process_function(lowerCamelCase			,     lowerCamelCase			,     output.predictions			,     "predict"		)
          __A	:       Tuple      =							self.compute_metrics(lowerCamelCase		)

          # Prefix all keys with metric_key_prefix + '_'
          for key in list(metrics.keys()		):
               if not key.startswith(f"{metric_key_prefix}_"		):
                    __A	:       Any      =							metrics.pop(lowerCamelCase		)
          metrics.update(output.metrics		)
          return PredictionOutput(predictions=predictions.predictions			,     label_ids=predictions.label_ids			,     metrics=lowerCamelCase		)

