

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


lowerCamelCase_					  =   logging.get_logger(__name__)







class 					__lowerCamelCase							(      a_      ):
      lowerCamelCase_      : PyTorchBenchmarkArguments
      lowerCamelCase_      : PretrainedConfig
      lowerCamelCase_      : str  			=   "PyTorch"







      @property
      def       lowerCAmelCase_       (	self  )		->			Any:
        return torch.__version__







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase  )		->			float:
        snake_case_					= self._prepare_inference_func(lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase  )
        return self._measure_speed(_inference  )







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase  )		->			[Memory, Optional[MemorySummary]]:
        snake_case_					= self._prepare_inference_func(lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase  )
        return self._measure_memory(_inference  )







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase  )		->			float:
        snake_case_					= self._prepare_train_func(lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase  )
        return self._measure_speed(_train  )







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase  )		->			[Memory, Optional[MemorySummary]]:
        snake_case_					= self._prepare_train_func(lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase  )
        return self._measure_memory(_train  )







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase  )		->			Callable[[], None]:
        snake_case_					= self.config_dict[model_name]

        if self.args.torchscript:
          snake_case_					= True

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

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

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

        def encoder_decoder_forward():
          with torch.no_grad():
            snake_case_					= inference_model(lowerCamelCase				,    decoder_input_ids=lowerCamelCase  )
          return outputs

        def encoder_forward():
          with torch.no_grad():
            snake_case_					= inference_model(lowerCamelCase  )
          return outputs

        snake_case_					= encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
        return _forward







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase  )		->			Callable[[], None]:
        snake_case_					= self.config_dict[model_name]

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

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

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

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

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

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







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







      def       lowerCAmelCase_       (	self				,    lowerCamelCase  )		->			[Memory, MemorySummary]:
        try:
          if self.args.trace_memory_line_by_line:
            snake_case_					= 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."""  )
              snake_case_					= '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()
              snake_case_					= nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx  )
              snake_case_					= nvml.nvmlDeviceGetMemoryInfo(lowerCamelCase  )
              snake_case_					= meminfo.used
              snake_case_					= Memory(lowerCamelCase  )
              # shutdown nvml
              nvml.nvmlShutdown()
          else:
            # cpu
            snake_case_					= measure_peak_memory_cpu(lowerCamelCase  )
            snake_case_					= Memory(lowerCamelCase  ) if isinstance(lowerCamelCase				,    lowerCamelCase  ) else memory_bytes

          if self.args.trace_memory_line_by_line:
            snake_case_					= stop_memory_tracing(lowerCamelCase  )
          else:
            snake_case_					= None

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

from ..utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_tf_available, is_torch_available


UpperCAmelCase      		=					{
    '''configuration_utils''': ['''GenerationConfig'''],
    '''streamers''': ['''TextIteratorStreamer''', '''TextStreamer'''],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 UpperCAmelCase      		=					[
     '''Constraint''',
     '''ConstraintListState''',
     '''DisjunctiveConstraint''',
     '''PhrasalConstraint''',
 ]
 UpperCAmelCase      		=					[
     '''BeamHypotheses''',
     '''BeamScorer''',
     '''BeamSearchScorer''',
     '''ConstrainedBeamSearchScorer''',
 ]
 UpperCAmelCase      		=					[
     '''EpsilonLogitsWarper''',
     '''EtaLogitsWarper''',
     '''ForcedBOSTokenLogitsProcessor''',
     '''ForcedEOSTokenLogitsProcessor''',
     '''HammingDiversityLogitsProcessor''',
     '''InfNanRemoveLogitsProcessor''',
     '''LogitsProcessor''',
     '''LogitsProcessorList''',
     '''LogitsWarper''',
     '''MinLengthLogitsProcessor''',
     '''MinNewTokensLengthLogitsProcessor''',
     '''NoBadWordsLogitsProcessor''',
     '''NoRepeatNGramLogitsProcessor''',
     '''PrefixConstrainedLogitsProcessor''',
     '''RepetitionPenaltyLogitsProcessor''',
     '''SequenceBiasLogitsProcessor''',
     '''EncoderRepetitionPenaltyLogitsProcessor''',
     '''TemperatureLogitsWarper''',
     '''TopKLogitsWarper''',
     '''TopPLogitsWarper''',
     '''TypicalLogitsWarper''',
     '''EncoderNoRepeatNGramLogitsProcessor''',
     '''ExponentialDecayLengthPenalty''',
     '''LogitNormalization''',
 ]
 UpperCAmelCase      		=					[
     '''MaxNewTokensCriteria''',
     '''MaxLengthCriteria''',
     '''MaxTimeCriteria''',
     '''StoppingCriteria''',
     '''StoppingCriteriaList''',
     '''validate_stopping_criteria''',
 ]
 UpperCAmelCase      		=					[
     '''GenerationMixin''',
     '''top_k_top_p_filtering''',
     '''GreedySearchEncoderDecoderOutput''',
     '''GreedySearchDecoderOnlyOutput''',
     '''SampleEncoderDecoderOutput''',
     '''SampleDecoderOnlyOutput''',
     '''BeamSearchEncoderDecoderOutput''',
     '''BeamSearchDecoderOnlyOutput''',
     '''BeamSampleEncoderDecoderOutput''',
     '''BeamSampleDecoderOnlyOutput''',
     '''ContrastiveSearchEncoderDecoderOutput''',
     '''ContrastiveSearchDecoderOnlyOutput''',
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 UpperCAmelCase      		=					[
     '''TFForcedBOSTokenLogitsProcessor''',
     '''TFForcedEOSTokenLogitsProcessor''',
     '''TFLogitsProcessor''',
     '''TFLogitsProcessorList''',
     '''TFLogitsWarper''',
     '''TFMinLengthLogitsProcessor''',
     '''TFNoBadWordsLogitsProcessor''',
     '''TFNoRepeatNGramLogitsProcessor''',
     '''TFRepetitionPenaltyLogitsProcessor''',
     '''TFTemperatureLogitsWarper''',
     '''TFTopKLogitsWarper''',
     '''TFTopPLogitsWarper''',
     '''TFForceTokensLogitsProcessor''',
     '''TFSuppressTokensAtBeginLogitsProcessor''',
     '''TFSuppressTokensLogitsProcessor''',
 ]
 UpperCAmelCase      		=					[
     '''TFGenerationMixin''',
     '''tf_top_k_top_p_filtering''',
     '''TFGreedySearchDecoderOnlyOutput''',
     '''TFGreedySearchEncoderDecoderOutput''',
     '''TFSampleEncoderDecoderOutput''',
     '''TFSampleDecoderOnlyOutput''',
     '''TFBeamSearchEncoderDecoderOutput''',
     '''TFBeamSearchDecoderOnlyOutput''',
     '''TFBeamSampleEncoderDecoderOutput''',
     '''TFBeamSampleDecoderOnlyOutput''',
     '''TFContrastiveSearchEncoderDecoderOutput''',
     '''TFContrastiveSearchDecoderOnlyOutput''',
 ]

try:
 if not is_flax_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 UpperCAmelCase      		=					[
     '''FlaxForcedBOSTokenLogitsProcessor''',
     '''FlaxForcedEOSTokenLogitsProcessor''',
     '''FlaxLogitsProcessor''',
     '''FlaxLogitsProcessorList''',
     '''FlaxLogitsWarper''',
     '''FlaxMinLengthLogitsProcessor''',
     '''FlaxTemperatureLogitsWarper''',
     '''FlaxTopKLogitsWarper''',
     '''FlaxTopPLogitsWarper''',
 ]
 UpperCAmelCase      		=					[
     '''FlaxGenerationMixin''',
     '''FlaxGreedySearchOutput''',
     '''FlaxSampleOutput''',
     '''FlaxBeamSearchOutput''',
 ]

if TYPE_CHECKING:
 from .configuration_utils import GenerationConfig
 from .streamers import TextIteratorStreamer, TextStreamer

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .beam_constraints import Constraint, ConstraintListState, DisjunctiveConstraint, PhrasalConstraint
  from .beam_search import BeamHypotheses, BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer
  from .logits_process import (
      EncoderNoRepeatNGramLogitsProcessor,
      EncoderRepetitionPenaltyLogitsProcessor,
      EpsilonLogitsWarper,
      EtaLogitsWarper,
      ExponentialDecayLengthPenalty,
      ForcedBOSTokenLogitsProcessor,
      ForcedEOSTokenLogitsProcessor,
      HammingDiversityLogitsProcessor,
      InfNanRemoveLogitsProcessor,
      LogitNormalization,
      LogitsProcessor,
      LogitsProcessorList,
      LogitsWarper,
      MinLengthLogitsProcessor,
      MinNewTokensLengthLogitsProcessor,
      NoBadWordsLogitsProcessor,
      NoRepeatNGramLogitsProcessor,
      PrefixConstrainedLogitsProcessor,
      RepetitionPenaltyLogitsProcessor,
      SequenceBiasLogitsProcessor,
      TemperatureLogitsWarper,
      TopKLogitsWarper,
      TopPLogitsWarper,
      TypicalLogitsWarper,
  )
  from .stopping_criteria import (
      MaxLengthCriteria,
      MaxNewTokensCriteria,
      MaxTimeCriteria,
      StoppingCriteria,
      StoppingCriteriaList,
      validate_stopping_criteria,
  )
  from .utils import (
      BeamSampleDecoderOnlyOutput,
      BeamSampleEncoderDecoderOutput,
      BeamSearchDecoderOnlyOutput,
      BeamSearchEncoderDecoderOutput,
      ContrastiveSearchDecoderOnlyOutput,
      ContrastiveSearchEncoderDecoderOutput,
      GenerationMixin,
      GreedySearchDecoderOnlyOutput,
      GreedySearchEncoderDecoderOutput,
      SampleDecoderOnlyOutput,
      SampleEncoderDecoderOutput,
      top_k_top_p_filtering,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tf_logits_process import (
      TFForcedBOSTokenLogitsProcessor,
      TFForcedEOSTokenLogitsProcessor,
      TFForceTokensLogitsProcessor,
      TFLogitsProcessor,
      TFLogitsProcessorList,
      TFLogitsWarper,
      TFMinLengthLogitsProcessor,
      TFNoBadWordsLogitsProcessor,
      TFNoRepeatNGramLogitsProcessor,
      TFRepetitionPenaltyLogitsProcessor,
      TFSuppressTokensAtBeginLogitsProcessor,
      TFSuppressTokensLogitsProcessor,
      TFTemperatureLogitsWarper,
      TFTopKLogitsWarper,
      TFTopPLogitsWarper,
  )
  from .tf_utils import (
      TFBeamSampleDecoderOnlyOutput,
      TFBeamSampleEncoderDecoderOutput,
      TFBeamSearchDecoderOnlyOutput,
      TFBeamSearchEncoderDecoderOutput,
      TFContrastiveSearchDecoderOnlyOutput,
      TFContrastiveSearchEncoderDecoderOutput,
      TFGenerationMixin,
      TFGreedySearchDecoderOnlyOutput,
      TFGreedySearchEncoderDecoderOutput,
      TFSampleDecoderOnlyOutput,
      TFSampleEncoderDecoderOutput,
      tf_top_k_top_p_filtering,
  )

 try:
  if not is_flax_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .flax_logits_process import (
      FlaxForcedBOSTokenLogitsProcessor,
      FlaxForcedEOSTokenLogitsProcessor,
      FlaxLogitsProcessor,
      FlaxLogitsProcessorList,
      FlaxLogitsWarper,
      FlaxMinLengthLogitsProcessor,
      FlaxTemperatureLogitsWarper,
      FlaxTopKLogitsWarper,
      FlaxTopPLogitsWarper,
  )
  from .flax_utils import FlaxBeamSearchOutput, FlaxGenerationMixin, FlaxGreedySearchOutput, FlaxSampleOutput
else:
 import sys

 UpperCAmelCase      		=					_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



