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


__lowerCamelCase      :			Dict							       =		16
__lowerCamelCase      :			List[Any]							       =		32



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	Accelerator			,					__UpperCamelCase   :	DatasetDict			,					__UpperCamelCase   :	List[int]			,					__UpperCamelCase   :	List[int]			,					__UpperCamelCase   :	int = 16     )							-> str:







  """simple docstring"""





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

  def tokenize_function(__UpperCamelCase   :	List[Any]     ):
    # max_length=None => use the model max length (it's actually the default)
    SCREAMING_SNAKE_CASE__      							=				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():
    SCREAMING_SNAKE_CASE__      							=				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
  SCREAMING_SNAKE_CASE__      							=				tokenized_datasets.rename_column("""label"""			,					"""labels"""     )

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

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

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

  SCREAMING_SNAKE_CASE__      							=				DataLoader(
      tokenized_datasets["""test"""]			,					shuffle=__UpperCamelCase			,					collate_fn=__UpperCamelCase			,					batch_size=__UpperCamelCase     )

  return train_dataloader, eval_dataloader, test_dataloader



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	int			,					__UpperCamelCase   :	List[Any]     )							-> Optional[int]:







  """simple docstring"""





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

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

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

  set_seed(__UpperCamelCase     )

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

    # Instantiate optimizer
    SCREAMING_SNAKE_CASE__      							=				AdamW(params=model.parameters()			,					lr=__UpperCamelCase     )

    # Instantiate scheduler
    SCREAMING_SNAKE_CASE__      							=				get_linear_schedule_with_warmup(
        optimizer=__UpperCamelCase			,					num_warmup_steps=1_00			,					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.
    SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__      							=				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     )
        SCREAMING_SNAKE_CASE__      							=				model(**__UpperCamelCase     )
        SCREAMING_SNAKE_CASE__      							=				outputs.loss
        SCREAMING_SNAKE_CASE__      							=				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():
          SCREAMING_SNAKE_CASE__      							=				model(**__UpperCamelCase     )
        SCREAMING_SNAKE_CASE__      							=				outputs.logits.argmax(dim=-1     )
        SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__      							=				accelerator.gather_for_metrics((predictions, batch["""labels"""])     )
        metric.add_batch(
            predictions=__UpperCamelCase			,					references=__UpperCamelCase			,					)

      SCREAMING_SNAKE_CASE__      							=				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
    SCREAMING_SNAKE_CASE__      							=				[]
    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():
        SCREAMING_SNAKE_CASE__      							=				model(**__UpperCamelCase     )
      SCREAMING_SNAKE_CASE__      							=				outputs.logits
      SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__      							=				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:
  SCREAMING_SNAKE_CASE__      							=				torch.cat(__UpperCamelCase			,					dim=0     )
  SCREAMING_SNAKE_CASE__      							=				torch.stack(__UpperCamelCase			,					dim=0     ).sum(dim=0     ).div(int(args.num_folds     )     ).argmax(dim=-1     )
  SCREAMING_SNAKE_CASE__      							=				metric.compute(predictions=__UpperCamelCase			,					references=__UpperCamelCase     )
  accelerator.print("""Average test metrics from all folds:"""			,					__UpperCamelCase     )



def      __SCREAMING_SNAKE_CASE  ( )							-> Any:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__      							=				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"""     )
  SCREAMING_SNAKE_CASE__      							=				parser.parse_args()
  SCREAMING_SNAKE_CASE__      							=				{"""lr""": 2E-5, """num_epochs""": 3, """seed""": 42, """batch_size""": 16}
  training_function(__UpperCamelCase			,					__UpperCamelCase     )


if __name__ == "__main__":
   main()

from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__lowerCamelCase      :			Optional[int]							       =		{
    '''configuration_whisper''': ['''WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''WhisperConfig''', '''WhisperOnnxConfig'''],
    '''feature_extraction_whisper''': ['''WhisperFeatureExtractor'''],
    '''processing_whisper''': ['''WhisperProcessor'''],
    '''tokenization_whisper''': ['''WhisperTokenizer'''],
}

try:
   if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __lowerCamelCase      :			Optional[Any]							       =		['''WhisperTokenizerFast''']

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __lowerCamelCase      :			Any							       =		[
       '''WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST''',
       '''WhisperForConditionalGeneration''',
       '''WhisperModel''',
       '''WhisperPreTrainedModel''',
       '''WhisperForAudioClassification''',
   ]

try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __lowerCamelCase      :			Tuple							       =		[
       '''TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST''',
       '''TFWhisperForConditionalGeneration''',
       '''TFWhisperModel''',
       '''TFWhisperPreTrainedModel''',
   ]

try:
   if not is_flax_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __lowerCamelCase      :			List[Any]							       =		[
       '''FlaxWhisperForConditionalGeneration''',
       '''FlaxWhisperModel''',
       '''FlaxWhisperPreTrainedModel''',
       '''FlaxWhisperForAudioClassification''',
   ]


if TYPE_CHECKING:
   from .configuration_whisper import WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP, WhisperConfig, WhisperOnnxConfig
   from .feature_extraction_whisper import WhisperFeatureExtractor
   from .processing_whisper import WhisperProcessor
   from .tokenization_whisper import WhisperTokenizer

   try:
      if not is_tokenizers_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_whisper_fast import WhisperTokenizerFast

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_whisper import (
          WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
          WhisperForAudioClassification,
          WhisperForConditionalGeneration,
          WhisperModel,
          WhisperPreTrainedModel,
      )

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_whisper import (
          TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFWhisperForConditionalGeneration,
          TFWhisperModel,
          TFWhisperPreTrainedModel,
      )

   try:
      if not is_flax_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_flax_whisper import (
          FlaxWhisperForAudioClassification,
          FlaxWhisperForConditionalGeneration,
          FlaxWhisperModel,
          FlaxWhisperPreTrainedModel,
      )

else:
   import sys

   __lowerCamelCase      :			Tuple							       =		_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)
