



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 							A			(  _UpperCAmelCase       :     Dict					)       ->				Any:







			'''simple docstring'''
			return int(x / 2**20					)







class     __lowerCAmelCase					:



			def __enter__(		self     :						Optional[int])     ->					Tuple:



						"""simple docstring"""






						gc.collect()
						torch.cuda.empty_cache()
						torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
						_UpperCAmelCase	 =     torch.cuda.memory_allocated()
						return self



			def __exit__(		self     :						int				,	*A     :						Any)     ->					Tuple:



						"""simple docstring"""






						gc.collect()
						torch.cuda.empty_cache()
						_UpperCAmelCase	 =     torch.cuda.memory_allocated()
						_UpperCAmelCase	 =     torch.cuda.max_memory_allocated()
						_UpperCAmelCase	 =     bamb(self.end - self.begin)
						_UpperCAmelCase	 =     bamb(self.peak - self.begin)
						# print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")





def 							A			(  _UpperCAmelCase       :     Accelerator							,     _UpperCAmelCase       :     int = 16							,     _UpperCAmelCase       :     str = "bert-base-cased"							,     _UpperCAmelCase       :     int = 320							,     _UpperCAmelCase       :     int = 160							,     )       ->				List[Any]:







			'''simple docstring'''
			_UpperCAmelCase	 =     AutoTokenizer.from_pretrained(_UpperCAmelCase					)
			_UpperCAmelCase	 =     load_dataset(
			    'glue'							,     'mrpc'							,     split={'train': F"train[:{n_train}]", 'validation': F"validation[:{n_val}]"}					)

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

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

			return train_dataloader, eval_dataloader





def 							A			(  _UpperCAmelCase       :     Any							,     _UpperCAmelCase       :     Optional[Any]					)       ->				Optional[int]:







			'''simple docstring'''
			# Initialize accelerator
			_UpperCAmelCase	 =     Accelerator()

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

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

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

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

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

			# Now we train the model
			_UpperCAmelCase	 =     {}
			for epoch in range(_UpperCAmelCase							,     _UpperCAmelCase					):
						with TorchTracemalloc() as tracemalloc:
									model.train()
									for step, batch in enumerate(_UpperCAmelCase					):
												_UpperCAmelCase	 =     model(**_UpperCAmelCase					)
												_UpperCAmelCase	 =     outputs.loss
												_UpperCAmelCase	 =     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					)					)					)
						_UpperCAmelCase	 =     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 							A			(  )       ->				Optional[Any]:







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


if __name__ == "__main__":
							main()





import os
from shutil import copyfile
from typing import List, Optional, Tuple

from ...tokenization_utils import AddedToken
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
							from .tokenization_xlnet import XLNetTokenizer
else:
							UpperCAmelCase__     = None


UpperCAmelCase__     = logging.get_logger(__name__)

UpperCAmelCase__     = {"vocab_file": "spiece.model", "tokenizer_file": "tokenizer.json"}

UpperCAmelCase__     = {
    "vocab_file": {
        "xlnet-base-cased": "https://huggingface.co/xlnet-base-cased/resolve/main/spiece.model",
        "xlnet-large-cased": "https://huggingface.co/xlnet-large-cased/resolve/main/spiece.model",
    },
    "tokenizer_file": {
        "xlnet-base-cased": "https://huggingface.co/xlnet-base-cased/resolve/main/tokenizer.json",
        "xlnet-large-cased": "https://huggingface.co/xlnet-large-cased/resolve/main/tokenizer.json",
    },
}

UpperCAmelCase__     = {
    "xlnet-base-cased": None,
    "xlnet-large-cased": None,
}

UpperCAmelCase__     = "▁"

# Segments (not really needed)
UpperCAmelCase__     = 0
UpperCAmelCase__     = 1
UpperCAmelCase__     = 2
UpperCAmelCase__     = 3
UpperCAmelCase__     = 4







class     __lowerCAmelCase					(    A		):
			UpperCamelCase           =  VOCAB_FILES_NAMES
			UpperCamelCase           =  PRETRAINED_VOCAB_FILES_MAP
			UpperCamelCase           =  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			UpperCamelCase           =  '''left'''
			UpperCamelCase           =  XLNetTokenizer



			def __init__(		self     :						Any				,	A     :						Union[str, Any]=None				,	A     :						str=None				,	A     :						Tuple=False				,	A     :						Tuple=True				,	A     :						Any=False				,	A     :						List[str]="<s>"				,	A     :						List[str]="</s>"				,	A     :						Optional[int]="<unk>"				,	A     :						Tuple="<sep>"				,	A     :						str="<pad>"				,	A     :						Dict="<cls>"				,	A     :						Dict="<mask>"				,	A     :						Optional[Any]=["<eop>", "<eod>"]				,	**A     :						Optional[Any]				,	)     ->					str:



						"""simple docstring"""






						_UpperCAmelCase	 =     AddedToken(A				,	lstrip=A				,	rstrip=A) if isinstance(A				,	A) else mask_token

						super().__init__(
						    vocab_file=A				,	tokenizer_file=A				,	do_lower_case=A				,	remove_space=A				,	keep_accents=A				,	bos_token=A				,	eos_token=A				,	unk_token=A				,	sep_token=A				,	pad_token=A				,	cls_token=A				,	mask_token=A				,	additional_special_tokens=A				,	**A				,	)

						_UpperCAmelCase	 =     3
						_UpperCAmelCase	 =     do_lower_case
						_UpperCAmelCase	 =     remove_space
						_UpperCAmelCase	 =     keep_accents
						_UpperCAmelCase	 =     vocab_file
						_UpperCAmelCase	 =     False if not self.vocab_file else True



			def 					_lowerCamelCase (		self     :						Tuple				,	A     :						List[int]				,	A     :						Optional[List[int]] = None)     ->					List[int]:



						"""simple docstring"""






						_UpperCAmelCase	 =     [self.sep_token_id]
						_UpperCAmelCase	 =     [self.cls_token_id]
						if token_ids_a is None:
									return token_ids_a + sep + cls
						return token_ids_a + sep + token_ids_a + sep + cls



			def 					_lowerCamelCase (		self     :						Tuple				,	A     :						List[int]				,	A     :						Optional[List[int]] = None)     ->					List[int]:



						"""simple docstring"""






						_UpperCAmelCase	 =     [self.sep_token_id]
						_UpperCAmelCase	 =     [2]

						if token_ids_a is None:
									return len(token_ids_a + sep) * [0] + cls_segment_id
						return len(token_ids_a + sep) * [0] + len(token_ids_a + sep) * [1] + cls_segment_id



			def 					_lowerCamelCase (		self     :						List[str]				,	A     :						str				,	A     :						Optional[str] = None)     ->					Tuple[str]:



						"""simple docstring"""






						if not self.can_save_slow_tokenizer:
									raise ValueError(
									    'Your fast tokenizer does not have the necessary information to save the vocabulary for a slow '
									    'tokenizer.')

						if not os.path.isdir(A):
									logger.error(F"Vocabulary path ({save_directory}) should be a directory")
									return
						_UpperCAmelCase	 =     os.path.join(
						    A				,	(filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file'])

						if os.path.abspath(self.vocab_file) != os.path.abspath(A):
									copyfile(self.vocab_file				,	A)

						return (out_vocab_file,)
