



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     =	16
__SCREAMING_SNAKE_CASE     =	32




def 						UpperCAmelCase (    _lowerCamelCase			,     _lowerCamelCase = 16     ):
			A				:       Tuple														=    AutoTokenizer.from_pretrained("bert-base-cased"     )
			A				:       Optional[Any]														=    load_dataset("glue"			,     "mrpc"     )

			def tokenize_function(_lowerCamelCase     ):
						# 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[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
			A				:       str														=    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.
						A				:       Optional[int]														=    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":
									A				:       Tuple														=    16
						elif accelerator.mixed_precision != "no":
									A				:       List[Any]														=    8
						else:
									A				:       Any														=    None

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

			# Instantiate dataloaders.
			A				:       Tuple														=    DataLoader(
			    tokenized_datasets["train"]			,     shuffle=_UpperCAmelCase			,     collate_fn=_UpperCAmelCase			,     batch_size=_UpperCAmelCase     )
			A				:       List[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

							__SCREAMING_SNAKE_CASE     =	mocked_dataloaders  # noqa: F811




def 						UpperCAmelCase (    _lowerCamelCase			,     _lowerCamelCase     ):





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

			set_seed(_UpperCAmelCase     )

			A				:       Tuple														=    get_dataloaders(_UpperCAmelCase			,     _UpperCAmelCase     )
			A				:       str														=    evaluate.load("glue"			,     "mrpc"     )

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

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

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

			# Instantiate scheduler
			A				:       Optional[int]														=    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.
			A				:       Optional[Any]														=    accelerator.prepare(
			    _UpperCAmelCase			,     _UpperCAmelCase			,     _UpperCAmelCase			,     _UpperCAmelCase			,     _UpperCAmelCase     )

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

						if "epoch" in training_difference:
									A				:       Union[str, Any]														=    int(training_difference.replace("epoch_"			,     ""     )     ) + 1
									A				:       List[str]														=    None
						else:
									A				:       str														=    int(training_difference.replace("step_"			,     ""     )     )
									A				:       Dict														=    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
									A				:       Any														=    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
									A				:       Optional[Any]														=    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     )
									A				:       List[str]														=    model(**_UpperCAmelCase     )
									A				:       str														=    outputs.loss
									A				:       Any														=    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     ):
												A				:       Any														=    f"""step_{overall_step}"""
												if overall_step % checkpointing_steps == 0:
															if args.output_dir is not None:
																		A				:       Optional[int]														=    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():
												A				:       str														=    model(**_UpperCAmelCase     )
									A				:       List[str]														=    outputs.logits.argmax(dim=-1     )
									A				:       Union[str, Any]														=    accelerator.gather_for_metrics((predictions, batch["labels"])     )
									metric.add_batch(
									    predictions=_UpperCAmelCase			,     references=_UpperCAmelCase			,     )

						A				:       str														=    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":
									A				:       List[Any]														=    f"""epoch_{epoch}"""
									if args.output_dir is not None:
												A				:       Dict														=    os.path.join(args.output_dir			,     _UpperCAmelCase     )
									accelerator.save_state(_UpperCAmelCase     )




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


if __name__ == "__main__":
							main()




import json
import os
import shutil
import tempfile
import unittest

from transformers import is_speech_available, is_torch_available
from transformers.models.speechta import SpeechTaTokenizer
from transformers.testing_utils import get_tests_dir, require_torch
from transformers.utils import FEATURE_EXTRACTOR_NAME


if is_speech_available() and is_torch_available():
				from transformers import SpeechTaFeatureExtractor, SpeechTaProcessor

				from .test_feature_extraction_speechta import floats_list


__A       =get_tests_dir("fixtures/test_sentencepiece_bpe_char.model")

@require_torch
class 						UpperCAmelCase__     (     unittest.TestCase  ):




			'''simple docstring'''






			def 					snake_case__	(  self				:					List[str]					):
					'''simple docstring'''





					__UpperCAmelCase       :	str						=      tempfile.mkdtemp()

					__UpperCAmelCase       :	Optional[int]						=      SpeechTaTokenizer(a_					)
					tokenizer.save_pretrained(self.tmpdirname					)

					__UpperCAmelCase       :	Union[str, Any]						=      {
					    '''feature_size''': 1,
					    '''padding_value''': 0.0,
					    '''sampling_rate''': 1_60_00,
					    '''do_normalize''': False,
					    '''num_mel_bins''': 80,
					    '''hop_length''': 16,
					    '''win_length''': 64,
					    '''win_function''': '''hann_window''',
					    '''fmin''': 80,
					    '''fmax''': 76_00,
					    '''mel_floor''': 1e-10,
					    '''reduction_factor''': 2,
					    '''return_attention_mask''': True,
					}

					__UpperCAmelCase       :	Tuple						=      os.path.join(self.tmpdirname			,			a_					)
					with open(self.feature_extraction_file			,			'''w'''			,			encoding='''utf-8'''					) as fp:
							fp.write(json.dumps(a_					) + '''\n'''					)




			def 					snake_case__	(  self				:					Any			,			**a_				:					Tuple					):
					'''simple docstring'''





					return SpeechTaTokenizer.from_pretrained(self.tmpdirname			,			**a_					)




			def 					snake_case__	(  self				:					Dict			,			**a_				:					str					):
					'''simple docstring'''





					return SpeechTaFeatureExtractor.from_pretrained(self.tmpdirname			,			**a_					)




			def 					snake_case__	(  self				:					List[str]					):
					'''simple docstring'''





					shutil.rmtree(self.tmpdirname					)




			def 					snake_case__	(  self				:					Dict					):
					'''simple docstring'''





					__UpperCAmelCase       :	int						=      self.get_tokenizer()
					__UpperCAmelCase       :	str						=      self.get_feature_extractor()

					__UpperCAmelCase       :	List[Any]						=      SpeechTaProcessor(tokenizer=a_			,			feature_extractor=a_					)

					processor.save_pretrained(self.tmpdirname					)
					__UpperCAmelCase       :	List[str]						=      SpeechTaProcessor.from_pretrained(self.tmpdirname					)

					self.assertEqual(processor.tokenizer.get_vocab()			,			tokenizer.get_vocab()					)
					self.assertIsInstance(processor.tokenizer			,			a_					)

					self.assertEqual(processor.feature_extractor.to_json_string()			,			feature_extractor.to_json_string()					)
					self.assertIsInstance(processor.feature_extractor			,			a_					)




			def 					snake_case__	(  self				:					int					):
					'''simple docstring'''





					__UpperCAmelCase       :	Optional[int]						=      SpeechTaProcessor(tokenizer=self.get_tokenizer()			,			feature_extractor=self.get_feature_extractor()					)
					processor.save_pretrained(self.tmpdirname					)

					__UpperCAmelCase       :	Optional[Any]						=      self.get_tokenizer(bos_token='''(BOS)'''			,			eos_token='''(EOS)'''					)
					__UpperCAmelCase       :	List[Any]						=      self.get_feature_extractor(do_normalize=a_			,			padding_value=1.0					)

					__UpperCAmelCase       :	List[str]						=      SpeechTaProcessor.from_pretrained(
					    self.tmpdirname			,			bos_token='''(BOS)'''			,			eos_token='''(EOS)'''			,			do_normalize=a_			,			padding_value=1.0					)

					self.assertEqual(processor.tokenizer.get_vocab()			,			tokenizer_add_kwargs.get_vocab()					)
					self.assertIsInstance(processor.tokenizer			,			a_					)

					self.assertEqual(processor.feature_extractor.to_json_string()			,			feature_extractor_add_kwargs.to_json_string()					)
					self.assertIsInstance(processor.feature_extractor			,			a_					)




			def 					snake_case__	(  self				:					Union[str, Any]					):
					'''simple docstring'''





					__UpperCAmelCase       :	Optional[Any]						=      self.get_feature_extractor()
					__UpperCAmelCase       :	Union[str, Any]						=      self.get_tokenizer()

					__UpperCAmelCase       :	List[Any]						=      SpeechTaProcessor(tokenizer=a_			,			feature_extractor=a_					)

					__UpperCAmelCase       :	List[Any]						=      floats_list((3, 10_00)					)

					__UpperCAmelCase       :	Tuple						=      feature_extractor(audio=a_			,			return_tensors='''np'''					)
					__UpperCAmelCase       :	List[str]						=      processor(audio=a_			,			return_tensors='''np'''					)

					for key in input_feat_extract.keys():
							self.assertAlmostEqual(input_feat_extract[key].sum()			,			input_processor[key].sum()			,			delta=1e-2					)




			def 					snake_case__	(  self				:					Optional[Any]					):
					'''simple docstring'''





					__UpperCAmelCase       :	List[Any]						=      self.get_feature_extractor()
					__UpperCAmelCase       :	Dict						=      self.get_tokenizer()

					__UpperCAmelCase       :	Optional[Any]						=      SpeechTaProcessor(tokenizer=a_			,			feature_extractor=a_					)

					__UpperCAmelCase       :	Union[str, Any]						=      floats_list((3, 10_00)					)

					__UpperCAmelCase       :	Dict						=      feature_extractor(audio_target=a_			,			return_tensors='''np'''					)
					__UpperCAmelCase       :	str						=      processor(audio_target=a_			,			return_tensors='''np'''					)

					for key in input_feat_extract.keys():
							self.assertAlmostEqual(input_feat_extract[key].sum()			,			input_processor[key].sum()			,			delta=1e-2					)




			def 					snake_case__	(  self				:					Union[str, Any]					):
					'''simple docstring'''





					__UpperCAmelCase       :	Union[str, Any]						=      self.get_feature_extractor()
					__UpperCAmelCase       :	Optional[int]						=      self.get_tokenizer()

					__UpperCAmelCase       :	Optional[Any]						=      SpeechTaProcessor(tokenizer=a_			,			feature_extractor=a_					)

					__UpperCAmelCase       :	List[str]						=      '''This is a test string'''

					__UpperCAmelCase       :	List[Any]						=      processor(text=a_					)
					__UpperCAmelCase       :	Optional[int]						=      tokenizer(a_					)

					for key in encoded_tok.keys():
							self.assertListEqual(encoded_tok[key]			,			encoded_processor[key]					)




			def 					snake_case__	(  self				:					List[str]					):
					'''simple docstring'''





					__UpperCAmelCase       :	Optional[Any]						=      self.get_feature_extractor()
					__UpperCAmelCase       :	Optional[Any]						=      self.get_tokenizer()

					__UpperCAmelCase       :	List[Any]						=      SpeechTaProcessor(tokenizer=a_			,			feature_extractor=a_					)

					__UpperCAmelCase       :	List[str]						=      '''This is a test string'''

					__UpperCAmelCase       :	Optional[Any]						=      processor(text_target=a_					)
					__UpperCAmelCase       :	Optional[int]						=      tokenizer(a_					)

					for key in encoded_tok.keys():
							self.assertListEqual(encoded_tok[key]			,			encoded_processor[key]					)




			def 					snake_case__	(  self				:					List[str]					):
					'''simple docstring'''





					__UpperCAmelCase       :	int						=      self.get_feature_extractor()
					__UpperCAmelCase       :	Union[str, Any]						=      self.get_tokenizer()

					__UpperCAmelCase       :	int						=      SpeechTaProcessor(tokenizer=a_			,			feature_extractor=a_					)

					__UpperCAmelCase       :	Optional[Any]						=      [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

					__UpperCAmelCase       :	Optional[Any]						=      processor.batch_decode(a_					)
					__UpperCAmelCase       :	Any						=      tokenizer.batch_decode(a_					)

					self.assertListEqual(a_			,			a_					)




			def 					snake_case__	(  self				:					str					):
					'''simple docstring'''





					__UpperCAmelCase       :	str						=      self.get_feature_extractor()
					__UpperCAmelCase       :	str						=      self.get_tokenizer()

					__UpperCAmelCase       :	Optional[Any]						=      SpeechTaProcessor(tokenizer=a_			,			feature_extractor=a_					)

					self.assertListEqual(
					    processor.model_input_names			,			feature_extractor.model_input_names			,			msg='''`processor` and `feature_extractor` model input names do not match'''			,			)





