



import argparse
import gc
import json
import os

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


A__  :  Any								=					16
A__  :  int								=					32


def 			UpperCamelCase(  __UpperCamelCase							:		int ):
		return int(x / 2**20 )



class 				__snake_case						:




							def __enter__(       self     :       List[str]):
									gc.collect()
									torch.cuda.empty_cache()
									torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
									lowerCAmelCase_			:     int									= torch.cuda.memory_allocated()
									return self







							def __exit__(       self     :       int				,				*A_     :       Optional[int]):
									gc.collect()
									torch.cuda.empty_cache()
									lowerCAmelCase_			:     List[str]									= torch.cuda.memory_allocated()
									lowerCAmelCase_			:     Union[str, Any]									= torch.cuda.max_memory_allocated()
									lowerCAmelCase_			:     Any									= bamb(self.end - self.begin)
									lowerCAmelCase_			:     str									= bamb(self.peak - self.begin)
									# print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")





def 			UpperCamelCase(  __UpperCamelCase							:		Accelerator     ,__UpperCamelCase							:		int = 16     ,__UpperCamelCase							:		str = "bert-base-cased"     ,__UpperCamelCase							:		int = 320     ,__UpperCamelCase							:		int = 160     ,):
		lowerCAmelCase_			:     List[Any]									= AutoTokenizer.from_pretrained(__UpperCamelCase )
		lowerCAmelCase_			:     int									= load_dataset(
		    '''glue'''     ,'''mrpc'''     ,split={'''train''': f"""train[:{n_train}]""", '''validation''': f"""validation[:{n_val}]"""} )

		def tokenize_function(__UpperCamelCase							:		Union[str, Any] ):
				# max_length=None => use the model max length (it's actually the default)
				lowerCAmelCase_			:     Any									= 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
		lowerCAmelCase_			:     Any									= datasets.map(
		    __UpperCamelCase     ,batched=__UpperCamelCase     ,remove_columns=['''idx''', '''sentence1''', '''sentence2''']     ,load_from_cache_file=__UpperCamelCase )

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

		def collate_fn(__UpperCamelCase							:		List[Any] ):
				# 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(__UpperCamelCase     ,padding='''max_length'''     ,max_length=128     ,return_tensors='''pt''' )
				return tokenizer.pad(__UpperCamelCase     ,padding='''longest'''     ,return_tensors='''pt''' )

		# Instantiate dataloaders.
		lowerCAmelCase_			:     List[Any]									= DataLoader(
		    tokenized_datasets['''train''']     ,shuffle=__UpperCamelCase     ,collate_fn=__UpperCamelCase     ,batch_size=__UpperCamelCase )
		lowerCAmelCase_			:     Dict									= DataLoader(
		    tokenized_datasets['''validation''']     ,shuffle=__UpperCamelCase     ,collate_fn=__UpperCamelCase     ,batch_size=__UpperCamelCase )

		return train_dataloader, eval_dataloader


def 			UpperCamelCase(  __UpperCamelCase							:		Any     ,__UpperCamelCase							:		int ):
		# Initialize accelerator
		lowerCAmelCase_			:     str									= Accelerator()

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

		set_seed(__UpperCamelCase )
		lowerCAmelCase_					,  lowerCAmelCase_			:     Optional[Any]									= get_dataloaders(__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase     ,args.n_train     ,args.n_val )

		# Instantiate the model (we build the model here so that the seed also control new weights initialization)
		lowerCAmelCase_			:     str									= AutoModelForSequenceClassification.from_pretrained(__UpperCamelCase     ,return_dict=__UpperCamelCase )

		# Instantiate optimizer
		lowerCAmelCase_			:     Optional[int]									= (
		    AdamW
		    if accelerator.state.deepspeed_plugin is None
		    or '''optimizer''' not in accelerator.state.deepspeed_plugin.deepspeed_config
		    else DummyOptim
		)
		lowerCAmelCase_			:     List[Any]									= optimizer_cls(params=model.parameters()     ,lr=__UpperCamelCase )

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

		# Instantiate scheduler
		if (
		    accelerator.state.deepspeed_plugin is None
		    or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
		):
				lowerCAmelCase_			:     Dict									= get_linear_schedule_with_warmup(
				    optimizer=__UpperCamelCase     ,num_warmup_steps=0     ,num_training_steps=__UpperCamelCase     ,)
		else:
				lowerCAmelCase_			:     Tuple									= DummyScheduler(__UpperCamelCase     ,total_num_steps=__UpperCamelCase     ,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.
		lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_			:     Dict									= accelerator.prepare(
		    __UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase )

		# We need to keep track of how many total steps we have iterated over
		lowerCAmelCase_			:     Union[str, Any]									= 0
		# We also need to keep track of the stating epoch so files are named properly
		lowerCAmelCase_			:     List[Any]									= 0

		# Now we train the model
		lowerCAmelCase_			:     Dict									= {}
		for epoch in range(__UpperCamelCase     ,__UpperCamelCase ):
				with TorchTracemalloc() as tracemalloc:
						model.train()
						for step, batch in enumerate(__UpperCamelCase ):
								lowerCAmelCase_			:     List[Any]									= model(**__UpperCamelCase )
								lowerCAmelCase_			:     int									= outputs.loss
								lowerCAmelCase_			:     List[str]									= loss / gradient_accumulation_steps
								accelerator.backward(__UpperCamelCase )
								if step % gradient_accumulation_steps == 0:
										optimizer.step()
										lr_scheduler.step()
										optimizer.zero_grad()

								overall_step += 1

        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
				accelerator.print('''Memory before entering the train : {}'''.format(bamb(tracemalloc.begin ) ) )
				accelerator.print('''Memory consumed at the end of the train (end-begin): {}'''.format(tracemalloc.used ) )
				accelerator.print('''Peak Memory consumed during the train (max-begin): {}'''.format(tracemalloc.peaked ) )
				accelerator.print(
				    '''Total Peak Memory consumed during the train (max): {}'''.format(
				        tracemalloc.peaked + bamb(tracemalloc.begin ) ) )
				lowerCAmelCase_			:     Tuple									= tracemalloc.peaked + bamb(tracemalloc.begin )
				if args.peak_memory_upper_bound is not None:
						assert (
						    train_total_peak_memory[f"""epoch-{epoch}"""] <= args.peak_memory_upper_bound
						), "Peak memory usage exceeded the upper bound"

		accelerator.wait_for_everyone()
		if accelerator.is_main_process:
				with open(os.path.join(args.output_dir     ,'''peak_memory_utilization.json''' )     ,'''w''' ) as f:
						json.dump(__UpperCamelCase     ,__UpperCamelCase )


def 			UpperCamelCase(  ):
		lowerCAmelCase_			:     List[str]									= argparse.ArgumentParser(description='''Simple example of training script tracking peak GPU memory usage.''' )
		parser.add_argument(
		    '''--model_name_or_path'''     ,type=__UpperCamelCase     ,default='''bert-base-cased'''     ,help='''Path to pretrained model or model identifier from huggingface.co/models.'''     ,required=__UpperCamelCase     ,)
		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(
		    '''--peak_memory_upper_bound'''     ,type=__UpperCamelCase     ,default=__UpperCamelCase     ,help='''The upper bound of peak memory usage in MB. If set, the training will throw an error if the peak memory usage exceeds this value.'''     ,)
		parser.add_argument(
		    '''--n_train'''     ,type=__UpperCamelCase     ,default=320     ,help='''Number of training examples to use.'''     ,)
		parser.add_argument(
		    '''--n_val'''     ,type=__UpperCamelCase     ,default=160     ,help='''Number of validation examples to use.'''     ,)
		parser.add_argument(
		    '''--num_epochs'''     ,type=__UpperCamelCase     ,default=1     ,help='''Number of train epochs.'''     ,)
		lowerCAmelCase_			:     str									= parser.parse_args()
		lowerCAmelCase_			:     List[str]									= {'''lr''': 2e-5, '''num_epochs''': args.num_epochs, '''seed''': 42, '''batch_size''': 16}
		training_function(__UpperCamelCase     ,__UpperCamelCase )


if __name__ == "__main__":
						main()










import argparse
import os
import re

import packaging.version


A__  :  Dict								=					'''examples/'''
A__  :  Any								=					{
    '''examples''': (re.compile(R'''^check_min_version\("[^"]+"\)\s*$''', re.MULTILINE), '''check_min_version("VERSION")\n'''),
    '''init''': (re.compile(R'''^__version__\s+=\s+"([^"]+)"\s*$''', re.MULTILINE), '''__version__ = "VERSION"\n'''),
    '''setup''': (re.compile(R'''^(\s*)version\s*=\s*"[^"]+",''', re.MULTILINE), R'''\1version="VERSION",'''),
    '''doc''': (re.compile(R'''^(\s*)release\s*=\s*"[^"]+"$''', re.MULTILINE), '''release = "VERSION"\n'''),
}
A__  :  Any								=					{
    '''init''': '''src/transformers/__init__.py''',
    '''setup''': '''setup.py''',
}
A__  :  Any								=					'''README.md'''


def 			UpperCamelCase(  __UpperCamelCase							:		int     ,__UpperCamelCase							:		List[Any]     ,__UpperCamelCase							:		List[Any] ):
		with open(__UpperCamelCase     ,'''r'''     ,encoding='''utf-8'''     ,newline='''\n''' ) as f:
				lowerCAmelCase_			:     Tuple									= f.read()
		lowerCAmelCase_					,  lowerCAmelCase_			:     Dict									= REPLACE_PATTERNS[pattern]
		lowerCAmelCase_			:     Tuple									= replace.replace('''VERSION'''     ,__UpperCamelCase )
		lowerCAmelCase_			:     Optional[int]									= re_pattern.sub(__UpperCamelCase     ,__UpperCamelCase )
		with open(__UpperCamelCase     ,'''w'''     ,encoding='''utf-8'''     ,newline='''\n''' ) as f:
				f.write(__UpperCamelCase )


def 			UpperCamelCase(  __UpperCamelCase							:		Union[str, Any] ):
		for folder, directories, fnames in os.walk(__UpperCamelCase ):
				# Removing some of the folders with non-actively maintained examples from the walk
				if "research_projects" in directories:
						directories.remove('''research_projects''' )
				if "legacy" in directories:
						directories.remove('''legacy''' )
				for fname in fnames:
						if fname.endswith('''.py''' ):
								update_version_in_file(os.path.join(__UpperCamelCase     ,__UpperCamelCase )     ,__UpperCamelCase     ,pattern='''examples''' )


def 			UpperCamelCase(  __UpperCamelCase							:		int     ,__UpperCamelCase							:		List[Any]=False ):
		for pattern, fname in REPLACE_FILES.items():
				update_version_in_file(__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase )
		if not patch:
				update_version_in_examples(__UpperCamelCase )


def 			UpperCamelCase(  ):
		lowerCAmelCase_			:     List[str]									= '''🤗 Transformers currently provides the following architectures'''
		lowerCAmelCase_			:     List[Any]									= '''1. Want to contribute a new model?'''
		with open(__UpperCamelCase     ,'''r'''     ,encoding='''utf-8'''     ,newline='''\n''' ) as f:
				lowerCAmelCase_			:     Union[str, Any]									= f.readlines()

		# Find the start of the list.
		lowerCAmelCase_			:     int									= 0
		while not lines[start_index].startswith(_start_prompt ):
				start_index += 1
		start_index += 1

		lowerCAmelCase_			:     str									= start_index
		# Update the lines in the model list.
		while not lines[index].startswith(_end_prompt ):
				if lines[index].startswith('''1.''' ):
						lowerCAmelCase_			:     int									= lines[index].replace(
						    '''https://huggingface.co/docs/transformers/main/model_doc'''     ,'''https://huggingface.co/docs/transformers/model_doc'''     ,)
				index += 1

		with open(__UpperCamelCase     ,'''w'''     ,encoding='''utf-8'''     ,newline='''\n''' ) as f:
				f.writelines(__UpperCamelCase )


def 			UpperCamelCase(  ):
		with open(REPLACE_FILES['''init''']     ,'''r''' ) as f:
				lowerCAmelCase_			:     Optional[Any]									= f.read()
		lowerCAmelCase_			:     Dict									= REPLACE_PATTERNS['''init'''][0].search(__UpperCamelCase ).groups()[0]
		return packaging.version.parse(__UpperCamelCase )


def 			UpperCamelCase(  __UpperCamelCase							:		Dict=False ):
		lowerCAmelCase_			:     Union[str, Any]									= get_version()
		if patch and default_version.is_devrelease:
				raise ValueError('''Can\'t create a patch version from the dev branch, checkout a released version!''' )
		if default_version.is_devrelease:
				lowerCAmelCase_			:     List[str]									= default_version.base_version
		elif patch:
				lowerCAmelCase_			:     int									= f"""{default_version.major}.{default_version.minor}.{default_version.micro + 1}"""
		else:
				lowerCAmelCase_			:     int									= f"""{default_version.major}.{default_version.minor + 1}.0"""

		# Now let's ask nicely if that's the right one.
		lowerCAmelCase_			:     Optional[Any]									= input(f"""Which version are you releasing? [{default_version}]""" )
		if len(__UpperCamelCase ) == 0:
				lowerCAmelCase_			:     List[str]									= default_version

		print(f"""Updating version to {version}.""" )
		global_version_update(__UpperCamelCase     ,patch=__UpperCamelCase )
		if not patch:
				print('''Cleaning main README, don\'t forget to run `make fix-copies`.''' )
				clean_main_ref_in_model_list()


def 			UpperCamelCase(  ):
		lowerCAmelCase_			:     Any									= get_version()
		lowerCAmelCase_			:     int									= f"""{current_version.major}.{current_version.minor + 1}.0.dev0"""
		lowerCAmelCase_			:     Optional[Any]									= current_version.base_version

		# Check with the user we got that right.
		lowerCAmelCase_			:     Optional[Any]									= input(f"""Which version are we developing now? [{dev_version}]""" )
		if len(__UpperCamelCase ) == 0:
				lowerCAmelCase_			:     int									= dev_version

		print(f"""Updating version to {version}.""" )
		global_version_update(__UpperCamelCase )
		print('''Cleaning main README, don\'t forget to run `make fix-copies`.''' )
		clean_main_ref_in_model_list()


if __name__ == "__main__":
						A__  :  Dict								=					argparse.ArgumentParser()
						parser.add_argument('''--post_release''', action='''store_true''', help='''Whether this is pre or post release.''')
						parser.add_argument('''--patch''', action='''store_true''', help='''Whether or not this is a patch release.''')
						A__  :  Optional[int]								=					parser.parse_args()
						if not args.post_release:
												pre_release_work(patch=args.patch)
						elif args.patch:
												print('''Nothing to do after a patch :-)''')
						else:
												post_release_work()





