




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, DistributedType
from accelerate.utils import find_executable_batch_size


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to ensure out-of-memory errors never
# interrupt training, 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)
#
# 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[Any]		=   1_6
_lowerCAmelCase						:   Union[str, Any]		=   3_2



def 						lowerCAmelCase		(		_lowerCAmelCase							:   Dict						,							_lowerCAmelCase							:   Union[str, Any] = 16	):




 """simple docstring"""







 UpperCAmelCase__							 =						AutoTokenizer.from_pretrained("bert-base-cased"	)
 UpperCAmelCase__							 =						load_dataset("glue"						,							"mrpc"	)

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

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

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

 # Instantiate dataloaders.
 UpperCAmelCase__							 =						DataLoader(
     tokenized_datasets["train"]						,							shuffle=_lowerCAmelCase						,							collate_fn=_lowerCAmelCase						,							batch_size=_lowerCAmelCase	)
 UpperCAmelCase__							 =						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						:   Optional[int]		=   mocked_dataloaders  # noqa: F811



def 						lowerCAmelCase		(		_lowerCAmelCase							:   Any						,							_lowerCAmelCase							:   List[str]	):




 """simple docstring"""







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

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

 # New Code #
 # We now can define an inner training loop function. It should take a batch size as the only parameter,
 # and build the dataloaders in there.
 # It also gets our decorator
 @find_executable_batch_size(starting_batch_size=_lowerCAmelCase	)
 def inner_training_loop(_lowerCAmelCase							:   Any	):
  # And now just move everything below under this function
  # We need to bring in the Accelerator object from earlier
  nonlocal accelerator
  # And reset all of its attributes that could hold onto any memory:
  accelerator.free_memory()

  # Then we can declare the model, optimizer, and everything else:
  set_seed(_lowerCAmelCase	)

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

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

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

   UpperCAmelCase__							 =						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		(		):




 """simple docstring"""







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


if __name__ == "__main__":
  main()









from typing import TYPE_CHECKING

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


_lowerCAmelCase						:   Union[str, Any]		=   {
    "configuration_bert": ["BERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "BertConfig", "BertOnnxConfig"],
    "tokenization_bert": ["BasicTokenizer", "BertTokenizer", "WordpieceTokenizer"],
}

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase						:   str		=   ["BertTokenizerFast"]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase						:   int		=   [
      "BERT_PRETRAINED_MODEL_ARCHIVE_LIST",
      "BertForMaskedLM",
      "BertForMultipleChoice",
      "BertForNextSentencePrediction",
      "BertForPreTraining",
      "BertForQuestionAnswering",
      "BertForSequenceClassification",
      "BertForTokenClassification",
      "BertLayer",
      "BertLMHeadModel",
      "BertModel",
      "BertPreTrainedModel",
      "load_tf_weights_in_bert",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase						:   Dict		=   [
      "TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST",
      "TFBertEmbeddings",
      "TFBertForMaskedLM",
      "TFBertForMultipleChoice",
      "TFBertForNextSentencePrediction",
      "TFBertForPreTraining",
      "TFBertForQuestionAnswering",
      "TFBertForSequenceClassification",
      "TFBertForTokenClassification",
      "TFBertLMHeadModel",
      "TFBertMainLayer",
      "TFBertModel",
      "TFBertPreTrainedModel",
  ]
try:
  if not is_tensorflow_text_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase						:   str		=   ["TFBertTokenizer"]

try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase						:   Optional[Any]		=   [
      "FlaxBertForCausalLM",
      "FlaxBertForMaskedLM",
      "FlaxBertForMultipleChoice",
      "FlaxBertForNextSentencePrediction",
      "FlaxBertForPreTraining",
      "FlaxBertForQuestionAnswering",
      "FlaxBertForSequenceClassification",
      "FlaxBertForTokenClassification",
      "FlaxBertModel",
      "FlaxBertPreTrainedModel",
  ]

if TYPE_CHECKING:
  from .configuration_bert import BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, BertConfig, BertOnnxConfig
  from .tokenization_bert import BasicTokenizer, BertTokenizer, WordpieceTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_bert_fast import BertTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_bert import (
        BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        BertForMaskedLM,
        BertForMultipleChoice,
        BertForNextSentencePrediction,
        BertForPreTraining,
        BertForQuestionAnswering,
        BertForSequenceClassification,
        BertForTokenClassification,
        BertLayer,
        BertLMHeadModel,
        BertModel,
        BertPreTrainedModel,
        load_tf_weights_in_bert,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_bert import (
        TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFBertEmbeddings,
        TFBertForMaskedLM,
        TFBertForMultipleChoice,
        TFBertForNextSentencePrediction,
        TFBertForPreTraining,
        TFBertForQuestionAnswering,
        TFBertForSequenceClassification,
        TFBertForTokenClassification,
        TFBertLMHeadModel,
        TFBertMainLayer,
        TFBertModel,
        TFBertPreTrainedModel,
    )

  try:
    if not is_tensorflow_text_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_bert_tf import TFBertTokenizer

  try:
    if not is_flax_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_flax_bert import (
        FlaxBertForCausalLM,
        FlaxBertForMaskedLM,
        FlaxBertForMultipleChoice,
        FlaxBertForNextSentencePrediction,
        FlaxBertForPreTraining,
        FlaxBertForQuestionAnswering,
        FlaxBertForSequenceClassification,
        FlaxBertForTokenClassification,
        FlaxBertModel,
        FlaxBertPreTrainedModel,
    )

else:
  import sys

  _lowerCAmelCase						:   Union[str, Any]		=   _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





