



'''simple docstring'''






import argparse
import os

# New Code #
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
from accelerate.utils import find_executable_batch_size


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to combine both the gradient accumulation
# and automatic batch size finder utilities of Accelerate to perfrom
# automatic gradient accumulation
#
# 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)
#
# 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
#
########################################################################

A_				:					int							=      3_2




def        snake_case_   (	lowerCAmelCase_			,  lowerCAmelCase_ = 16						)->						str:





       '''simple docstring'''


       _UpperCAmelCase    :		Optional[int]     	=					AutoTokenizer.from_pretrained("""bert-base-cased"""						)
       _UpperCAmelCase    :		List[str]     	=					load_dataset("""glue"""			,  """mrpc"""						)

       def tokenize_function(lowerCAmelCase_						):
              # max_length=None => use the model max length (it's actually the default)
              _UpperCAmelCase    :		Any     	=					tokenizer(examples["""sentence1"""]			,  examples["""sentence2"""]			,  truncation=lowerCAmelCase_			,  max_length=lowerCAmelCase_						)
              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():
              _UpperCAmelCase    :		int     	=					datasets.map(
                  lowerCAmelCase_			,  batched=lowerCAmelCase_			,  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
       _UpperCAmelCase    :		Union[str, Any]     	=					tokenized_datasets.rename_column("""label"""			,  """labels"""						)

       def collate_fn(lowerCAmelCase_						):
              # When using mixed precision we want round multiples of 8/16
              if accelerator.mixed_precision == "fp8":
                     _UpperCAmelCase    :		int     	=					16
              elif accelerator.mixed_precision != "no":
                     _UpperCAmelCase    :		Tuple     	=					8
              else:
                     _UpperCAmelCase    :		str     	=					None

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

       # Instantiate dataloaders.
       _UpperCAmelCase    :		Any     	=					DataLoader(
           tokenized_datasets["""train"""]			,  shuffle=lowerCAmelCase_			,  collate_fn=lowerCAmelCase_			,  batch_size=lowerCAmelCase_						)
       _UpperCAmelCase    :		int     	=					DataLoader(
           tokenized_datasets["""validation"""]			,  shuffle=lowerCAmelCase_			,  collate_fn=lowerCAmelCase_			,  batch_size=lowerCAmelCase_						)

       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

   A_				:					Union[str, Any]							=      mocked_dataloaders  # noqa: F811




def        snake_case_   (	lowerCAmelCase_			,  lowerCAmelCase_						)->						int:





       '''simple docstring'''


       if os.environ.get("""TESTING_MOCKED_DATALOADERS"""			,  lowerCAmelCase_						) == "1":
              _UpperCAmelCase    :		Optional[int]     	=					2
       # Initialize accelerator
       _UpperCAmelCase    :		Tuple     	=					Accelerator(cpu=args.cpu			,  mixed_precision=args.mixed_precision						)
       # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
       _UpperCAmelCase    :		Optional[Any]     	=					config["""lr"""]
       _UpperCAmelCase    :		List[str]     	=					int(config["""num_epochs"""]						)
       _UpperCAmelCase    :		Any     	=					int(config["""seed"""]						)
       _UpperCAmelCase    :		str     	=					int(config["""batch_size"""]						)

       _UpperCAmelCase    :		List[str]     	=					evaluate.load("""glue"""			,  """mrpc"""						)

       # New Code #
       # We use the `find_executable_batch_size` decorator, passing in the desired observed batch size
       # to train on. If a CUDA OOM error occurs, it will retry this loop cutting the batch size in
       # half each time. From this, we can calculate the number of gradient accumulation steps needed
       # and modify the Accelerator object as a result
       @find_executable_batch_size(starting_batch_size=int(lowerCAmelCase_						)						)
       def inner_training_loop(lowerCAmelCase_						):
              # Since we need to modify the outside accelerator object, we need to bring it
              # to the local scope
              nonlocal accelerator

              # We can calculate the number of gradient accumulation steps based on the current
              # batch size vs the starting batch size
              _UpperCAmelCase    :		Tuple     	=					observed_batch_size // batch_size

              # And then set it in the Accelerator directly:
              _UpperCAmelCase    :		Optional[int]     	=					num_gradient_accumulation_steps

              # Next we need to free all of the stored model references in the Accelerator each time
              accelerator.free_memory()

              # And set the seed so our results are reproducable each reset
              set_seed(lowerCAmelCase_						)

              # Instantiate the model (we build the model here so that the seed also control new weights initialization)
              _UpperCAmelCase    :		str     	=					AutoModelForSequenceClassification.from_pretrained("""bert-base-cased"""			,  return_dict=lowerCAmelCase_						)

              # 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).
              _UpperCAmelCase    :		int     	=					model.to(accelerator.device						)

              # Instantiate optimizer
              _UpperCAmelCase    :		List[str]     	=					AdamW(params=model.parameters()			,  lr=lowerCAmelCase_						)
              _UpperCAmelCase	,_UpperCAmelCase    :		Optional[Any]     	=					get_dataloaders(lowerCAmelCase_			,  lowerCAmelCase_						)

              # Instantiate scheduler
              _UpperCAmelCase    :		Optional[Any]     	=					get_linear_schedule_with_warmup(
                  optimizer=lowerCAmelCase_			,  num_warmup_steps=100			,  num_training_steps=(len(lowerCAmelCase_						) * num_epochs)			,  )

              # 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    :		Dict     	=					accelerator.prepare(
                  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_						)

              # Now we train the model
              for epoch in range(lowerCAmelCase_						):
                     model.train()
                     for step, batch in enumerate(lowerCAmelCase_						):
                            # And perform gradient accumulation
                            with accelerator.accumulate(lowerCAmelCase_						):
                                   # We could avoid this line since we set the accelerator with `device_placement=True`.
                                   batch.to(accelerator.device						)
                                   _UpperCAmelCase    :		List[Any]     	=					model(**lowerCAmelCase_						)
                                   _UpperCAmelCase    :		str     	=					outputs.loss
                                   accelerator.backward(lowerCAmelCase_						)
                                   optimizer.step()
                                   lr_scheduler.step()
                                   optimizer.zero_grad()

                     model.eval()
                     for step, batch in enumerate(lowerCAmelCase_						):
                            # We could avoid this line since we set the accelerator with `device_placement=True`.
                            batch.to(accelerator.device						)
                            with torch.no_grad():
                                   _UpperCAmelCase    :		Optional[Any]     	=					model(**lowerCAmelCase_						)
                            _UpperCAmelCase    :		List[Any]     	=					outputs.logits.argmax(dim=-1						)
                            _UpperCAmelCase	,_UpperCAmelCase    :		Tuple     	=					accelerator.gather_for_metrics((predictions, batch["""labels"""])						)
                            metric.add_batch(
                                predictions=lowerCAmelCase_			,  references=lowerCAmelCase_			,  )

                     _UpperCAmelCase    :		Union[str, Any]     	=					metric.compute()
                     # Use accelerator.print to print only on the main process.
                     accelerator.print(F'''epoch {epoch}:'''			,  lowerCAmelCase_						)

    # New Code #
    # And call it at the end with no arguments
    # Note: You could also refactor this outside of your training loop function
       inner_training_loop()




def        snake_case_   (	)->						int:





       '''simple docstring'''


       _UpperCAmelCase    :		Optional[int]     	=					argparse.ArgumentParser(description="""Simple example of training script."""						)
       parser.add_argument(
           """--mixed_precision"""			,  type=lowerCAmelCase_			,  default=lowerCAmelCase_			,  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."""						)
       _UpperCAmelCase    :		Union[str, Any]     	=					parser.parse_args()
       # New Code #
       # We modify the starting batch size to be an observed batch size of 256, to guarentee an initial CUDA OOM
       _UpperCAmelCase    :		List[Any]     	=					{"""lr""": 2e-5, """num_epochs""": 3, """seed""": 42, """batch_size""": 256}
       training_function(lowerCAmelCase_			,  lowerCAmelCase_						)


if __name__ == "__main__":
   main()





import argparse
import logging
import math
import time

import torch

from transformers import TransfoXLCorpus, TransfoXLLMHeadModel


logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO
)
_UpperCAmelCase   :					List[str]    =			logging.getLogger(__name__)


def  A			(   )  ->    List[Any]:
		'''simple docstring'''

		UpperCamelCase            =	argparse.ArgumentParser(description='PyTorch Transformer Language Model'					)
		parser.add_argument('--model_name'       ,			type=lowercase       ,			default='transfo-xl-wt103'       ,			help='pretrained model name'					)
		parser.add_argument(
		    '--split'       ,			type=lowercase       ,			default='test'       ,			choices=['all', 'valid', 'test']       ,			help='which split to evaluate'					)
		parser.add_argument('--batch_size'       ,			type=lowercase       ,			default=10       ,			help='batch size'					)
		parser.add_argument('--tgt_len'       ,			type=lowercase       ,			default=128       ,			help='number of tokens to predict'					)
		parser.add_argument('--ext_len'       ,			type=lowercase       ,			default=0       ,			help='length of the extended context'					)
		parser.add_argument('--mem_len'       ,			type=lowercase       ,			default=1_600       ,			help='length of the retained previous heads'					)
		parser.add_argument('--clamp_len'       ,			type=lowercase       ,			default=1_000       ,			help='max positional embedding index'					)
		parser.add_argument('--no_cuda'       ,			action='store_true'       ,			help='Do not use CUDA even though CUA is available'					)
		parser.add_argument('--work_dir'       ,			type=lowercase       ,			required=lowercase       ,			help='path to the work_dir'					)
		parser.add_argument('--no_log'       ,			action='store_true'       ,			help='do not log the eval result'					)
		parser.add_argument('--same_length'       ,			action='store_true'       ,			help='set same length attention with masking'					)
		parser.add_argument('--server_ip'       ,			type=lowercase       ,			default=''       ,			help='Can be used for distant debugging.'					)
		parser.add_argument('--server_port'       ,			type=lowercase       ,			default=''       ,			help='Can be used for distant debugging.'					)
		UpperCamelCase            =	parser.parse_args()
		assert args.ext_len >= 0, "extended context length must be non-negative"

		if args.server_ip and args.server_port:
				# Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
				import ptvsd

				print('Waiting for debugger attach'					)
				ptvsd.enable_attach(address=(args.server_ip, args.server_port)       ,			redirect_output=lowercase					)
				ptvsd.wait_for_attach()

		UpperCamelCase            =	torch.device('cuda' if torch.cuda.is_available() and not args.no_cuda else 'cpu'					)
		logger.info('device: {}'.format(lowercase					)					)

		# Load a pre-processed dataset
		# You can also build the corpus yourself using TransfoXLCorpus methods
		# The pre-processing involve computing word frequencies to prepare the Adaptive input and SoftMax
		# and tokenizing the dataset
		# The pre-processed corpus is a convertion (using the conversion script )
		UpperCamelCase            =	TransfoXLCorpus.from_pretrained(args.model_name					)

		UpperCamelCase            =	corpus.get_iterator('valid'       ,			args.batch_size       ,			args.tgt_len       ,			device=lowercase       ,			ext_len=args.ext_len					)
		UpperCamelCase            =	corpus.get_iterator('test'       ,			args.batch_size       ,			args.tgt_len       ,			device=lowercase       ,			ext_len=args.ext_len					)

		# Load a pre-trained model
		UpperCamelCase            =	TransfoXLLMHeadModel.from_pretrained(args.model_name					)
		model.to(lowercase					)

		logger.info(
		    'Evaluating with bsz {} tgt_len {} ext_len {} mem_len {} clamp_len {}'.format(
		        args.batch_size       ,			args.tgt_len       ,			args.ext_len       ,			args.mem_len       ,			args.clamp_len					)					)

		model.reset_memory_length(args.mem_len					)
		if args.clamp_len > 0:
				UpperCamelCase            =	args.clamp_len
		if args.same_length:
				UpperCamelCase            =	True

		###############################################################################
		# Evaluation code
		###############################################################################
		def evaluate(lowercase					):
				# Turn on evaluation mode which disables dropout.
				model.eval()
				UpperCamelCase      ,    UpperCamelCase            =	0, 0.0
				UpperCamelCase            =	time.time()
				with torch.no_grad():
						UpperCamelCase            =	None
						for idx, (data, target, seq_len) in enumerate(lowercase					):
								UpperCamelCase            =	model(lowercase       ,			lm_labels=lowercase       ,			mems=lowercase					)
								UpperCamelCase      ,    UpperCamelCase      ,    UpperCamelCase            =	ret
								UpperCamelCase            =	loss.mean()
								total_loss += seq_len * loss.item()
								total_len += seq_len
						UpperCamelCase            =	time.time() - start_time
				logger.info('Time : {:.2f}s, {:.2f}ms/segment'.format(lowercase       ,			1_000 * total_time / (idx + 1)					)					)
				return total_loss / total_len

		# Run on test data.
		if args.split == "all":
				UpperCamelCase            =	evaluate(lowercase					)
				UpperCamelCase            =	evaluate(lowercase					)
		elif args.split == "valid":
				UpperCamelCase            =	evaluate(lowercase					)
				UpperCamelCase            =	None
		elif args.split == "test":
				UpperCamelCase            =	evaluate(lowercase					)
				UpperCamelCase            =	None

		def format_log(lowercase       ,			lowercase					):
				UpperCamelCase            =	'| {0} loss {1:5.2f} | {0} ppl {2:9.3f} '.format(lowercase       ,			lowercase       ,			math.exp(lowercase					)					)
				return log_str

		UpperCamelCase            =	''
		if valid_loss is not None:
				log_str += format_log(lowercase       ,			'valid'					)
		if test_loss is not None:
				log_str += format_log(lowercase       ,			'test'					)

		logger.info('=' * 100					)
		logger.info(lowercase					)
		logger.info('=' * 100					)


if __name__ == "__main__":
					main()
