





'''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 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__      :   Union[str, Any]      	=			16
lowerCAmelCase__      :   str      	=			32


def 						__UpperCamelCase						(			_UpperCAmelCase,       _UpperCAmelCase = 16  ):
		__UpperCAmelCase	:							Tuple     =     AutoTokenizer.from_pretrained("bert-base-cased"  )
		__UpperCAmelCase	:							Optional[int]     =     load_dataset("glue",       "mrpc"  )

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

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

				return tokenizer.pad(
				    _UpperCAmelCase,       padding="longest",       max_length=_UpperCAmelCase,       pad_to_multiple_of=_UpperCAmelCase,       return_tensors="pt",       )

		# Instantiate dataloaders.
		__UpperCAmelCase	:							Dict     =     DataLoader(
		    tokenized_datasets["train"],       shuffle=_UpperCAmelCase,       collate_fn=_UpperCAmelCase,       batch_size=_UpperCAmelCase  )
		__UpperCAmelCase	:							List[str]     =     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__      :   str      	=			mocked_dataloaders  # noqa: F811


def 						__UpperCamelCase						(			_UpperCAmelCase,       _UpperCAmelCase  ):



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

		set_seed(_UpperCAmelCase  )

		__UpperCAmelCase	:							Union[str, Any]     =     get_dataloaders(_UpperCAmelCase,       _UpperCAmelCase  )
		__UpperCAmelCase	:							List[Any]     =     evaluate.load("glue",       "mrpc"  )

		# If the batch size is too big we use gradient accumulation
		__UpperCAmelCase	:							Optional[Any]     =     1
		if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
				__UpperCAmelCase	:							List[Any]     =     batch_size // MAX_GPU_BATCH_SIZE
				__UpperCAmelCase	:							int     =     MAX_GPU_BATCH_SIZE

		# Instantiate the model (we build the model here so that the seed also control new weights initialization)
		__UpperCAmelCase	:							Dict     =     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).
		__UpperCAmelCase	:							Union[str, Any]     =     model.to(accelerator.device  )

		# Instantiate optimizer
		__UpperCAmelCase	:							List[Any]     =     AdamW(params=model.parameters(),       lr=_UpperCAmelCase  )

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

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

				if "epoch" in training_difference:
						__UpperCAmelCase	:							Optional[int]     =     int(training_difference.replace("epoch_",       ""  )  ) + 1
						__UpperCAmelCase	:							str     =     None
				else:
						__UpperCAmelCase	:							Dict     =     int(training_difference.replace("step_",       ""  )  )
						__UpperCAmelCase	:							Any     =     resume_step // len(_UpperCAmelCase  )
						resume_step -= starting_epoch * len(_UpperCAmelCase  )

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

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

				__UpperCAmelCase	:							List[Any]     =     metric.compute()
				# Use accelerator.print to print only on the main process.
				accelerator.print(F"epoch {epoch}:",       _UpperCAmelCase  )

				# 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":
						__UpperCAmelCase	:							Optional[Any]     =     F"epoch_{epoch}"
						if args.output_dir is not None:
								__UpperCAmelCase	:							Union[str, Any]     =     os.path.join(args.output_dir,       _UpperCAmelCase  )
						accelerator.save_state(_UpperCAmelCase  )


def 						__UpperCamelCase						(			):
		__UpperCAmelCase	:							Union[str, Any]     =     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."  )
		parser.add_argument(
		    "--checkpointing_steps",       type=_UpperCAmelCase,       default=_UpperCAmelCase,       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=_UpperCAmelCase,       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=_UpperCAmelCase,       default=_UpperCAmelCase,       help="If the training should continue from a checkpoint folder.",       )
		__UpperCAmelCase	:							Any     =     parser.parse_args()
		__UpperCAmelCase	:							Tuple     =     {'''lr''': 2E-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
		training_function(_UpperCAmelCase,       _UpperCAmelCase  )


if __name__ == "__main__":
			main()








import unittest
import warnings
from dataclasses import dataclass

from transformers.convert_slow_tokenizer import SpmConverter
from transformers.testing_utils import get_tests_dir

@dataclass
class    __magic_name__						:
     A:		str



class    __magic_name__						(  unittest.TestCase):


     def   UpperCAmelCase__						(   self				:  Any      )	->   List[Any]:





      '''simple docstring'''






      UpperCamelCase__       :     Optional[Any] 							= get_tests_dir('''fixtures/test_sentencepiece.model'''      )
      UpperCamelCase__       :     Any 							= get_tests_dir('''fixtures/test_sentencepiece_with_bytefallback.model'''      )

      UpperCamelCase__       :     List[str] 							= FakeOriginalTokenizer(vocab_file=lowerCamelCase__      )

      with warnings.catch_warnings(record=lowerCamelCase__      ) as w:
       UpperCamelCase__       :     Optional[int] 							= SpmConverter(lowerCamelCase__      )
      self.assertEqual(len(lowerCamelCase__      )				,						0      )

      UpperCamelCase__       :     Tuple 							= FakeOriginalTokenizer(vocab_file=lowerCamelCase__      )

      with warnings.catch_warnings(record=lowerCamelCase__      ) as w:
       UpperCamelCase__       :     str 							= SpmConverter(lowerCamelCase__      )
      self.assertEqual(len(lowerCamelCase__      )				,						1      )

      self.assertIn(
          '''The sentencepiece tokenizer that you are converting to a fast tokenizer uses the byte fallback option'''
          ''' which is not implemented in the fast tokenizers.'''				,						str(w[0].message      )				,						)


