





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


UpperCAmelCase_      							=     16
UpperCAmelCase_      							=     32





def 						lowerCamelCase__     (     UpperCamelCase__			:       Union[str, Any]						)			-> Optional[int]:



					'''simple docstring'''




					return int(x / 2**20						)
class 			UpperCamelCase_		:




					def __enter__(		self  )		->							Union[str, Any]:
										gc.collect()
										torch.cuda.empty_cache()
										torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
										_snake_case									=				torch.cuda.memory_allocated()
										return self






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





def 						lowerCamelCase__     (     UpperCamelCase__			:       Accelerator       ,    UpperCamelCase__			:       int = 16       ,    UpperCamelCase__			:       str = "bert-base-cased"       ,    UpperCamelCase__			:       int = 320       ,    UpperCamelCase__			:       int = 160       ,    )			-> Optional[int]:



					'''simple docstring'''




					_snake_case									=				AutoTokenizer.from_pretrained(UpperCamelCase__						)
					_snake_case									=				load_dataset(
					    'glue'       ,    'mrpc'       ,    split={'train': F'''train[:{n_train}]''', 'validation': F'''validation[:{n_val}]'''}						)

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

					def collate_fn(UpperCamelCase__			:       Optional[int]						):
										# 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.
					_snake_case									=				DataLoader(
					    tokenized_datasets['train']       ,    shuffle=UpperCamelCase__       ,    collate_fn=UpperCamelCase__       ,    batch_size=UpperCamelCase__						)
					_snake_case									=				DataLoader(
					    tokenized_datasets['validation']       ,    shuffle=UpperCamelCase__       ,    collate_fn=UpperCamelCase__       ,    batch_size=UpperCamelCase__						)

					return train_dataloader, eval_dataloader





def 						lowerCamelCase__     (     UpperCamelCase__			:       Optional[Any]       ,    UpperCamelCase__			:       Optional[Any]						)			-> Optional[int]:



					'''simple docstring'''




					_snake_case									=				Accelerator()

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

					set_seed(UpperCamelCase__						)
					_snake_case							,     _snake_case									=				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)
					_snake_case									=				AutoModelForSequenceClassification.from_pretrained(UpperCamelCase__       ,    return_dict=UpperCamelCase__						)

					# Instantiate optimizer
					_snake_case									=				(
					    AdamW
					    if accelerator.state.deepspeed_plugin is None
					    or 'optimizer' not in accelerator.state.deepspeed_plugin.deepspeed_config
					    else DummyOptim
					)
					_snake_case									=				optimizer_cls(params=model.parameters()       ,    lr=UpperCamelCase__						)

					if accelerator.state.deepspeed_plugin is not None:
										_snake_case									=				accelerator.state.deepspeed_plugin.deepspeed_config[
										    'gradient_accumulation_steps'
										]
					else:
										_snake_case									=				1
					_snake_case									=				(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
					):
										_snake_case									=				get_linear_schedule_with_warmup(
										    optimizer=UpperCamelCase__       ,    num_warmup_steps=0       ,    num_training_steps=UpperCamelCase__       ,    )
					else:
										_snake_case									=				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.
					_snake_case							,     _snake_case							,     _snake_case							,     _snake_case							,     _snake_case									=				accelerator.prepare(
					    UpperCamelCase__       ,    UpperCamelCase__       ,    UpperCamelCase__       ,    UpperCamelCase__       ,    UpperCamelCase__						)

					# We need to keep track of how many total steps we have iterated over
					_snake_case									=				0
					# We also need to keep track of the stating epoch so files are named properly
					_snake_case									=				0

					# Now we train the model
					_snake_case									=				{}
					for epoch in range(UpperCamelCase__       ,    UpperCamelCase__						):
										with TorchTracemalloc() as tracemalloc:
															model.train()
															for step, batch in enumerate(UpperCamelCase__						):
																				_snake_case									=				model(**UpperCamelCase__						)
																				_snake_case									=				outputs.loss
																				_snake_case									=				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						)						)						)
										_snake_case									=				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 						lowerCamelCase__     (     )			-> Dict:



					'''simple docstring'''




					_snake_case									=				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.'       ,    )
					_snake_case									=				parser.parse_args()
					_snake_case									=				{'lr': 2e-5, 'num_epochs': args.num_epochs, 'seed': 42, 'batch_size': 16}
					training_function(UpperCamelCase__       ,    UpperCamelCase__						)


if __name__ == "__main__":
					main()










import copy
import tempfile
import unittest

from huggingface_hub import HfFolder, delete_repo
from parameterized import parameterized
from requests.exceptions import HTTPError

from transformers import AutoConfig, GenerationConfig
from transformers.testing_utils import TOKEN, USER, is_staging_test
class 			UpperCamelCase_		(  unittest.TestCase					):




					@parameterized.expand([(None,), ('foo.json',)]  )
					def 	lowerCAmelCase			(		self				,						lowerCAmelCase_  )		->							str:
										_snake_case									=				GenerationConfig(
										    do_sample=lowerCAmelCase_				,						temperature=0.7				,						length_penalty=1.0				,						bad_words_ids=[[1, 2, 3], [4, 5]]				,						)
										with tempfile.TemporaryDirectory() as tmp_dir:
															config.save_pretrained(lowerCAmelCase_				,						config_name=lowerCAmelCase_  )
															_snake_case									=				GenerationConfig.from_pretrained(lowerCAmelCase_				,						config_name=lowerCAmelCase_  )

										# Checks parameters that were specified
										self.assertEqual(loaded_config.do_sample				,						lowerCAmelCase_  )
										self.assertEqual(loaded_config.temperature				,						0.7  )
										self.assertEqual(loaded_config.length_penalty				,						1.0  )
										self.assertEqual(loaded_config.bad_words_ids				,						[[1, 2, 3], [4, 5]]  )

										# Checks parameters that were not specified (defaults)
										self.assertEqual(loaded_config.top_k				,						50  )
										self.assertEqual(loaded_config.max_length				,						20  )
										self.assertEqual(loaded_config.max_time				,						lowerCAmelCase_  )




					def 	lowerCAmelCase			(		self  )		->							Optional[int]:
										_snake_case									=				AutoConfig.from_pretrained('gpt2'  )
										_snake_case									=				GenerationConfig.from_model_config(lowerCAmelCase_  )
										_snake_case									=				GenerationConfig()

										# The generation config has loaded a few non-default parameters from the model config
										self.assertNotEqual(lowerCAmelCase_				,						lowerCAmelCase_  )

										# One of those parameters is eos_token_id -- check if it matches
										self.assertNotEqual(generation_config_from_model.eos_token_id				,						default_generation_config.eos_token_id  )
										self.assertEqual(generation_config_from_model.eos_token_id				,						model_config.eos_token_id  )




					def 	lowerCAmelCase			(		self  )		->							Tuple:
										_snake_case									=				GenerationConfig()
										_snake_case									=				{
										    'max_new_tokens': 1024,
										    'foo': 'bar',
										}
										_snake_case									=				copy.deepcopy(lowerCAmelCase_  )
										_snake_case									=				generation_config.update(**lowerCAmelCase_  )

										# update_kwargs was not modified (no side effects)
										self.assertEqual(lowerCAmelCase_				,						lowerCAmelCase_  )

										# update_kwargs was used to update the config on valid attributes
										self.assertEqual(generation_config.max_new_tokens				,						1024  )

										# `.update()` returns a dictionary of unused kwargs
										self.assertEqual(lowerCAmelCase_				,						{'foo': 'bar'}  )




					def 	lowerCAmelCase			(		self  )		->							Optional[int]:
										_snake_case									=				GenerationConfig()
										_snake_case									=				'bar'

										with tempfile.TemporaryDirectory('test-generation-config'  ) as tmp_dir:
															generation_config.save_pretrained(lowerCAmelCase_  )

															_snake_case									=				GenerationConfig.from_pretrained(lowerCAmelCase_  )
										# update_kwargs was used to update the config on valid attributes
										self.assertEqual(new_config.foo				,						'bar'  )

										_snake_case									=				GenerationConfig.from_model_config(lowerCAmelCase_  )
										assert not hasattr(lowerCAmelCase_				,						'foo'  )  # no new kwargs should be initialized if from config






					def 	lowerCAmelCase			(		self  )		->							List[Any]:
										_snake_case									=				GenerationConfig()
										self.assertEqual(default_config.temperature				,						1.0  )
										self.assertEqual(default_config.do_sample				,						lowerCAmelCase_  )
										self.assertEqual(default_config.num_beams				,						1  )

										_snake_case									=				GenerationConfig(
										    do_sample=lowerCAmelCase_				,						temperature=0.7				,						length_penalty=1.0				,						bad_words_ids=[[1, 2, 3], [4, 5]]				,						)
										self.assertEqual(config.temperature				,						0.7  )
										self.assertEqual(config.do_sample				,						lowerCAmelCase_  )
										self.assertEqual(config.num_beams				,						1  )

										with tempfile.TemporaryDirectory() as tmp_dir:
															config.save_pretrained(lowerCAmelCase_  )
															_snake_case									=				GenerationConfig.from_pretrained(lowerCAmelCase_				,						temperature=1.0  )

										self.assertEqual(loaded_config.temperature				,						1.0  )
										self.assertEqual(loaded_config.do_sample				,						lowerCAmelCase_  )
										self.assertEqual(loaded_config.num_beams				,						1  )  # default value





@is_staging_test
class 			UpperCamelCase_		(  unittest.TestCase					):




					@classmethod
					def 	lowerCAmelCase			(		cls  )		->							List[str]:
										_snake_case									=				TOKEN
										HfFolder.save_token(lowerCAmelCase_  )




					@classmethod
					def 	lowerCAmelCase			(		cls  )		->							Union[str, Any]:
										try:
															delete_repo(token=cls._token				,						repo_id='test-generation-config'  )
										except HTTPError:
															pass

										try:
															delete_repo(token=cls._token				,						repo_id='valid_org/test-generation-config-org'  )
										except HTTPError:
															pass




					def 	lowerCAmelCase			(		self  )		->							int:
										_snake_case									=				GenerationConfig(
										    do_sample=lowerCAmelCase_				,						temperature=0.7				,						length_penalty=1.0				,						)
										config.push_to_hub('test-generation-config'				,						use_auth_token=self._token  )

										_snake_case									=				GenerationConfig.from_pretrained(F'''{USER}/test-generation-config'''  )
										for k, v in config.to_dict().items():
															if k != "transformers_version":
																				self.assertEqual(lowerCAmelCase_				,						getattr(lowerCAmelCase_				,						lowerCAmelCase_  )  )

        # Reset repo
										delete_repo(token=self._token				,						repo_id='test-generation-config'  )

										# Push to hub via save_pretrained
										with tempfile.TemporaryDirectory() as tmp_dir:
															config.save_pretrained(
															    lowerCAmelCase_				,						repo_id='test-generation-config'				,						push_to_hub=lowerCAmelCase_				,						use_auth_token=self._token  )

										_snake_case									=				GenerationConfig.from_pretrained(F'''{USER}/test-generation-config'''  )
										for k, v in config.to_dict().items():
															if k != "transformers_version":
																				self.assertEqual(lowerCAmelCase_				,						getattr(lowerCAmelCase_				,						lowerCAmelCase_  )  )






					def 	lowerCAmelCase			(		self  )		->							List[str]:
										_snake_case									=				GenerationConfig(
										    do_sample=lowerCAmelCase_				,						temperature=0.7				,						length_penalty=1.0				,						)
										config.push_to_hub('valid_org/test-generation-config-org'				,						use_auth_token=self._token  )

										_snake_case									=				GenerationConfig.from_pretrained('valid_org/test-generation-config-org'  )
										for k, v in config.to_dict().items():
															if k != "transformers_version":
																				self.assertEqual(lowerCAmelCase_				,						getattr(lowerCAmelCase_				,						lowerCAmelCase_  )  )

        # Reset repo
										delete_repo(token=self._token				,						repo_id='valid_org/test-generation-config-org'  )

										# Push to hub via save_pretrained
										with tempfile.TemporaryDirectory() as tmp_dir:
															config.save_pretrained(
															    lowerCAmelCase_				,						repo_id='valid_org/test-generation-config-org'				,						push_to_hub=lowerCAmelCase_				,						use_auth_token=self._token  )

										_snake_case									=				GenerationConfig.from_pretrained('valid_org/test-generation-config-org'  )
										for k, v in config.to_dict().items():
															if k != "transformers_version":
																				self.assertEqual(lowerCAmelCase_				,						getattr(lowerCAmelCase_				,						lowerCAmelCase_  )  )



