


'''simple docstring'''





from sklearn.metrics import fa_score

import datasets


__UpperCamelCase				:  Dict      					=							'\nThe F1 score is the harmonic mean of the precision and recall. It can be computed with the equation:\nF1 = 2 * (precision * recall) / (precision + recall)\n'


__UpperCamelCase				:  int      					=							'\nArgs:\n    predictions (`list` of `int`): Predicted labels.\n    references (`list` of `int`): Ground truth labels.\n    labels (`list` of `int`): The set of labels to include when `average` is not set to `\'binary\'`, and the order of the labels if `average` is `None`. Labels present in the data can be excluded, for example to calculate a multiclass average ignoring a majority negative class. Labels not present in the data will result in 0 components in a macro average. For multilabel targets, labels are column indices. By default, all labels in `predictions` and `references` are used in sorted order. Defaults to None.\n    pos_label (`int`): The class to be considered the positive class, in the case where `average` is set to `binary`. Defaults to 1.\n    average (`string`): This parameter is required for multiclass/multilabel targets. If set to `None`, the scores for each class are returned. Otherwise, this determines the type of averaging performed on the data. Defaults to `\'binary\'`.\n\n        - \'binary\': Only report results for the class specified by `pos_label`. This is applicable only if the classes found in `predictions` and `references` are binary.\n        - \'micro\': Calculate metrics globally by counting the total true positives, false negatives and false positives.\n        - \'macro\': Calculate metrics for each label, and find their unweighted mean. This does not take label imbalance into account.\n        - \'weighted\': Calculate metrics for each label, and find their average weighted by support (the number of true instances for each label). This alters `\'macro\'` to account for label imbalance. This option can result in an F-score that is not between precision and recall.\n        - \'samples\': Calculate metrics for each instance, and find their average (only meaningful for multilabel classification).\n    sample_weight (`list` of `float`): Sample weights Defaults to None.\n\nReturns:\n    f1 (`float` or `array` of `float`): F1 score or list of f1 scores, depending on the value passed to `average`. Minimum possible value is 0. Maximum possible value is 1. Higher f1 scores are better.\n\nExamples:\n\n    Example 1-A simple binary example\n        >>> f1_metric = datasets.load_metric("f1")\n        >>> results = f1_metric.compute(references=[0, 1, 0, 1, 0], predictions=[0, 0, 1, 1, 0])\n        >>> print(results)\n        {\'f1\': 0.5}\n\n    Example 2-The same simple binary example as in Example 1, but with `pos_label` set to `0`.\n        >>> f1_metric = datasets.load_metric("f1")\n        >>> results = f1_metric.compute(references=[0, 1, 0, 1, 0], predictions=[0, 0, 1, 1, 0], pos_label=0)\n        >>> print(round(results[\'f1\'], 2))\n        0.67\n\n    Example 3-The same simple binary example as in Example 1, but with `sample_weight` included.\n        >>> f1_metric = datasets.load_metric("f1")\n        >>> results = f1_metric.compute(references=[0, 1, 0, 1, 0], predictions=[0, 0, 1, 1, 0], sample_weight=[0.9, 0.5, 3.9, 1.2, 0.3])\n        >>> print(round(results[\'f1\'], 2))\n        0.35\n\n    Example 4-A multiclass example, with different values for the `average` input.\n        >>> predictions = [0, 2, 1, 0, 0, 1]\n        >>> references = [0, 1, 2, 0, 1, 2]\n        >>> results = f1_metric.compute(predictions=predictions, references=references, average="macro")\n        >>> print(round(results[\'f1\'], 2))\n        0.27\n        >>> results = f1_metric.compute(predictions=predictions, references=references, average="micro")\n        >>> print(round(results[\'f1\'], 2))\n        0.33\n        >>> results = f1_metric.compute(predictions=predictions, references=references, average="weighted")\n        >>> print(round(results[\'f1\'], 2))\n        0.27\n        >>> results = f1_metric.compute(predictions=predictions, references=references, average=None)\n        >>> print(results)\n        {\'f1\': array([0.8, 0. , 0. ])}\n'


__UpperCamelCase				:  List[str]      					=							'\n@article{scikit-learn,\n    title={Scikit-learn: Machine Learning in {P}ython},\n    author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V.\n           and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P.\n           and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and\n           Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},\n    journal={Journal of Machine Learning Research},\n    volume={12},\n    pages={2825--2830},\n    year={2011}\n}\n'





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






    '''simple docstring'''





    def 					_lowerCamelCase (	self				:   Union[str, Any]	):





       '''simple docstring'''
       return datasets.MetricInfo(
           description=_DESCRIPTION       ,   citation=_CITATION       ,   inputs_description=_KWARGS_DESCRIPTION       ,   features=datasets.Features(
               {
                   '''predictions''': datasets.Sequence(datasets.Value('''int32'''	)	),
                   '''references''': datasets.Sequence(datasets.Value('''int32'''	)	),
               }
               if self.config_name == '''multilabel'''
               else {
                   '''predictions''': datasets.Value('''int32'''	),
                   '''references''': datasets.Value('''int32'''	),
               }	)       ,   reference_urls=['''https://scikit-learn.org/stable/modules/generated/sklearn.metrics.f1_score.html''']       ,   )





    def 					_lowerCamelCase (	self				:   int       ,   __snake_case				:   Any       ,   __snake_case				:   List[str]       ,   __snake_case				:   Any=None       ,   __snake_case				:   int=1       ,   __snake_case				:   Any="binary"       ,   __snake_case				:   int=None	):





       '''simple docstring'''
       UpperCAmelCase_							:      Dict						=   fa_score(
           __snake_case       ,   __snake_case       ,   labels=__snake_case       ,   pos_label=__snake_case       ,   average=__snake_case       ,   sample_weight=__snake_case	)
       return {"f1": float(__snake_case	) if score.size == 1 else score}



import math
import os
from copy import deepcopy

import datasets
import evaluate
import torch
import transformers
from datasets import load_dataset
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer

from accelerate import Accelerator
from accelerate.test_utils import RegressionDataset, RegressionModel
from accelerate.utils import is_tpu_available, set_seed


__UpperCamelCase				:  str      					=							'true'
def      snake_case_						( __lowercase			,       __lowercase=8_2			,       __lowercase=1_6       ):
   set_seed(4_2       )
   UpperCAmelCase_							:      Optional[int]						=   RegressionModel()
   UpperCAmelCase_							:      Optional[int]						=   deepcopy(__lowercase       )
   UpperCAmelCase_							:      Union[str, Any]						=   RegressionDataset(length=__lowercase       )
   UpperCAmelCase_							:      Any						=   DataLoader(__lowercase			,       batch_size=__lowercase       )
   model.to(accelerator.device       )
   UpperCAmelCase_		,      UpperCAmelCase_							:      Dict						=   accelerator.prepare(__lowercase			,       __lowercase       )
   return model, ddp_model, dataloader
def      snake_case_						( __lowercase			,       __lowercase=False       ):
   UpperCAmelCase_							:      Optional[int]						=   AutoTokenizer.from_pretrained('''hf-internal-testing/mrpc-bert-base-cased'''       )
   UpperCAmelCase_							:      List[Any]						=   load_dataset('''glue'''			,       '''mrpc'''			,       split='''validation'''       )

   def tokenize_function(__lowercase       ):
      UpperCAmelCase_							:      int						=   tokenizer(examples['''sentence1''']			,       examples['''sentence2''']			,       truncation=__lowercase			,       max_length=__lowercase       )
      return outputs

   with accelerator.main_process_first():
      UpperCAmelCase_							:      List[str]						=   dataset.map(
          __lowercase			,       batched=__lowercase			,       remove_columns=['''idx''', '''sentence1''', '''sentence2''']			,       )

   UpperCAmelCase_							:      Any						=   tokenized_datasets.rename_column('''label'''			,       '''labels'''       )

   def collate_fn(__lowercase       ):
      if use_longest:
         return tokenizer.pad(__lowercase			,       padding='''longest'''			,       return_tensors='''pt'''       )
      return tokenizer.pad(__lowercase			,       padding='''max_length'''			,       max_length=1_2_8			,       return_tensors='''pt'''       )

   return DataLoader(__lowercase			,       shuffle=__lowercase			,       collate_fn=__lowercase			,       batch_size=1_6       )
def      snake_case_						( __lowercase			,       __lowercase       ):
   UpperCAmelCase_							:      Optional[int]						=   Accelerator(dispatch_batches=__lowercase			,       split_batches=__lowercase       )
   UpperCAmelCase_							:      int						=   get_dataloader(__lowercase			,       not dispatch_batches       )
   UpperCAmelCase_							:      Optional[int]						=   AutoModelForSequenceClassification.from_pretrained(
       '''hf-internal-testing/mrpc-bert-base-cased'''			,       return_dict=__lowercase       )
   UpperCAmelCase_		,      UpperCAmelCase_							:      Any						=   accelerator.prepare(__lowercase			,       __lowercase       )
   return {"ddp": [ddp_model, ddp_dataloader, "cuda:0"], "no": [model, dataloader, accelerator.device]}, accelerator
def      snake_case_						( __lowercase			,       __lowercase			,       __lowercase       ):
   UpperCAmelCase_							:      Dict						=   []
   for batch in dataloader:
      UpperCAmelCase_		,      UpperCAmelCase_							:      Optional[int]						=   batch.values()
      with torch.no_grad():
         UpperCAmelCase_							:      List[Any]						=   model(__lowercase       )
         UpperCAmelCase_		,      UpperCAmelCase_							:      Dict						=   accelerator.gather_for_metrics((logit, target)       )
         logits_and_targets.append((logit, target)       )
   UpperCAmelCase_		,      UpperCAmelCase_							:      Any						=   [], []
   for logit, targ in logits_and_targets:
      logits.append(__lowercase       )
      targs.append(__lowercase       )
   UpperCAmelCase_		,      UpperCAmelCase_							:      Optional[int]						=   torch.cat(__lowercase       ), torch.cat(__lowercase       )
   return logits, targs
def      snake_case_						( __lowercase			,       __lowercase=8_2			,       __lowercase=False			,       __lowercase=False			,       __lowercase=1_6       ):
   UpperCAmelCase_		,      UpperCAmelCase_		,      UpperCAmelCase_							:      Any						=   get_basic_setup(__lowercase			,       __lowercase			,       __lowercase       )
   UpperCAmelCase_		,      UpperCAmelCase_							:      Optional[int]						=   generate_predictions(__lowercase			,       __lowercase			,       __lowercase       )
   assert (
       len(__lowercase       ) == num_samples
   ), F'''Unexpected number of inputs:\n    Expected: {num_samples}\n    Actual: {len(__lowercase       )}'''
def      snake_case_						( __lowercase = False			,       __lowercase = False       ):
   UpperCAmelCase_							:      Optional[Any]						=   evaluate.load('''glue'''			,       '''mrpc'''       )
   UpperCAmelCase_		,      UpperCAmelCase_							:      Tuple						=   get_mrpc_setup(__lowercase			,       __lowercase       )
   # First do baseline
   UpperCAmelCase_		,      UpperCAmelCase_		,      UpperCAmelCase_							:      Union[str, Any]						=   setup['''no''']
   model.to(__lowercase       )
   model.eval()
   for batch in dataloader:
      batch.to(__lowercase       )
      with torch.inference_mode():
         UpperCAmelCase_							:      str						=   model(**__lowercase       )
      UpperCAmelCase_							:      Dict						=   outputs.logits.argmax(dim=-1       )
      metric.add_batch(predictions=__lowercase			,       references=batch['''labels''']       )
   UpperCAmelCase_							:      Optional[int]						=   metric.compute()

   # Then do distributed
   UpperCAmelCase_		,      UpperCAmelCase_		,      UpperCAmelCase_							:      Optional[Any]						=   setup['''ddp''']
   model.eval()
   for batch in dataloader:
      with torch.inference_mode():
         UpperCAmelCase_							:      Optional[int]						=   model(**__lowercase       )
      UpperCAmelCase_							:      int						=   outputs.logits.argmax(dim=-1       )
      UpperCAmelCase_							:      Optional[int]						=   batch['''labels''']
      UpperCAmelCase_		,      UpperCAmelCase_							:      Tuple						=   accelerator.gather_for_metrics((preds, references)       )
      metric.add_batch(predictions=__lowercase			,       references=__lowercase       )
   UpperCAmelCase_							:      Dict						=   metric.compute()

   for key in "accuracy f1".split():
      assert math.isclose(
          baseline[key]			,       distributed[key]       ), F'''Baseline and Distributed are not the same for key {key}:\n\tBaseline: {baseline[key]}\n\tDistributed: {distributed[key]}\n'''
def      snake_case_						( ):
   UpperCAmelCase_							:      str						=   Accelerator(split_batches=__lowercase			,       dispatch_batches=__lowercase       )
   if accelerator.is_local_main_process:
      datasets.utils.logging.set_verbosity_warning()
      transformers.utils.logging.set_verbosity_warning()
   else:
      datasets.utils.logging.set_verbosity_error()
      transformers.utils.logging.set_verbosity_error()
   # These are a bit slower so they should only be ran on the GPU or TPU
   if torch.cuda.is_available() or is_tpu_available():
      if accelerator.is_local_main_process:
         print('''**Testing gather_for_metrics**'''       )
      for split_batches in [True, False]:
         for dispatch_batches in [True, False]:
            if accelerator.is_local_main_process:
               print(F'''With: `split_batches={split_batches}`, `dispatch_batches={dispatch_batches}`'''       )
            test_mrpc(__lowercase			,       __lowercase       )
            accelerator.state._reset_state()
   if accelerator.is_local_main_process:
      print('''**Test torch metrics**'''       )
   for split_batches in [True, False]:
      for dispatch_batches in [True, False]:
         UpperCAmelCase_							:      Optional[Any]						=   Accelerator(split_batches=__lowercase			,       dispatch_batches=__lowercase       )
         if accelerator.is_local_main_process:
            print(F'''With: `split_batches={split_batches}`, `dispatch_batches={dispatch_batches}`, length=99'''       )
         test_torch_metrics(__lowercase			,       9_9       )
         accelerator.state._reset_state()
   if accelerator.is_local_main_process:
      print('''**Test last batch is not dropped when perfectly divisible**'''       )
   UpperCAmelCase_							:      List[Any]						=   Accelerator()
   test_torch_metrics(__lowercase			,       5_1_2       )
   accelerator.state._reset_state()
def      snake_case_						( __lowercase       ):
   # For xla_spawn (TPUs)
   main()


if __name__ == "__main__":
       main()