





'''simple docstring'''


import argparse
from typing import List

import evaluate
import numpy as np
import torch
from datasets import DatasetDict, load_dataset

# New Code #
# We'll be using StratifiedKFold for this example
from sklearn.model_selection import StratifiedKFold
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


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to perform Cross Validation,
# and builds off the `nlp_example.py` script.
#
# 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)
#
# To help focus on the differences in the code, building `DataLoaders`
# was refactored into its own function.
# 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__			:					Dict       						=						16
a__			:					int       						=						32
def       __snake_case		( SCREAMING_SNAKE_CASE_ :			Dict			,	SCREAMING_SNAKE_CASE_ :			List[str]			,	SCREAMING_SNAKE_CASE_ :			Tuple			,	SCREAMING_SNAKE_CASE_ :			Optional[int]			,	SCREAMING_SNAKE_CASE_ :			int = 16 )							->     Optional[Any]:

       """simple docstring"""


       UpperCAmelCase  =						AutoTokenizer.from_pretrained('''bert-base-cased''' )
       UpperCAmelCase  =						DatasetDict(
           {
               '''train''': dataset['''train'''].select(SCREAMING_SNAKE_CASE_ ),
               '''validation''': dataset['''train'''].select(SCREAMING_SNAKE_CASE_ ),
               '''test''': dataset['''validation'''],
           } )

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

       def collate_fn(SCREAMING_SNAKE_CASE_ :			Tuple ):
              # On TPU it's best to pad everything to the same length or training will be very slow.
              UpperCAmelCase  =						128 if accelerator.distributed_type == DistributedType.TPU else None
              # When using mixed precision we want round multiples of 8/16
              if accelerator.mixed_precision == "fp8":
                     UpperCAmelCase  =						16
              elif accelerator.mixed_precision != "no":
                     UpperCAmelCase  =						8
              else:
                     UpperCAmelCase  =						None

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

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

       UpperCAmelCase  =						DataLoader(
           tokenized_datasets['''test''']			,	shuffle=SCREAMING_SNAKE_CASE_			,	collate_fn=SCREAMING_SNAKE_CASE_			,	batch_size=SCREAMING_SNAKE_CASE_ )

       return train_dataloader, eval_dataloader, test_dataloader
def       __snake_case		( SCREAMING_SNAKE_CASE_ :			Any			,	SCREAMING_SNAKE_CASE_ :			Optional[Any] )							->     Any:

       """simple docstring"""


       UpperCAmelCase  =						[]
       # Download the dataset
       UpperCAmelCase  =						load_dataset('''glue'''			,	'''mrpc''' )
       # Create our splits
       UpperCAmelCase  =						StratifiedKFold(n_splits=int(args.num_folds ) )
       # Initialize accelerator
       UpperCAmelCase  =						Accelerator(cpu=args.cpu			,	mixed_precision=args.mixed_precision )
       # 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  =						evaluate.load('''glue'''			,	'''mrpc''' )

       # If the batch size is too big we use gradient accumulation
       UpperCAmelCase  =						1
       if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
              UpperCAmelCase  =						batch_size // MAX_GPU_BATCH_SIZE
              UpperCAmelCase  =						MAX_GPU_BATCH_SIZE

       set_seed(SCREAMING_SNAKE_CASE_ )

       # New Code #
       # Create our folds:
       UpperCAmelCase  =						kfold.split(np.zeros(datasets['''train'''].num_rows )			,	datasets['''train''']['''label'''] )
       UpperCAmelCase  =						[]
       # Iterate over them
       for i, (train_idxs, valid_idxs) in enumerate(SCREAMING_SNAKE_CASE_ ):
              UpperCAmelCase,					UpperCAmelCase,					UpperCAmelCase  =						get_fold_dataloaders(
                  SCREAMING_SNAKE_CASE_			,	SCREAMING_SNAKE_CASE_			,	SCREAMING_SNAKE_CASE_			,	SCREAMING_SNAKE_CASE_			,	)
              # Instantiate the model (we build the model here so that the seed also control new weights initialization)
              UpperCAmelCase  =						AutoModelForSequenceClassification.from_pretrained('''bert-base-cased'''			,	return_dict=SCREAMING_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).
              UpperCAmelCase  =						model.to(accelerator.device )

              # Instantiate optimizer
              UpperCAmelCase  =						AdamW(params=model.parameters()			,	lr=SCREAMING_SNAKE_CASE_ )

              # Instantiate scheduler
              UpperCAmelCase  =						get_linear_schedule_with_warmup(
                  optimizer=SCREAMING_SNAKE_CASE_			,	num_warmup_steps=100			,	num_training_steps=(len(SCREAMING_SNAKE_CASE_ ) * num_epochs) // gradient_accumulation_steps			,	)

              # 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(
                  SCREAMING_SNAKE_CASE_			,	SCREAMING_SNAKE_CASE_			,	SCREAMING_SNAKE_CASE_			,	SCREAMING_SNAKE_CASE_			,	SCREAMING_SNAKE_CASE_ )

              # Now we train the model
              for epoch in range(SCREAMING_SNAKE_CASE_ ):
                     model.train()
                     for step, batch in enumerate(SCREAMING_SNAKE_CASE_ ):
                            # We could avoid this line since we set the accelerator with `device_placement=True`.
                            batch.to(accelerator.device )
                            UpperCAmelCase  =						model(**SCREAMING_SNAKE_CASE_ )
                            UpperCAmelCase  =						outputs.loss
                            UpperCAmelCase  =						loss / gradient_accumulation_steps
                            accelerator.backward(SCREAMING_SNAKE_CASE_ )
                            if step % gradient_accumulation_steps == 0:
                                   optimizer.step()
                                   lr_scheduler.step()
                                   optimizer.zero_grad()

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

                     UpperCAmelCase  =						metric.compute()
                     # Use accelerator.print to print only on the main process.
                     accelerator.print(f"epoch {epoch}:"			,	SCREAMING_SNAKE_CASE_ )

              # New Code #
              # We also run predictions on the test set at the very end
              UpperCAmelCase  =						[]
              for step, batch in enumerate(SCREAMING_SNAKE_CASE_ ):
                     # We could avoid this line since we set the accelerator with `device_placement=True`.
                     batch.to(accelerator.device )
                     with torch.no_grad():
                            UpperCAmelCase  =						model(**SCREAMING_SNAKE_CASE_ )
                     UpperCAmelCase  =						outputs.logits
                     UpperCAmelCase,					UpperCAmelCase  =						accelerator.gather_for_metrics((predictions, batch['''labels''']) )
                     fold_predictions.append(predictions.cpu() )
                     if i == 0:
                            # We need all of the test predictions
                            test_references.append(references.cpu() )
        # Use accelerator.print to print only on the main process.
              test_predictions.append(torch.cat(SCREAMING_SNAKE_CASE_			,	dim=0 ) )
              # We now need to release all our memory and get rid of the current model, optimizer, etc
              accelerator.free_memory()
       # New Code #
       # Finally we check the accuracy of our folded results:
       UpperCAmelCase  =						torch.cat(SCREAMING_SNAKE_CASE_			,	dim=0 )
       UpperCAmelCase  =						torch.stack(SCREAMING_SNAKE_CASE_			,	dim=0 ).sum(dim=0 ).div(int(args.num_folds ) ).argmax(dim=-1 )
       UpperCAmelCase  =						metric.compute(predictions=SCREAMING_SNAKE_CASE_			,	references=SCREAMING_SNAKE_CASE_ )
       accelerator.print('''Average test metrics from all folds:'''			,	SCREAMING_SNAKE_CASE_ )
def       __snake_case		( )							->     Optional[Any]:

       """simple docstring"""


       UpperCAmelCase  =						argparse.ArgumentParser(description='''Simple example of training script.''' )
       parser.add_argument(
           '''--mixed_precision'''			,	type=SCREAMING_SNAKE_CASE_			,	default=SCREAMING_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.''' )
       # New Code #
       parser.add_argument('''--num_folds'''			,	type=SCREAMING_SNAKE_CASE_			,	default=3			,	help='''The number of splits to perform across the dataset''' )
       UpperCAmelCase  =						parser.parse_args()
       UpperCAmelCase  =						{'''lr''': 2e-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
       training_function(SCREAMING_SNAKE_CASE_			,	SCREAMING_SNAKE_CASE_ )


if __name__ == "__main__":
    main()








'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available


lowerCAmelCase_     						=				{}

try:
    if not is_sentencepiece_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    lowerCAmelCase_     						=				['''GPTSw3Tokenizer''']


if TYPE_CHECKING:
    try:
        if not is_sentencepiece_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .tokenization_gpt_swa import GPTSwaTokenizer

else:
    import sys

    lowerCAmelCase_     						=				_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)


