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


UpperCAmelCase_    		=				logging.get_logger(__name__)





class   lowerCAmelCase    (       lowercase_      ):
     _SCREAMING_SNAKE_CASE  :  str		=42
     _SCREAMING_SNAKE_CASE  :  str		=42
     _SCREAMING_SNAKE_CASE  :  Dict		="""PyTorch"""



     @property
     def       a__	(       self							):
            return torch.__version__



     def       a__	(       self					,   lowerCAmelCase__					,   lowerCAmelCase__					,   lowerCAmelCase__							):
            _A=						self._prepare_inference_func(snake_case__					,   snake_case__					,   snake_case__							)
            return self._measure_speed(_inference							)



     def       a__	(       self					,   lowerCAmelCase__					,   lowerCAmelCase__					,   lowerCAmelCase__							):
            _A=						self._prepare_inference_func(snake_case__					,   snake_case__					,   snake_case__							)
            return self._measure_memory(_inference							)



     def       a__	(       self					,   lowerCAmelCase__					,   lowerCAmelCase__					,   lowerCAmelCase__							):
            _A=						self._prepare_train_func(snake_case__					,   snake_case__					,   snake_case__							)
            return self._measure_speed(_train							)



     def       a__	(       self					,   lowerCAmelCase__					,   lowerCAmelCase__					,   lowerCAmelCase__							):
            _A=						self._prepare_train_func(snake_case__					,   snake_case__					,   snake_case__							)
            return self._measure_memory(_train							)



     def       a__	(       self					,   lowerCAmelCase__					,   lowerCAmelCase__					,   lowerCAmelCase__							):
            _A=						self.config_dict[model_name]

            if self.args.torchscript:
                   _A=						True

            _A=						(
                hasattr(snake_case__					,   'architectures'							)
                and isinstance(config.architectures					,   snake_case__							)
                and len(config.architectures							) > 0
            )
            if not self.args.only_pretrain_model and has_model_class_in_config:
                   try:
                          _A=						config.architectures[0]
                          _A=						__import__('transformers'					,   fromlist=[model_class]							)
                          _A=						getattr(snake_case__					,   snake_case__							)
                          _A=						model_cls(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:
                   _A=						MODEL_MAPPING[config.__class__](snake_case__							)

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

            # encoder-decoder has vocab size saved differently
            _A=						config.vocab_size if hasattr(snake_case__					,   'vocab_size'							) else config.encoder.vocab_size
            _A=						torch.randint(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():
                          _A=						torch.jit.trace(snake_case__					,   snake_case__							)
            else:
                   _A=						model

            def encoder_decoder_forward():
                   with torch.no_grad():
                          _A=						inference_model(snake_case__					,   decoder_input_ids=snake_case__							)
                   return outputs

            def encoder_forward():
                   with torch.no_grad():
                          _A=						inference_model(snake_case__							)
                   return outputs

            _A=						encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
            return _forward



     def       a__	(       self					,   lowerCAmelCase__					,   lowerCAmelCase__					,   lowerCAmelCase__							):
            _A=						self.config_dict[model_name]

            _A=						(
                hasattr(snake_case__					,   'architectures'							)
                and isinstance(config.architectures					,   snake_case__							)
                and len(config.architectures							) > 0
            )
            if not self.args.only_pretrain_model and has_model_class_in_config:
                   try:
                          _A=						config.architectures[0]
                          _A=						__import__('transformers'					,   fromlist=[model_class]							)
                          _A=						getattr(snake_case__					,   snake_case__							)
                          _A=						model_cls(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:
                   _A=						MODEL_WITH_LM_HEAD_MAPPING[config.__class__](snake_case__							)

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

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

            # encoder-decoder has vocab size saved differently
            _A=						config.vocab_size if hasattr(snake_case__					,   'vocab_size'							) else config.encoder.vocab_size
            _A=						torch.randint(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():
                   _A=						train_model(snake_case__					,   labels=snake_case__							)[0]
                   loss.backward()
                   return loss

            def compute_loss_and_backprob_encoder_decoder():
                   _A=						train_model(snake_case__					,   decoder_input_ids=snake_case__					,   labels=snake_case__							)[0]
                   loss.backward()
                   return loss

            _A=						(
                compute_loss_and_backprob_encoder_decoder
                if config.is_encoder_decoder
                else compute_loss_and_backprob_encoder
            )
            return _train



     def       a__	(       self					,   lowerCAmelCase__							):



            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(
                              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
                   _A=						timeit.repeat(
                       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(snake_case__							) / 10.0
            except RuntimeError as e:
                   self.print_fn(f"Doesn\'t fit on GPU. {e}"							)
                   return "N/A"





     def       a__	(       self					,   lowerCAmelCase__							):



            try:
                   if self.args.trace_memory_line_by_line:
                          _A=						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.'							)
                                 _A=						'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()
                                 _A=						nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx							)
                                 _A=						nvml.nvmlDeviceGetMemoryInfo(snake_case__							)
                                 _A=						meminfo.used
                                 _A=						Memory(snake_case__							)
                                 # shutdown nvml
                                 nvml.nvmlShutdown()
                   else:
                          # cpu
                          _A=						measure_peak_memory_cpu(snake_case__							)
                          _A=						Memory(snake_case__							) if isinstance(snake_case__					,   snake_case__							) else memory_bytes

                   if self.args.trace_memory_line_by_line:
                          _A=						stop_memory_tracing(snake_case__							)
                   else:
                          _A=						None

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







import logging
import os
import sys
import threading
from logging import (
    CRITICAL,  # NOQA
    DEBUG,  # NOQA
    ERROR,  # NOQA
    FATAL,  # NOQA
    INFO,  # NOQA
    NOTSET,  # NOQA
    WARN,  # NOQA
    WARNING,  # NOQA
)
from typing import Optional

from tqdm import auto as tqdm_lib


lowerCamelCase					    =	threading.Lock()
lowerCamelCase					    =	None

lowerCamelCase					    =	{
    '''debug''': logging.DEBUG,
    '''info''': logging.INFO,
    '''warning''': logging.WARNING,
    '''error''': logging.ERROR,
    '''critical''': logging.CRITICAL,
}

lowerCamelCase					    =	logging.WARNING

lowerCamelCase					    =	True
def 							UpperCAmelCase   ( ):


    '''simple docstring'''



    _lowercase     			=						os.getenv('DIFFUSERS_VERBOSITY',   _lowercase   )
    if env_level_str:
        if env_level_str in log_levels:
            return log_levels[env_level_str]
        else:
            logging.getLogger().warning(
                f'''Unknown option DIFFUSERS_VERBOSITY={env_level_str}, '''
                f'''has to be one of: { ", ".join(log_levels.keys()   ) }'''   )
    return _default_log_level
def 							UpperCAmelCase   ( ):


    '''simple docstring'''



    return __name__.split('.'   )[0]
def 							UpperCAmelCase   ( ):


    '''simple docstring'''



    return logging.getLogger(_get_library_name()   )
def 							UpperCAmelCase   ( ):


    '''simple docstring'''



    global _default_handler

    with _lock:
        if _default_handler:
            # This library has already configured the library root logger.
            return
        _lowercase     			=						logging.StreamHandler()  # Set sys.stderr as stream.
        _lowercase     			=						sys.stderr.flush

        # Apply our default configuration to the library root logger.
        _lowercase     			=						_get_library_root_logger()
        library_root_logger.addHandler(_default_handler   )
        library_root_logger.setLevel(_get_default_logging_level()   )
        _lowercase     			=						False
def 							UpperCAmelCase   ( ):


    '''simple docstring'''



    global _default_handler

    with _lock:
        if not _default_handler:
            return

        _lowercase     			=						_get_library_root_logger()
        library_root_logger.removeHandler(_default_handler   )
        library_root_logger.setLevel(logging.NOTSET   )
        _lowercase     			=						None
def 							UpperCAmelCase   ( ):


    '''simple docstring'''



    return log_levels
def 							UpperCAmelCase   ( _lowercase = None   ):


    '''simple docstring'''



    if name is None:
        _lowercase     			=						_get_library_name()

    _configure_library_root_logger()
    return logging.getLogger(_lowercase   )
def 							UpperCAmelCase   ( ):


    '''simple docstring'''



    _configure_library_root_logger()
    return _get_library_root_logger().getEffectiveLevel()
def 							UpperCAmelCase   ( _lowercase   ):


    '''simple docstring'''



    _configure_library_root_logger()
    _get_library_root_logger().setLevel(_lowercase   )
def 							UpperCAmelCase   ( ):


    '''simple docstring'''



    return set_verbosity(_lowercase   )
def 							UpperCAmelCase   ( ):


    '''simple docstring'''



    return set_verbosity(_lowercase   )
def 							UpperCAmelCase   ( ):


    '''simple docstring'''



    return set_verbosity(_lowercase   )
def 							UpperCAmelCase   ( ):


    '''simple docstring'''



    return set_verbosity(_lowercase   )
def 							UpperCAmelCase   ( ):


    '''simple docstring'''



    _configure_library_root_logger()

    assert _default_handler is not None
    _get_library_root_logger().removeHandler(_default_handler   )
def 							UpperCAmelCase   ( ):


    '''simple docstring'''



    _configure_library_root_logger()

    assert _default_handler is not None
    _get_library_root_logger().addHandler(_default_handler   )
def 							UpperCAmelCase   ( _lowercase   ):


    '''simple docstring'''



    _configure_library_root_logger()

    assert handler is not None
    _get_library_root_logger().addHandler(_lowercase   )
def 							UpperCAmelCase   ( _lowercase   ):


    '''simple docstring'''



    _configure_library_root_logger()

    assert handler is not None and handler not in _get_library_root_logger().handlers
    _get_library_root_logger().removeHandler(_lowercase   )
def 							UpperCAmelCase   ( ):


    '''simple docstring'''



    _configure_library_root_logger()
    _lowercase     			=						False
def 							UpperCAmelCase   ( ):


    '''simple docstring'''



    _configure_library_root_logger()
    _lowercase     			=						True
def 							UpperCAmelCase   ( ):


    '''simple docstring'''



    _lowercase     			=						_get_library_root_logger().handlers

    for handler in handlers:
        _lowercase     			=						logging.Formatter('[%(levelname)s|%(filename)s:%(lineno)s] %(asctime)s >> %(message)s'   )
        handler.setFormatter(_lowercase   )
def 							UpperCAmelCase   ( ):


    '''simple docstring'''



    _lowercase     			=						_get_library_root_logger().handlers

    for handler in handlers:
        handler.setFormatter(_lowercase   )
def 							UpperCAmelCase   ( self,   *_lowercase,   **_lowercase   ):


    '''simple docstring'''



    _lowercase     			=						os.getenv('DIFFUSERS_NO_ADVISORY_WARNINGS',   _lowercase   )
    if no_advisory_warnings:
        return
    self.warning(*_lowercase,   **_lowercase   )


lowerCamelCase					    =	warning_advice

class        snake_case_				:







   """simple docstring"""



   def __init__(     self:						int ,							*snake_case__:						str ,							**snake_case__:						Any						)	->		Any:  # pylint: disable=unused-argument


       """simple docstring"""



       _lowercase     			=						args[0] if args else None

   def __iter__(     self:						Union[str, Any]						)	->		str:


       """simple docstring"""



       return iter(self._iterator						)

   def __getattr__(     self:						int ,							snake_case__:						Any						)	->		List[Any]:


       """simple docstring"""



       def empty_fn(*snake_case__:						Optional[int] ,							**snake_case__:						Union[str, Any]						):  # pylint: disable=unused-argument
           return

       return empty_fn

   def __enter__(     self:						List[Any]						)	->		Tuple:


       """simple docstring"""



       return self

   def __exit__(     self:						List[str] ,							snake_case__:						Optional[Any] ,							snake_case__:						Dict ,							snake_case__:						List[str]						)	->		Optional[Any]:


       """simple docstring"""



       return






class        snake_case_				:







   """simple docstring"""



   def __call__(     self:						List[str] ,							*snake_case__:						Tuple ,							**snake_case__:						Union[str, Any]						)	->		List[str]:


       """simple docstring"""



       if _tqdm_active:
           return tqdm_lib.tqdm(*snake_case__ ,							**snake_case__						)
       else:
           return EmptyTqdm(*snake_case__ ,							**snake_case__						)

   def       SCREAMING_SNAKE_CASE__					(     self:						Tuple ,							*snake_case__:						Union[str, Any] ,							**snake_case__:						str						)	->		Optional[int]:


       """simple docstring"""



       _lowercase     			=						None
       if _tqdm_active:
           return tqdm_lib.tqdm.set_lock(*snake_case__ ,							**snake_case__						)

   def       SCREAMING_SNAKE_CASE__					(     self:						int						)	->		int:


       """simple docstring"""



       if _tqdm_active:
           return tqdm_lib.tqdm.get_lock()






lowerCamelCase					    =	_tqdm_cls()
def 							UpperCAmelCase   ( ):


    '''simple docstring'''



    global _tqdm_active
    return bool(_tqdm_active   )
def 							UpperCAmelCase   ( ):


    '''simple docstring'''



    global _tqdm_active
    _lowercase     			=						True
def 							UpperCAmelCase   ( ):


    '''simple docstring'''



    global _tqdm_active
    _lowercase     			=						False
