




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          =			16
__a          =			32







def lowerCamelCase__       (      _lowercase			,						_lowercase			,						_lowercase			,						_lowercase			,						_lowercase = 16				):

     '''simple docstring'''


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

     def tokenize_function(_lowercase				):
          # max_length=None => use the model max length (it's actually the default)
          UpperCAmelCase_       =      tokenizer(examples['''sentence1''']			,						examples['''sentence2''']			,						truncation=__a			,						max_length=__a				)
          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(
              __a			,						batched=__a			,						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(_lowercase				):
          # 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(
              __a			,						padding='''longest'''			,						max_length=__a			,						pad_to_multiple_of=__a			,						return_tensors='''pt'''			,						)

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

     UpperCAmelCase_       =      DataLoader(
         tokenized_datasets['''test''']			,						shuffle=__a			,						collate_fn=__a			,						batch_size=__a				)

     return train_dataloader, eval_dataloader, test_dataloader







def lowerCamelCase__       (      _lowercase			,						_lowercase				):

     '''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(__a				)

     # 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(__a				):
          UpperCAmelCase_       =      get_fold_dataloaders(
              __a			,						__a			,						__a			,						__a			,						)
          # 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=__a				)

          # 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=__a				)

          # Instantiate scheduler
          UpperCAmelCase_       =      get_linear_schedule_with_warmup(
              optimizer=__a			,						num_warmup_steps=100			,						num_training_steps=(len(__a				) * 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_       =      accelerator.prepare(
              __a			,						__a			,						__a			,						__a			,						__a				)

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

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

               UpperCAmelCase_       =      metric.compute()
               # Use accelerator.print to print only on the main process.
               accelerator.print(f'''epoch {epoch}:'''			,						__a				)

          # New Code #
          # We also run predictions on the test set at the very end
          UpperCAmelCase_       =      []
          for step, batch in enumerate(__a				):
               # We could avoid this line since we set the accelerator with `device_placement=True`.
               batch.to(accelerator.device				)
               with torch.no_grad():
                    UpperCAmelCase_       =      model(**__a				)
               UpperCAmelCase_       =      outputs.logits
               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(__a			,						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(__a			,						dim=0				)
     UpperCAmelCase_       =      torch.stack(__a			,						dim=0				).sum(dim=0				).div(int(args.num_folds				)				).argmax(dim=-1				)
     UpperCAmelCase_       =      metric.compute(predictions=__a			,						references=__a				)
     accelerator.print('''Average test metrics from all folds:'''			,						__a				)







def lowerCamelCase__       (      ):

     '''simple docstring'''


     UpperCAmelCase_       =      argparse.ArgumentParser(description='''Simple example of training script.'''				)
     parser.add_argument(
         '''--mixed_precision'''			,						type=__a			,						default=__a			,						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=__a			,						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(__a			,						__a				)


if __name__ == "__main__":
     main()





import argparse
import logging
import sys
from unittest.mock import patch

import run_glue_deebert

from transformers.testing_utils import TestCasePlus, get_gpu_count, require_torch_non_multi_gpu, slow


logging.basicConfig(level=logging.DEBUG)

__a          =			logging.getLogger()







def lowerCamelCase__       (      ):

     '''simple docstring'''


     UpperCAmelCase_      :							List[str]       =      argparse.ArgumentParser()
     parser.add_argument('''-f'''				)
     UpperCAmelCase_      :							Dict       =      parser.parse_args()
     return args.f



class   __a(					_a       ):






    """simple docstring"""





    def 							a__				(				self  )			->		None:
         UpperCAmelCase_      :							int       =      logging.StreamHandler(sys.stdout  )
         logger.addHandler(_SCREAMING_SNAKE_CASE  )





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE  )			->		str:
         UpperCAmelCase_      :							int       =      get_gpu_count()

         if n_gpu > 1:
              pass
              # XXX: doesn't quite work with n_gpu > 1 https://github.com/huggingface/transformers/issues/10560
              # script = f"{self.examples_dir_str}/research_projects/deebert/run_glue_deebert.py"
              # distributed_args = f"-m torch.distributed.launch --nproc_per_node={n_gpu} {script}".split()
              # cmd = [sys.executable] + distributed_args + args
              # execute_subprocess_async(cmd, env=self.get_env())
              # XXX: test the results - need to save them first into .json file
         else:
              args.insert(0       ,'''run_glue_deebert.py'''  )
              with patch.object(_SCREAMING_SNAKE_CASE       ,'''argv'''       ,_SCREAMING_SNAKE_CASE  ):
                   UpperCAmelCase_      :							List[str]       =      run_glue_deebert.main()
                   for value in result.values():
                        self.assertGreaterEqual(_SCREAMING_SNAKE_CASE       ,0.6_66  )





    @slow
    @require_torch_non_multi_gpu
    def 							a__				(				self  )			->		List[str]:
         UpperCAmelCase_      :							List[Any]       =      '''
            --model_type roberta
            --model_name_or_path roberta-base
            --task_name MRPC
            --do_train
            --do_eval
            --do_lower_case
            --data_dir ./tests/fixtures/tests_samples/MRPC/
            --max_seq_length 128
            --per_gpu_eval_batch_size=1
            --per_gpu_train_batch_size=8
            --learning_rate 2e-4
            --num_train_epochs 3
            --overwrite_output_dir
            --seed 42
            --output_dir ./examples/deebert/saved_models/roberta-base/MRPC/two_stage
            --plot_data_dir ./examples/deebert/results/
            --save_steps 0
            --overwrite_cache
            --eval_after_first_stage
            '''.split()
         self.run_and_check(_SCREAMING_SNAKE_CASE  )

         UpperCAmelCase_      :							Union[str, Any]       =      '''
            --model_type roberta
            --model_name_or_path ./examples/deebert/saved_models/roberta-base/MRPC/two_stage
            --task_name MRPC
            --do_eval
            --do_lower_case
            --data_dir ./tests/fixtures/tests_samples/MRPC/
            --output_dir ./examples/deebert/saved_models/roberta-base/MRPC/two_stage
            --plot_data_dir ./examples/deebert/results/
            --max_seq_length 128
            --eval_each_highway
            --eval_highway
            --overwrite_cache
            --per_gpu_eval_batch_size=1
            '''.split()
         self.run_and_check(_SCREAMING_SNAKE_CASE  )

         UpperCAmelCase_      :							Any       =      '''
            --model_type roberta
            --model_name_or_path ./examples/deebert/saved_models/roberta-base/MRPC/two_stage
            --task_name MRPC
            --do_eval
            --do_lower_case
            --data_dir ./tests/fixtures/tests_samples/MRPC/
            --output_dir ./examples/deebert/saved_models/roberta-base/MRPC/two_stage
            --plot_data_dir ./examples/deebert/results/
            --max_seq_length 128
            --early_exit_entropy 0.1
            --eval_highway
            --overwrite_cache
            --per_gpu_eval_batch_size=1
            '''.split()
         self.run_and_check(_SCREAMING_SNAKE_CASE  )