






"""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


UpperCamelCase											=     logging.get_logger(__name__)




class  lowercase_     (lowerCAmelCase_  ):
      A__						:		str							      =					42
      A__						:		Dict							      =					42
      A__						:		Optional[int]							      =					'''PyTorch'''






      @property
      def       lowerCamelCase__     (     self						)						->Optional[Any]:







        '''simple docstring'''


        return torch.__version__






      def       lowerCamelCase__     (     self  ,  a_  ,  a_  ,  a_						)						->float:







        '''simple docstring'''


        _a					      =						self._prepare_inference_func(__UpperCamelCase  ,  __UpperCamelCase  ,  __UpperCamelCase						)
        return self._measure_speed(_inference						)






      def       lowerCamelCase__     (     self  ,  a_  ,  a_  ,  a_						)						->[Memory, Optional[MemorySummary]]:







        '''simple docstring'''


        _a					      =						self._prepare_inference_func(__UpperCamelCase  ,  __UpperCamelCase  ,  __UpperCamelCase						)
        return self._measure_memory(_inference						)






      def       lowerCamelCase__     (     self  ,  a_  ,  a_  ,  a_						)						->float:







        '''simple docstring'''


        _a					      =						self._prepare_train_func(__UpperCamelCase  ,  __UpperCamelCase  ,  __UpperCamelCase						)
        return self._measure_speed(_train						)






      def       lowerCamelCase__     (     self  ,  a_  ,  a_  ,  a_						)						->[Memory, Optional[MemorySummary]]:







        '''simple docstring'''


        _a					      =						self._prepare_train_func(__UpperCamelCase  ,  __UpperCamelCase  ,  __UpperCamelCase						)
        return self._measure_memory(_train						)






      def       lowerCamelCase__     (     self  ,  a_  ,  a_  ,  a_						)						->Callable[[], None]:







        '''simple docstring'''


        _a					      =						self.config_dict[model_name]

        if self.args.torchscript:
          _a					      =						True

        _a					      =						(
            hasattr(__UpperCamelCase  ,  "architectures"						)
            and isinstance(config.architectures  ,  __UpperCamelCase						)
            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(__UpperCamelCase  ,  __UpperCamelCase						)
            _a					      =						model_cls(__UpperCamelCase						)
          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__](__UpperCamelCase						)

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

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

        def encoder_decoder_forward():
          with torch.no_grad():
            _a					      =						inference_model(__UpperCamelCase  ,  decoder_input_ids=__UpperCamelCase						)
          return outputs

        def encoder_forward():
          with torch.no_grad():
            _a					      =						inference_model(__UpperCamelCase						)
          return outputs

        _a					      =						encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
        return _forward






      def       lowerCamelCase__     (     self  ,  a_  ,  a_  ,  a_						)						->Callable[[], None]:







        '''simple docstring'''


        _a					      =						self.config_dict[model_name]

        _a					      =						(
            hasattr(__UpperCamelCase  ,  "architectures"						)
            and isinstance(config.architectures  ,  __UpperCamelCase						)
            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(__UpperCamelCase  ,  __UpperCamelCase						)
            _a					      =						model_cls(__UpperCamelCase						)
          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__](__UpperCamelCase						)

        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(__UpperCamelCase  ,  "vocab_size"						) else config.encoder.vocab_size
        _a					      =						torch.randint(__UpperCamelCase  ,  (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(__UpperCamelCase  ,  labels=__UpperCamelCase						)[0]
          loss.backward()
          return loss

        def compute_loss_and_backprob_encoder_decoder():
          _a					      =						train_model(__UpperCamelCase  ,  decoder_input_ids=__UpperCamelCase  ,  labels=__UpperCamelCase						)[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       lowerCamelCase__     (     self  ,  a_						)						->float:







        '''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(
                __UpperCamelCase  ,  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(
              __UpperCamelCase  ,  repeat=self.args.repeat  ,  number=1_0  ,  )

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






      def       lowerCamelCase__     (     self  ,  a_						)						->[Memory, MemorySummary]:







        '''simple docstring'''


        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(__UpperCamelCase						)
              _a					      =						meminfo.used
              _a					      =						Memory(__UpperCamelCase						)
              # shutdown nvml
              nvml.nvmlShutdown()
          else:
            # cpu
            _a					      =						measure_peak_memory_cpu(__UpperCamelCase						)
            _a					      =						Memory(__UpperCamelCase						) if isinstance(__UpperCamelCase  ,  __UpperCamelCase						) else memory_bytes

          if self.args.trace_memory_line_by_line:
            _a					      =						stop_memory_tracing(__UpperCamelCase						)
          else:
            _a					      =						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 argparse
import os


A__   :						Any 							=     {
    """pytorch""": {
        """1.13""": {
            """torch""": """1.13.1""",
            """torchvision""": """0.14.1""",
            """torchaudio""": """0.13.1""",
            """python""": 3.9,
            """cuda""": """cu116""",
            """install""": (
                """python3 -m pip install --no-cache-dir -U torch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1"""
                """ --extra-index-url https://download.pytorch.org/whl/cu116"""
            ),
            """base_image""": """nvidia/cuda:11.6.2-cudnn8-devel-ubuntu20.04""",
        },
        """1.12""": {
            """torch""": """1.12.1""",
            """torchvision""": """0.13.1""",
            """torchaudio""": """0.12.1""",
            """python""": 3.9,
            """cuda""": """cu113""",
            """install""": (
                """python3 -m pip install --no-cache-dir -U torch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1"""
                """ --extra-index-url https://download.pytorch.org/whl/cu113"""
            ),
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
        """1.11""": {
            """torch""": """1.11.0""",
            """torchvision""": """0.12.0""",
            """torchaudio""": """0.11.0""",
            """python""": 3.9,
            """cuda""": """cu113""",
            """install""": (
                """python3 -m pip install --no-cache-dir -U torch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0"""
                """ --extra-index-url https://download.pytorch.org/whl/cu113"""
            ),
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
        """1.10""": {
            """torch""": """1.10.2""",
            """torchvision""": """0.11.3""",
            """torchaudio""": """0.10.2""",
            """python""": 3.9,
            """cuda""": """cu113""",
            """install""": (
                """python3 -m pip install --no-cache-dir -U torch==1.10.2 torchvision==0.11.3 torchaudio==0.10.2"""
                """ --extra-index-url https://download.pytorch.org/whl/cu113"""
            ),
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
        # torchaudio < 0.10 has no CUDA-enabled binary distributions
        """1.9""": {
            """torch""": """1.9.1""",
            """torchvision""": """0.10.1""",
            """torchaudio""": """0.9.1""",
            """python""": 3.9,
            """cuda""": """cu111""",
            """install""": (
                """python3 -m pip install --no-cache-dir -U torch==1.9.1 torchvision==0.10.1 torchaudio==0.9.1"""
                """ --extra-index-url https://download.pytorch.org/whl/cu111"""
            ),
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
    },
    """tensorflow""": {
        """2.11""": {
            """tensorflow""": """2.11.1""",
            """install""": """python3 -m pip install --no-cache-dir -U tensorflow==2.11.1""",
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
        """2.10""": {
            """tensorflow""": """2.10.1""",
            """install""": """python3 -m pip install --no-cache-dir -U tensorflow==2.10.1""",
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
        """2.9""": {
            """tensorflow""": """2.9.3""",
            """install""": """python3 -m pip install --no-cache-dir -U tensorflow==2.9.3""",
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
        """2.8""": {
            """tensorflow""": """2.8.2""",
            """install""": """python3 -m pip install --no-cache-dir -U tensorflow==2.8.2""",
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
        """2.7""": {
            """tensorflow""": """2.7.3""",
            """install""": """python3 -m pip install --no-cache-dir -U tensorflow==2.7.3""",
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
        """2.6""": {
            """tensorflow""": """2.6.5""",
            """install""": """python3 -m pip install --no-cache-dir -U tensorflow==2.6.5""",
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
        """2.5""": {
            """tensorflow""": """2.5.3""",
            """install""": """python3 -m pip install --no-cache-dir -U tensorflow==2.5.3""",
            """base_image""": """nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04""",
        },
    },
}


if __name__ == "__main__":
    A__   :						str 							=     argparse.ArgumentParser("""Choose the framework and version to install""")
    parser.add_argument(
        """--framework""", help="""The framework to install. Should be `torch` or `tensorflow`""", type=str, required=True
    )
    parser.add_argument("""--version""", help="""The version of the framework to install.""", type=str, required=True)
    A__   :						Any 							=     parser.parse_args()

    A__   :						List[Any] 							=     past_versions_testing[args.framework][args.version]

    os.system(f"""echo \"export INSTALL_CMD=\'{info["install"]}\'\" >> ~/.profile""")
    print(f"""echo \"export INSTALL_CMD=\'{info["install"]}\'\" >> ~/.profile""")

    A__   :						Any 							=     """"""
    if args.framework == "pytorch":
        A__   :						str 							=     info["""cuda"""]
    os.system(f"""echo \"export CUDA='{cuda}'\" >> ~/.profile""")
    print(f"""echo \"export CUDA='{cuda}'\" >> ~/.profile""")



