


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
#
########################################################################

UpperCamelCase			:					str         =							32
def 				UpperCamelCase_	(					__snake_case      :    Accelerator						,      __snake_case      :    int = 16  )->   List[Any]:




						'''simple docstring'''
						__magic_name__					=      AutoTokenizer.from_pretrained('bert-base-cased'  )
						__magic_name__					=      load_dataset('glue'						,      'mrpc'  )

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

						def collate_fn(__snake_case      :    List[Any]  ):
												# When using mixed precision we want round multiples of 8/16
												if accelerator.mixed_precision == "fp8":
																		__magic_name__					=      16
												elif accelerator.mixed_precision != "no":
																		__magic_name__					=      8
												else:
																		__magic_name__					=      None

												return tokenizer.pad(
												    __snake_case						,      padding='longest'						,      pad_to_multiple_of=__snake_case						,      return_tensors='pt'						,      )

						# Instantiate dataloaders.
						__magic_name__					=      DataLoader(
						    tokenized_datasets['train']						,      shuffle=__snake_case						,      collate_fn=__snake_case						,      batch_size=__snake_case  )
						__magic_name__					=      DataLoader(
						    tokenized_datasets['validation']						,      shuffle=__snake_case						,      collate_fn=__snake_case						,      batch_size=__snake_case  )

						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

					UpperCamelCase			:					List[str]         =							mocked_dataloaders  # noqa: F811
def 				UpperCamelCase_	(					__snake_case      :    Dict						,      __snake_case      :    Optional[int]  )->   int:




						'''simple docstring'''
						# For testing only
						if os.environ.get('TESTING_MOCKED_DATALOADERS'						,      __snake_case  ) == "1":
												__magic_name__					=      2
						# Initialize accelerator
						__magic_name__					=      Accelerator(cpu=args.cpu						,      mixed_precision=args.mixed_precision  )
						# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
						__magic_name__					=      config['lr']
						__magic_name__					=      int(config['num_epochs']  )
						__magic_name__					=      int(config['seed']  )
						__magic_name__					=      int(config['batch_size']  )

						__magic_name__					=      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(__snake_case  )  )
						def inner_training_loop(__snake_case      :    List[Any]  ):
												# 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
												__magic_name__					=      observed_batch_size // batch_size

												# And then set it in the Accelerator directly:
												__magic_name__					=      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(__snake_case  )

												# Instantiate the model (we build the model here so that the seed also control new weights initialization)
												__magic_name__					=      AutoModelForSequenceClassification.from_pretrained('bert-base-cased'						,      return_dict=__snake_case  )

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

												# Instantiate optimizer
												__magic_name__					=      AdamW(params=model.parameters()						,      lr=__snake_case  )
												__magic_name__		,      __magic_name__					=      get_dataloaders(__snake_case						,      __snake_case  )

												# Instantiate scheduler
												__magic_name__					=      get_linear_schedule_with_warmup(
												    optimizer=__snake_case						,      num_warmup_steps=1_00						,      num_training_steps=(len(__snake_case  ) * 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.
												__magic_name__		,      __magic_name__		,      __magic_name__		,      __magic_name__		,      __magic_name__					=      accelerator.prepare(
												    __snake_case						,      __snake_case						,      __snake_case						,      __snake_case						,      __snake_case  )

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

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

																		__magic_name__					=      metric.compute()
																		# Use accelerator.print to print only on the main process.
																		accelerator.print(F'''epoch {epoch}:'''						,      __snake_case  )

    # 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 				UpperCamelCase_	(					)->   Dict:




						'''simple docstring'''
						__magic_name__					=      argparse.ArgumentParser(description='Simple example of training script.'  )
						parser.add_argument(
						    '--mixed_precision'						,      type=__snake_case						,      default=__snake_case						,      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.'  )
						__magic_name__					=      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
						__magic_name__					=      {'lr': 2E-5, 'num_epochs': 3, 'seed': 42, 'batch_size': 2_56}
						training_function(__snake_case						,      __snake_case  )


if __name__ == "__main__":
					main()





import argparse
import logging

import datasets
from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers, processors
from tokenizers.models import Unigram
from tokenizers.trainers import UnigramTrainer

from transformers import AlbertTokenizerFast


UpperCamelCase			:					List[str]         =							logging.getLogger(__name__)
def 				UpperCamelCase_	(					)->   Dict:




						'''simple docstring'''
						__magic_name__					=      argparse.ArgumentParser(description='Train a unigram tokenizer on the wikitext dataset.'  )
						parser.add_argument(
						    '--dataset_name'						,      type=__snake_case						,      default='wikitext'						,      help='Name of the training. Explore datasets at: hf.co/datasets.'						,      )
						parser.add_argument(
						    '--dataset_config'						,      type=__snake_case						,      default='wikitext-103-raw-v1'						,      help='Configuration name of the dataset.'  )
						parser.add_argument(
						    '--batch_size'						,      type=__snake_case						,      default=10_00						,      help='Batch size during training.'						,      )
						parser.add_argument(
						    '--vocab_size'						,      type=__snake_case						,      default=1_00_48						,      help='Size of the desired vocabulary.'						,      )
						parser.add_argument(
						    '--limit'						,      default=__snake_case						,      type=__snake_case						,      help='Limit the number of shards (used for debugging).'						,      )
						parser.add_argument(
						    '--export_to_hub'						,      action='store_true'						,      )

						__magic_name__					=      parser.parse_args()
						return args
def 				UpperCamelCase_	(					__snake_case      :    Any  )->   Optional[int]:




						'''simple docstring'''
						__magic_name__					=      datasets.load_dataset(args.dataset_name						,      args.dataset_config						,      split='train'  )

						if args.limit is not None:
												__magic_name__					=      min(len(__snake_case  )						,      args.limit  )
												__magic_name__					=      dataset.select(range(__snake_case  )  )
												logger.info(F'''Limiting the dataset to {args.limit} entries.'''  )

						def batch_iterator():
												for i in range(0						,      len(__snake_case  )						,      args.batch_size  ):
																		yield dataset[i : i + args.batch_size]["text"]

    # Prepare the tokenizer.
						__magic_name__					=      Tokenizer(Unigram()  )
						__magic_name__					=      normalizers.Sequence([normalizers.Replace('``'						,      '"'  ), normalizers.Replace('\'\''						,      '"'  )]  )
						__magic_name__					=      pre_tokenizers.Metaspace()

						# Prepare the trainer.
						__magic_name__					=      UnigramTrainer(
						    unk_token='<unk>'						,      special_tokens=['[CLS]', '[SEP]', '<unk>', '<pad>', '[MASK]']						,      vocab_size=args.vocab_size						,      )

						logger.info('Training the tokenizer.'  )
						tokenizer.train_from_iterator(batch_iterator()						,      trainer=__snake_case  )
						logger.info('Tokenizer training complete!'  )

						__magic_name__					=      tokenizer.token_to_id('[CLS]'  )
						__magic_name__					=      tokenizer.token_to_id('[SEP]'  )
						__magic_name__					=      processors.TemplateProcessing(
						    single='[CLS]:0 $A:0 [SEP]:0'						,      pair='[CLS]:0 $A:0 [SEP]:0 $B:1 [SEP]:1'						,      special_tokens=[
						        ('[CLS]', cls_token_id),
						        ('[SEP]', sep_token_id),
						    ]						,      )
						__magic_name__					=      decoders.Metaspace()

						if args.export_to_hub:
												logger.info('Exporting the trained tokenzier to Hub.'  )
												__magic_name__					=      AlbertTokenizerFast(tokenizer_object=__snake_case  )
												new_tokenizer.push_to_hub('unigram-tokenizer-dataset'  )


if __name__ == "__main__":
					UpperCamelCase			:					Dict         =							parse_args()
					main(args)

