





"""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


lowerCamelCase_       =		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(       A  ):




         def __init__(						self		,SCREAMING_SNAKE_CASE__ = 1_01							)      ->				Optional[int]:





               """simple docstring"""






               __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  length




         def __len__(						self							)      ->				Any:





               """simple docstring"""






               return self.length




         def __getitem__(						self		,SCREAMING_SNAKE_CASE__							)      ->				int:





               """simple docstring"""






               return i






  class  _SCREAMING_SNAKE_CASE:




         def __call__(						self		,SCREAMING_SNAKE_CASE__							)      ->				Dict:





               """simple docstring"""






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






  class  _SCREAMING_SNAKE_CASE(       nn.Module  ):




         def __init__(						self							)      ->				int:





               """simple docstring"""






               super().__init__()
               # Add some (unused) params otherwise DDP will complain.
               __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  nn.Linear(1_20		,80							)




         def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__=None							)      ->				Union[str, Any]:





               """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(       A  ):




       @require_torch_neuroncore
       def    _UpperCamelCase					(						self							)      ->				List[str]:





             """simple docstring"""






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






class  _SCREAMING_SNAKE_CASE(       A  ):




       @require_torch_multi_gpu
       def    _UpperCamelCase					(						self							)      ->				Tuple:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :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()
             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  self.get_auto_remove_tmp_dir()
             __SCREAMING_SNAKE_CASE       :List[Any]		    =  f'''--output_dir {output_dir}'''.split()
             __SCREAMING_SNAKE_CASE       :str		    =  ['''torchrun'''] + distributed_args + args
             execute_subprocess_async(SCREAMING_SNAKE_CASE__		,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

  lowerCamelCase_       =		HfArgumentParser((TrainingArguments,))
  lowerCamelCase_       =		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 [1_0_1, 4_0, 7]:
    lowerCamelCase_       =		DummyDataset(dataset_length)



    def 					__lowerCamelCase       (							a_					:		EvalPrediction							)      ->				Dict:
          __SCREAMING_SNAKE_CASE       :Any		    =  list(range(len(a_							)							)							)
          __SCREAMING_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}

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

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

    lowerCamelCase_       =		2

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

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







    lowerCamelCase_       =		None






"""simple docstring"""
def 					__lowerCamelCase       (							a_					:		str							)      ->				list:
      return [
          txt[:a] + txt[a].upper() + txt[a + 1 :]
          for a in range(len(a_							)							)
          if txt[a].isalpha()
      ]


if __name__ == "__main__":
  __import__("doctest").testmod()