


import argparse
import json
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.utils.deepspeed import DummyOptim, DummyScheduler


__UpperCAmelCase    :Optional[Any]				       =     1_6
__UpperCAmelCase    :List[str]				       =     3_2




def 			_a			(		_lowercase      :			Accelerator    ,      _lowercase      :			int = 16    ,      _lowercase      :			str = "bert-base-cased"						):





     '''simple docstring'''






     __UpperCAmelCase			:				Union[str, Any]													=    AutoTokenizer.from_pretrained(_lowercase						)
     __UpperCAmelCase			:				List[str]													=    load_dataset('''glue'''    ,      '''mrpc'''						)

     def tokenize_function(_lowercase      :			Union[str, Any]						):
          # max_length=None => use the model max length (it's actually the default)
          __UpperCAmelCase			:				List[str]													=    tokenizer(examples['''sentence1''']    ,      examples['''sentence2''']    ,      truncation=_lowercase    ,      max_length=_lowercase						)
          return outputs

     # Apply the method we just defined to all the examples in all the splits of the dataset
     __UpperCAmelCase			:				List[str]													=    datasets.map(
         _lowercase    ,      batched=_lowercase    ,      remove_columns=['''idx''', '''sentence1''', '''sentence2''']    ,      load_from_cache_file=_lowercase						)

     # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
     # transformers library
     __UpperCAmelCase			:				str													=    tokenized_datasets.rename_column('''label'''    ,      '''labels'''						)

     def collate_fn(_lowercase      :			Any						):
          # On TPU it's best to pad everything to the same length or training will be very slow.
          if accelerator.distributed_type == DistributedType.TPU:
               return tokenizer.pad(_lowercase    ,      padding='''max_length'''    ,      max_length=128    ,      return_tensors='''pt'''						)
          return tokenizer.pad(_lowercase    ,      padding='''longest'''    ,      return_tensors='''pt'''						)

     # Instantiate dataloaders.
     __UpperCAmelCase			:				int													=    DataLoader(
         tokenized_datasets['''train''']    ,      shuffle=_lowercase    ,      collate_fn=_lowercase    ,      batch_size=_lowercase						)
     __UpperCAmelCase			:				Dict													=    DataLoader(
         tokenized_datasets['''validation''']    ,      shuffle=_lowercase    ,      collate_fn=_lowercase    ,      batch_size=_lowercase						)

     return train_dataloader, eval_dataloader




def 			_a			(		_lowercase      :			str    ,      _lowercase      :			Any    ,      _lowercase      :			Any    ,      _lowercase      :			Any						):





     '''simple docstring'''






     model.eval()
     __UpperCAmelCase			:				Union[str, Any]													=    0
     for step, batch in enumerate(_lowercase						):
          # 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(**_lowercase						)
          __UpperCAmelCase			:				int													=    outputs.logits.argmax(dim=-1						)
          # It is slightly faster to call this once, than multiple times
          __UpperCAmelCase			:				Optional[int]													=    accelerator.gather(
              (predictions, batch['''labels'''])						)  # If we are in a multiprocess environment, the last batch has duplicates
          if accelerator.use_distributed:
               if step == len(_lowercase						) - 1:
                    __UpperCAmelCase			:				List[str]													=    predictions[: len(eval_dataloader.dataset						) - samples_seen]
                    __UpperCAmelCase			:				Tuple													=    references[: len(eval_dataloader.dataset						) - samples_seen]
               else:
                    samples_seen += references.shape[0]
          metric.add_batch(
              predictions=_lowercase    ,      references=_lowercase    ,      )

     __UpperCAmelCase			:				Union[str, Any]													=    metric.compute()
     return eval_metric["accuracy"]




def 			_a			(		_lowercase      :			Optional[Any]    ,      _lowercase      :			List[Any]						):





     '''simple docstring'''






     __UpperCAmelCase			:				List[str]													=    Accelerator()

     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     __UpperCAmelCase			:				Tuple													=    config['''lr''']
     __UpperCAmelCase			:				Tuple													=    int(config['''num_epochs''']						)
     __UpperCAmelCase			:				Tuple													=    int(config['''seed''']						)
     __UpperCAmelCase			:				Dict													=    int(config['''batch_size''']						)
     __UpperCAmelCase			:				Tuple													=    args.model_name_or_path

     set_seed(_lowercase						)
     __UpperCAmelCase			:				Optional[Any]													=    get_dataloaders(_lowercase    ,      _lowercase    ,      _lowercase						)

     # Instantiate the model (we build the model here so that the seed also control new weights initialization)
     __UpperCAmelCase			:				Dict													=    AutoModelForSequenceClassification.from_pretrained(_lowercase    ,      return_dict=_lowercase						)

     # Instantiate optimizer
     __UpperCAmelCase			:				List[Any]													=    (
         AdamW
         if accelerator.state.deepspeed_plugin is None
         or '''optimizer''' not in accelerator.state.deepspeed_plugin.deepspeed_config
         else DummyOptim
     )
     __UpperCAmelCase			:				Optional[Any]													=    optimizer_cls(params=model.parameters()    ,      lr=_lowercase						)

     if accelerator.state.deepspeed_plugin is not None:
          __UpperCAmelCase			:				Union[str, Any]													=    accelerator.state.deepspeed_plugin.deepspeed_config[
              '''gradient_accumulation_steps'''
          ]
     else:
          __UpperCAmelCase			:				Union[str, Any]													=    1
     __UpperCAmelCase			:				str													=    (len(_lowercase						) * num_epochs) // gradient_accumulation_steps

     # Instantiate scheduler
     if (
         accelerator.state.deepspeed_plugin is None
         or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
     ):
          __UpperCAmelCase			:				Any													=    get_linear_schedule_with_warmup(
              optimizer=_lowercase    ,      num_warmup_steps=0    ,      num_training_steps=_lowercase    ,      )
     else:
          __UpperCAmelCase			:				Optional[Any]													=    DummyScheduler(_lowercase    ,      total_num_steps=_lowercase    ,      warmup_num_steps=0						)

     # 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			:				Optional[Any]													=    accelerator.prepare(
         _lowercase    ,      _lowercase    ,      _lowercase    ,      _lowercase    ,      _lowercase						)

     # We need to keep track of how many total steps we have iterated over
     __UpperCAmelCase			:				Optional[int]													=    0
     # We also need to keep track of the stating epoch so files are named properly
     __UpperCAmelCase			:				Any													=    0
     __UpperCAmelCase			:				int													=    evaluate.load('''glue'''    ,      '''mrpc'''						)
     __UpperCAmelCase			:				List[str]													=    num_epochs

     if args.partial_train_epoch is not None:
          __UpperCAmelCase			:				List[Any]													=    args.partial_train_epoch

     if args.resume_from_checkpoint:
          accelerator.load_state(args.resume_from_checkpoint						)
          __UpperCAmelCase			:				List[Any]													=    args.resume_from_checkpoint.split('''epoch_'''						)[1]
          __UpperCAmelCase			:				str													=    ''''''
          for char in epoch_string:
               if char.isdigit():
                    state_epoch_num += char
               else:
                    break
          __UpperCAmelCase			:				List[str]													=    int(_lowercase						) + 1
          __UpperCAmelCase			:				Any													=    evaluation_loop(_lowercase    ,      _lowercase    ,      _lowercase    ,      _lowercase						)
          accelerator.print('''resumed checkpoint performance:'''    ,      _lowercase						)
          accelerator.print('''resumed checkpoint\'s scheduler\'s lr:'''    ,      lr_scheduler.get_lr()[0]						)
          accelerator.print('''resumed optimizers\'s lr:'''    ,      optimizer.param_groups[0]['''lr''']						)
          with open(os.path.join(args.output_dir    ,      F'state_{starting_epoch-1}.json'						)    ,      '''r'''						) as f:
               __UpperCAmelCase			:				Union[str, Any]													=    json.load(_lowercase						)
               assert resumed_state["accuracy"] == accuracy, "Accuracy mismatch, loading from checkpoint failed"
               assert (
                   resumed_state["lr"] == lr_scheduler.get_lr()[0]
               ), "Scheduler learning rate mismatch, loading from checkpoint failed"
               assert (
                   resumed_state["optimizer_lr"] == optimizer.param_groups[0]["lr"]
               ), "Optimizer learning rate mismatch, loading from checkpoint failed"
               assert resumed_state["epoch"] == starting_epoch - 1, "Epoch mismatch, loading from checkpoint failed"
               return

    # Now we train the model
     __UpperCAmelCase			:				str													=    {}
     for epoch in range(_lowercase    ,      _lowercase						):
          model.train()
          for step, batch in enumerate(_lowercase						):
               __UpperCAmelCase			:				Union[str, Any]													=    model(**_lowercase						)
               __UpperCAmelCase			:				Optional[Any]													=    outputs.loss
               __UpperCAmelCase			:				Tuple													=    loss / gradient_accumulation_steps
               accelerator.backward(_lowercase						)
               if step % gradient_accumulation_steps == 0:
                    optimizer.step()
                    lr_scheduler.step()
                    optimizer.zero_grad()

               overall_step += 1
          __UpperCAmelCase			:				Optional[int]													=    F'epoch_{epoch}'
          __UpperCAmelCase			:				Tuple													=    os.path.join(args.output_dir    ,      _lowercase						)
          accelerator.save_state(_lowercase						)
          __UpperCAmelCase			:				Any													=    evaluation_loop(_lowercase    ,      _lowercase    ,      _lowercase    ,      _lowercase						)
          __UpperCAmelCase			:				Any													=    accuracy
          __UpperCAmelCase			:				Union[str, Any]													=    lr_scheduler.get_lr()[0]
          __UpperCAmelCase			:				Optional[int]													=    optimizer.param_groups[0]['''lr''']
          __UpperCAmelCase			:				Dict													=    epoch
          __UpperCAmelCase			:				Tuple													=    overall_step
          accelerator.print(F'epoch {epoch}:'    ,      _lowercase						)

          accelerator.wait_for_everyone()
          if accelerator.is_main_process:
               with open(os.path.join(args.output_dir    ,      F'state_{epoch}.json'						)    ,      '''w'''						) as f:
                    json.dump(_lowercase    ,      _lowercase						)




def 			_a			(		):





     '''simple docstring'''






     __UpperCAmelCase			:				Any													=    argparse.ArgumentParser(description='''Simple example of training script tracking peak GPU memory usage.'''						)
     parser.add_argument(
         '''--model_name_or_path'''    ,      type=_lowercase    ,      default='''bert-base-cased'''    ,      help='''Path to pretrained model or model identifier from huggingface.co/models.'''    ,      required=_lowercase    ,      )
     parser.add_argument(
         '''--output_dir'''    ,      type=_lowercase    ,      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=_lowercase    ,      default=_lowercase    ,      help='''If the training should continue from a checkpoint folder.'''    ,      )
     parser.add_argument(
         '''--partial_train_epoch'''    ,      type=_lowercase    ,      default=_lowercase    ,      help='''If passed, the training will stop after this number of epochs.'''    ,      )
     parser.add_argument(
         '''--num_epochs'''    ,      type=_lowercase    ,      default=2    ,      help='''Number of train epochs.'''    ,      )
     __UpperCAmelCase			:				str													=    parser.parse_args()
     __UpperCAmelCase			:				Tuple													=    {'''lr''': 2E-5, '''num_epochs''': args.num_epochs, '''seed''': 42, '''batch_size''': 16}

     training_function(_lowercase    ,      _lowercase						)


if __name__ == "__main__":
      main()



'''simple docstring'''



from scipy.stats import pearsonr

import datasets


__UpperCAmelCase    :Tuple				       =     "\nPearson correlation coefficient and p-value for testing non-correlation.\nThe Pearson correlation coefficient measures the linear relationship between two datasets. The calculation of the p-value relies on the assumption that each dataset is normally distributed. Like other correlation coefficients, this one varies between -1 and +1 with 0 implying no correlation. Correlations of -1 or +1 imply an exact linear relationship. Positive correlations imply that as x increases, so does y. Negative correlations imply that as x increases, y decreases.\nThe p-value roughly indicates the probability of an uncorrelated system producing datasets that have a Pearson correlation at least as extreme as the one computed from these datasets.\n"


__UpperCAmelCase    :List[str]				       =     "\nArgs:\n    predictions (`list` of `int`): Predicted class labels, as returned by a model.\n    references (`list` of `int`): Ground truth labels.\n    return_pvalue (`boolean`): If `True`, returns the p-value, along with the correlation coefficient. If `False`, returns only the correlation coefficient. Defaults to `False`.\n\nReturns:\n    pearsonr (`float`): Pearson correlation coefficient. Minimum possible value is -1. Maximum possible value is 1. Values of 1 and -1 indicate exact linear positive and negative relationships, respectively. A value of 0 implies no correlation.\n    p-value (`float`): P-value, which roughly indicates the probability of an The p-value roughly indicates the probability of an uncorrelated system producing datasets that have a Pearson correlation at least as extreme as the one computed from these datasets. Minimum possible value is 0. Maximum possible value is 1. Higher values indicate higher probabilities.\n\nExamples:\n\n    Example 1-A simple example using only predictions and references.\n        >>> pearsonr_metric = datasets.load_metric(\"pearsonr\")\n        >>> results = pearsonr_metric.compute(predictions=[10, 9, 2.5, 6, 4], references=[1, 2, 3, 4, 5])\n        >>> print(round(results['pearsonr'], 2))\n        -0.74\n\n    Example 2-The same as Example 1, but that also returns the `p-value`.\n        >>> pearsonr_metric = datasets.load_metric(\"pearsonr\")\n        >>> results = pearsonr_metric.compute(predictions=[10, 9, 2.5, 6, 4], references=[1, 2, 3, 4, 5], return_pvalue=True)\n        >>> print(sorted(list(results.keys())))\n        ['p-value', 'pearsonr']\n        >>> print(round(results['pearsonr'], 2))\n        -0.74\n        >>> print(round(results['p-value'], 2))\n        0.15\n"


__UpperCAmelCase    :List[Any]				       =     "\n@article{2020SciPy-NMeth,\nauthor  = {Virtanen, Pauli and Gommers, Ralf and Oliphant, Travis E. and\n      Haberland, Matt and Reddy, Tyler and Cournapeau, David and\n      Burovski, Evgeni and Peterson, Pearu and Weckesser, Warren and\n      Bright, Jonathan and {van der Walt}, St{\'e}fan J. and\n      Brett, Matthew and Wilson, Joshua and Millman, K. Jarrod and\n      Mayorov, Nikolay and Nelson, Andrew R. J. and Jones, Eric and\n      Kern, Robert and Larson, Eric and Carey, C J and\n      Polat, Ilhan and Feng, Yu and Moore, Eric W. and\n      {VanderPlas}, Jake and Laxalde, Denis and Perktold, Josef and\n      Cimrman, Robert and Henriksen, Ian and Quintero, E. A. and\n      Harris, Charles R. and Archibald, Anne M. and\n      Ribeiro, Antonio H. and Pedregosa, Fabian and\n      {van Mulbregt}, Paul and {SciPy 1.0 Contributors}},\ntitle   = {{{SciPy} 1.0: Fundamental Algorithms for Scientific\n      Computing in Python}},\njournal = {Nature Methods},\nyear    = {2020},\nvolume  = {17},\npages   = {261--272},\nadsurl  = {https://rdcu.be/b08Wh},\ndoi = {10.1038/s41592-019-0686-2},\n}\n"



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


     """simple docstring"""






     def 				lowerCamelCase__				(       self  :	List[Any]    )		->							Union[str, Any]:
          return datasets.MetricInfo(
              description=_DESCRIPTION							,				citation=_CITATION							,				inputs_description=_KWARGS_DESCRIPTION							,				features=datasets.Features(
                  {
                      '''predictions''': datasets.Value('''float'''    ),
                      '''references''': datasets.Value('''float'''    ),
                  }    )							,				reference_urls=['''https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.pearsonr.html''']							,				)






     def 				lowerCamelCase__				(       self  :	Any							,				snake_case  :	Union[str, Any]							,				snake_case  :	int							,				snake_case  :	Dict=False    )		->							List[Any]:
          if return_pvalue:
               __UpperCAmelCase			:				Tuple													=    pearsonr(snake_case							,				snake_case    )
               return {"pearsonr": results[0], "p-value": results[1]}
          else:
               return {"pearsonr": float(pearsonr(snake_case							,				snake_case    )[0]    )}