





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       :     Optional[int]   =     16
__lowerCamelCase       :     Union[str, Any]   =     32



def 		A_				(     _lowerCAmelCase							,   _lowerCAmelCase							,   _lowerCAmelCase							,   _lowerCAmelCase							,   _lowerCAmelCase = 16     )						-> Any:
 UpperCamelCase	:					Optional[int]       	=   AutoTokenizer.from_pretrained("bert-base-cased"     )
 UpperCamelCase	:					Dict       	=   DatasetDict(
     {
         "train": dataset["train"].select(_lowerCAmelCase     ),
         "validation": dataset["train"].select(_lowerCAmelCase     ),
         "test": dataset["validation"],
     }     )

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

 def collate_fn(_lowerCAmelCase     ):
  # On TPU it's best to pad everything to the same length or training will be very slow.
  UpperCamelCase	:					Dict       	=   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	:					Optional[Any]       	=   16
  elif accelerator.mixed_precision != "no":
   UpperCamelCase	:					List[Any]       	=   8
  else:
   UpperCamelCase	:					Tuple       	=   None

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

 # Instantiate dataloaders.
 UpperCamelCase	:					List[Any]       	=   DataLoader(
     tokenized_datasets["train"]							,   shuffle=_lowerCAmelCase							,   collate_fn=_lowerCAmelCase							,   batch_size=_lowerCAmelCase     )
 UpperCamelCase	:					str       	=   DataLoader(
     tokenized_datasets["validation"]							,   shuffle=_lowerCAmelCase							,   collate_fn=_lowerCAmelCase							,   batch_size=_lowerCAmelCase     )

 UpperCamelCase	:					Dict       	=   DataLoader(
     tokenized_datasets["test"]							,   shuffle=_lowerCAmelCase							,   collate_fn=_lowerCAmelCase							,   batch_size=_lowerCAmelCase     )

 return train_dataloader, eval_dataloader, test_dataloader



def 		A_				(     _lowerCAmelCase							,   _lowerCAmelCase     )						-> Tuple:
 # New Code #
 UpperCamelCase	:					Dict       	=   []
 # Download the dataset
 UpperCamelCase	:					Optional[Any]       	=   load_dataset("glue"							,   "mrpc"     )
 # Create our splits
 UpperCamelCase	:					Any       	=   StratifiedKFold(n_splits=int(args.num_folds     )     )
 # Initialize accelerator
 UpperCamelCase	:					int       	=   Accelerator(cpu=args.cpu							,   mixed_precision=args.mixed_precision     )
 # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
 UpperCamelCase	:					int       	=   config["lr"]
 UpperCamelCase	:					Union[str, Any]       	=   int(config["num_epochs"]     )
 UpperCamelCase	:					List[Any]       	=   int(config["seed"]     )
 UpperCamelCase	:					Optional[int]       	=   int(config["batch_size"]     )

 UpperCamelCase	:					int       	=   evaluate.load("glue"							,   "mrpc"     )

 # If the batch size is too big we use gradient accumulation
 UpperCamelCase	:					List[str]       	=   1
 if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
  UpperCamelCase	:					Union[str, Any]       	=   batch_size // MAX_GPU_BATCH_SIZE
  UpperCamelCase	:					Optional[int]       	=   MAX_GPU_BATCH_SIZE

 set_seed(_lowerCAmelCase     )

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

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

  # Instantiate optimizer
  UpperCamelCase	:					Tuple       	=   AdamW(params=model.parameters()							,   lr=_lowerCAmelCase     )

  # Instantiate scheduler
  UpperCamelCase	:					Union[str, Any]       	=   get_linear_schedule_with_warmup(
      optimizer=_lowerCAmelCase							,   num_warmup_steps=100							,   num_training_steps=(len(_lowerCAmelCase     ) * 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	:					Any       	=   accelerator.prepare(
      _lowerCAmelCase							,   _lowerCAmelCase							,   _lowerCAmelCase							,   _lowerCAmelCase							,   _lowerCAmelCase     )

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

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

   UpperCamelCase	:					Optional[Any]       	=   metric.compute()
   # Use accelerator.print to print only on the main process.
   accelerator.print(F"""epoch {epoch}:"""							,   _lowerCAmelCase     )

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



def 		A_				(     )						-> int:
 UpperCamelCase	:					Optional[Any]       	=   argparse.ArgumentParser(description="Simple example of training script."     )
 parser.add_argument(
     "--mixed_precision"							,   type=_lowerCAmelCase							,   default=_lowerCAmelCase							,   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=_lowerCAmelCase							,   default=3							,   help="The number of splits to perform across the dataset"     )
 UpperCamelCase	:					int       	=   parser.parse_args()
 UpperCamelCase	:					Union[str, Any]       	=   {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
 training_function(_lowerCAmelCase							,   _lowerCAmelCase     )


if __name__ == "__main__":
    main()



import math
from typing import Dict, Iterable, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    get_image_size,
    is_torch_available,
    is_torch_tensor,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_torch_available():
 import torch

if is_vision_available():
 import PIL


_UpperCamelCase 			=		logging.get_logger(__name__)







def 		_lowercase (     lowercase__   ,	lowercase__   ,	lowercase__   ,	lowercase__						):
      def constraint_to_multiple_of(lowercase__   ,	lowercase__   ,	lowercase__=0   ,	lowercase__=None						):
            __lowerCAmelCase					:						int					=					round(val / multiple						) * multiple

            if max_val is not None and x > max_val:
                  __lowerCAmelCase					:						Optional[int]					=					math.floor(val / multiple						) * multiple

            if x < min_val:
                  __lowerCAmelCase					:						Any					=					math.ceil(val / multiple						) * multiple

            return x

      __lowerCAmelCase					:						Dict					=					(output_size, output_size) if isinstance(lowercase__   ,	lowercase__						) else output_size

      __lowerCAmelCase,							__lowerCAmelCase					:						Optional[Any]					=					get_image_size(lowercase__						)
      __lowerCAmelCase,							__lowerCAmelCase					:						int					=					output_size

      # determine new height and width
      __lowerCAmelCase					:						Optional[Any]					=					output_height / input_height
      __lowerCAmelCase					:						List[Any]					=					output_width / input_width

      if keep_aspect_ratio:
            # scale as little as possible
            if abs(1 - scale_width						) < abs(1 - scale_height						):
                  # fit width
                  __lowerCAmelCase					:						str					=					scale_width
            else:
                  # fit height
                  __lowerCAmelCase					:						str					=					scale_height

      __lowerCAmelCase					:						Any					=					constraint_to_multiple_of(scale_height * input_height   ,	multiple=lowercase__						)
      __lowerCAmelCase					:						Union[str, Any]					=					constraint_to_multiple_of(scale_width * input_width   ,	multiple=lowercase__						)

      return (new_height, new_width)







class   __lowercase      (_UpperCAmelCase ):
 _UpperCamelCase			  =  ["""pixel_values"""]
 def __init__(							self      ,							A_ = True      ,							A_ = None      ,							A_ = PILImageResampling.BILINEAR      ,							A_ = False      ,							A_ = 1      ,							A_ = True      ,							A_ = 1 / 255      ,							A_ = True      ,							A_ = None      ,							A_ = None      ,							**A_      ,							)					->None:



       '''simple docstring'''





       super().__init__(**A_					)
       __lowerCAmelCase					:						Union[str, Any]					=					size if size is not None else {'''height''': 384, '''width''': 384}
       __lowerCAmelCase					:						Dict					=					get_size_dict(A_					)
       __lowerCAmelCase					:						Optional[Any]					=					do_resize
       __lowerCAmelCase					:						int					=					size
       __lowerCAmelCase					:						Dict					=					keep_aspect_ratio
       __lowerCAmelCase					:						List[Any]					=					ensure_multiple_of
       __lowerCAmelCase					:						Tuple					=					resample
       __lowerCAmelCase					:						Dict					=					do_rescale
       __lowerCAmelCase					:						Any					=					rescale_factor
       __lowerCAmelCase					:						List[Any]					=					do_normalize
       __lowerCAmelCase					:						Optional[int]					=					image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
       __lowerCAmelCase					:						Optional[int]					=					image_std if image_std is not None else IMAGENET_STANDARD_STD
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_      ,							A_ = False      ,							A_ = 1      ,							A_ = PILImageResampling.BICUBIC      ,							A_ = None      ,							**A_      ,							)					->np.ndarray:



       '''simple docstring'''





       __lowerCAmelCase					:						int					=					get_size_dict(A_					)
       if "height" not in size or "width" not in size:
             raise ValueError(f"""The size dictionary must contain the keys 'height' and 'width'. Got {size.keys()}"""					)
       __lowerCAmelCase					:						Union[str, Any]					=					get_resize_output_image_size(
           A_      ,							output_size=(size['''height'''], size['''width'''])      ,							keep_aspect_ratio=A_      ,							multiple=A_      ,							)
       return resize(A_      ,							size=A_      ,							resample=A_      ,							data_format=A_      ,							**A_					)
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_      ,							A_ = None      ,							**A_      ,							)					->Dict:



       '''simple docstring'''





       return rescale(A_      ,							scale=A_      ,							data_format=A_      ,							**A_					)
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_      ,							A_      ,							A_ = None      ,							**A_      ,							)					->np.ndarray:



       '''simple docstring'''





       return normalize(A_      ,							mean=A_      ,							std=A_      ,							data_format=A_      ,							**A_					)
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_ = None      ,							A_ = None      ,							A_ = None      ,							A_ = None      ,							A_ = None      ,							A_ = None      ,							A_ = None      ,							A_ = None      ,							A_ = None      ,							A_ = None      ,							A_ = None      ,							A_ = ChannelDimension.FIRST      ,							**A_      ,							)					->PIL.Image.Image:



       '''simple docstring'''





       __lowerCAmelCase					:						int					=					do_resize if do_resize is not None else self.do_resize
       __lowerCAmelCase					:						Optional[int]					=					size if size is not None else self.size
       __lowerCAmelCase					:						Union[str, Any]					=					get_size_dict(A_					)
       __lowerCAmelCase					:						List[Any]					=					keep_aspect_ratio if keep_aspect_ratio is not None else self.keep_aspect_ratio
       __lowerCAmelCase					:						Optional[int]					=					ensure_multiple_of if ensure_multiple_of is not None else self.ensure_multiple_of
       __lowerCAmelCase					:						Tuple					=					resample if resample is not None else self.resample
       __lowerCAmelCase					:						Optional[int]					=					do_rescale if do_rescale is not None else self.do_rescale
       __lowerCAmelCase					:						Union[str, Any]					=					rescale_factor if rescale_factor is not None else self.rescale_factor
       __lowerCAmelCase					:						Tuple					=					do_normalize if do_normalize is not None else self.do_normalize
       __lowerCAmelCase					:						str					=					image_mean if image_mean is not None else self.image_mean
       __lowerCAmelCase					:						Optional[Any]					=					image_std if image_std is not None else self.image_std

       __lowerCAmelCase					:						Optional[Any]					=					make_list_of_images(A_					)

       if not valid_images(A_					):
             raise ValueError(
                 '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
                 '''torch.Tensor, tf.Tensor or jax.ndarray.'''					)

       if do_resize and size is None or resample is None:
             raise ValueError('''Size and resample must be specified if do_resize is True.'''					)

       if do_rescale and rescale_factor is None:
             raise ValueError('''Rescale factor must be specified if do_rescale is True.'''					)

       if do_normalize and (image_mean is None or image_std is None):
             raise ValueError('''Image mean and std must be specified if do_normalize is True.'''					)

       # All transformations expect numpy arrays.
       __lowerCAmelCase					:						Any					=					[to_numpy_array(A_					) for image in images]

       if do_resize:
             __lowerCAmelCase					:						Optional[Any]					=					[self.resize(image=A_      ,							size=A_      ,							resample=A_					) for image in images]

       if do_rescale:
             __lowerCAmelCase					:						Tuple					=					[self.rescale(image=A_      ,							scale=A_					) for image in images]

       if do_normalize:
             __lowerCAmelCase					:						str					=					[self.normalize(image=A_      ,							mean=A_      ,							std=A_					) for image in images]

       __lowerCAmelCase					:						Union[str, Any]					=					[to_channel_dimension_format(A_      ,							A_					) for image in images]

       __lowerCAmelCase					:						Dict					=					{'''pixel_values''': images}
       return BatchFeature(data=A_      ,							tensor_type=A_					)





 def 	UpperCamelCase__  (							self      ,							A_      ,							A_ = None					)					->Any:



       '''simple docstring'''





       __lowerCAmelCase					:						Any					=					outputs.logits

       # Resize logits and compute semantic segmentation maps
       if target_sizes is not None:
             if len(A_					) != len(A_					):
                   raise ValueError(
                       '''Make sure that you pass in as many target sizes as the batch dimension of the logits'''					)

             if is_torch_tensor(A_					):
                   __lowerCAmelCase					:						Optional[int]					=					target_sizes.numpy()

             __lowerCAmelCase					:						List[str]					=					[]

             for idx in range(len(A_					)					):
                   __lowerCAmelCase					:						Any					=					torch.nn.functional.interpolate(
                       logits[idx].unsqueeze(dim=0					)      ,							size=target_sizes[idx]      ,							mode='''bilinear'''      ,							align_corners=A_					)
                   __lowerCAmelCase					:						str					=					resized_logits[0].argmax(dim=0					)
                   semantic_segmentation.append(A_					)
       else:
             __lowerCAmelCase					:						Any					=					logits.argmax(dim=1					)
             __lowerCAmelCase					:						List[Any]					=					[semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]					)]

       return semantic_segmentation


