




'''simple docstring'''






import argparse
import json
import logging
import os
import sys
from unittest import skip
from unittest.mock import patch

import tensorflow as tf

from transformers.testing_utils import TestCasePlus, get_gpu_count, slow


lowercase						:   List[str]      						=							[
    os.path.join(os.path.dirname(__file__), dirname)
    for dirname in [
        '''text-generation''',
        '''text-classification''',
        '''token-classification''',
        '''language-modeling''',
        '''multiple-choice''',
        '''question-answering''',
        '''summarization''',
        '''translation''',
        '''image-classification''',
    ]
]
sys.path.extend(SRC_DIRS)


if SRC_DIRS is not None:
      import run_clm
      import run_image_classification
      import run_mlm
      import run_ner
      import run_qa as run_squad
      import run_summarization
      import run_swag
      import run_text_classification
      import run_translation


logging.basicConfig(level=logging.DEBUG)

lowercase						:   Union[str, Any]      						=							logging.getLogger()




def 		a_(			)					->Tuple:






       '''simple docstring'''




       lowercase	:		List[Any]    =   argparse.ArgumentParser()
       parser.add_argument('''-f'''   )
       lowercase	:		Union[str, Any]    =   parser.parse_args()
       return args.f




def 		a_(			lowerCamelCase   )					->str:






       '''simple docstring'''




       lowercase	:		Union[str, Any]    =   {}
       lowercase	:		Optional[int]    =   os.path.join(lowerCamelCase ,					'''all_results.json'''   )
       if os.path.exists(lowerCamelCase   ):
              with open(lowerCamelCase ,					'''r'''   ) as f:
                     lowercase	:		Dict    =   json.load(lowerCamelCase   )
       else:
              raise ValueError(f"""can't find {path}"""   )
       return results




def 		a_(			)					->Optional[int]:






       '''simple docstring'''




       return bool(tf.config.list_physical_devices('''GPU'''   )   )


lowercase						:   Optional[int]      						=							logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)





class 	SCREAMING_SNAKE_CASE_		(							lowerCamelCase_					):



 @skip('''Skipping until shape inference for to_tf_dataset PR is merged.'''			)
 def  SCREAMING_SNAKE_CASE(	self			):






        '''simple docstring'''




        lowercase	:		Union[str, Any]    =   self.get_auto_remove_tmp_dir()
        lowercase	:		Any    =   f"""
            run_text_classification.py
            --model_name_or_path distilbert-base-uncased
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --train_file ./tests/fixtures/tests_samples/MRPC/train.csv
            --validation_file ./tests/fixtures/tests_samples/MRPC/dev.csv
            --do_train
            --do_eval
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --learning_rate=1e-4
            --max_steps=10
            --warmup_steps=2
            --seed=42
            --max_seq_length=128
            """.split()

        if is_cuda_available():
               testargs.append('''--fp16'''			)

        with patch.object(A ,		'''argv''' ,		A			):
               run_text_classification.main()
               # Reset the mixed precision policy so we don't break other tests
               tf.keras.mixed_precision.set_global_policy('''float32'''			)
               lowercase	:		str    =   get_results(A			)
               self.assertGreaterEqual(result['''eval_accuracy'''] ,		0.7_5			)



 def  SCREAMING_SNAKE_CASE(	self			):






        '''simple docstring'''




        lowercase	:		Tuple    =   self.get_auto_remove_tmp_dir()
        lowercase	:		str    =   f"""
            run_clm.py
            --model_name_or_path distilgpt2
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --do_train
            --do_eval
            --block_size 128
            --per_device_train_batch_size 2
            --per_device_eval_batch_size 1
            --num_train_epochs 2
            --output_dir {tmp_dir}
            --overwrite_output_dir
            """.split()

        if len(tf.config.list_physical_devices('''GPU'''			)			) > 1:
               # Skipping because there are not enough batches to train the model + would need a drop_last to work.
               return

        with patch.object(A ,		'''argv''' ,		A			):
               run_clm.main()
               lowercase	:		Dict    =   get_results(A			)
               self.assertLess(result['''eval_perplexity'''] ,		100			)



 def  SCREAMING_SNAKE_CASE(	self			):






        '''simple docstring'''




        lowercase	:		int    =   self.get_auto_remove_tmp_dir()
        lowercase	:		int    =   f"""
            run_mlm.py
            --model_name_or_path distilroberta-base
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --max_seq_length 64
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --do_train
            --do_eval
            --prediction_loss_only
            --num_train_epochs=1
            --learning_rate=1e-4
        """.split()

        with patch.object(A ,		'''argv''' ,		A			):
               run_mlm.main()
               lowercase	:		Dict    =   get_results(A			)
               self.assertLess(result['''eval_perplexity'''] ,		42			)



 def  SCREAMING_SNAKE_CASE(	self			):






        '''simple docstring'''




        lowercase	:		List[str]    =   7 if get_gpu_count() > 1 else 2

        lowercase	:		List[Any]    =   self.get_auto_remove_tmp_dir()
        lowercase	:		Dict    =   f"""
            run_ner.py
            --model_name_or_path bert-base-uncased
            --train_file tests/fixtures/tests_samples/conll/sample.json
            --validation_file tests/fixtures/tests_samples/conll/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --do_train
            --do_eval
            --warmup_steps=2
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=2
            --num_train_epochs={epochs}
            --seed 7
        """.split()

        with patch.object(A ,		'''argv''' ,		A			):
               run_ner.main()
               lowercase	:		Union[str, Any]    =   get_results(A			)
               self.assertGreaterEqual(result['''accuracy'''] ,		0.7_5			)



 def  SCREAMING_SNAKE_CASE(	self			):






        '''simple docstring'''




        lowercase	:		Any    =   self.get_auto_remove_tmp_dir()
        lowercase	:		Dict    =   f"""
            run_qa.py
            --model_name_or_path bert-base-uncased
            --version_2_with_negative
            --train_file tests/fixtures/tests_samples/SQUAD/sample.json
            --validation_file tests/fixtures/tests_samples/SQUAD/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --max_steps=10
            --warmup_steps=2
            --do_train
            --do_eval
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
        """.split()

        with patch.object(A ,		'''argv''' ,		A			):
               run_squad.main()
               lowercase	:		List[str]    =   get_results(A			)
               self.assertGreaterEqual(result['''f1'''] ,		30			)
               self.assertGreaterEqual(result['''exact'''] ,		30			)



 def  SCREAMING_SNAKE_CASE(	self			):






        '''simple docstring'''




        lowercase	:		Tuple    =   self.get_auto_remove_tmp_dir()
        lowercase	:		List[Any]    =   f"""
            run_swag.py
            --model_name_or_path bert-base-uncased
            --train_file tests/fixtures/tests_samples/swag/sample.json
            --validation_file tests/fixtures/tests_samples/swag/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --max_steps=20
            --warmup_steps=2
            --do_train
            --do_eval
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
        """.split()

        with patch.object(A ,		'''argv''' ,		A			):
               run_swag.main()
               lowercase	:		Any    =   get_results(A			)
               self.assertGreaterEqual(result['''val_accuracy'''] ,		0.8			)



 @slow
 def  SCREAMING_SNAKE_CASE(	self			):






        '''simple docstring'''




        lowercase	:		Tuple    =   self.get_auto_remove_tmp_dir()
        lowercase	:		List[str]    =   f"""
            run_summarization.py
            --model_name_or_path t5-small
            --train_file tests/fixtures/tests_samples/xsum/sample.json
            --validation_file tests/fixtures/tests_samples/xsum/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --max_steps=50
            --warmup_steps=8
            --do_train
            --do_eval
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
        """.split()

        with patch.object(A ,		'''argv''' ,		A			):
               run_summarization.main()
               lowercase	:		str    =   get_results(A			)
               self.assertGreaterEqual(result['''rouge1'''] ,		10			)
               self.assertGreaterEqual(result['''rouge2'''] ,		2			)
               self.assertGreaterEqual(result['''rougeL'''] ,		7			)
               self.assertGreaterEqual(result['''rougeLsum'''] ,		7			)



 @slow
 def  SCREAMING_SNAKE_CASE(	self			):






        '''simple docstring'''




        lowercase	:		str    =   self.get_auto_remove_tmp_dir()
        lowercase	:		str    =   f"""
            run_translation.py
            --model_name_or_path Rocketknight1/student_marian_en_ro_6_1
            --source_lang en
            --target_lang ro
            --train_file tests/fixtures/tests_samples/wmt16/sample.json
            --validation_file tests/fixtures/tests_samples/wmt16/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --warmup_steps=8
            --do_train
            --do_eval
            --learning_rate=3e-3
            --num_train_epochs 12
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --source_lang en_XX
            --target_lang ro_RO
        """.split()

        with patch.object(A ,		'''argv''' ,		A			):
               run_translation.main()
               lowercase	:		Tuple    =   get_results(A			)
               self.assertGreaterEqual(result['''bleu'''] ,		30			)





 def  SCREAMING_SNAKE_CASE(	self			):






        '''simple docstring'''




        lowercase	:		int    =   self.get_auto_remove_tmp_dir()
        lowercase	:		Tuple    =   f"""
            run_image_classification.py
            --dataset_name hf-internal-testing/cats_vs_dogs_sample
            --model_name_or_path microsoft/resnet-18
            --do_train
            --do_eval
            --learning_rate 1e-4
            --per_device_train_batch_size 2
            --per_device_eval_batch_size 1
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --dataloader_num_workers 16
            --num_train_epochs 2
            --train_val_split 0.1
            --seed 42
            --ignore_mismatched_sizes True
            """.split()

        with patch.object(A ,		'''argv''' ,		A			):
               run_image_classification.main()
               lowercase	:		Dict    =   get_results(A			)
               self.assertGreaterEqual(result['''accuracy'''] ,		0.7			)










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


lowercase						:   List[Any]      						=							logging.get_logger(__name__)

class 	SCREAMING_SNAKE_CASE_		(							lowerCamelCase_					):
 __snake_case						:  PyTorchBenchmarkArguments
 __snake_case						:  PretrainedConfig
 __snake_case						:  str       =   "PyTorch"



 @property
 def  SCREAMING_SNAKE_CASE(	self			):






        '''simple docstring'''




        return torch.__version__



 def  SCREAMING_SNAKE_CASE(	self ,		A ,		A ,		A			):






        '''simple docstring'''




        lowercase	:		str    =   self._prepare_inference_func(A ,		A ,		A			)
        return self._measure_speed(_inference			)



 def  SCREAMING_SNAKE_CASE(	self ,		A ,		A ,		A			):






        '''simple docstring'''




        lowercase	:		int    =   self._prepare_inference_func(A ,		A ,		A			)
        return self._measure_memory(_inference			)



 def  SCREAMING_SNAKE_CASE(	self ,		A ,		A ,		A			):






        '''simple docstring'''




        lowercase	:		Union[str, Any]    =   self._prepare_train_func(A ,		A ,		A			)
        return self._measure_speed(_train			)



 def  SCREAMING_SNAKE_CASE(	self ,		A ,		A ,		A			):






        '''simple docstring'''




        lowercase	:		Any    =   self._prepare_train_func(A ,		A ,		A			)
        return self._measure_memory(_train			)



 def  SCREAMING_SNAKE_CASE(	self ,		A ,		A ,		A			):






        '''simple docstring'''




        lowercase	:		Any    =   self.config_dict[model_name]

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

        lowercase	:		Union[str, Any]    =   (
            hasattr(A ,		'''architectures'''			)
            and isinstance(config.architectures ,		A			)
            and len(config.architectures			) > 0
        )
        if not self.args.only_pretrain_model and has_model_class_in_config:
               try:
                      lowercase	:		Any    =   config.architectures[0]
                      lowercase	:		Dict    =   __import__('''transformers''' ,		fromlist=[model_class]			)
                      lowercase	:		List[str]    =   getattr(A ,		A			)
                      lowercase	:		Dict    =   model_cls(A			)
               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:
               lowercase	:		Optional[int]    =   MODEL_MAPPING[config.__class__](A			)

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

        # encoder-decoder has vocab size saved differently
        lowercase	:		str    =   config.vocab_size if hasattr(A ,		'''vocab_size'''			) else config.encoder.vocab_size
        lowercase	:		Union[str, Any]    =   torch.randint(A ,		(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():
                      lowercase	:		int    =   torch.jit.trace(A ,		A			)
        else:
               lowercase	:		Any    =   model

        def encoder_decoder_forward():
               with torch.no_grad():
                      lowercase	:		int    =   inference_model(A ,		decoder_input_ids=A			)
               return outputs

        def encoder_forward():
               with torch.no_grad():
                      lowercase	:		Tuple    =   inference_model(A			)
               return outputs

        lowercase	:		int    =   encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
        return _forward



 def  SCREAMING_SNAKE_CASE(	self ,		A ,		A ,		A			):






        '''simple docstring'''




        lowercase	:		Optional[int]    =   self.config_dict[model_name]

        lowercase	:		Optional[Any]    =   (
            hasattr(A ,		'''architectures'''			)
            and isinstance(config.architectures ,		A			)
            and len(config.architectures			) > 0
        )
        if not self.args.only_pretrain_model and has_model_class_in_config:
               try:
                      lowercase	:		str    =   config.architectures[0]
                      lowercase	:		List[str]    =   __import__('''transformers''' ,		fromlist=[model_class]			)
                      lowercase	:		int    =   getattr(A ,		A			)
                      lowercase	:		Any    =   model_cls(A			)
               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:
               lowercase	:		Union[str, Any]    =   MODEL_WITH_LM_HEAD_MAPPING[config.__class__](A			)

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

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

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

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

        lowercase	:		Any    =   (
            compute_loss_and_backprob_encoder_decoder
            if config.is_encoder_decoder
            else compute_loss_and_backprob_encoder
        )
        return _train



 def  SCREAMING_SNAKE_CASE(	self ,		A			):






        '''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(
                          A ,		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
               lowercase	:		str    =   timeit.repeat(
                   A ,		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(A			) / 1_0.0
        except RuntimeError as e:
               self.print_fn(f"""Doesn't fit on GPU. {e}"""			)
               return "N/A"





 def  SCREAMING_SNAKE_CASE(	self ,		A			):






        '''simple docstring'''




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

               if self.args.trace_memory_line_by_line:
                      lowercase	:		Tuple    =   stop_memory_tracing(A			)
               else:
                      lowercase	:		int    =   None

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




