"""simple docstring"""


import os
from argparse import ArgumentParser
from pathlib import Path
from shutil import copyfile
from typing import List

from datasets import config
from datasets.builder import DatasetBuilder
from datasets.commands import BaseDatasetsCLICommand
from datasets.download.download_config import DownloadConfig
from datasets.download.download_manager import DownloadMode
from datasets.load import dataset_module_factory, import_main_class
from datasets.utils.info_utils import VerificationMode







def       UpperCAmelCase					(	UpperCAmelCase		,      **UpperCAmelCase      )       ->   Tuple:
     return RunBeamCommand(
         args.dataset		,      args.name		,      args.cache_dir		,      args.beam_pipeline_options		,      args.data_dir		,      args.all_configs		,      args.save_info or args.save_infos		,      args.ignore_verifications		,      args.force_redownload		,      **UpperCAmelCase		,      )




class 				UpperCamelCase      (       lowerCAmelCase__     ):







       @staticmethod
       def 	a_		(				lowerCAmelCase__)   ->   List[str]:
            snake_case_       =      parser.add_parser('run_beam',							help='Run a Beam dataset processing pipeline')
            run_beam_parser.add_argument('dataset',							type=lowerCAmelCase_,							help='Name of the dataset to download')
            run_beam_parser.add_argument('--name',							type=lowerCAmelCase_,							default=lowerCAmelCase_,							help='Dataset config name')
            run_beam_parser.add_argument(
                '--cache_dir',							type=lowerCAmelCase_,							default=lowerCAmelCase_,							help='Cache directory where the datasets are stored',							)
            run_beam_parser.add_argument(
                '--beam_pipeline_options',							type=lowerCAmelCase_,							default='',							help='Beam pipeline options, separated by commas. Example:: `--beam_pipeline_options=job_name=my-job,project=my-project`',							)
            run_beam_parser.add_argument(
                '--data_dir',							type=lowerCAmelCase_,							default=lowerCAmelCase_,							help='Can be used to specify a manual directory to get the files from',							)
            run_beam_parser.add_argument('--all_configs',							action='store_true',							help='Test all dataset configurations')
            run_beam_parser.add_argument('--save_info',							action='store_true',							help='Save the dataset infos file')
            run_beam_parser.add_argument(
                '--ignore_verifications',							action='store_true',							help='Run the test without checksums and splits checks')
            run_beam_parser.add_argument('--force_redownload',							action='store_true',							help='Force dataset redownload')
            # aliases
            run_beam_parser.add_argument('--save_infos',							action='store_true',							help='alias for save_info')
            run_beam_parser.set_defaults(func=lowerCAmelCase_)







       def __init__(				self,							lowerCAmelCase__,							lowerCAmelCase__,							lowerCAmelCase__,							lowerCAmelCase__,							lowerCAmelCase__,							lowerCAmelCase__,							lowerCAmelCase__,							lowerCAmelCase__,							lowerCAmelCase__,							**lowerCAmelCase__,							)   ->   Tuple:
            snake_case_       =      dataset
            snake_case_       =      name
            snake_case_       =      cache_dir
            snake_case_       =      beam_pipeline_options
            snake_case_       =      data_dir
            snake_case_       =      all_configs
            snake_case_       =      save_infos
            snake_case_       =      ignore_verifications
            snake_case_       =      force_redownload
            snake_case_       =      config_kwargs







       def 	a_		(				self)   ->   int:
            import apache_beam as beam

            if self._name is not None and self._all_configs:
                 print('Both parameters `name` and `all_configs` can\'t be used at once.')
                 exit(1)
            snake_case_ ,			snake_case_       =      self._dataset, self._name
            snake_case_       =      dataset_module_factory(lowerCAmelCase_)
            snake_case_       =      import_main_class(dataset_module.module_path)
            snake_case_       =      []
            if self._beam_pipeline_options:
                 snake_case_       =      beam.options.pipeline_options.PipelineOptions(
                     flags=[f'--{opt.strip()}' for opt in self._beam_pipeline_options.split(',') if opt])
            else:
                 snake_case_       =      None
            if self._all_configs and len(builder_cls.BUILDER_CONFIGS) > 0:
                 for builder_config in builder_cls.BUILDER_CONFIGS:
                      builders.append(
                          builder_cls(
                              config_name=builder_config.name,							data_dir=self._data_dir,							hash=dataset_module.hash,							beam_options=lowerCAmelCase_,							cache_dir=self._cache_dir,							base_path=dataset_module.builder_kwargs.get('base_path'),							))
            else:
                 builders.append(
                     builder_cls(
                         config_name=lowerCAmelCase_,							data_dir=self._data_dir,							beam_options=lowerCAmelCase_,							cache_dir=self._cache_dir,							base_path=dataset_module.builder_kwargs.get('base_path'),							**self._config_kwargs,							))

            for builder in builders:
                 builder.download_and_prepare(
                     download_mode=DownloadMode.REUSE_CACHE_IF_EXISTS
                     if not self._force_redownload
                     else DownloadMode.FORCE_REDOWNLOAD,							download_config=DownloadConfig(cache_dir=config.DOWNLOADED_DATASETS_PATH),							verification_mode=VerificationMode.NO_CHECKS
                     if self._ignore_verifications
                     else VerificationMode.ALL_CHECKS,							try_from_hf_gcs=lowerCAmelCase_,							)
                 if self._save_infos:
                      builder._save_infos()

            print('Apache beam run successful.')

            # If save_infos=True, the dataset infos file is created next to the loaded module file.
            # Let's move it to the original directory of the dataset script, to allow the user to
            # upload them on S3 at the same time afterwards.
            if self._save_infos:
                 snake_case_       =      os.path.join(builder_cls.get_imported_module_dir(),							config.DATASETDICT_INFOS_FILENAME)

                 snake_case_       =      Path(lowerCAmelCase_).name + '.py'

                 snake_case_       =      os.path.join(lowerCAmelCase_,							lowerCAmelCase_)
                 if os.path.isfile(lowerCAmelCase_):
                      snake_case_       =      os.path.dirname(lowerCAmelCase_)
                 elif os.path.isfile(lowerCAmelCase_):
                      snake_case_       =      path
                 else:  # in case of a remote dataset
                      print(f'Dataset Infos file saved at {dataset_infos_path}')
                      exit(1)

                 # Move datasetinfo back to the user
                 snake_case_       =      os.path.join(lowerCAmelCase_,							config.DATASETDICT_INFOS_FILENAME)
                 copyfile(lowerCAmelCase_,							lowerCAmelCase_)
                 print(f'Dataset Infos file saved at {user_dataset_infos_path}')

"""simple docstring"""


import collections

from .utils import ExplicitEnum, is_torch_available, logging


if is_torch_available():
    import torch


__UpperCamelCase								=    logging.get_logger(__name__)




class 				UpperCamelCase      :







       def __init__(				self,							lowerCAmelCase__,							lowerCAmelCase__=21,							lowerCAmelCase__=[],							lowerCAmelCase__=None)   ->   Optional[Any]:
            snake_case_       =      model
            snake_case_       =      trace_batch_nums
            snake_case_       =      abort_after_batch_num

            # keep a LIFO buffer of frames to dump as soon as inf/nan is encountered to give context to the problem emergence
            snake_case_       =      collections.deque([],							lowerCAmelCase__)
            snake_case_       =      []
            snake_case_       =      0
            snake_case_       =      0
            snake_case_       =      False
            snake_case_       =      '                 '

            self.analyse_model()

            self.register_forward_hook()







       def 	a_		(				self,							lowerCAmelCase__=None)   ->   Optional[int]:
            if frame is not None:
                 self.expand_frame(lowerCAmelCase__)
            self.frames.append('\n'.join(self.frame))
            snake_case_       =      []  # start a new frame







       def 	a_		(				self,							lowerCAmelCase__)   ->   List[str]:
            self.frame.append(lowerCAmelCase__)







       def 	a_		(				self)   ->   List[str]:
            print('\n'.join(self.frames))
            snake_case_       =      []







       def 	a_		(				self)   ->   Tuple:
            snake_case_       =      []







       def 	a_		(				self)   ->   List[str]:
            print(f'\nDetected inf/nan during batch_number={self.batch_number}')
            print(f'Last {len(self.frames)} forward frames:')
            print(f'{"abs min":8} {"abs max":8} metadata')
            print('\n'.join(self.frames))
            print('\n\n')
            snake_case_       =      []







       def 	a_		(				self)   ->   Any:
            # extract the fully qualified module names, to be able to report at run time. e.g.:
            # encoder.block.2.layer.0.SelfAttention.o
            #
            # for shared weights only the first shared module name will be registered
            snake_case_       =      {m: name for name, m in self.model.named_modules()}
            # self.longest_module_name = max(len(v) for v in self.module_names.values())







       def 	a_		(				self,							lowerCAmelCase__,							lowerCAmelCase__)   ->   str:
            if torch.is_tensor(lowerCAmelCase__):
                 self.expand_frame(get_abs_min_max(lowerCAmelCase__,							lowerCAmelCase__))
                 if detect_overflow(lowerCAmelCase__,							lowerCAmelCase__):
                      snake_case_       =      True
            elif var is None:
                 self.expand_frame(f'{"None":>17} {ctx}')
            else:
                 self.expand_frame(f'{"not a tensor":>17} {ctx}')







       def 	a_		(				self)   ->   List[str]:
            self.expand_frame(f'\n\n{self.prefix} *** Starting batch number={self.batch_number} ***')
            self.expand_frame(f'{"abs min":8} {"abs max":8} metadata')







       def 	a_		(				self)   ->   Dict:
            self.expand_frame(f'{self.prefix} *** Finished batch number={self.batch_number-1} ***\n\n')







       def 	a_		(				self,							lowerCAmelCase__,							lowerCAmelCase__,							lowerCAmelCase__)   ->   List[str]:
            self.expand_frame(f'{self.prefix} {self.module_names[module]} {module.__class__.__name__}')

            # params
            for name, p in module.named_parameters(recurse=lowerCAmelCase__):
                 self.analyse_variable(lowerCAmelCase__,							lowerCAmelCase__)

            # inputs
            if isinstance(lowerCAmelCase__,							lowerCAmelCase__):
                 for i, x in enumerate(lowerCAmelCase__):
                      self.analyse_variable(lowerCAmelCase__,							f'input[{i}]')
            else:
                 self.analyse_variable(lowerCAmelCase__,							'input')

            # outputs
            if isinstance(lowerCAmelCase__,							lowerCAmelCase__):
                 for i, x in enumerate(lowerCAmelCase__):
                      # possibly a tuple of tuples
                      if isinstance(lowerCAmelCase__,							lowerCAmelCase__):
                           for j, y in enumerate(lowerCAmelCase__):
                                self.analyse_variable(lowerCAmelCase__,							f'output[{i}][{j}]')
                      else:
                           self.analyse_variable(lowerCAmelCase__,							f'output[{i}]')
            else:
                 self.analyse_variable(lowerCAmelCase__,							'output')

            self.save_frame()







       def 	a_		(				self)   ->   List[Any]:
            self.model.apply(self._register_forward_hook)







       def 	a_		(				self,							lowerCAmelCase__)   ->   int:
            module.register_forward_hook(self.forward_hook)







       def 	a_		(				self,							lowerCAmelCase__,							lowerCAmelCase__,							lowerCAmelCase__)   ->   Optional[int]:
            # - input is a tuple of packed inputs (could be non-Tensors)
            # - output could be a Tensor or a tuple of Tensors and non-Tensors

            snake_case_       =      False

            snake_case_       =      True if self.batch_number in self.trace_batch_nums else False
            if trace_mode:
                 self.reset_saved_frames()

            if self.total_calls == 0:
                 self.batch_start_frame()
            self.total_calls += 1

            # count batch numbers - the very first forward hook of the batch will be called when the
            # batch completes - i.e. it gets called very last - we know this batch has finished
            if module == self.model:
                 self.batch_number += 1
                 snake_case_       =      True

            self.create_frame(lowerCAmelCase__,							lowerCAmelCase__,							lowerCAmelCase__)

            # if last_frame_of_batch:
            #     self.batch_end_frame()

            if trace_mode:
                 self.trace_frames()

            if last_frame_of_batch:
                 self.batch_start_frame()

            if self.detected_overflow and not trace_mode:
                 self.dump_saved_frames()

                 # now we can abort, as it's pointless to continue running
                 raise ValueError(
                     'DebugUnderflowOverflow: inf/nan detected, aborting as there is no point running further. '
                     'Please scroll up above this traceback to see the activation values prior to this event.')

            # abort after certain batch if requested to do so
            if self.abort_after_batch_num is not None and self.batch_number > self.abort_after_batch_num:
                 raise ValueError(
                     f'DebugUnderflowOverflow: aborting after {self.batch_number} batches due to'
                     f' `abort_after_batch_num={self.abort_after_batch_num}` arg')







def       UpperCAmelCase					(	UpperCAmelCase		,      UpperCAmelCase      )       ->   Optional[int]:
     snake_case_       =      var.abs()
     return f'{abs_var.min():8.2e} {abs_var.max():8.2e} {ctx}'







def       UpperCAmelCase					(	UpperCAmelCase		,      UpperCAmelCase      )       ->   Union[str, Any]:
     snake_case_       =      False
     if torch.isnan(UpperCAmelCase      ).any().item():
          snake_case_       =      True
          print(f'{ctx} has nans'      )
     if torch.isinf(UpperCAmelCase      ).any().item():
          snake_case_       =      True
          print(f'{ctx} has infs'      )

     # if needed to monitor large elements can enable the following
     if 0:  # and detected:
          snake_case_       =      var[torch.ge(var.abs()		,      100      )]
          if naaa.numel() > 0:
               print(f'{ctx}:  n100={naaa.numel()}'      )
          snake_case_       =      var[torch.ge(var.abs()		,      1000      )]
          if naaaa.numel() > 0:
               print(f'{ctx}: n1000={naaaa.numel()}'      )
          snake_case_       =      var[torch.ge(var.abs()		,      10000      )]
          if naaaaa.numel() > 0:
               print(f'{ctx}: n10000={naaaaa.numel()}'      )

     if 0:
          print(f'min={var.min():9.2e} max={var.max():9.2e}'      )

     if 0:
          print(f'min={var.min():9.2e} max={var.max():9.2e} var={var.var():9.2e} mean={var.mean():9.2e} ({ctx})'      )

     return detected




class 				UpperCamelCase      (       lowerCAmelCase__     ):
       SCREAMING_SNAKE_CASE_				=   "underflow_overflow"
       SCREAMING_SNAKE_CASE_				=   "tpu_metrics_debug"

