




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

SCREAMING_SNAKE_CASE:    List[str]											=   1_6
SCREAMING_SNAKE_CASE:    str											=   3_2


def 				__lowerCAmelCase       (				lowerCAmelCase						,	lowerCAmelCase = 16   )->  List[str]:
		SCREAMING_SNAKE_CASE_							=       AutoTokenizer.from_pretrained('bert-base-cased'   )
		SCREAMING_SNAKE_CASE_							=       load_dataset('glue'						,	'mrpc'   )

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

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

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

		# Instantiate dataloaders.
		SCREAMING_SNAKE_CASE_							=       DataLoader(
		    tokenized_datasets['train']						,	shuffle=lowerCAmelCase						,	collate_fn=lowerCAmelCase						,	batch_size=lowerCAmelCase   )
		SCREAMING_SNAKE_CASE_							=       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:    Optional[Any]											=   mocked_dataloaders  # noqa: F811


def 				__lowerCAmelCase       (				lowerCAmelCase						,	lowerCAmelCase   )->  Optional[int]:
		# For testing only
		if os.environ.get('TESTING_MOCKED_DATALOADERS'						,	lowerCAmelCase   ) == "1":
				SCREAMING_SNAKE_CASE_							=       2
		# Initialize accelerator
		SCREAMING_SNAKE_CASE_							=       Accelerator(cpu=args.cpu						,	mixed_precision=args.mixed_precision   )
		# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
		SCREAMING_SNAKE_CASE_							=       config['lr']
		SCREAMING_SNAKE_CASE_							=       int(config['num_epochs']   )
		SCREAMING_SNAKE_CASE_							=       int(config['seed']   )
		SCREAMING_SNAKE_CASE_							=       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":
						SCREAMING_SNAKE_CASE_							=       args.checkpointing_steps
				elif args.checkpointing_steps.isdigit():
						SCREAMING_SNAKE_CASE_							=       int(args.checkpointing_steps   )
				else:
						raise ValueError(
						    F'''Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed.'''   )
		else:
				SCREAMING_SNAKE_CASE_							=       None

		set_seed(lowerCAmelCase   )

		SCREAMING_SNAKE_CASE_			,							SCREAMING_SNAKE_CASE_							=       get_dataloaders(lowerCAmelCase						,	lowerCAmelCase   )
		SCREAMING_SNAKE_CASE_							=       evaluate.load('glue'						,	'mrpc'   )

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

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

		# Instantiate optimizer
		SCREAMING_SNAKE_CASE_							=       AdamW(params=model.parameters()						,	lr=lowerCAmelCase   )

		# Instantiate scheduler
		SCREAMING_SNAKE_CASE_							=       get_linear_schedule_with_warmup(
		    optimizer=lowerCAmelCase						,	num_warmup_steps=100						,	num_training_steps=(len(lowerCAmelCase   ) * 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.
		SCREAMING_SNAKE_CASE_			,							SCREAMING_SNAKE_CASE_			,							SCREAMING_SNAKE_CASE_			,							SCREAMING_SNAKE_CASE_			,							SCREAMING_SNAKE_CASE_							=       accelerator.prepare(
		    lowerCAmelCase						,	lowerCAmelCase						,	lowerCAmelCase						,	lowerCAmelCase						,	lowerCAmelCase   )

		# New Code #
		# We need to keep track of how many total steps we have iterated over
		SCREAMING_SNAKE_CASE_							=       0
		# We also need to keep track of the stating epoch so files are named properly
		SCREAMING_SNAKE_CASE_							=       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   )
						SCREAMING_SNAKE_CASE_							=       os.path.basename(args.resume_from_checkpoint   )
				else:
						# Get the most recent checkpoint
						SCREAMING_SNAKE_CASE_							=       [f.name for f in os.scandir(os.getcwd()   ) if f.is_dir()]
						dirs.sort(key=os.path.getctime   )
						SCREAMING_SNAKE_CASE_							=       dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
				# Extract `epoch_{i}` or `step_{i}`
				SCREAMING_SNAKE_CASE_							=       os.path.splitext(lowerCAmelCase   )[0]

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

    # Now we train the model
		for epoch in range(lowerCAmelCase						,	lowerCAmelCase   ):
				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
						SCREAMING_SNAKE_CASE_							=       accelerator.skip_first_batches(lowerCAmelCase						,	lowerCAmelCase   )
						overall_step += resume_step
				else:
						# After the first iteration though, we need to go back to the original dataloader
						SCREAMING_SNAKE_CASE_							=       train_dataloader
				for step, batch in enumerate(lowerCAmelCase   ):
						# We could avoid this line since we set the accelerator with `device_placement=True`.
						batch.to(accelerator.device   )
						SCREAMING_SNAKE_CASE_							=       model(**lowerCAmelCase   )
						SCREAMING_SNAKE_CASE_							=       outputs.loss
						SCREAMING_SNAKE_CASE_							=       loss / gradient_accumulation_steps
						accelerator.backward(lowerCAmelCase   )
						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(lowerCAmelCase						,	lowerCAmelCase   ):
								SCREAMING_SNAKE_CASE_							=       F'''step_{overall_step}'''
								if overall_step % checkpointing_steps == 0:
										if args.output_dir is not None:
												SCREAMING_SNAKE_CASE_							=       os.path.join(args.output_dir						,	lowerCAmelCase   )
										accelerator.save_state(lowerCAmelCase   )

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

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

				# 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":
						SCREAMING_SNAKE_CASE_							=       F'''epoch_{epoch}'''
						if args.output_dir is not None:
								SCREAMING_SNAKE_CASE_							=       os.path.join(args.output_dir						,	lowerCAmelCase   )
						accelerator.save_state(lowerCAmelCase   )


def 				__lowerCAmelCase       (				)->  Optional[int]:
		SCREAMING_SNAKE_CASE_							=       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.'						,	)
		parser.add_argument('--cpu'						,	action='store_true'						,	help='If passed, will train on the CPU.'   )
		parser.add_argument(
		    '--checkpointing_steps'						,	type=lowerCAmelCase						,	default=lowerCAmelCase						,	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=lowerCAmelCase						,	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=lowerCAmelCase						,	default=lowerCAmelCase						,	help='If the training should continue from a checkpoint folder.'						,	)
		SCREAMING_SNAKE_CASE_							=       parser.parse_args()
		SCREAMING_SNAKE_CASE_							=       {'lr': 2E-5, 'num_epochs': 3, 'seed': 42, 'batch_size': 16}
		training_function(lowerCAmelCase						,	lowerCAmelCase   )


if __name__ == "__main__":
						main()





import unittest

from transformers import (
    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    SummarizationPipeline,
    TFPreTrainedModel,
    pipeline,
)
from transformers.testing_utils import get_gpu_count, is_pipeline_test, require_tf, require_torch, slow, torch_device
from transformers.tokenization_utils import TruncationStrategy

from .test_pipelines_common import ANY


SCREAMING_SNAKE_CASE:    Optional[Any]											=   -1 if torch_device == '''cpu''' else 0



@is_pipeline_test
class       lowercase_		(unittest.TestCase   ):
						lowerCAmelCase__ 			=MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
						lowerCAmelCase__ 			=TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING



						def   __a				( self  :      Optional[Any]				,       snake_case__  :      Tuple				,       snake_case__  :      Tuple				,       snake_case__  :      Optional[Any]     ):

								"""simple docstring"""

								SCREAMING_SNAKE_CASE_							=       SummarizationPipeline(model=snake_case__				,       tokenizer=snake_case__     )
								return summarizer, ["(CNN)The Palestinian Authority officially became", "Some other text"]



						def   __a				( self  :      Tuple				,       snake_case__  :      str				,       snake_case__  :      Union[str, Any]     ):

								"""simple docstring"""

								SCREAMING_SNAKE_CASE_							=       summarizer.model

								SCREAMING_SNAKE_CASE_							=       summarizer('(CNN)The Palestinian Authority officially became'     )
								self.assertEqual(snake_case__				,       [{'summary_text': ANY(snake_case__     )}]     )

								SCREAMING_SNAKE_CASE_							=       summarizer(
								    '(CNN)The Palestinian Authority officially became '				,       num_beams=2				,       min_length=2				,       max_length=5				,       )
								self.assertEqual(snake_case__				,       [{'summary_text': ANY(snake_case__     )}]     )

								# Some models (Switch Transformers, LED, T5, LongT5, etc) can handle long sequences.
								SCREAMING_SNAKE_CASE_							=       [
								    'SwitchTransformersConfig',
								    'T5Config',
								    'LongT5Config',
								    'LEDConfig',
								    'PegasusXConfig',
								    'FSMTConfig',
								    'M2M100Config',
								    'ProphetNetConfig',  # positional embeddings up to a fixed maximum size (otherwise clamping the values)
								]
								if model.config.__class__.__name__ not in model_can_handle_longer_seq:
										# Too long and exception is expected.
										# For TF models, if the weights are initialized in GPU context, we won't get expected index error from
										# the embedding layer.
										if not (
										    isinstance(snake_case__				,       snake_case__     )
										    and get_gpu_count() > 0
										    and len(summarizer.model.trainable_weights     ) > 0
										):
												with self.assertRaises(snake_case__     ):
														SCREAMING_SNAKE_CASE_							=       summarizer('This ' * 10_00     )
								SCREAMING_SNAKE_CASE_							=       summarizer('This ' * 10_00				,       truncation=TruncationStrategy.ONLY_FIRST     )



						@require_torch
						def   __a				( self  :      Dict     ):

								"""simple docstring"""

								SCREAMING_SNAKE_CASE_							=       pipeline(task='summarization'				,       model='sshleifer/tiny-mbart'				,       framework='pt'     )
								SCREAMING_SNAKE_CASE_							=       summarizer('This is a small test'     )
								self.assertEqual(
								    snake_case__				,       [
								        {
								            'summary_text': 'เข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไป'
								        }
								    ]				,       )



						@require_tf
						def   __a				( self  :      Any     ):

								"""simple docstring"""

								SCREAMING_SNAKE_CASE_							=       pipeline(task='summarization'				,       model='sshleifer/tiny-mbart'				,       framework='tf'     )
								SCREAMING_SNAKE_CASE_							=       summarizer('This is a small test'     )
								self.assertEqual(
								    snake_case__				,       [
								        {
								            'summary_text': 'เข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไป'
								        }
								    ]				,       )






						@require_torch
						@slow
						def   __a				( self  :      Any     ):

								"""simple docstring"""

								SCREAMING_SNAKE_CASE_							=       pipeline(task='summarization'				,       device=snake_case__     )
								SCREAMING_SNAKE_CASE_							=       (
								    ' (CNN)The Palestinian Authority officially became the 123rd member of the International Criminal Court on'
								    ' Wednesday, a step that gives the court jurisdiction over alleged crimes in Palestinian territories. The'
								    ' formal accession was marked with a ceremony at The Hague, in the Netherlands, where the court is based.'
								    ' The Palestinians signed the ICC\'s founding Rome Statute in January, when they also accepted its'
								    ' jurisdiction over alleged crimes committed "in the occupied Palestinian territory, including East'
								    ' Jerusalem, since June 13, 2014." Later that month, the ICC opened a preliminary examination into the'
								    ' situation in Palestinian territories, paving the way for possible war crimes investigations against'
								    ' Israelis. As members of the court, Palestinians may be subject to counter-charges as well. Israel and'
								    ' the United States, neither of which is an ICC member, opposed the Palestinians\' efforts to join the'
								    ' body. But Palestinian Foreign Minister Riad al-Malki, speaking at Wednesday\'s ceremony, said it was a'
								    ' move toward greater justice. "As Palestine formally becomes a State Party to the Rome Statute today, the'
								    ' world is also a step closer to ending a long era of impunity and injustice," he said, according to an'
								    ' ICC news release. "Indeed, today brings us closer to our shared goals of justice and peace." Judge'
								    ' Kuniko Ozaki, a vice president of the ICC, said acceding to the treaty was just the first step for the'
								    ' Palestinians. "As the Rome Statute today enters into force for the State of Palestine, Palestine'
								    ' acquires all the rights as well as responsibilities that come with being a State Party to the Statute.'
								    ' These are substantive commitments, which cannot be taken lightly," she said. Rights group Human Rights'
								    ' Watch welcomed the development. "Governments seeking to penalize Palestine for joining the ICC should'
								    ' immediately end their pressure, and countries that support universal acceptance of the court\'s treaty'
								    ' should speak out to welcome its membership," said Balkees Jarrah, international justice counsel for the'
								    ' group. \"What\'s objectionable is the attempts to undermine international justice, not Palestine\'s'
								    ' decision to join a treaty to which over 100 countries around the world are members." In January, when'
								    ' the preliminary ICC examination was opened, Israeli Prime Minister Benjamin Netanyahu described it as an'
								    ' outrage, saying the court was overstepping its boundaries. The United States also said it "strongly"'
								    ' disagreed with the court\'s decision. \"As we have said repeatedly, we do not believe that Palestine is a'
								    ' state and therefore we do not believe that it is eligible to join the ICC," the State Department said in'
								    ' a statement. It urged the warring sides to resolve their differences through direct negotiations. "We'
								    ' will continue to oppose actions against Israel at the ICC as counterproductive to the cause of peace,"'
								    ' it said. But the ICC begs to differ with the definition of a state for its purposes and refers to the'
								    ' territories as "Palestine." While a preliminary examination is not a formal investigation, it allows the'
								    ' court to review evidence and determine whether to investigate suspects on both sides. Prosecutor Fatou'
								    ' Bensouda said her office would "conduct its analysis in full independence and impartiality." The war'
								    ' between Israel and Hamas militants in Gaza last summer left more than 2,000 people dead. The inquiry'
								    ' will include alleged war crimes committed since June. The International Criminal Court was set up in'
								    ' 2002 to prosecute genocide, crimes against humanity and war crimes. CNN\'s Vasco Cotovio, Kareem Khadder'
								    ' and Faith Karimi contributed to this report.'
								)
								SCREAMING_SNAKE_CASE_							=       (
								    ' The Palestinian Authority becomes the 123rd member of the International Criminal Court . The move gives'
								    ' the court jurisdiction over alleged crimes in Palestinian territories . Israel and the United States'
								    ' opposed the Palestinians\' efforts to join the court . Rights group Human Rights Watch welcomes the move,'
								    ' says governments seeking to penalize Palestine should end pressure .'
								)
								SCREAMING_SNAKE_CASE_							=       summarizer(snake_case__     )
								self.assertEqual(result[0]['summary_text']				,       snake_case__     )