


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


UpperCAmelCase__					=16
UpperCAmelCase__					=32


def lowerCAmelCase_   (					UpperCamelCase__						:			Accelerator					, UpperCamelCase__						:			DatasetDict					, UpperCamelCase__						:			List[int]					, UpperCamelCase__						:			List[int]					, UpperCamelCase__						:			int = 16			):






 """simple docstring"""







 __lowercase						=						AutoTokenizer.from_pretrained("""bert-base-cased"""			)
 __lowercase						=						DatasetDict(
     {
         """train""": dataset["""train"""].select(UpperCamelCase__			),
         """validation""": dataset["""train"""].select(UpperCamelCase__			),
         """test""": dataset["""validation"""],
     }			)

 def tokenize_function(UpperCamelCase__						:			Optional[int]			):
  # max_length=None => use the model max length (it's actually the default)
  __lowercase						=						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
 # starting with the main process first:
 with accelerator.main_process_first():
  __lowercase						=						datasets.map(
      UpperCamelCase__					, batched=UpperCamelCase__					, 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
 __lowercase						=						tokenized_datasets.rename_column("""label"""					, """labels"""			)

 def collate_fn(UpperCamelCase__						:			str			):
  # On TPU it's best to pad everything to the same length or training will be very slow.
  __lowercase						=						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":
   __lowercase						=						16
  elif accelerator.mixed_precision != "no":
   __lowercase						=						8
  else:
   __lowercase						=						None

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

 # Instantiate dataloaders.
 __lowercase						=						DataLoader(
     tokenized_datasets["""train"""]					, shuffle=UpperCamelCase__					, collate_fn=UpperCamelCase__					, batch_size=UpperCamelCase__			)
 __lowercase						=						DataLoader(
     tokenized_datasets["""validation"""]					, shuffle=UpperCamelCase__					, collate_fn=UpperCamelCase__					, batch_size=UpperCamelCase__			)

 __lowercase						=						DataLoader(
     tokenized_datasets["""test"""]					, shuffle=UpperCamelCase__					, collate_fn=UpperCamelCase__					, batch_size=UpperCamelCase__			)

 return train_dataloader, eval_dataloader, test_dataloader


def lowerCAmelCase_   (					UpperCamelCase__						:			int					, UpperCamelCase__						:			List[Any]			):






 """simple docstring"""







 __lowercase						=						[]
 # Download the dataset
 __lowercase						=						load_dataset("""glue"""					, """mrpc"""			)
 # Create our splits
 __lowercase						=						StratifiedKFold(n_splits=int(args.num_folds			)			)
 # Initialize accelerator
 __lowercase						=						Accelerator(cpu=args.cpu					, mixed_precision=args.mixed_precision			)
 # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
 __lowercase						=						config["""lr"""]
 __lowercase						=						int(config["""num_epochs"""]			)
 __lowercase						=						int(config["""seed"""]			)
 __lowercase						=						int(config["""batch_size"""]			)

 __lowercase						=						evaluate.load("""glue"""					, """mrpc"""			)

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

 set_seed(UpperCamelCase__			)

 # New Code #
 # Create our folds:
 __lowercase						=						kfold.split(np.zeros(datasets["""train"""].num_rows			)					, datasets["""train"""]["""label"""]			)
 __lowercase						=						[]
 # Iterate over them
 for i, (train_idxs, valid_idxs) in enumerate(UpperCamelCase__			):
  __lowercase				,		__lowercase				,		__lowercase						=						get_fold_dataloaders(
      UpperCamelCase__					, UpperCamelCase__					, UpperCamelCase__					, UpperCamelCase__					, )
  # Instantiate the model (we build the model here so that the seed also control new weights initialization)
  __lowercase						=						AutoModelForSequenceClassification.from_pretrained("""bert-base-cased"""					, return_dict=UpperCamelCase__			)

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

  # Instantiate optimizer
  __lowercase						=						AdamW(params=model.parameters()					, lr=UpperCamelCase__			)

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

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

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

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

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


def lowerCAmelCase_   (					):






 """simple docstring"""







 __lowercase						=						argparse.ArgumentParser(description="""Simple example of training script."""			)
 parser.add_argument(
     """--mixed_precision"""					, type=UpperCamelCase__					, default=UpperCamelCase__					, 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=UpperCamelCase__					, default=3					, help="""The number of splits to perform across the dataset"""			)
 __lowercase						=						parser.parse_args()
 __lowercase						=						{"""lr""": 2E-5, """num_epochs""": 3, """seed""": 42, """batch_size""": 16}
 training_function(UpperCamelCase__					, UpperCamelCase__			)


if __name__ == "__main__":
       main()




"""simple docstring"""

import copy
from typing import TYPE_CHECKING, Any, Mapping, Optional, OrderedDict

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging
from ..auto.configuration_auto import AutoConfig


if TYPE_CHECKING:
       from ... import PreTrainedTokenizerBase, TensorType

UpperCAmelCase__					=logging.get_logger(__name__)





class  lowerCamelCase__    (		_a     ):
      a			:  Optional[int]           =		"""vision-encoder-decoder"""
      a			:  Dict           =		True
      def __init__(		self						:						Dict   ,							**A_						:						List[Any]				):






       '''simple docstring'''


       super().__init__(**A_				)
       if "encoder" not in kwargs or "decoder" not in kwargs:
        raise ValueError(
            F'''A configuraton of type {self.model_type} cannot be instantiated because '''
            F'''not both `encoder` and `decoder` sub-configurations are passed, but only {kwargs}'''				)

       __lowercase						=						kwargs.pop("""encoder"""				)
       __lowercase						=						encoder_config.pop("""model_type"""				)
       __lowercase						=						kwargs.pop("""decoder"""				)
       __lowercase						=						decoder_config.pop("""model_type"""				)

       __lowercase						=						AutoConfig.for_model(A_   ,							**A_				)
       __lowercase						=						AutoConfig.for_model(A_   ,							**A_				)
       __lowercase						=						True
      @classmethod
      def SCREAMING_SNAKE_CASE_						(		cls						:						int   ,							A_						:						PretrainedConfig   ,							A_						:						PretrainedConfig   ,							**A_						:						Any				):






       '''simple docstring'''


       logger.info("""Setting `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config"""				)
       __lowercase						=						True
       __lowercase						=						True

       return cls(encoder=encoder_config.to_dict()   ,							decoder=decoder_config.to_dict()   ,							**A_				)






      def SCREAMING_SNAKE_CASE_						(		self						:						Union[str, Any]				):






       '''simple docstring'''


       __lowercase						=						copy.deepcopy(self.__dict__				)
       __lowercase						=						self.encoder.to_dict()
       __lowercase						=						self.decoder.to_dict()
       __lowercase						=						self.__class__.model_type
       return output






class  lowerCamelCase__    (		_a     ):
      a			:  Dict           =		version.parse("""1.11"""     )
      @property
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[int]				):






       '''simple docstring'''


       return OrderedDict(
           [
               ("""pixel_values""", {0: """batch""", 1: """num_channels""", 2: """height""", 3: """width"""}),
           ]				)
      @property
      def SCREAMING_SNAKE_CASE_						(		self						:						Tuple				):






       '''simple docstring'''


       return 1e-4






      @property
      def SCREAMING_SNAKE_CASE_						(		self						:						Any				):






       '''simple docstring'''


       return OrderedDict({"""last_hidden_state""": {0: """batch""", 1: """encoder_sequence"""}}				)






class  lowerCamelCase__    (		_a     ):
      @property
      def SCREAMING_SNAKE_CASE_						(		self						:						Dict				):






       '''simple docstring'''


       __lowercase						=						OrderedDict()
       __lowercase						=						{0: """batch""", 1: """past_decoder_sequence + sequence"""}
       __lowercase						=						{0: """batch""", 1: """past_decoder_sequence + sequence"""}
       __lowercase						=						{0: """batch""", 1: """encoder_sequence"""}

       return common_inputs






      def SCREAMING_SNAKE_CASE_						(		self						:						str   ,							A_						:						"PreTrainedTokenizerBase"   ,							A_						:						int = -1   ,							A_						:						int = -1   ,							A_						:						bool = False   ,							A_						:						Optional["TensorType"] = None   ,							):






       '''simple docstring'''


       import torch

       __lowercase						=						OrderedDict()

       __lowercase						=						super().generate_dummy_inputs(
           A_   ,							batch_size=A_   ,							seq_length=A_   ,							is_pair=A_   ,							framework=A_				)

       __lowercase				,		__lowercase						=						dummy_input["""input_ids"""].shape
       __lowercase						=						(batch, encoder_sequence, self._config.encoder_hidden_size)
       __lowercase						=						dummy_input.pop("""input_ids"""				)
       __lowercase						=						dummy_input.pop("""attention_mask"""				)
       __lowercase						=						torch.zeros(A_				)

       return common_inputs






class  lowerCamelCase__    (		_a     ):
      @property
      def SCREAMING_SNAKE_CASE_						(		self						:						str				):






       '''simple docstring'''


       pass
      def SCREAMING_SNAKE_CASE_						(		self						:						str   ,							A_						:						PretrainedConfig				):






       '''simple docstring'''


       return VisionEncoderDecoderEncoderOnnxConfig(A_				)






      def SCREAMING_SNAKE_CASE_						(		self						:						Tuple   ,							A_						:						PretrainedConfig   ,							A_						:						PretrainedConfig   ,							A_						:						str = "default"				):






       '''simple docstring'''


       __lowercase						=						encoder_config.hidden_size
       return VisionEncoderDecoderDecoderOnnxConfig(A_   ,							A_				)
