





'''simple docstring'''


import unittest

from transformers import load_tool

from .test_tools_common import ToolTesterMixin


A_				     =							"\nHugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf originally as a company that developed a chatbot app targeted at teenagers.[2] After open-sourcing the model behind the chatbot, the company pivoted to focus on being a platform for machine learning.\n\nIn March 2021, Hugging Face raised $40 million in a Series B funding round.[3]\n\nOn April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model.[4] In 2022, the workshop concluded with the announcement of BLOOM, a multilingual large language model with 176 billion parameters.[5]\n"


class 		UpperCamelCase__     (       unittest.TestCase	,	a							):
      '''simple docstring'''






      def 	snake_case  (     self )				->  Optional[int]:
        __lowerCAmelCase		:    int			     =  load_tool('summarization' )
        self.tool.setup()
        __lowerCAmelCase		:    Optional[Any]			     =  load_tool('summarization'     , remote=SCREAMING_SNAKE_CASE )




      def 	snake_case  (     self )				->  Any:
        __lowerCAmelCase		:    Any			     =  self.tool(SCREAMING_SNAKE_CASE )
        self.assertEqual(
            SCREAMING_SNAKE_CASE     , 'Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'     , )




      def 	snake_case  (     self )				->  Tuple:
        __lowerCAmelCase		:    Union[str, Any]			     =  self.remote_tool(SCREAMING_SNAKE_CASE )
        self.assertEqual(
            SCREAMING_SNAKE_CASE     , 'Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'     , )




      def 	snake_case  (     self )				->  Dict:
        __lowerCAmelCase		:    Any			     =  self.tool(text=SCREAMING_SNAKE_CASE )
        self.assertEqual(
            SCREAMING_SNAKE_CASE     , 'Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'     , )







      def 	snake_case  (     self )				->  Optional[int]:
        __lowerCAmelCase		:    str			     =  self.remote_tool(text=SCREAMING_SNAKE_CASE )
        self.assertEqual(
            SCREAMING_SNAKE_CASE     , 'Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'     , )












'''simple docstring'''


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


A_				     =							logging.get_logger(__name__)


class 		UpperCamelCase__     (       a							):
      '''simple docstring'''






      _snake_case           =				42
      _snake_case           =				42
      _snake_case           =				"PyTorch"




      @property
      def 	snake_case  (     self )				->  Dict:
        return torch.__version__




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )				->  float:
        __lowerCAmelCase		:    Any			     =  self._prepare_inference_func(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )
        return self._measure_speed(_inference )




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )				->  [Memory, Optional[MemorySummary]]:
        __lowerCAmelCase		:    Dict			     =  self._prepare_inference_func(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )
        return self._measure_memory(_inference )




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )				->  float:
        __lowerCAmelCase		:    Optional[Any]			     =  self._prepare_train_func(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )
        return self._measure_speed(_train )




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )				->  [Memory, Optional[MemorySummary]]:
        __lowerCAmelCase		:    Optional[int]			     =  self._prepare_train_func(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )
        return self._measure_memory(_train )




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )				->  Callable[[], None]:
        __lowerCAmelCase		:    str			     =  self.config_dict[model_name]

        if self.args.torchscript:
          __lowerCAmelCase		:    Tuple			     =  True

        __lowerCAmelCase		:    Union[str, Any]			     =  (
            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:
            __lowerCAmelCase		:    Dict			     =  config.architectures[0]
            __lowerCAmelCase		:    List[str]			     =  __import__('transformers'     , fromlist=[model_class] )
            __lowerCAmelCase		:    Tuple			     =  getattr(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )
            __lowerCAmelCase		:    Dict			     =  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:
          __lowerCAmelCase		:    List[Any]			     =  MODEL_MAPPING[config.__class__](SCREAMING_SNAKE_CASE )

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

        # encoder-decoder has vocab size saved differently
        __lowerCAmelCase		:    List[str]			     =  config.vocab_size if hasattr(SCREAMING_SNAKE_CASE     , 'vocab_size' ) else config.encoder.vocab_size
        __lowerCAmelCase		:    Optional[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()

        if self.args.torchscript:
          with torch.no_grad():
            __lowerCAmelCase		:    List[Any]			     =  torch.jit.trace(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )
        else:
          __lowerCAmelCase		:    str			     =  model

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

        def encoder_forward():
          with torch.no_grad():
            __lowerCAmelCase		:    str			     =  inference_model(SCREAMING_SNAKE_CASE )
          return outputs

        __lowerCAmelCase		:    Optional[Any]			     =  encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
        return _forward




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )				->  Callable[[], None]:
        __lowerCAmelCase		:    int			     =  self.config_dict[model_name]

        __lowerCAmelCase		:    Optional[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:
            __lowerCAmelCase		:    List[Any]			     =  config.architectures[0]
            __lowerCAmelCase		:    List[str]			     =  __import__('transformers'     , fromlist=[model_class] )
            __lowerCAmelCase		:    Tuple			     =  getattr(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )
            __lowerCAmelCase		:    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:
          __lowerCAmelCase		:    List[str]			     =  MODEL_WITH_LM_HEAD_MAPPING[config.__class__](SCREAMING_SNAKE_CASE )

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

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

        # encoder-decoder has vocab size saved differently
        __lowerCAmelCase		:    Tuple			     =  config.vocab_size if hasattr(SCREAMING_SNAKE_CASE     , 'vocab_size' ) else config.encoder.vocab_size
        __lowerCAmelCase		:    Optional[int]			     =  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():
          __lowerCAmelCase		:    List[str]			     =  train_model(SCREAMING_SNAKE_CASE     , labels=SCREAMING_SNAKE_CASE )[0]
          loss.backward()
          return loss

        def compute_loss_and_backprob_encoder_decoder():
          __lowerCAmelCase		:    List[Any]			     =  train_model(SCREAMING_SNAKE_CASE     , decoder_input_ids=SCREAMING_SNAKE_CASE     , labels=SCREAMING_SNAKE_CASE )[0]
          loss.backward()
          return loss

        __lowerCAmelCase		:    Dict			     =  (
            compute_loss_and_backprob_encoder_decoder
            if config.is_encoder_decoder
            else compute_loss_and_backprob_encoder
        )
        return _train




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE )				->  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(
                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
          __lowerCAmelCase		:    Tuple			     =  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 ) / 1_0.0
        except RuntimeError as e:
          self.print_fn(F"""Doesn't fit on GPU. {e}""" )
          return "N/A"







      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE )				->  [Memory, MemorySummary]:
        try:
          if self.args.trace_memory_line_by_line:
            __lowerCAmelCase		:    Tuple			     =  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.' )
              __lowerCAmelCase		:    Union[str, 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()
              __lowerCAmelCase		:    int			     =  nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx )
              __lowerCAmelCase		:    List[str]			     =  nvml.nvmlDeviceGetMemoryInfo(SCREAMING_SNAKE_CASE )
              __lowerCAmelCase		:    Tuple			     =  meminfo.used
              __lowerCAmelCase		:    Dict			     =  Memory(SCREAMING_SNAKE_CASE )
              # shutdown nvml
              nvml.nvmlShutdown()
          else:
            # cpu
            __lowerCAmelCase		:    Tuple			     =  measure_peak_memory_cpu(SCREAMING_SNAKE_CASE )
            __lowerCAmelCase		:    str			     =  Memory(SCREAMING_SNAKE_CASE ) if isinstance(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE ) else memory_bytes

          if self.args.trace_memory_line_by_line:
            __lowerCAmelCase		:    int			     =  stop_memory_tracing(SCREAMING_SNAKE_CASE )
          else:
            __lowerCAmelCase		:    Any			     =  None

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





