


import timeit
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_auto import MODEL_MAPPING, MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_torch_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_torch_available():
			import torch

			from .benchmark_args import PyTorchBenchmarkArguments


if is_pyanvml_available():
			import pyanvml.pyanvml as nvml


snake_case__  				=			logging.get_logger(__name__)



class      snake_case__	(   _A	):
				__lowerCAmelCase		=   42
				__lowerCAmelCase		=   42
				__lowerCAmelCase		=   '''PyTorch'''
				@property
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										return torch.__version__
				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase  ,   _UpperCAmelCase  ,   _UpperCAmelCase   ):




										"""simple docstring"""




										__magic_name__					:				Optional[Any]				=   self._prepare_inference_func(__SCREAMING_SNAKE_CASE  ,   __SCREAMING_SNAKE_CASE  ,   __SCREAMING_SNAKE_CASE   )
										return self._measure_speed(_inference   )
				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase  ,   _UpperCAmelCase  ,   _UpperCAmelCase   ):




										"""simple docstring"""




										__magic_name__					:				List[str]				=   self._prepare_inference_func(__SCREAMING_SNAKE_CASE  ,   __SCREAMING_SNAKE_CASE  ,   __SCREAMING_SNAKE_CASE   )
										return self._measure_memory(_inference   )
				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase  ,   _UpperCAmelCase  ,   _UpperCAmelCase   ):




										"""simple docstring"""




										__magic_name__					:				str				=   self._prepare_train_func(__SCREAMING_SNAKE_CASE  ,   __SCREAMING_SNAKE_CASE  ,   __SCREAMING_SNAKE_CASE   )
										return self._measure_speed(_train   )
				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase  ,   _UpperCAmelCase  ,   _UpperCAmelCase   ):




										"""simple docstring"""




										__magic_name__					:				str				=   self._prepare_train_func(__SCREAMING_SNAKE_CASE  ,   __SCREAMING_SNAKE_CASE  ,   __SCREAMING_SNAKE_CASE   )
										return self._measure_memory(_train   )
				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase  ,   _UpperCAmelCase  ,   _UpperCAmelCase   ):




										"""simple docstring"""




										__magic_name__					:				Union[str, Any]				=   self.config_dict[model_name]

										if self.args.torchscript:
																__magic_name__					:				Any				=   True

										__magic_name__					:				str				=   (
										    hasattr(__SCREAMING_SNAKE_CASE  ,   'architectures'   )
										    and isinstance(config.architectures  ,   __SCREAMING_SNAKE_CASE   )
										    and len(config.architectures   ) > 0
										)
										if not self.args.only_pretrain_model and has_model_class_in_config:
																try:
																						__magic_name__					:				Optional[int]				=   config.architectures[0]
																						__magic_name__					:				Dict				=   __import__('transformers'  ,   fromlist=[model_class]   )
																						__magic_name__					:				Union[str, Any]				=   getattr(__SCREAMING_SNAKE_CASE  ,   __SCREAMING_SNAKE_CASE   )
																						__magic_name__					:				Any				=   model_cls(__SCREAMING_SNAKE_CASE   )
																except ImportError:
																						raise ImportError(
																						    F"{model_class} does not exist. If you just want to test the pretrained model, you might want to"
																						    ' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'   )
										else:
																__magic_name__					:				Any				=   MODEL_MAPPING[config.__class__](__SCREAMING_SNAKE_CASE   )

										model.eval()
										model.to(self.args.device   )

										# encoder-decoder has vocab size saved differently
										__magic_name__					:				Optional[int]				=   config.vocab_size if hasattr(__SCREAMING_SNAKE_CASE  ,   'vocab_size'   ) else config.encoder.vocab_size
										__magic_name__					:				List[str]				=   torch.randint(__SCREAMING_SNAKE_CASE  ,   (batch_size, sequence_length)  ,   dtype=torch.long  ,   device=self.args.device   )

										if self.args.fpaa:
																logger.info('Running training in Mixed Precision...'   )
																if not self.args.is_gpu:
																						raise ValueError('Mixed precision is possible only for GPU.'   )
																# amp seems to have memory leaks so that memory usage
																# is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
																model.half()

										if self.args.torchscript:
																with torch.no_grad():
																						__magic_name__					:				Dict				=   torch.jit.trace(__SCREAMING_SNAKE_CASE  ,   __SCREAMING_SNAKE_CASE   )
										else:
																__magic_name__					:				Optional[int]				=   model

										def encoder_decoder_forward():
																with torch.no_grad():
																						__magic_name__					:				str				=   inference_model(__SCREAMING_SNAKE_CASE  ,   decoder_input_ids=__SCREAMING_SNAKE_CASE   )
																return outputs

										def encoder_forward():
																with torch.no_grad():
																						__magic_name__					:				List[str]				=   inference_model(__SCREAMING_SNAKE_CASE   )
																return outputs

										__magic_name__					:				int				=   encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
										return _forward
				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase  ,   _UpperCAmelCase  ,   _UpperCAmelCase   ):




										"""simple docstring"""




										__magic_name__					:				Optional[Any]				=   self.config_dict[model_name]

										__magic_name__					:				int				=   (
										    hasattr(__SCREAMING_SNAKE_CASE  ,   'architectures'   )
										    and isinstance(config.architectures  ,   __SCREAMING_SNAKE_CASE   )
										    and len(config.architectures   ) > 0
										)
										if not self.args.only_pretrain_model and has_model_class_in_config:
																try:
																						__magic_name__					:				List[str]				=   config.architectures[0]
																						__magic_name__					:				List[str]				=   __import__('transformers'  ,   fromlist=[model_class]   )
																						__magic_name__					:				Any				=   getattr(__SCREAMING_SNAKE_CASE  ,   __SCREAMING_SNAKE_CASE   )
																						__magic_name__					:				str				=   model_cls(__SCREAMING_SNAKE_CASE   )
																except ImportError:
																						raise ImportError(
																						    F"{model_class} does not exist. If you just want to test the pretrained model, you might want to"
																						    ' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'   )
										else:
																__magic_name__					:				Dict				=   MODEL_WITH_LM_HEAD_MAPPING[config.__class__](__SCREAMING_SNAKE_CASE   )

										if self.args.torchscript:
																raise NotImplementedError('Training for torchscript is currently not implemented'   )
										else:
																__magic_name__					:				Dict				=   model

										model.train()
										model.to(self.args.device   )

										# encoder-decoder has vocab size saved differently
										__magic_name__					:				Dict				=   config.vocab_size if hasattr(__SCREAMING_SNAKE_CASE  ,   'vocab_size'   ) else config.encoder.vocab_size
										__magic_name__					:				List[Any]				=   torch.randint(__SCREAMING_SNAKE_CASE  ,   (batch_size, sequence_length)  ,   dtype=torch.long  ,   device=self.args.device   )

										if self.args.fpaa:
																logger.info('Running training in Mixed Precision...'   )
																if not self.args.is_gpu:
																						raise ValueError('Mixed precision is possible only for GPU.'   )

																# amp seems to have memory leaks so that memory usage
																# is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
																model.half()

										def compute_loss_and_backprob_encoder():
																__magic_name__					:				List[Any]				=   train_model(__SCREAMING_SNAKE_CASE  ,   labels=__SCREAMING_SNAKE_CASE   )[0]
																loss.backward()
																return loss

										def compute_loss_and_backprob_encoder_decoder():
																__magic_name__					:				Optional[int]				=   train_model(__SCREAMING_SNAKE_CASE  ,   decoder_input_ids=__SCREAMING_SNAKE_CASE  ,   labels=__SCREAMING_SNAKE_CASE   )[0]
																loss.backward()
																return loss

										__magic_name__					:				Optional[int]				=   (
										    compute_loss_and_backprob_encoder_decoder
										    if config.is_encoder_decoder
										    else compute_loss_and_backprob_encoder
										)
										return _train
				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase   ):




										"""simple docstring"""




										try:
																if self.args.is_tpu or self.args.torchscript:
																						# run additional 10 times to stabilize compilation for tpu and torchscript
																						logger.info('Do inference on TPU or torchscript. Running model 5 times to stabilize compilation'   )
																						timeit.repeat(
																						    __SCREAMING_SNAKE_CASE  ,   repeat=1  ,   number=5  ,   )

																# as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
																__magic_name__					:				List[Any]				=   timeit.repeat(
																    __SCREAMING_SNAKE_CASE  ,   repeat=self.args.repeat  ,   number=10  ,   )

																if self.args.is_tpu and self.args.torch_xla_tpu_print_metrics:
																						import torch_xla.debug.metrics as met

																						self.print_fn(met.metrics_report()   )

																return min(__SCREAMING_SNAKE_CASE   ) / 10.0
										except RuntimeError as e:
																self.print_fn(F"Doesn't fit on GPU. {e}"   )
																return "N/A"





				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase   ):




										"""simple docstring"""




										try:
																if self.args.trace_memory_line_by_line:
																						__magic_name__					:				Union[str, Any]				=   start_memory_tracing('transformers'   )

																if self.args.is_tpu:
																						# tpu
																						raise NotImplementedError(
																						    'Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking with'
																						    ' `--no-memory` or `args.memory=False`'   )
																elif self.args.is_gpu:
																						if not is_pyanvml_available():
																												logger.warning(
																												    'py3nvml not installed, we won\'t log GPU memory usage. '
																												    'Install py3nvml (pip install py3nvml) to log information about GPU.'   )
																												__magic_name__					:				Optional[Any]				=   'N/A'
																						else:
																												logger.info(
																												    'Measuring total GPU usage on GPU device. Make sure to not have additional processes running'
																												    ' on the same GPU.'   )
																												# init nvml
																												nvml.nvmlInit()
																												func()
																												__magic_name__					:				Any				=   nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx   )
																												__magic_name__					:				Optional[Any]				=   nvml.nvmlDeviceGetMemoryInfo(__SCREAMING_SNAKE_CASE   )
																												__magic_name__					:				Dict				=   meminfo.used
																												__magic_name__					:				str				=   Memory(__SCREAMING_SNAKE_CASE   )
																												# shutdown nvml
																												nvml.nvmlShutdown()
																else:
																						# cpu
																						__magic_name__					:				List[Any]				=   measure_peak_memory_cpu(__SCREAMING_SNAKE_CASE   )
																						__magic_name__					:				int				=   Memory(__SCREAMING_SNAKE_CASE   ) if isinstance(__SCREAMING_SNAKE_CASE  ,   __SCREAMING_SNAKE_CASE   ) else memory_bytes

																if self.args.trace_memory_line_by_line:
																						__magic_name__					:				int				=   stop_memory_tracing(__SCREAMING_SNAKE_CASE   )
																else:
																						__magic_name__					:				str				=   None

																return memory, summary
										except RuntimeError as e:
																self.print_fn(F"Doesn't fit on GPU. {e}"   )
																return "N/A", None




'''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        SCREAMING_SNAKE_CASE_  (	snake_case_				:	List[str]   ,     **snake_case_				:	str							)       ->	Optional[int]:
 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   ,     **snake_case_   ,     )





class 		_a     (							SCREAMING_SNAKE_CASE  ):







   """simple docstring"""




   @staticmethod
   def 					__a			(       __SCREAMING_SNAKE_CASE      ):
    SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		parser.add_parser('run_beam'    ,help='Run a Beam dataset processing pipeline'      )
    run_beam_parser.add_argument('dataset'    ,type=__SCREAMING_SNAKE_CASE    ,help='Name of the dataset to download'      )
    run_beam_parser.add_argument('--name'    ,type=__SCREAMING_SNAKE_CASE    ,default=__SCREAMING_SNAKE_CASE    ,help='Dataset config name'      )
    run_beam_parser.add_argument(
        '--cache_dir'    ,type=__SCREAMING_SNAKE_CASE    ,default=__SCREAMING_SNAKE_CASE    ,help='Cache directory where the datasets are stored'    ,)
    run_beam_parser.add_argument(
        '--beam_pipeline_options'    ,type=__SCREAMING_SNAKE_CASE    ,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=__SCREAMING_SNAKE_CASE    ,default=__SCREAMING_SNAKE_CASE    ,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=__SCREAMING_SNAKE_CASE      )






   def __init__(       self    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,**__SCREAMING_SNAKE_CASE    ,):
    SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		dataset
    SCREAMING_SNAKE_CASE  :  List[Any]								=		name
    SCREAMING_SNAKE_CASE  :  int								=		cache_dir
    SCREAMING_SNAKE_CASE  :  str								=		beam_pipeline_options
    SCREAMING_SNAKE_CASE  :  List[str]								=		data_dir
    SCREAMING_SNAKE_CASE  :  int								=		all_configs
    SCREAMING_SNAKE_CASE  :  Any								=		save_infos
    SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		ignore_verifications
    SCREAMING_SNAKE_CASE  :  Optional[int]								=		force_redownload
    SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		config_kwargs






   def 					__a			(       self      ):
    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      )
    SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE  :  List[Any]								=		self._dataset, self._name
    SCREAMING_SNAKE_CASE  :  str								=		dataset_module_factory(__SCREAMING_SNAKE_CASE      )
    SCREAMING_SNAKE_CASE  :  str								=		import_main_class(dataset_module.module_path      )
    SCREAMING_SNAKE_CASE  :  List[DatasetBuilder]								=		[]
    if self._beam_pipeline_options:
     SCREAMING_SNAKE_CASE  :  Optional[int]								=		beam.options.pipeline_options.PipelineOptions(
         flags=[f"""--{opt.strip()}""" for opt in self._beam_pipeline_options.split(','      ) if opt]      )
    else:
     SCREAMING_SNAKE_CASE  :  List[Any]								=		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=__SCREAMING_SNAKE_CASE    ,cache_dir=self._cache_dir    ,base_path=dataset_module.builder_kwargs.get('base_path'      )    ,)      )
    else:
     builders.append(
         builder_cls(
             config_name=__SCREAMING_SNAKE_CASE    ,data_dir=self._data_dir    ,beam_options=__SCREAMING_SNAKE_CASE    ,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=__SCREAMING_SNAKE_CASE    ,)
     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:
     SCREAMING_SNAKE_CASE  :  Dict								=		os.path.join(builder_cls.get_imported_module_dir()    ,config.DATASETDICT_INFOS_FILENAME      )

     SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		Path(__SCREAMING_SNAKE_CASE      ).name + '.py'

     SCREAMING_SNAKE_CASE  :  Dict								=		os.path.join(__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE      )
     if os.path.isfile(__SCREAMING_SNAKE_CASE      ):
      SCREAMING_SNAKE_CASE  :  Tuple								=		os.path.dirname(__SCREAMING_SNAKE_CASE      )
     elif os.path.isfile(__SCREAMING_SNAKE_CASE      ):
      SCREAMING_SNAKE_CASE  :  Optional[Any]								=		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
     SCREAMING_SNAKE_CASE  :  List[Any]								=		os.path.join(__SCREAMING_SNAKE_CASE    ,config.DATASETDICT_INFOS_FILENAME      )
     copyfile(__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE      )
     print(f"""Dataset Infos file saved at {user_dataset_infos_path}"""      )


