import argparse
import os

# New Code #
import evaluate
import torch
from datasets import load_dataset
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
from accelerate.utils import find_executable_batch_size


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to combine both the gradient accumulation
# and automatic batch size finder utilities of Accelerate to perfrom
# automatic gradient accumulation
#
# 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)
#
# 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__							= 32






def       __lowerCamelCase		(						lowerCAmelCase__	,  lowerCAmelCase__ = 1_6	):
       lowerCAmelCase__     	=					AutoTokenizer.from_pretrained('bert-base-cased'	)
       lowerCAmelCase__     	=					load_dataset('glue'	,  'mrpc'	)

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

       def collate_fn(lowerCAmelCase__	):
              # When using mixed precision we want round multiples of 8/16
              if accelerator.mixed_precision == "fp8":
                     lowerCAmelCase__     	=					1_6
              elif accelerator.mixed_precision != "no":
                     lowerCAmelCase__     	=					8
              else:
                     lowerCAmelCase__     	=					None

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

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

       return train_dataloader, eval_dataloader


# For testing only
if os.environ.get('TESTING_MOCKED_DATALOADERS', None) == "1":
    from accelerate.test_utils.training import mocked_dataloaders

    lowerCAmelCase__							= mocked_dataloaders  # noqa: F811






def       __lowerCamelCase		(						lowerCAmelCase__	,  lowerCAmelCase__	):
       # For testing only
       if os.environ.get('TESTING_MOCKED_DATALOADERS'	,  lowerCAmelCase__	) == "1":
              lowerCAmelCase__     	=					2
       # Initialize accelerator
       lowerCAmelCase__     	=					Accelerator(cpu=args.cpu	,  mixed_precision=args.mixed_precision	)
       # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
       lowerCAmelCase__     	=					config['lr']
       lowerCAmelCase__     	=					int(config['num_epochs']	)
       lowerCAmelCase__     	=					int(config['seed']	)
       lowerCAmelCase__     	=					int(config['batch_size']	)

       lowerCAmelCase__     	=					evaluate.load('glue'	,  'mrpc'	)

       # New Code #
       # We use the `find_executable_batch_size` decorator, passing in the desired observed batch size
       # to train on. If a CUDA OOM error occurs, it will retry this loop cutting the batch size in
       # half each time. From this, we can calculate the number of gradient accumulation steps needed
       # and modify the Accelerator object as a result
       @find_executable_batch_size(starting_batch_size=int(lowerCAmelCase__	)	)
       def inner_training_loop(lowerCAmelCase__	):
              # Since we need to modify the outside accelerator object, we need to bring it
              # to the local scope
              nonlocal accelerator

              # We can calculate the number of gradient accumulation steps based on the current
              # batch size vs the starting batch size
              lowerCAmelCase__     	=					observed_batch_size // batch_size

              # And then set it in the Accelerator directly:
              lowerCAmelCase__     	=					num_gradient_accumulation_steps

              # Next we need to free all of the stored model references in the Accelerator each time
              accelerator.free_memory()

              # And set the seed so our results are reproducable each reset
              set_seed(lowerCAmelCase__	)

              # Instantiate the model (we build the model here so that the seed also control new weights initialization)
              lowerCAmelCase__     	=					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).
              lowerCAmelCase__     	=					model.to(accelerator.device	)

              # Instantiate optimizer
              lowerCAmelCase__     	=					AdamW(params=model.parameters()	,  lr=lowerCAmelCase__	)
              lowerCAmelCase__	,					lowerCAmelCase__     	=					get_dataloaders(lowerCAmelCase__	,  lowerCAmelCase__	)

              # Instantiate scheduler
              lowerCAmelCase__     	=					get_linear_schedule_with_warmup(
                  optimizer=lowerCAmelCase__	,  num_warmup_steps=1_0_0	,  num_training_steps=(len(lowerCAmelCase__	) * num_epochs)	,  )

              # 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.
              lowerCAmelCase__	,					lowerCAmelCase__	,					lowerCAmelCase__	,					lowerCAmelCase__	,					lowerCAmelCase__     	=					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__	):
                            # And perform gradient accumulation
                            with accelerator.accumulate(lowerCAmelCase__	):
                                   # We could avoid this line since we set the accelerator with `device_placement=True`.
                                   batch.to(accelerator.device	)
                                   lowerCAmelCase__     	=					model(**lowerCAmelCase__	)
                                   lowerCAmelCase__     	=					outputs.loss
                                   accelerator.backward(lowerCAmelCase__	)
                                   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():
                                   lowerCAmelCase__     	=					model(**lowerCAmelCase__	)
                            lowerCAmelCase__     	=					outputs.logits.argmax(dim=-1	)
                            lowerCAmelCase__	,					lowerCAmelCase__     	=					accelerator.gather_for_metrics((predictions, batch['labels'])	)
                            metric.add_batch(
                                predictions=lowerCAmelCase__	,  references=lowerCAmelCase__	,  )

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

    # New Code #
    # And call it at the end with no arguments
    # Note: You could also refactor this outside of your training loop function
       inner_training_loop()






def       __lowerCamelCase		(						):
       lowerCAmelCase__     	=					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.'	)
       lowerCAmelCase__     	=					parser.parse_args()
       # New Code #
       # We modify the starting batch size to be an observed batch size of 256, to guarentee an initial CUDA OOM
       lowerCAmelCase__     	=					{'lr': 2e-5, 'num_epochs': 3, 'seed': 4_2, 'batch_size': 2_5_6}
       training_function(lowerCAmelCase__	,  lowerCAmelCase__	)


if __name__ == "__main__":
    main()

import itertools
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa

import datasets
from datasets.table import table_cast


lowerCAmelCase__							= datasets.utils.logging.get_logger(__name__)
@dataclass
class 				a_ (						datasets.BuilderConfig						):






 '''simple docstring'''






 UpperCAmelCase_	 =					None






class 				a_ (						datasets.ArrowBasedBuilder						):






 '''simple docstring'''






 UpperCAmelCase_	 =					ArrowConfig




 def 		__snake_case						(			self					:  Any):







        '''simple docstring'''



        return datasets.DatasetInfo(features=self.config.features)




 def 		__snake_case						(			self					:  Dict			,    lowercase__					:  str):







        '''simple docstring'''



        if not self.config.data_files:
               raise ValueError(F"""At least one data file must be specified, but got data_files={self.config.data_files}""")
        lowerCAmelCase__     	=					dl_manager.download_and_extract(self.config.data_files)
        if isinstance(lowercase__			,    (str, list, tuple)):
               lowerCAmelCase__     	=					data_files
               if isinstance(lowercase__			,    lowercase__):
                      lowerCAmelCase__     	=					[files]
               # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
               lowerCAmelCase__     	=					[dl_manager.iter_files(lowercase__) for file in files]
               return [datasets.SplitGenerator(name=datasets.Split.TRAIN			,    gen_kwargs={'files': files})]
        lowerCAmelCase__     	=					[]
        for split_name, files in data_files.items():
               if isinstance(lowercase__			,    lowercase__):
                      lowerCAmelCase__     	=					[files]
               # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
               lowerCAmelCase__     	=					[dl_manager.iter_files(lowercase__) for file in files]
               # Infer features is they are stoed in the arrow schema
               if self.info.features is None:
                      for file in itertools.chain.from_iterable(lowercase__):
                             with open(lowercase__			,    'rb') as f:
                                    lowerCAmelCase__     	=					datasets.Features.from_arrow_schema(pa.ipc.open_stream(lowercase__).schema)
                             break
               splits.append(datasets.SplitGenerator(name=lowercase__			,    gen_kwargs={'files': files}))
        return splits




 def 		__snake_case						(			self					:  Dict			,    lowercase__					:  pa.Table):







        '''simple docstring'''



        if self.info.features is not None:
               # more expensive cast to support nested features with keys in a different order
               # allows str <-> int/float or str to Audio for example
               lowerCAmelCase__     	=					table_cast(lowercase__			,    self.info.features.arrow_schema)
        return pa_table




 def 		__snake_case						(			self					:  List[Any]			,    lowercase__					:  str):







        '''simple docstring'''



        for file_idx, file in enumerate(itertools.chain.from_iterable(lowercase__)):
               with open(lowercase__			,    'rb') as f:
                      try:
                             for batch_idx, record_batch in enumerate(pa.ipc.open_stream(lowercase__)):
                                    lowerCAmelCase__     	=					pa.Table.from_batches([record_batch])
                                    # Uncomment for debugging (will print the Arrow table size and elements)
                                    # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                                    # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
                                    yield F"""{file_idx}_{batch_idx}""", self._cast_table(lowercase__)
                      except ValueError as e:
                             logger.error(F"""Failed to read file '{file}' with error {type(lowercase__)}: {e}""")
                             raise
