




import argparse
import os

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, DistributedType


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing the checkpointing capability,
# 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__   	=				16
lowerCAmelCase__   	=				32
def 							_UpperCAmelCase       (UpperCamelCase__  :						Accelerator       ,							UpperCamelCase__  :						int = 16		):
       _A   : Any       					=				AutoTokenizer.from_pretrained("bert-base-cased"		)
       _A   : List[str]       					=				load_dataset("glue"       ,							"mrpc"		)

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

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

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

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

       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 							_UpperCAmelCase       (UpperCamelCase__  :						Dict       ,							UpperCamelCase__  :						Dict		):
       # For testing only
       if os.environ.get("TESTING_MOCKED_DATALOADERS"       ,							UpperCamelCase__		) == "1":
              _A   : Union[str, Any]       					=				2
       # Initialize accelerator
       _A   : Dict       					=				Accelerator(cpu=args.cpu       ,							mixed_precision=args.mixed_precision		)
       # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
       _A   : List[str]       					=				config["lr"]
       _A   : int       					=				int(config["num_epochs"]		)
       _A   : str       					=				int(config["seed"]		)
       _A   : Dict       					=				int(config["batch_size"]		)

       # New Code #
       # Parse out whether we are saving every epoch or after a certain number of batches
       if hasattr(args.checkpointing_steps       ,							"isdigit"		):
              if args.checkpointing_steps == "epoch":
                     _A   : int       					=				args.checkpointing_steps
              elif args.checkpointing_steps.isdigit():
                     _A   : List[Any]       					=				int(args.checkpointing_steps		)
              else:
                     raise ValueError(
                         f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."		)
       else:
              _A   : List[str]       					=				None

       set_seed(UpperCamelCase__		)

       _A   : Tuple       					=				get_dataloaders(UpperCamelCase__       ,							UpperCamelCase__		)
       _A   : Any       					=				evaluate.load("glue"       ,							"mrpc"		)

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

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

       # Instantiate optimizer
       _A   : Union[str, Any]       					=				AdamW(params=model.parameters()       ,							lr=UpperCamelCase__		)

       # Instantiate scheduler
       _A   : List[str]       					=				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.
       _A   : Union[str, Any]       					=				accelerator.prepare(
           UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__		)

       # New Code #
       # We need to keep track of how many total steps we have iterated over
       _A   : Tuple       					=				0
       # We also need to keep track of the stating epoch so files are named properly
       _A   : Optional[Any]       					=				0

       # We need to load the checkpoint back in before training here with `load_state`
       # The total number of epochs is adjusted based on where the state is being loaded from,
       # as we assume continuation of the same training script
       if args.resume_from_checkpoint:
              if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
                     accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}"		)
                     accelerator.load_state(args.resume_from_checkpoint		)
                     _A   : int       					=				os.path.basename(args.resume_from_checkpoint		)
              else:
                     # Get the most recent checkpoint
                     _A   : List[str]       					=				[f.name for f in os.scandir(os.getcwd()		) if f.is_dir()]
                     dirs.sort(key=os.path.getctime		)
                     _A   : str       					=				dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
              # Extract `epoch_{i}` or `step_{i}`
              _A   : Optional[int]       					=				os.path.splitext(UpperCamelCase__		)[0]

              if "epoch" in training_difference:
                     _A   : List[Any]       					=				int(training_difference.replace("epoch_"       ,							""		)		) + 1
                     _A   : Dict       					=				None
              else:
                     _A   : Tuple       					=				int(training_difference.replace("step_"       ,							""		)		)
                     _A   : Any       					=				resume_step // len(UpperCamelCase__		)
                     resume_step -= starting_epoch * len(UpperCamelCase__		)

    # Now we train the model
       for epoch in range(UpperCamelCase__       ,							UpperCamelCase__		):
              model.train()
              # New Code #
              if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:
                     # We need to skip steps until we reach the resumed step
                     _A   : int       					=				accelerator.skip_first_batches(UpperCamelCase__       ,							UpperCamelCase__		)
                     overall_step += resume_step
              else:
                     # After the first iteration though, we need to go back to the original dataloader
                     _A   : Any       					=				train_dataloader
              for step, batch in enumerate(UpperCamelCase__		):
                     # We could avoid this line since we set the accelerator with `device_placement=True`.
                     batch.to(accelerator.device		)
                     _A   : int       					=				model(**UpperCamelCase__		)
                     _A   : Union[str, Any]       					=				outputs.loss
                     _A   : Union[str, Any]       					=				loss / gradient_accumulation_steps
                     accelerator.backward(UpperCamelCase__		)
                     if step % gradient_accumulation_steps == 0:
                            optimizer.step()
                            lr_scheduler.step()
                            optimizer.zero_grad()
                     # New Code #
                     overall_step += 1

                     # New Code #
                     # We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`
                     # These are saved to folders named `step_{overall_step}`
                     # Will contain files: "pytorch_model.bin", "optimizer.bin", "scheduler.bin", and "random_states.pkl"
                     # If mixed precision was used, will also save a "scalar.bin" file
                     if isinstance(UpperCamelCase__       ,							UpperCamelCase__		):
                            _A   : Optional[Any]       					=				f"step_{overall_step}"
                            if overall_step % checkpointing_steps == 0:
                                   if args.output_dir is not None:
                                          _A   : Optional[Any]       					=				os.path.join(args.output_dir       ,							UpperCamelCase__		)
                                   accelerator.save_state(UpperCamelCase__		)

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

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

              # New Code #
              # We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`
              # These are saved to folders named `epoch_{epoch}`
              # Will contain files: "pytorch_model.bin", "optimizer.bin", "scheduler.bin", and "random_states.pkl"
              # If mixed precision was used, will also save a "scalar.bin" file
              if checkpointing_steps == "epoch":
                     _A   : Any       					=				f"epoch_{epoch}"
                     if args.output_dir is not None:
                            _A   : Tuple       					=				os.path.join(args.output_dir       ,							UpperCamelCase__		)
                     accelerator.save_state(UpperCamelCase__		)
def 							_UpperCAmelCase       ():
       _A   : Dict       					=				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."		)
       parser.add_argument(
           "--checkpointing_steps"       ,							type=UpperCamelCase__       ,							default=UpperCamelCase__       ,							help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch."       ,							)
       parser.add_argument(
           "--output_dir"       ,							type=UpperCamelCase__       ,							default="."       ,							help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory."       ,							)
       parser.add_argument(
           "--resume_from_checkpoint"       ,							type=UpperCamelCase__       ,							default=UpperCamelCase__       ,							help="If the training should continue from a checkpoint folder."       ,							)
       _A   : Optional[Any]       					=				parser.parse_args()
       _A   : Optional[int]       					=				{"lr": 2E-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
       training_function(UpperCamelCase__       ,							UpperCamelCase__		)


if __name__ == "__main__":
    main()




import importlib
import inspect
from functools import wraps
from typing import TYPE_CHECKING, Optional, Union

from .download.streaming_download_manager import (
    xbasename,
    xdirname,
    xet_parse,
    xexists,
    xgetsize,
    xglob,
    xgzip_open,
    xisdir,
    xisfile,
    xjoin,
    xlistdir,
    xnumpy_load,
    xopen,
    xpandas_read_csv,
    xpandas_read_excel,
    xPath,
    xrelpath,
    xsio_loadmat,
    xsplit,
    xsplitext,
    xwalk,
    xxml_dom_minidom_parse,
)
from .utils.logging import get_logger
from .utils.patching import patch_submodule
from .utils.py_utils import get_imports


lowerCAmelCase__   	=				get_logger(__name__)


if TYPE_CHECKING:
    from .builder import DatasetBuilder
def 							_UpperCAmelCase       (UpperCamelCase__  :						Optional[Any]       ,							UpperCamelCase__  :						Optional[Union[str, bool]] = None		):
       _A   : Dict       					=				importlib.import_module(UpperCamelCase__		)

       # TODO(QL): always update the module to add subsequent new authentication
       if hasattr(UpperCamelCase__       ,							"_patched_for_streaming"		) and module._patched_for_streaming:
              return

       def wrap_auth(UpperCamelCase__  :						Optional[int]		):
              @wraps(UpperCamelCase__		)
              def wrapper(*UpperCamelCase__  :						str       ,							**UpperCamelCase__  :						Optional[int]		):
                     return function(*UpperCamelCase__       ,							use_auth_token=UpperCamelCase__       ,							**UpperCamelCase__		)

              _A   : List[Any]       					=				"wrap_auth"
              return wrapper

       # open files in a streaming fashion
       patch_submodule(UpperCamelCase__       ,							"open"       ,							wrap_auth(UpperCamelCase__		)		).start()
       patch_submodule(UpperCamelCase__       ,							"os.listdir"       ,							wrap_auth(UpperCamelCase__		)		).start()
       patch_submodule(UpperCamelCase__       ,							"os.walk"       ,							wrap_auth(UpperCamelCase__		)		).start()
       patch_submodule(UpperCamelCase__       ,							"glob.glob"       ,							wrap_auth(UpperCamelCase__		)		).start()
       # allow to navigate in remote zip files
       patch_submodule(UpperCamelCase__       ,							"os.path.join"       ,							UpperCamelCase__		).start()
       patch_submodule(UpperCamelCase__       ,							"os.path.dirname"       ,							UpperCamelCase__		).start()
       patch_submodule(UpperCamelCase__       ,							"os.path.basename"       ,							UpperCamelCase__		).start()
       patch_submodule(UpperCamelCase__       ,							"os.path.relpath"       ,							UpperCamelCase__		).start()
       patch_submodule(UpperCamelCase__       ,							"os.path.split"       ,							UpperCamelCase__		).start()
       patch_submodule(UpperCamelCase__       ,							"os.path.splitext"       ,							UpperCamelCase__		).start()
       # allow checks on paths
       patch_submodule(UpperCamelCase__       ,							"os.path.exists"       ,							wrap_auth(UpperCamelCase__		)		).start()
       patch_submodule(UpperCamelCase__       ,							"os.path.isdir"       ,							wrap_auth(UpperCamelCase__		)		).start()
       patch_submodule(UpperCamelCase__       ,							"os.path.isfile"       ,							wrap_auth(UpperCamelCase__		)		).start()
       patch_submodule(UpperCamelCase__       ,							"os.path.getsize"       ,							wrap_auth(UpperCamelCase__		)		).start()
       patch_submodule(UpperCamelCase__       ,							"pathlib.Path"       ,							UpperCamelCase__		).start()
       # file readers
       patch_submodule(UpperCamelCase__       ,							"gzip.open"       ,							wrap_auth(UpperCamelCase__		)		).start()
       patch_submodule(UpperCamelCase__       ,							"numpy.load"       ,							wrap_auth(UpperCamelCase__		)		).start()
       patch_submodule(UpperCamelCase__       ,							"pandas.read_csv"       ,							wrap_auth(UpperCamelCase__		)       ,							attrs=["__version__"]		).start()
       patch_submodule(UpperCamelCase__       ,							"pandas.read_excel"       ,							wrap_auth(UpperCamelCase__		)       ,							attrs=["__version__"]		).start()
       patch_submodule(UpperCamelCase__       ,							"scipy.io.loadmat"       ,							wrap_auth(UpperCamelCase__		)       ,							attrs=["__version__"]		).start()
       patch_submodule(UpperCamelCase__       ,							"xml.etree.ElementTree.parse"       ,							wrap_auth(UpperCamelCase__		)		).start()
       patch_submodule(UpperCamelCase__       ,							"xml.dom.minidom.parse"       ,							wrap_auth(UpperCamelCase__		)		).start()
       _A   : Any       					=				True






def 							_UpperCAmelCase       (UpperCamelCase__  :						"DatasetBuilder"		):
       extend_module_for_streaming(builder.__module__       ,							use_auth_token=builder.use_auth_token		)
       # if needed, we also have to extend additional internal imports (like wmt14 -> wmt_utils)
       if not builder.__module__.startswith("datasets."		):  # check that it's not a packaged builder like csv
              for imports in get_imports(inspect.getfile(builder.__class__		)		):
                     if imports[0] == "internal":
                            _A   : Optional[Any]       					=				imports[1]
                            _A   : int       					=				".".join(builder.__module__.split("."		)[:-1] + [internal_import_name]		)
                            extend_module_for_streaming(UpperCamelCase__       ,							use_auth_token=builder.use_auth_token		)

    # builders can inherit from other builders that might use streaming functionality
    # (for example, ImageFolder and AudioFolder inherit from FolderBuilder which implements examples generation)
    # but these parents builders are not patched automatically as they are not instantiated, so we patch them here
       from .builder import DatasetBuilder

       _A   : Tuple       					=				[
           cls.__module__
           for cls in type(UpperCamelCase__		).__mro__[1:]  # make sure it's not the same module we've already patched
           if issubclass(cls       ,							UpperCamelCase__		) and cls.__module__ != DatasetBuilder.__module__
       ]  # check it's not a standard builder from datasets.builder
       for module in parent_builder_modules:
              extend_module_for_streaming(UpperCamelCase__       ,							use_auth_token=builder.use_auth_token		)
