





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


_SCREAMING_SNAKE_CASE					=    1_6
_SCREAMING_SNAKE_CASE					=    3_2




def    _lowerCAmelCase						(				lowerCamelCase_			:  Accelerator		,       lowerCamelCase_			:  int = 1_6				):
   __lowercase										=							AutoTokenizer.from_pretrained('''bert-base-cased'''				)
   __lowercase										=							load_dataset('''glue'''		,       '''mrpc'''				)

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

   def collate_fn(lowerCamelCase_			:  int				):
      # On TPU it's best to pad everything to the same length or training will be very slow.
      __lowercase										=							1_2_8 if accelerator.distributed_type == DistributedType.TPU else None
      # When using mixed precision we want round multiples of 8/16
      if accelerator.mixed_precision == "fp8":
         __lowercase										=							1_6
      elif accelerator.mixed_precision != "no":
         __lowercase										=							8
      else:
         __lowercase										=							None

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

   # Instantiate dataloaders.
   __lowercase										=							DataLoader(
       tokenized_datasets['''train''']		,       shuffle=lowerCAmelCase__		,       collate_fn=lowerCAmelCase__		,       batch_size=lowerCAmelCase__				)
   __lowercase										=							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					=    mocked_dataloaders  # noqa: F811




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

   __lowercase										=							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_			:  List[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)
      __lowercase										=							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).
      __lowercase										=							model.to(accelerator.device				)

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

      # Instantiate scheduler
      __lowercase										=							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.
      __lowercase						,		__lowercase						,		__lowercase						,		__lowercase						,		__lowercase										=							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				)
            __lowercase										=							model(**lowerCAmelCase__				)
            __lowercase										=							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():
               __lowercase										=							model(**lowerCAmelCase__				)
            __lowercase										=							outputs.logits.argmax(dim=-1				)
            __lowercase						,		__lowercase										=							accelerator.gather_for_metrics((predictions, batch['''labels'''])				)
            metric.add_batch(
                predictions=lowerCAmelCase__		,       references=lowerCAmelCase__		,       )

         __lowercase										=							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						(				):
   __lowercase										=							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.'''				)
   __lowercase										=							parser.parse_args()
   __lowercase										=							{'''lr''': 2E-5, '''num_epochs''': 3, '''seed''': 4_2, '''batch_size''': 1_6}
   training_function(lowerCAmelCase__		,       lowerCAmelCase__				)


if __name__ == "__main__":
 main()












'''simple docstring'''



from scipy.stats import pearsonr, spearmanr
from sklearn.metrics import fa_score, matthews_corrcoef

import datasets


_SCREAMING_SNAKE_CASE					=    '''\
@inproceedings{wang2019glue,
  title={{GLUE}: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding},
  author={Wang, Alex and Singh, Amanpreet and Michael, Julian and Hill, Felix and Levy, Omer and Bowman, Samuel R.},
  note={In the Proceedings of ICLR.},
  year={2019}
}
'''

_SCREAMING_SNAKE_CASE					=    '''\
GLUE, the General Language Understanding Evaluation benchmark
(https://gluebenchmark.com/) is a collection of resources for training,
evaluating, and analyzing natural language understanding systems.
'''

_SCREAMING_SNAKE_CASE					=    '''
Compute GLUE evaluation metric associated to each GLUE dataset.
Args:
    predictions: list of predictions to score.
        Each translation should be tokenized into a list of tokens.
    references: list of lists of references for each translation.
        Each reference should be tokenized into a list of tokens.
Returns: depending on the GLUE subset, one or several of:
    "accuracy": Accuracy
    "f1": F1 score
    "pearson": Pearson Correlation
    "spearmanr": Spearman Correlation
    "matthews_correlation": Matthew Correlation
Examples:

    >>> glue_metric = datasets.load_metric(\'glue\', \'sst2\')  # \'sst2\' or any of ["mnli", "mnli_mismatched", "mnli_matched", "qnli", "rte", "wnli", "hans"]
    >>> references = [0, 1]
    >>> predictions = [0, 1]
    >>> results = glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'accuracy\': 1.0}

    >>> glue_metric = datasets.load_metric(\'glue\', \'mrpc\')  # \'mrpc\' or \'qqp\'
    >>> references = [0, 1]
    >>> predictions = [0, 1]
    >>> results = glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'accuracy\': 1.0, \'f1\': 1.0}

    >>> glue_metric = datasets.load_metric(\'glue\', \'stsb\')
    >>> references = [0., 1., 2., 3., 4., 5.]
    >>> predictions = [0., 1., 2., 3., 4., 5.]
    >>> results = glue_metric.compute(predictions=predictions, references=references)
    >>> print({"pearson": round(results["pearson"], 2), "spearmanr": round(results["spearmanr"], 2)})
    {\'pearson\': 1.0, \'spearmanr\': 1.0}

    >>> glue_metric = datasets.load_metric(\'glue\', \'cola\')
    >>> references = [0, 1]
    >>> predictions = [0, 1]
    >>> results = glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'matthews_correlation\': 1.0}
'''




def    _lowerCAmelCase						(				lowerCamelCase_			:  Any		,       lowerCamelCase_			:  int				):
			return float((preds == labels).mean()				)




def    _lowerCAmelCase						(				lowerCamelCase_			:  List[str]		,       lowerCamelCase_			:  str				):
			__lowercase										=							simple_accuracy(lowerCamelCase_		,       lowerCamelCase_				)
			__lowercase										=							float(fa_score(y_true=lowerCamelCase_		,       y_pred=lowerCamelCase_				)				)
			return {
			    "accuracy": acc,
			    "f1": fa,
			}




def    _lowerCAmelCase						(				lowerCamelCase_			:  List[Any]		,       lowerCamelCase_			:  Any				):
			__lowercase										=							float(pearsonr(lowerCamelCase_		,       lowerCamelCase_				)[0]				)
			__lowercase										=							float(spearmanr(lowerCamelCase_		,       lowerCamelCase_				)[0]				)
			return {
			    "pearson": pearson_corr,
			    "spearmanr": spearman_corr,
			}






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







	'''simple docstring'''

	def 						_UpperCAmelCase     (self )							->							str:







				'''simple docstring'''





				if self.config_name not in [
				    "sst2",
				    "mnli",
				    "mnli_mismatched",
				    "mnli_matched",
				    "cola",
				    "stsb",
				    "mrpc",
				    "qqp",
				    "qnli",
				    "rte",
				    "wnli",
				    "hans",
				]:
							raise KeyError(
							    '''You should supply a configuration name selected in '''
							    '''["sst2", "mnli", "mnli_mismatched", "mnli_matched", '''
							    '''"cola", "stsb", "mrpc", "qqp", "qnli", "rte", "wnli", "hans"]''' )
				return datasets.MetricInfo(
				    description=_DESCRIPTION     ,citation=_CITATION     ,inputs_description=_KWARGS_DESCRIPTION     ,features=datasets.Features(
				        {
				            '''predictions''': datasets.Value('''int64''' if self.config_name != '''stsb''' else '''float32''' ),
				            '''references''': datasets.Value('''int64''' if self.config_name != '''stsb''' else '''float32''' ),
				        } )     ,codebase_urls=[]     ,reference_urls=[]     ,format='''numpy'''     ,)


	def 						_UpperCAmelCase     (self     ,_lowerCamelCase     ,_lowerCamelCase )							->							Dict:







				'''simple docstring'''





				if self.config_name == "cola":
							return {"matthews_correlation": matthews_corrcoef(_lowerCamelCase     ,_lowerCamelCase )}
				elif self.config_name == "stsb":
							return pearson_and_spearman(_lowerCamelCase     ,_lowerCamelCase )
				elif self.config_name in ["mrpc", "qqp"]:
							return acc_and_fa(_lowerCamelCase     ,_lowerCamelCase )
				elif self.config_name in ["sst2", "mnli", "mnli_mismatched", "mnli_matched", "qnli", "rte", "wnli", "hans"]:
							return {"accuracy": simple_accuracy(_lowerCamelCase     ,_lowerCamelCase )}
				else:
							raise KeyError(
							    '''You should supply a configuration name selected in '''
							    '''["sst2", "mnli", "mnli_mismatched", "mnli_matched", '''
							    '''"cola", "stsb", "mrpc", "qqp", "qnli", "rte", "wnli", "hans"]''' )





