





'''simple docstring'''




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
from accelerate.local_sgd import LocalSGD


########################################################################
# This is a fully working simple example to use Accelerate
# with LocalSGD, which is a method to synchronize model
# parameters every K batches. It is different, but complementary
# to 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)
#
# 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
#
########################################################################


__SCREAMING_SNAKE_CASE     :List[str]       =				16
__SCREAMING_SNAKE_CASE     :Optional[Any]       =				32
def 	UpperCAmelCase_							(					__lowercase							:      Accelerator			,     __lowercase							:      int = 16       )		->       Dict:


     '''simple docstring'''



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

     def tokenize_function(__lowercase							:      Optional[int]       ):
          # 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(__lowercase							:      Dict       ):
          # 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

  __SCREAMING_SNAKE_CASE     :Union[str, Any]       =				mocked_dataloaders  # noqa: F811
def 	UpperCAmelCase_							(					__lowercase							:      Dict			,     __lowercase							:      Any       )		->       Tuple:


     '''simple docstring'''



     if os.environ.get("TESTING_MOCKED_DATALOADERS"			,     _lowerCamelCase       ) == "1":
          _UpperCAmelCase     							=   2
     # New Code #
     _UpperCAmelCase     							=   int(args.gradient_accumulation_steps       )
     _UpperCAmelCase     							=   int(args.local_sgd_steps       )
     # Initialize accelerator
     _UpperCAmelCase     							=   Accelerator(
         cpu=args.cpu			,     mixed_precision=args.mixed_precision			,     gradient_accumulation_steps=_lowerCamelCase       )
     if accelerator.distributed_type not in [DistributedType.NO, DistributedType.MULTI_CPU, DistributedType.MULTI_GPU]:
          raise NotImplementedError("LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)"       )
     # 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"       )

     set_seed(_lowerCamelCase       )
     _UpperCAmelCase							,  _UpperCAmelCase     							=   get_dataloaders(_lowerCamelCase			,     _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       )

     # 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							,  _UpperCAmelCase							,  _UpperCAmelCase							,  _UpperCAmelCase							,  _UpperCAmelCase     							=   accelerator.prepare(
         _lowerCamelCase			,     _lowerCamelCase			,     _lowerCamelCase			,     _lowerCamelCase			,     _lowerCamelCase       )

     # Now we train the model
     for epoch in range(_lowerCamelCase       ):
          model.train()
          with LocalSGD(
              accelerator=_lowerCamelCase			,     model=_lowerCamelCase			,     local_sgd_steps=_lowerCamelCase			,     enabled=local_sgd_steps is not None       ) as local_sgd:
               for step, batch in enumerate(_lowerCamelCase       ):
                    # We could avoid this line since we set the accelerator with `device_placement=True`.
                    batch.to(accelerator.device       )
                    # New code #
                    # We use the new `accumulate` context manager to perform gradient accumulation
                    # We also currently do not support TPUs nor advise it as bugs were found on the XLA side when running our tests.
                    with accelerator.accumulate(_lowerCamelCase       ):
                         _UpperCAmelCase     							=   model(**_lowerCamelCase       )
                         _UpperCAmelCase     							=   output.loss
                         accelerator.backward(_lowerCamelCase       )
                         optimizer.step()
                         lr_scheduler.step()
                         optimizer.zero_grad()
                         # LocalSGD-specific line
                         local_sgd.step()

          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							,  _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       )
def 	UpperCAmelCase_							(					)		->       Optional[int]:


     '''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."			,     )
     # New Code #
     parser.add_argument(
         "--gradient_accumulation_steps"			,     type=_lowerCamelCase			,     default=1			,     help="The number of minibatches to be ran before gradients are accumulated."			,     )
     parser.add_argument(
         "--local_sgd_steps"			,     type=_lowerCamelCase			,     default=8			,     help="Number of local SGD steps or None to disable local SGD"       )
     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()












"""simple docstring"""





import numpy as np
from nltk.translate import meteor_score

import datasets
from datasets.config import importlib_metadata, version


_SCREAMING_SNAKE_CASE					:				Tuple					    =						version.parse(importlib_metadata.version('''nltk'''))
if NLTK_VERSION >= version.Version('''3.6.4'''):
  from nltk import word_tokenize


_SCREAMING_SNAKE_CASE					:				Tuple					    =						'''\
@inproceedings{banarjee2005,
  title     = {{METEOR}: An Automatic Metric for {MT} Evaluation with Improved Correlation with Human Judgments},
  author    = {Banerjee, Satanjeev  and Lavie, Alon},
  booktitle = {Proceedings of the {ACL} Workshop on Intrinsic and Extrinsic Evaluation Measures for Machine Translation and/or Summarization},
  month     = jun,
  year      = {2005},
  address   = {Ann Arbor, Michigan},
  publisher = {Association for Computational Linguistics},
  url       = {https://www.aclweb.org/anthology/W05-0909},
  pages     = {65--72},
}
'''

_SCREAMING_SNAKE_CASE					:				int					    =						'''\
METEOR, an automatic metric for machine translation evaluation
that is based on a generalized concept of unigram matching between the
machine-produced translation and human-produced reference translations.
Unigrams can be matched based on their surface forms, stemmed forms,
and meanings; furthermore, METEOR can be easily extended to include more
advanced matching strategies. Once all generalized unigram matches
between the two strings have been found, METEOR computes a score for
this matching using a combination of unigram-precision, unigram-recall, and
a measure of fragmentation that is designed to directly capture how
well-ordered the matched words in the machine translation are in relation
to the reference.

METEOR gets an R correlation value of 0.347 with human evaluation on the Arabic
data and 0.331 on the Chinese data. This is shown to be an improvement on
using simply unigram-precision, unigram-recall and their harmonic F1
combination.
'''

_SCREAMING_SNAKE_CASE					:				List[Any]					    =						'''
Computes METEOR score of translated segments against one or more references.
Args:
    predictions: list of predictions to score. Each prediction
        should be a string with tokens separated by spaces.
    references: list of reference for each prediction. Each
        reference should be a string with tokens separated by spaces.
    alpha: Parameter for controlling relative weights of precision and recall. default: 0.9
    beta: Parameter for controlling shape of penalty as a function of fragmentation. default: 3
    gamma: Relative weight assigned to fragmentation penalty. default: 0.5
Returns:
    \'meteor\': meteor score.
Examples:

    >>> meteor = datasets.load_metric(\'meteor\')
    >>> predictions = ["It is a guide to action which ensures that the military always obeys the commands of the party"]
    >>> references = ["It is a guide to action that ensures that the military will forever heed Party commands"]
    >>> results = meteor.compute(predictions=predictions, references=references)
    >>> print(round(results["meteor"], 4))
    0.6944
'''






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






    def    UpperCamelCase       (       self     :							List[str]					)     ->		List[str]:
      return datasets.MetricInfo(
          description=_DESCRIPTION       ,						citation=_CITATION       ,						inputs_description=_KWARGS_DESCRIPTION       ,						features=datasets.Features(
              {
                  'predictions': datasets.Value('string'       ,						id='sequence'					),
                  'references': datasets.Value('string'       ,						id='sequence'					),
              }					)       ,						codebase_urls=['https://github.com/nltk/nltk/blob/develop/nltk/translate/meteor_score.py']       ,						reference_urls=[
              'https://www.nltk.org/api/nltk.translate.html#module-nltk.translate.meteor_score',
              'https://en.wikipedia.org/wiki/METEOR',
          ]       ,						)






    def    UpperCamelCase       (       self     :							Optional[int]       ,						__SCREAMING_SNAKE_CASE     :							Union[str, Any]					)     ->		Dict:
      import nltk

      nltk.download('wordnet'					)
      if NLTK_VERSION >= version.Version('3.6.5'					):
        nltk.download('punkt'					)
      if NLTK_VERSION >= version.Version('3.6.6'					):
        nltk.download('omw-1.4'					)






    def    UpperCamelCase       (       self     :							Optional[int]       ,						__SCREAMING_SNAKE_CASE     :							Any       ,						__SCREAMING_SNAKE_CASE     :							Any       ,						__SCREAMING_SNAKE_CASE     :							List[Any]=0.9       ,						__SCREAMING_SNAKE_CASE     :							Dict=3       ,						__SCREAMING_SNAKE_CASE     :							Union[str, Any]=0.5					)     ->		int:
      if NLTK_VERSION >= version.Version('3.6.5'					):
        lowerCamelCase_         =						[
            meteor_score.single_meteor_score(
                word_tokenize(__SCREAMING_SNAKE_CASE					)       ,						word_tokenize(__SCREAMING_SNAKE_CASE					)       ,						alpha=__SCREAMING_SNAKE_CASE       ,						beta=__SCREAMING_SNAKE_CASE       ,						gamma=__SCREAMING_SNAKE_CASE					)
            for ref, pred in zip(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE					)
        ]
      else:
        lowerCamelCase_         =						[
            meteor_score.single_meteor_score(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE       ,						alpha=__SCREAMING_SNAKE_CASE       ,						beta=__SCREAMING_SNAKE_CASE       ,						gamma=__SCREAMING_SNAKE_CASE					)
            for ref, pred in zip(__SCREAMING_SNAKE_CASE       ,						__SCREAMING_SNAKE_CASE					)
        ]

      return {"meteor": np.mean(__SCREAMING_SNAKE_CASE					)}




