






# XXX: we want transformers master here - in the absense of conftest manipulating sys.path:
# hack it in for now:
import sys
from pathlib import Path


__UpperCamelCase   :   Optional[Any]      =    Path(__file__).resolve().parents[3] / """src"""
sys.path.insert(1, str(git_repo_path))

import dataclasses  # noqa
import io  # noqa
import itertools  # noqa
import json  # noqa
import os  # noqa
import unittest  # noqa
from copy import deepcopy  # noqa

from parameterized import parameterized  # noqa
from transformers import TrainingArguments, is_torch_available  # noqa
from transformers.deepspeed import is_deepspeed_available  # noqa
from transformers.file_utils import WEIGHTS_NAME  # noqa
from transformers.testing_utils import (  # noqa
    CaptureLogger,
    ExtendSysPath,
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    mockenv_context,
    require_deepspeed,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)
from transformers.trainer_utils import set_seed  # noqa


set_seed(42)

__UpperCamelCase   :   str      =    {"""base""": """patrickvonplaten/wav2vec2_tiny_random""", """robust""": """patrickvonplaten/wav2vec2_tiny_random_robust"""}

__UpperCamelCase   :   List[str]      =    """zero2"""
__UpperCamelCase   :   List[Any]      =    """zero3"""
__UpperCamelCase   :   Optional[int]      =    [ZEROa, ZEROa]





def 					snake_case	(  lowerCamelCase					,	lowerCamelCase					,	lowerCamelCase       ):




	'''simple docstring'''







	__lowercase   							=	parameterized.to_safe_name("""_""".join(str(lowerCamelCase       ) for x in param.args       )       )
	return F'{func.__name__}_{param_based_name}'


# Cartesian-product of zero stages with models to test
__UpperCamelCase   :   Union[str, Any]      =    list(itertools.product(stages, models.keys()))




@slow
@require_deepspeed
@require_torch_gpu
class 							__UpperCamelCase     (  _lowerCAmelCase ):







	@parameterized.expand(_lowerCAmelCase							,					name_func=_lowerCAmelCase      )
	def 			_a (	self   :				Dict							,					_lowerCAmelCase   :				int							,					_lowerCAmelCase   :				Dict      )		->			Union[str, Any]:
		"""simple docstring"""







		self.run_and_check(
		    stage=_lowerCAmelCase							,					model=_lowerCAmelCase							,					distributed=_lowerCAmelCase							,					fpaa=_lowerCAmelCase							,					)







	@require_torch_multi_gpu
	@parameterized.expand(_lowerCAmelCase							,					name_func=_lowerCAmelCase      )
	def 			_a (	self   :				Optional[int]							,					_lowerCAmelCase   :				Dict							,					_lowerCAmelCase   :				List[Any]      )		->			Tuple:
		"""simple docstring"""







		self.run_and_check(
		    stage=_lowerCAmelCase							,					model=_lowerCAmelCase							,					distributed=_lowerCAmelCase							,					fpaa=_lowerCAmelCase							,					)







	@parameterized.expand(_lowerCAmelCase							,					name_func=_lowerCAmelCase      )
	def 			_a (	self   :				Tuple							,					_lowerCAmelCase   :				Dict							,					_lowerCAmelCase   :				Tuple      )		->			Optional[Any]:
		"""simple docstring"""







		self.run_and_check(
		    stage=_lowerCAmelCase							,					model=_lowerCAmelCase							,					distributed=_lowerCAmelCase							,					fpaa=_lowerCAmelCase							,					)







	@require_torch_multi_gpu
	@parameterized.expand(_lowerCAmelCase							,					name_func=_lowerCAmelCase      )
	def 			_a (	self   :				str							,					_lowerCAmelCase   :				Optional[Any]							,					_lowerCAmelCase   :				Dict      )		->			str:
		"""simple docstring"""







		self.run_and_check(
		    stage=_lowerCAmelCase							,					model=_lowerCAmelCase							,					distributed=_lowerCAmelCase							,					fpaa=_lowerCAmelCase							,					)







	def 			_a (	self   :				Union[str, Any]							,					_lowerCAmelCase   :				Optional[Any]      )		->			Any:
		"""simple docstring"""







		pass







	def 			_a (	self   :				str							,					_lowerCAmelCase   :				str							,					_lowerCAmelCase   :				str							,					_lowerCAmelCase   :				int = 10							,					_lowerCAmelCase   :				bool = True							,					_lowerCAmelCase   :				bool = True							,					_lowerCAmelCase   :				bool = True							,					)		->			str:
		"""simple docstring"""







		__lowercase   							=	models[model]

		__lowercase   							=	self.run_trainer(
		    stage=_lowerCAmelCase							,					model_name=_lowerCAmelCase							,					eval_steps=_lowerCAmelCase							,					num_train_epochs=1							,					distributed=_lowerCAmelCase							,					fpaa=_lowerCAmelCase							,					)

		self.do_checks(_lowerCAmelCase      )

		return output_dir







	def 			_a (	self   :				List[str]							,					_lowerCAmelCase   :				str							,					_lowerCAmelCase   :				str							,					_lowerCAmelCase   :				int = 10							,					_lowerCAmelCase   :				int = 1							,					_lowerCAmelCase   :				bool = True							,					_lowerCAmelCase   :				bool = True							,					)		->			Dict:
		"""simple docstring"""







		__lowercase   							=	self.get_auto_remove_tmp_dir("""./xxx"""							,					after=_lowerCAmelCase      )
		__lowercase   							=	F'\n            --model_name_or_path {model_name}\n            --dataset_name hf-internal-testing/librispeech_asr_dummy\n            --dataset_config_name clean\n            --train_split_name validation\n            --validation_split_name validation\n            --output_dir {output_dir}\n            --num_train_epochs {str(_lowerCAmelCase      )}\n            --per_device_train_batch_size 2\n            --per_device_eval_batch_size 2\n            --evaluation_strategy steps\n            --learning_rate 5e-4\n            --warmup_steps 8\n            --orthography timit\n            --preprocessing_num_workers 1\n            --group_by_length\n            --freeze_feature_extractor\n            --report_to none\n            --save_steps 0\n            --eval_steps {eval_steps}\n            --report_to none\n        '.split()

		if fpaa:
			args.extend(["""--fp16"""]      )

		# currently ds_config_wav2vec2_zero.json requires "zero_optimization.find_unused_parameters": true,
		# hence the separate config files
		__lowercase   							=	F'--deepspeed {self.test_file_dir_str}/ds_config_wav2vec2_{stage}.json'.split()
		__lowercase   							=	[F'{self.examples_dir_str}/research_projects/wav2vec2/run_asr.py']
		__lowercase   							=	self.get_launcher(_lowerCAmelCase      )

		__lowercase   							=	launcher + script + args + ds_args
		# keep for quick debug
		# print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
		execute_subprocess_async(_lowerCAmelCase							,					env=self.get_env()      )

		return output_dir







	def 			_a (	self   :				List[str]							,					_lowerCAmelCase   :				str=False      )		->			Dict:
		"""simple docstring"""







		__lowercase   							=	min(2							,					get_gpu_count()      ) if distributed else 1
		return F'deepspeed --num_nodes 1 --num_gpus {num_gpus}'.split()





import argparse
import json
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.utils.deepspeed import DummyOptim, DummyScheduler


lowerCAmelCase       		=    16
lowerCAmelCase       		=    32




def     __SCREAMING_SNAKE_CASE							(       lowercase_      ,       lowercase_ = 16      ,       lowercase_ = "bert-base-cased"      )		->     Union[str, Any]:




						'''simple docstring'''







						__UpperCAmelCase			:							List[str]	    =	AutoTokenizer.from_pretrained(lowercase_      )
						__UpperCAmelCase			:							int	    =	load_dataset('''glue'''      ,       '''mrpc'''      )

						def tokenize_function(lowercase_      ):
												# max_length=None => use the model max length (it's actually the default)
												__UpperCAmelCase			:							List[Any]	    =	tokenizer(examples['''sentence1''']      ,       examples['''sentence2''']      ,       truncation=lowercase_      ,       max_length=lowercase_      )
												return outputs

						# Apply the method we just defined to all the examples in all the splits of the dataset
						__UpperCAmelCase			:							Dict	    =	datasets.map(
						    lowercase_      ,       batched=lowercase_      ,       remove_columns=['''idx''', '''sentence1''', '''sentence2''']      ,       load_from_cache_file=lowercase_      )

						# We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
						# transformers library
						__UpperCAmelCase			:							List[str]	    =	tokenized_datasets.rename_column('''label'''      ,       '''labels'''      )

						def collate_fn(lowercase_      ):
												# On TPU it's best to pad everything to the same length or training will be very slow.
												if accelerator.distributed_type == DistributedType.TPU:
																		return tokenizer.pad(lowercase_      ,       padding='''max_length'''      ,       max_length=128      ,       return_tensors='''pt'''      )
												return tokenizer.pad(lowercase_      ,       padding='''longest'''      ,       return_tensors='''pt'''      )

						# Instantiate dataloaders.
						__UpperCAmelCase			:							Optional[Any]	    =	DataLoader(
						    tokenized_datasets['''train''']      ,       shuffle=lowercase_      ,       collate_fn=lowercase_      ,       batch_size=lowercase_      )
						__UpperCAmelCase			:							Any	    =	DataLoader(
						    tokenized_datasets['''validation''']      ,       shuffle=lowercase_      ,       collate_fn=lowercase_      ,       batch_size=lowercase_      )

						return train_dataloader, eval_dataloader




def     __SCREAMING_SNAKE_CASE							(       lowercase_      ,       lowercase_      ,       lowercase_      ,       lowercase_      )		->     str:




						'''simple docstring'''







						model.eval()
						__UpperCAmelCase			:							Optional[int]	    =	0
						for step, batch in enumerate(lowercase_      ):
												# We could avoid this line since we set the accelerator with `device_placement=True`.
												batch.to(accelerator.device      )
												with torch.no_grad():
																		__UpperCAmelCase			:							Any	    =	model(**lowercase_      )
												__UpperCAmelCase			:							Optional[int]	    =	outputs.logits.argmax(dim=-1      )
												# It is slightly faster to call this once, than multiple times
												__UpperCAmelCase    ,      __UpperCAmelCase			:							Union[str, Any]	    =	accelerator.gather(
												    (predictions, batch['''labels'''])      )  # If we are in a multiprocess environment, the last batch has duplicates
												if accelerator.use_distributed:
																		if step == len(lowercase_      ) - 1:
																								__UpperCAmelCase			:							Optional[int]	    =	predictions[: len(eval_dataloader.dataset      ) - samples_seen]
																								__UpperCAmelCase			:							List[Any]	    =	references[: len(eval_dataloader.dataset      ) - samples_seen]
																		else:
																								samples_seen += references.shape[0]
												metric.add_batch(
												    predictions=lowercase_      ,       references=lowercase_      ,       )

						__UpperCAmelCase			:							Union[str, Any]	    =	metric.compute()
						return eval_metric["accuracy"]




def     __SCREAMING_SNAKE_CASE							(       lowercase_      ,       lowercase_      )		->     Dict:




						'''simple docstring'''







						__UpperCAmelCase			:							List[Any]	    =	Accelerator()

						# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
						__UpperCAmelCase			:							Optional[Any]	    =	config['''lr''']
						__UpperCAmelCase			:							List[Any]	    =	int(config['''num_epochs''']      )
						__UpperCAmelCase			:							Dict	    =	int(config['''seed''']      )
						__UpperCAmelCase			:							Tuple	    =	int(config['''batch_size''']      )
						__UpperCAmelCase			:							Any	    =	args.model_name_or_path

						set_seed(lowercase_      )
						__UpperCAmelCase    ,      __UpperCAmelCase			:							Tuple	    =	get_dataloaders(lowercase_      ,       lowercase_      ,       lowercase_      )

						# Instantiate the model (we build the model here so that the seed also control new weights initialization)
						__UpperCAmelCase			:							List[str]	    =	AutoModelForSequenceClassification.from_pretrained(lowercase_      ,       return_dict=lowercase_      )

						# Instantiate optimizer
						__UpperCAmelCase			:							str	    =	(
						    AdamW
						    if accelerator.state.deepspeed_plugin is None
						    or '''optimizer''' not in accelerator.state.deepspeed_plugin.deepspeed_config
						    else DummyOptim
						)
						__UpperCAmelCase			:							int	    =	optimizer_cls(params=model.parameters()      ,       lr=lowercase_      )

						if accelerator.state.deepspeed_plugin is not None:
												__UpperCAmelCase			:							int	    =	accelerator.state.deepspeed_plugin.deepspeed_config[
												    '''gradient_accumulation_steps'''
												]
						else:
												__UpperCAmelCase			:							List[Any]	    =	1
						__UpperCAmelCase			:							str	    =	(len(lowercase_      ) * num_epochs) // gradient_accumulation_steps

						# Instantiate scheduler
						if (
						    accelerator.state.deepspeed_plugin is None
						    or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
						):
												__UpperCAmelCase			:							List[str]	    =	get_linear_schedule_with_warmup(
												    optimizer=lowercase_      ,       num_warmup_steps=0      ,       num_training_steps=lowercase_      ,       )
						else:
												__UpperCAmelCase			:							Union[str, Any]	    =	DummyScheduler(lowercase_      ,       total_num_steps=lowercase_      ,       warmup_num_steps=0      )

						# 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    ,      __UpperCAmelCase    ,      __UpperCAmelCase    ,      __UpperCAmelCase    ,      __UpperCAmelCase			:							int	    =	accelerator.prepare(
						    lowercase_      ,       lowercase_      ,       lowercase_      ,       lowercase_      ,       lowercase_      )

						# We need to keep track of how many total steps we have iterated over
						__UpperCAmelCase			:							Tuple	    =	0
						# We also need to keep track of the stating epoch so files are named properly
						__UpperCAmelCase			:							Any	    =	0
						__UpperCAmelCase			:							Tuple	    =	evaluate.load('''glue'''      ,       '''mrpc'''      )
						__UpperCAmelCase			:							Any	    =	num_epochs

						if args.partial_train_epoch is not None:
												__UpperCAmelCase			:							Any	    =	args.partial_train_epoch

						if args.resume_from_checkpoint:
												accelerator.load_state(args.resume_from_checkpoint      )
												__UpperCAmelCase			:							List[Any]	    =	args.resume_from_checkpoint.split('''epoch_'''      )[1]
												__UpperCAmelCase			:							Union[str, Any]	    =	''''''
												for char in epoch_string:
																		if char.isdigit():
																								state_epoch_num += char
																		else:
																								break
												__UpperCAmelCase			:							Tuple	    =	int(lowercase_      ) + 1
												__UpperCAmelCase			:							Optional[int]	    =	evaluation_loop(lowercase_      ,       lowercase_      ,       lowercase_      ,       lowercase_      )
												accelerator.print('''resumed checkpoint performance:'''      ,       lowercase_      )
												accelerator.print('''resumed checkpoint\'s scheduler\'s lr:'''      ,       lr_scheduler.get_lr()[0]      )
												accelerator.print('''resumed optimizers\'s lr:'''      ,       optimizer.param_groups[0]['''lr''']      )
												with open(os.path.join(args.output_dir      ,       f"state_{starting_epoch-1}.json"      )      ,       '''r'''      ) as f:
																		__UpperCAmelCase			:							str	    =	json.load(lowercase_      )
																		assert resumed_state["accuracy"] == accuracy, "Accuracy mismatch, loading from checkpoint failed"
																		assert (
																		    resumed_state["lr"] == lr_scheduler.get_lr()[0]
																		), "Scheduler learning rate mismatch, loading from checkpoint failed"
																		assert (
																		    resumed_state["optimizer_lr"] == optimizer.param_groups[0]["lr"]
																		), "Optimizer learning rate mismatch, loading from checkpoint failed"
																		assert resumed_state["epoch"] == starting_epoch - 1, "Epoch mismatch, loading from checkpoint failed"
																		return

    # Now we train the model
						__UpperCAmelCase			:							Tuple	    =	{}
						for epoch in range(lowercase_      ,       lowercase_      ):
												model.train()
												for step, batch in enumerate(lowercase_      ):
																		__UpperCAmelCase			:							Dict	    =	model(**lowercase_      )
																		__UpperCAmelCase			:							Union[str, Any]	    =	outputs.loss
																		__UpperCAmelCase			:							Tuple	    =	loss / gradient_accumulation_steps
																		accelerator.backward(lowercase_      )
																		if step % gradient_accumulation_steps == 0:
																								optimizer.step()
																								lr_scheduler.step()
																								optimizer.zero_grad()

																		overall_step += 1
												__UpperCAmelCase			:							Optional[Any]	    =	f"epoch_{epoch}"
												__UpperCAmelCase			:							Optional[Any]	    =	os.path.join(args.output_dir      ,       lowercase_      )
												accelerator.save_state(lowercase_      )
												__UpperCAmelCase			:							Any	    =	evaluation_loop(lowercase_      ,       lowercase_      ,       lowercase_      ,       lowercase_      )
												__UpperCAmelCase			:							List[str]	    =	accuracy
												__UpperCAmelCase			:							Dict	    =	lr_scheduler.get_lr()[0]
												__UpperCAmelCase			:							List[str]	    =	optimizer.param_groups[0]['''lr''']
												__UpperCAmelCase			:							List[Any]	    =	epoch
												__UpperCAmelCase			:							Any	    =	overall_step
												accelerator.print(f"epoch {epoch}:"      ,       lowercase_      )

												accelerator.wait_for_everyone()
												if accelerator.is_main_process:
																		with open(os.path.join(args.output_dir      ,       f"state_{epoch}.json"      )      ,       '''w'''      ) as f:
																								json.dump(lowercase_      ,       lowercase_      )




def     __SCREAMING_SNAKE_CASE							(       )		->     List[Any]:




						'''simple docstring'''







						__UpperCAmelCase			:							Any	    =	argparse.ArgumentParser(description='''Simple example of training script tracking peak GPU memory usage.'''      )
						parser.add_argument(
						    '''--model_name_or_path'''      ,       type=lowercase_      ,       default='''bert-base-cased'''      ,       help='''Path to pretrained model or model identifier from huggingface.co/models.'''      ,       required=lowercase_      ,       )
						parser.add_argument(
						    '''--output_dir'''      ,       type=lowercase_      ,       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=lowercase_      ,       default=lowercase_      ,       help='''If the training should continue from a checkpoint folder.'''      ,       )
						parser.add_argument(
						    '''--partial_train_epoch'''      ,       type=lowercase_      ,       default=lowercase_      ,       help='''If passed, the training will stop after this number of epochs.'''      ,       )
						parser.add_argument(
						    '''--num_epochs'''      ,       type=lowercase_      ,       default=2      ,       help='''Number of train epochs.'''      ,       )
						__UpperCAmelCase			:							Optional[int]	    =	parser.parse_args()
						__UpperCAmelCase			:							str	    =	{'''lr''': 2e-5, '''num_epochs''': args.num_epochs, '''seed''': 42, '''batch_size''': 16}

						training_function(lowercase_      ,       lowercase_      )


if __name__ == "__main__":
				main()

