

'''simple docstring'''



from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    get_torch_dist_unique_port,
    require_torch_multi_gpu,
    require_torch_neuroncore,
)
from transformers.training_args import ParallelMode
from transformers.utils import logging


SCREAMING_SNAKE_CASE_     = logging.get_logger(__name__)


if is_torch_available():
   import torch
   from torch import nn
   from torch.utils.data import Dataset

   from transformers import Trainer

   class       SCREAMING_SNAKE_CASE       (				_SCREAMING_SNAKE_CASE       ):






         '''simple docstring'''






         def __init__(       self			:		Any  ,       snake_case			:		Dict = 101 ):



             """simple docstring"""





             _snake_case			:			Dict				       =		length




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



             """simple docstring"""





             return self.length




         def __getitem__(       self			:		Optional[Any]  ,       snake_case			:		Any ):



             """simple docstring"""





             return i



   class       SCREAMING_SNAKE_CASE       :






         '''simple docstring'''






         def __call__(       self			:		Any  ,       snake_case			:		Optional[int] ):



             """simple docstring"""





             return {"input_ids": torch.tensor(lowerCAmelCase_ ), "labels": torch.tensor(lowerCAmelCase_ )}



   class       SCREAMING_SNAKE_CASE       (				nn.Module       ):






         '''simple docstring'''






         def __init__(       self			:		Dict ):



             """simple docstring"""





             super().__init__()
             # Add some (unused) params otherwise DDP will complain.
             _snake_case			:			List[Any]				       =		nn.Linear(120  ,       80 )




         def 						__UpperCAmelCase      (       self			:		str  ,       snake_case			:		Any  ,       snake_case			:		List[str]=None ):



             """simple docstring"""





             if labels is not None:
                 return torch.tensor(0.0  ,       device=input_ids.device ), input_ids
             else:
                 return input_ids

class       SCREAMING_SNAKE_CASE       (				_SCREAMING_SNAKE_CASE       ):






      '''simple docstring'''






      @require_torch_neuroncore
      def 						__UpperCAmelCase      (       self			:		Tuple ):



          """simple docstring"""





          _snake_case			:			int				       =		F"""--nproc_per_node=2
            --master_port={get_torch_dist_unique_port()}
            {self.test_file_dir}/test_trainer_distributed.py
        """.split()
          _snake_case			:			Union[str, Any]				       =		self.get_auto_remove_tmp_dir()
          _snake_case			:			int				       =		F"""--output_dir {output_dir}""".split()
          _snake_case			:			Optional[int]				       =		["""torchrun"""] + distributed_args + args
          execute_subprocess_async(lowerCAmelCase_  ,       env=self.get_env() )
          # successful return here == success - any errors would have caused an error in the sub-call



class       SCREAMING_SNAKE_CASE       (				_SCREAMING_SNAKE_CASE       ):






      '''simple docstring'''






      @require_torch_multi_gpu
      def 						__UpperCAmelCase      (       self			:		int ):



          """simple docstring"""





          _snake_case			:			Optional[Any]				       =		F"""--nproc_per_node={torch.cuda.device_count()}
            --master_port={get_torch_dist_unique_port()}
            {self.test_file_dir}/test_trainer_distributed.py
        """.split()
          _snake_case			:			Optional[Any]				       =		self.get_auto_remove_tmp_dir()
          _snake_case			:			int				       =		F"""--output_dir {output_dir}""".split()
          _snake_case			:			Dict				       =		["""torchrun"""] + distributed_args + args
          execute_subprocess_async(lowerCAmelCase_  ,       env=self.get_env() )
          # successful return here == success - any errors would have caused an error in the sub-call


if __name__ == "__main__":
   # The script below is meant to be run under torch.distributed, on a machine with multiple GPUs:
   #
   # PYTHONPATH="src" python -m torch.distributed.run --nproc_per_node 2 --output_dir output_dir ./tests/test_trainer_distributed.py

   SCREAMING_SNAKE_CASE_     = HfArgumentParser((TrainingArguments,))
   SCREAMING_SNAKE_CASE_     = parser.parse_args_into_dataclasses()[0]

   logger.warning(
       F'''Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, '''
       F'''distributed training: {training_args.parallel_mode != ParallelMode.NOT_DISTRIBUTED}'''
   )

   # Essentially, what we want to verify in the distributed case is that we get all samples back,
   # in the right order. (this is crucial for prediction for instance)
   for dataset_length in [101, 40, 7]:
      SCREAMING_SNAKE_CASE_     = DummyDataset(dataset_length)

      def 	lowerCamelCase__							(   a__)		->     str:



          """simple docstring"""







          _snake_case			:			Union[str, Any]				       =		list(range(len(snake_case__)))
          _snake_case			:			int				       =		p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
          if not success and training_args.local_rank == 0:
              logger.warning(
                  'Predictions and/or labels do not match expected results:\n  - predictions: '
                  F"""{p.predictions.tolist()}\n  - labels: {p.label_ids.tolist()}\n  - expected: {sequential}""")
          return {"success": success}

      SCREAMING_SNAKE_CASE_     = Trainer(
          model=DummyModel(),
          args=training_args,
          data_collator=DummyDataCollator(),
          eval_dataset=dataset,
          compute_metrics=compute_metrics,
      )
      SCREAMING_SNAKE_CASE_     = trainer.evaluate()
      logger.info(metrics)
      if metrics["eval_success"] is not True:
         logger.error(metrics)
         exit(1)

      SCREAMING_SNAKE_CASE_     = trainer.predict(dataset)
      logger.info(p.metrics)
      if p.metrics["test_success"] is not True:
         logger.error(p.metrics)
         exit(1)

      SCREAMING_SNAKE_CASE_     = 2

      SCREAMING_SNAKE_CASE_     = trainer.evaluate()
      logger.info(metrics)
      if metrics["eval_success"] is not True:
         logger.error(metrics)
         exit(1)

      SCREAMING_SNAKE_CASE_     = trainer.predict(dataset)
      logger.info(p.metrics)
      if p.metrics["test_success"] is not True:
         logger.error(p.metrics)
         exit(1)




      SCREAMING_SNAKE_CASE_     = None











"""simple docstring"""
import math_equivalence  # From: git+https://github.com/hendrycks/math.git

import datasets


_lowerCAmelCase       =				"\\n@article{hendrycksmath2021,\n  title={Measuring Mathematical Problem Solving With the MATH Dataset},\n  author={Dan Hendrycks\n    and Collin Burns\n    and Saurav Kadavath\n    and Akul Arora\n    and Steven Basart\n    and Eric Tang\n    and Dawn Song\n    and Jacob Steinhardt},\n  journal={arXiv preprint arXiv:2103.03874},\n  year={2021}\n}\n"


_lowerCAmelCase       =				"\\nThis metric is used to assess performance on the Mathematics Aptitude Test of Heuristics (MATH) dataset.\nIt first canonicalizes the inputs (e.g., converting \"1/2\" to \"\\frac{1}{2}\") and then computes accuracy.\n"


_lowerCAmelCase       =				r"\nCalculates accuracy after canonicalizing inputs.\n\nArgs:\n    predictions: list of predictions to score. Each prediction\n        is a string that contains natural language and LaTex.\n    references: list of reference for each prediction. Each\n        reference is a string that contains natural language\n        and LaTex.\nReturns:\n    accuracy: accuracy after canonicalizing inputs\n        (e.g., converting \"1/2\" to \"\\frac{1}{2}\")\n\nExamples:\n    >>> metric = datasets.load_metric(\"competition_math\")\n    >>> results = metric.compute(references=[\"\\frac{1}{2}\"], predictions=[\"1/2\"])\n    >>> print(results)\n    {'accuracy': 1.0}\n"



@datasets.utils.file_utils.add_end_docstrings(_DESCRIPTION							,							_KWARGS_DESCRIPTION    )
class   SCREAMING_SNAKE_CASE	(					datasets.Metric    ):





	"""simple docstring"""




	def  lowerCamelCase(self				):
		return datasets.MetricInfo(
		    description=_DESCRIPTION  ,  citation=_CITATION  ,  inputs_description=_KWARGS_DESCRIPTION  ,  features=datasets.Features(
		        {
		            """predictions""": datasets.Value("""string"""				),
		            """references""": datasets.Value("""string"""				),
		        }				)  ,  homepage="""https://github.com/hendrycks/math"""  ,  codebase_urls=["""https://github.com/hendrycks/math"""]  ,  )





	def  lowerCamelCase(self  ,  lowerCAmelCase_  ,  lowerCAmelCase_				):
		A_		:    Optional[Any]						     = 0.0
		for i, j in zip(lowerCAmelCase_  ,  lowerCAmelCase_				):
			n_correct += 1.0 if math_equivalence.is_equiv(lowerCAmelCase_  ,  lowerCAmelCase_				) else 0.0
		A_		:    List[str]						     = n_correct / len(lowerCAmelCase_				)
		return {
		    "accuracy": accuracy,
		}

