


import builtins
import sys

from ...utils.imports import _is_package_available
from . import cursor, input
from .helpers import Direction, clear_line, forceWrite, linebreak, move_cursor, reset_cursor, writeColor
from .keymap import KEYMAP


SCREAMING_SNAKE_CASE       :		Tuple       =				False
try:
   SCREAMING_SNAKE_CASE       :		int       =				_is_package_available("google.colab")
except ModuleNotFoundError:
   pass






@input.register
class      UpperCamelCase					:


   '''simple docstring'''







   def __init__(     self   ,							UpperCamelCase_ = None   ,							UpperCamelCase_ = []		):
        lowercase_       :Union[str, Any]     =       0
        lowercase_       :Tuple     =       choices
        lowercase_       :Optional[int]     =       prompt
        if sys.platform == "win32":
             lowercase_       :int     =       '''*'''
        else:
             lowercase_       :Dict     =       '''➔ '''
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_ = ""		):
        if sys.platform != "win32":
             writeColor(self.choices[index]   ,							32   ,							UpperCamelCase_		)
        else:
             forceWrite(self.choices[index]   ,							UpperCamelCase_		)
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_		):
        if index == self.position:
             forceWrite(f" {self.arrow_char} "		)
             self.write_choice(UpperCamelCase_		)
        else:
             forceWrite(f"    {self.choices[index]}"		)
        reset_cursor()
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_ = 1		):
        lowercase_       :Optional[Any]     =       self.position
        if direction == Direction.DOWN:
             if self.position + 1 >= len(self.choices		):
                  return
             self.position += num_spaces
        else:
             if self.position - 1 < 0:
                  return
             self.position -= num_spaces
        clear_line()
        self.print_choice(UpperCamelCase_		)
        move_cursor(UpperCamelCase_   ,							direction.name		)
        self.print_choice(self.position		)
   @input.mark(KEYMAP['''up''']		)
   def 		UpperCamelCase  (     self		):
        self.move_direction(Direction.UP		)
   @input.mark(KEYMAP['''down''']		)
   def 		UpperCamelCase  (     self		):
        self.move_direction(Direction.DOWN		)
   @input.mark(KEYMAP['''newline''']		)
   def 		UpperCamelCase  (     self		):
        move_cursor(len(self.choices		) - self.position   ,							'''DOWN'''		)
        return self.position
   @input.mark(KEYMAP['''interrupt''']		)
   def 		UpperCamelCase  (     self		):
        move_cursor(len(self.choices		) - self.position   ,							'''DOWN'''		)
        raise KeyboardInterrupt
   @input.mark_multiple(*[KEYMAP[str(UpperCamelCase_		)] for number in range(10		)]		)
   def 		UpperCamelCase  (     self		):
        lowercase_       :List[Any]     =       int(chr(self.current_selection		)		)
        lowercase_       :Tuple     =       index - self.position
        if index == self.position:
             return
        if index < len(self.choices		):
             if self.position > index:
                  self.move_direction(Direction.UP   ,							-movement		)
             elif self.position < index:
                  self.move_direction(Direction.DOWN   ,							UpperCamelCase_		)
             else:
                  return
        else:
             return




   def 		UpperCamelCase  (     self   ,							UpperCamelCase_ = 0		):
        if self.prompt:
             linebreak()
             forceWrite(self.prompt   ,							'''\n'''		)
             if in_colab:
                  forceWrite('''Please input a choice index (starting from 0), and press enter'''   ,							'''\n'''		)
             else:
                  forceWrite('''Please select a choice using the arrow or number keys, and selecting with enter'''   ,							'''\n'''		)
        lowercase_       :Dict     =       default_choice
        for i in range(len(self.choices		)		):
             self.print_choice(UpperCamelCase_		)
             forceWrite('''\n'''		)
        move_cursor(len(self.choices		) - self.position   ,							'''UP'''		)
        with cursor.hide():
             while True:
                  if in_colab:
                       try:
                            lowercase_       :Union[str, Any]     =       int(builtins.input()		)
                       except ValueError:
                            lowercase_       :Union[str, Any]     =       default_choice
                  else:
                       lowercase_       :List[str]     =       self.handle_input()
                  if choice is not None:
                       reset_cursor()
                       for _ in range(len(self.choices		) + 1		):
                            move_cursor(1   ,							'''UP'''		)
                            clear_line()
                       self.write_choice(UpperCamelCase_   ,							'''\n'''		)
                       return choice





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       :		List[str]         =				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      UpperCamelCase					(		lowercase__     ):


      '''simple docstring'''







      def __init__(     self   ,							UpperCamelCase_ = 101		):
           lowercase_       :Tuple     =       length
      def __len__(     self		):
           return self.length




      def __getitem__(     self   ,							UpperCamelCase_		):
           return i






   class      UpperCamelCase					:


      '''simple docstring'''







      def __call__(     self   ,							UpperCamelCase_		):
           return {"input_ids": torch.tensor(UpperCamelCase_		), "labels": torch.tensor(UpperCamelCase_		)}






   class      UpperCamelCase					(		nn.Module     ):


      '''simple docstring'''







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




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






class      UpperCamelCase					(		lowercase__     ):


   '''simple docstring'''







   @require_torch_neuroncore
   def 		UpperCamelCase  (     self		):
        lowercase_       :int     =       f"--nproc_per_node=2\n            --master_port={get_torch_dist_unique_port()}\n            {self.test_file_dir}/test_trainer_distributed.py\n        ".split()
        lowercase_       :List[str]     =       self.get_auto_remove_tmp_dir()
        lowercase_       :List[str]     =       f"--output_dir {output_dir}".split()
        lowercase_       :Any     =       ['''torchrun'''] + distributed_args + args
        execute_subprocess_async(UpperCamelCase_   ,							env=self.get_env()		)
        # successful return here == success - any errors would have caused an error in the sub-call






class      UpperCamelCase					(		lowercase__     ):


   '''simple docstring'''







   @require_torch_multi_gpu
   def 		UpperCamelCase  (     self		):
        lowercase_       :Any     =       f"--nproc_per_node={torch.cuda.device_count()}\n            --master_port={get_torch_dist_unique_port()}\n            {self.test_file_dir}/test_trainer_distributed.py\n        ".split()
        lowercase_       :str     =       self.get_auto_remove_tmp_dir()
        lowercase_       :List[str]     =       f"--output_dir {output_dir}".split()
        lowercase_       :Dict     =       ['''torchrun'''] + distributed_args + args
        execute_subprocess_async(UpperCamelCase_   ,							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       :		Optional[int]         =				HfArgumentParser((TrainingArguments,))
   SCREAMING_SNAKE_CASE       :		Optional[int]         =				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       :		Dict         =				DummyDataset(dataset_length)




      def     UpperCamelCase    (      _a		)					->							Dict:
           '''simple docstring'''





           lowercase_       :Optional[Any]     =       list(range(len(_a		)		)		)
           lowercase_       :Optional[Any]     =       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       :		Optional[Any]         =				Trainer(
          model=DummyModel(),
          args=training_args,
          data_collator=DummyDataCollator(),
          eval_dataset=dataset,
          compute_metrics=compute_metrics,
      )
      SCREAMING_SNAKE_CASE       :		Any         =				trainer.evaluate()
      logger.info(metrics)
      if metrics["eval_success"] is not True:
         logger.error(metrics)
         exit(1)

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

      SCREAMING_SNAKE_CASE       :		Union[str, Any]         =				2

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

      SCREAMING_SNAKE_CASE       :		List[str]         =				trainer.predict(dataset)
      logger.info(p.metrics)
      if p.metrics["test_success"] is not True:
         logger.error(p.metrics)
         exit(1)
      SCREAMING_SNAKE_CASE       :		Optional[int]         =				None

