

"""simple docstring"""




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


__UpperCAmelCase			=					16
__UpperCAmelCase			=					32
def 	lowerCAmelCase    (       __UpperCamelCase				,     __UpperCamelCase = 16   ):
 '''simple docstring'''





 UpperCAmelCase__       :	Optional[int]        =       AutoTokenizer.from_pretrained("""bert-base-cased"""   )
 UpperCAmelCase__       :	int        =       load_dataset("""glue"""				,     """mrpc"""   )

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

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

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

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

      __UpperCAmelCase			=					mocked_dataloaders  # noqa: F811
def 	lowerCAmelCase    (       __UpperCamelCase				,     __UpperCamelCase   ):
 '''simple docstring'''





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

 UpperCAmelCase__       :	Dict        =       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=__UpperCamelCase   )
 def inner_training_loop(__UpperCamelCase   ):
  # 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(__UpperCamelCase   )

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

  # Instantiate optimizer
  UpperCAmelCase__       :	Tuple        =       AdamW(params=model.parameters()				,     lr=__UpperCamelCase   )
  UpperCAmelCase__ ,   UpperCAmelCase__       :	Any        =       get_dataloaders(__UpperCamelCase				,     __UpperCamelCase   )

  # Instantiate scheduler
  UpperCAmelCase__       :	Optional[int]        =       get_linear_schedule_with_warmup(
      optimizer=__UpperCamelCase				,     num_warmup_steps=100				,     num_training_steps=(len(__UpperCamelCase   ) * 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__ ,   UpperCAmelCase__ ,   UpperCAmelCase__ ,   UpperCAmelCase__ ,   UpperCAmelCase__       :	Optional[Any]        =       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   )
    UpperCAmelCase__       :	Dict        =       model(**__UpperCamelCase   )
    UpperCAmelCase__       :	Optional[int]        =       outputs.loss
    accelerator.backward(__UpperCamelCase   )
    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():
     UpperCAmelCase__       :	Dict        =       model(**__UpperCamelCase   )
    UpperCAmelCase__       :	List[Any]        =       outputs.logits.argmax(dim=-1   )
    UpperCAmelCase__ ,   UpperCAmelCase__       :	List[Any]        =       accelerator.gather_for_metrics((predictions, batch["""labels"""])   )
    metric.add_batch(
        predictions=__UpperCamelCase				,     references=__UpperCamelCase				,     )

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

    # 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__       :	Optional[Any]        =       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."""   )
 UpperCAmelCase__       :	Optional[Any]        =       parser.parse_args()
 UpperCAmelCase__       :	Tuple        =       {"""lr""": 2e-5, """num_epochs""": 3, """seed""": 42, """batch_size""": 16}
 training_function(__UpperCamelCase				,     __UpperCamelCase   )


if __name__ == "__main__":
      main()









import datasets


snake_case							:       int										=	'''\
@InProceedings{conneau2018xnli,
  author = "Conneau, Alexis
                 and Rinott, Ruty
                 and Lample, Guillaume
                 and Williams, Adina
                 and Bowman, Samuel R.
                 and Schwenk, Holger
                 and Stoyanov, Veselin",
  title = "XNLI: Evaluating Cross-lingual Sentence Representations",
  booktitle = "Proceedings of the 2018 Conference on Empirical Methods
               in Natural Language Processing",
  year = "2018",
  publisher = "Association for Computational Linguistics",
  location = "Brussels, Belgium",
}
'''

snake_case							:       Tuple										=	'''\
XNLI is a subset of a few thousand examples from MNLI which has been translated
into a 14 different languages (some low-ish resource). As with MNLI, the goal is
to predict textual entailment (does sentence A imply/contradict/neither sentence
B) and is a classification task (given two sentences, predict one of three
labels).
'''

snake_case							:       Union[str, Any]										=	'''
Computes XNLI score which is just simple accuracy.
Args:
    predictions: Predicted labels.
    references: Ground truth labels.
Returns:
    \'accuracy\': accuracy
Examples:

    >>> predictions = [0, 1]
    >>> references = [0, 1]
    >>> xnli_metric = datasets.load_metric("xnli")
    >>> results = xnli_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'accuracy\': 1.0}
'''
def 							__lowerCamelCase		(					UpperCAmelCase_		:						List[Any]					,    UpperCAmelCase_		:						List[str]		):




      """simple docstring"""





      return (preds == labels).mean()







@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION		,       _KWARGS_DESCRIPTION				)
class 						_snake_case    ( datasets.Metric				):




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       return datasets.MetricInfo(
           description=_DESCRIPTION   ,    citation=_CITATION   ,    inputs_description=_KWARGS_DESCRIPTION   ,    features=datasets.Features(
               {
                   '''predictions''': datasets.Value('''int64''' if self.config_name != '''sts-b''' else '''float32'''       ),
                   '''references''': datasets.Value('''int64''' if self.config_name != '''sts-b''' else '''float32'''       ),
               }       )   ,    codebase_urls=[]   ,    reference_urls=[]   ,    format='''numpy'''   ,    )




 def 			SCREAMING_SNAKE_CASE__							(						self   ,    _lowerCamelCase   ,    _lowerCamelCase       ):
       return {"accuracy": simple_accuracy(_lowerCamelCase   ,    _lowerCamelCase       )}





