





import unittest

from transformers import MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING, AutoTokenizer, is_vision_available
from transformers.pipelines import pipeline
from transformers.pipelines.document_question_answering import apply_tesseract
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_detectrona,
    require_pytesseract,
    require_tf,
    require_torch,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_vision_available():
       from PIL import Image

       from transformers.image_utils import load_image
else:


       class        __SCREAMING_SNAKE_CASE      :







              @staticmethod
              def 						__lowerCamelCase   (				*A					:     Optional[Any]		, **A					:     str      )       ->List[str]:
                 pass



       def 	_a						(   UpperCAmelCase					)    ->    Union[str, Any]:







          """simple docstring"""
          return None


# This is a pinned image from a specific revision of a document question answering space, hosted by HuggingFace,
# so we can expect it to be available.
_A :   Dict  			=  (
    'https://huggingface.co/spaces/impira/docquery/resolve/2f6c96314dc84dfda62d40de9da55f2f5165d403/invoice.png'
)


@is_pipeline_test
@require_torch
@require_vision
class        __SCREAMING_SNAKE_CASE      (			unittest.TestCase ):
       _UpperCAmelCase				:  Dict			   =   MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING







       @require_pytesseract
       @require_vision
       def 						__lowerCamelCase   (				self					:     List[str]		, A					:     Optional[Any]		, A					:     Optional[int]		, A					:     str      )       ->Any:
          lowerCamelCase__       :						Any        =							pipeline(
              '''document-question-answering'''		, model=A		, tokenizer=A		, image_processor=A      )

          lowerCamelCase__       :						int        =							INVOICE_URL
          lowerCamelCase__       :						str        =							list(zip(*apply_tesseract(load_image(A      )		, A		, ''''''      )      )      )
          lowerCamelCase__       :						Dict        =							'''What is the placebo?'''
          lowerCamelCase__       :						List[Any]        =							[
              {
                  '''image''': load_image(A      ),
                  '''question''': question,
              },
              {
                  '''image''': image,
                  '''question''': question,
              },
              {
                  '''image''': image,
                  '''question''': question,
                  '''word_boxes''': word_boxes,
              },
          ]
          return dqa_pipeline, examples







       def 						__lowerCamelCase   (				self					:     Dict		, A					:     Any		, A					:     Dict      )       ->Optional[int]:
          lowerCamelCase__       :						Union[str, Any]        =							dqa_pipeline(A		, top_k=2      )
          self.assertEqual(
              A		, [
                  [
                      {'''score''': ANY(A      ), '''answer''': ANY(A      ), '''start''': ANY(A      ), '''end''': ANY(A      )},
                      {'''score''': ANY(A      ), '''answer''': ANY(A      ), '''start''': ANY(A      ), '''end''': ANY(A      )},
                  ]
              ]
              * 3		, )







       @require_torch
       @require_detectrona
       @require_pytesseract
       def 						__lowerCamelCase   (				self					:     Optional[Any]      )       ->List[Any]:
          lowerCamelCase__       :						Tuple        =							pipeline('''document-question-answering'''		, model='''hf-internal-testing/tiny-random-layoutlmv2'''      )
          lowerCamelCase__       :						List[Any]        =							INVOICE_URL
          lowerCamelCase__       :						Optional[int]        =							'''How many cats are there?'''

          lowerCamelCase__       :						Dict        =							[
              {'''score''': 0.00_01, '''answer''': '''oy 2312/2019''', '''start''': 3_8, '''end''': 3_9},
              {'''score''': 0.00_01, '''answer''': '''oy 2312/2019 DUE''', '''start''': 3_8, '''end''': 4_0},
          ]
          lowerCamelCase__       :						Dict        =							dqa_pipeline(image=A		, question=A		, top_k=2      )
          self.assertEqual(nested_simplify(A		, decimals=4      )		, A      )

          lowerCamelCase__       :						List[Any]        =							dqa_pipeline({'''image''': image, '''question''': question}		, top_k=2      )
          self.assertEqual(nested_simplify(A		, decimals=4      )		, A      )

          # This image does not detect ANY text in it, meaning layoutlmv2 should fail.
          # Empty answer probably
          lowerCamelCase__       :						Optional[Any]        =							'''./tests/fixtures/tests_samples/COCO/000000039769.png'''
          lowerCamelCase__       :						int        =							dqa_pipeline(image=A		, question=A		, top_k=2      )
          self.assertEqual(A		, []      )

          # We can optionnally pass directly the words and bounding boxes
          lowerCamelCase__       :						Tuple        =							'''./tests/fixtures/tests_samples/COCO/000000039769.png'''
          lowerCamelCase__       :						Optional[Any]        =							[]
          lowerCamelCase__       :						Dict        =							[]
          lowerCamelCase__       :						Dict        =							dqa_pipeline(image=A		, question=A		, words=A		, boxes=A		, top_k=2      )
          self.assertEqual(A		, []      )







       @slow
       @require_torch
       @require_detectrona
       @require_pytesseract
       def 						__lowerCamelCase   (				self					:     Optional[Any]      )       ->str:
          lowerCamelCase__       :						List[Any]        =							pipeline(
              '''document-question-answering'''		, model='''tiennvcs/layoutlmv2-base-uncased-finetuned-docvqa'''		, revision='''9977165'''		, )
          lowerCamelCase__       :						Union[str, Any]        =							INVOICE_URL
          lowerCamelCase__       :						int        =							'''What is the invoice number?'''

          lowerCamelCase__       :						Optional[int]        =							dqa_pipeline(image=A		, question=A		, top_k=2      )
          self.assertEqual(
              nested_simplify(A		, decimals=4      )		, [
                  {'''score''': 0.99_44, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
                  {'''score''': 0.00_09, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
              ]		, )

          lowerCamelCase__       :						Union[str, Any]        =							dqa_pipeline({'''image''': image, '''question''': question}		, top_k=2      )
          self.assertEqual(
              nested_simplify(A		, decimals=4      )		, [
                  {'''score''': 0.99_44, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
                  {'''score''': 0.00_09, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
              ]		, )

          lowerCamelCase__       :						int        =							dqa_pipeline(
              [{'''image''': image, '''question''': question}, {'''image''': image, '''question''': question}]		, top_k=2      )
          self.assertEqual(
              nested_simplify(A		, decimals=4      )		, [
                  [
                      {'''score''': 0.99_44, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
                      {'''score''': 0.00_09, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
                  ],
              ]
              * 2		, )







       @slow
       @require_torch
       @require_detectrona
       @require_pytesseract
       def 						__lowerCamelCase   (				self					:     Optional[int]      )       ->Tuple:
          lowerCamelCase__       :						Any        =							pipeline(
              '''document-question-answering'''		, model='''tiennvcs/layoutlmv2-base-uncased-finetuned-docvqa'''		, revision='''9977165'''		, max_seq_len=5_0		, )
          lowerCamelCase__       :						Optional[int]        =							INVOICE_URL
          lowerCamelCase__       :						Union[str, Any]        =							'''What is the invoice number?'''

          lowerCamelCase__       :						Tuple        =							dqa_pipeline(image=A		, question=A		, top_k=2      )
          self.assertEqual(
              nested_simplify(A		, decimals=4      )		, [
                  {'''score''': 0.99_74, '''answer''': '''1110212019''', '''start''': 2_3, '''end''': 2_3},
                  {'''score''': 0.99_48, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
              ]		, )

          lowerCamelCase__       :						Optional[int]        =							dqa_pipeline({'''image''': image, '''question''': question}		, top_k=2      )
          self.assertEqual(
              nested_simplify(A		, decimals=4      )		, [
                  {'''score''': 0.99_74, '''answer''': '''1110212019''', '''start''': 2_3, '''end''': 2_3},
                  {'''score''': 0.99_48, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
              ]		, )

          lowerCamelCase__       :						Optional[Any]        =							dqa_pipeline(
              [{'''image''': image, '''question''': question}, {'''image''': image, '''question''': question}]		, top_k=2      )
          self.assertEqual(
              nested_simplify(A		, decimals=4      )		, [
                  [
                      {'''score''': 0.99_74, '''answer''': '''1110212019''', '''start''': 2_3, '''end''': 2_3},
                      {'''score''': 0.99_48, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
                  ]
              ]
              * 2		, )







       @slow
       @require_torch
       @require_pytesseract
       @require_vision
       def 						__lowerCamelCase   (				self					:     List[str]      )       ->Union[str, Any]:
          lowerCamelCase__       :						int        =							AutoTokenizer.from_pretrained(
              '''impira/layoutlm-document-qa'''		, revision='''3dc6de3'''		, add_prefix_space=A      )
          lowerCamelCase__       :						str        =							pipeline(
              '''document-question-answering'''		, model='''impira/layoutlm-document-qa'''		, tokenizer=A		, revision='''3dc6de3'''		, )
          lowerCamelCase__       :						Union[str, Any]        =							INVOICE_URL
          lowerCamelCase__       :						str        =							'''What is the invoice number?'''

          lowerCamelCase__       :						Dict        =							dqa_pipeline(image=A		, question=A		, top_k=2      )
          self.assertEqual(
              nested_simplify(A		, decimals=4      )		, [
                  {'''score''': 0.42_51, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
                  {'''score''': 0.08_19, '''answer''': '''1110212019''', '''start''': 2_3, '''end''': 2_3},
              ]		, )

          lowerCamelCase__       :						int        =							dqa_pipeline({'''image''': image, '''question''': question}		, top_k=2      )
          self.assertEqual(
              nested_simplify(A		, decimals=4      )		, [
                  {'''score''': 0.42_51, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
                  {'''score''': 0.08_19, '''answer''': '''1110212019''', '''start''': 2_3, '''end''': 2_3},
              ]		, )

          lowerCamelCase__       :						Tuple        =							dqa_pipeline(
              [{'''image''': image, '''question''': question}, {'''image''': image, '''question''': question}]		, top_k=2      )
          self.assertEqual(
              nested_simplify(A		, decimals=4      )		, [
                  [
                      {'''score''': 0.42_51, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
                      {'''score''': 0.08_19, '''answer''': '''1110212019''', '''start''': 2_3, '''end''': 2_3},
                  ]
              ]
              * 2		, )

          lowerCamelCase__       :						Optional[int]        =							list(zip(*apply_tesseract(load_image(A      )		, A		, ''''''      )      )      )

          # This model should also work if `image` is set to None
          lowerCamelCase__       :						Any        =							dqa_pipeline({'''image''': None, '''word_boxes''': word_boxes, '''question''': question}		, top_k=2      )
          self.assertEqual(
              nested_simplify(A		, decimals=4      )		, [
                  {'''score''': 0.42_51, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
                  {'''score''': 0.08_19, '''answer''': '''1110212019''', '''start''': 2_3, '''end''': 2_3},
              ]		, )







       @slow
       @require_torch
       @require_pytesseract
       @require_vision
       def 						__lowerCamelCase   (				self					:     Optional[int]      )       ->List[Any]:
          lowerCamelCase__       :						str        =							AutoTokenizer.from_pretrained(
              '''impira/layoutlm-document-qa'''		, revision='''3dc6de3'''		, add_prefix_space=A      )
          lowerCamelCase__       :						Union[str, Any]        =							pipeline(
              '''document-question-answering'''		, model='''impira/layoutlm-document-qa'''		, tokenizer=A		, revision='''3dc6de3'''		, max_seq_len=5_0		, )
          lowerCamelCase__       :						str        =							INVOICE_URL
          lowerCamelCase__       :						Any        =							'''What is the invoice number?'''

          lowerCamelCase__       :						Optional[Any]        =							dqa_pipeline(image=A		, question=A		, top_k=2      )
          self.assertEqual(
              nested_simplify(A		, decimals=4      )		, [
                  {'''score''': 0.99_99, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
                  {'''score''': 0.99_98, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
              ]		, )

          lowerCamelCase__       :						Any        =							dqa_pipeline(
              [{'''image''': image, '''question''': question}, {'''image''': image, '''question''': question}]		, top_k=2      )
          self.assertEqual(
              nested_simplify(A		, decimals=4      )		, [
                  [
                      {'''score''': 0.99_99, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
                      {'''score''': 0.99_98, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
                  ]
              ]
              * 2		, )

          lowerCamelCase__       :						str        =							list(zip(*apply_tesseract(load_image(A      )		, A		, ''''''      )      )      )

          # This model should also work if `image` is set to None
          lowerCamelCase__       :						str        =							dqa_pipeline({'''image''': None, '''word_boxes''': word_boxes, '''question''': question}		, top_k=2      )
          self.assertEqual(
              nested_simplify(A		, decimals=4      )		, [
                  {'''score''': 0.99_99, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
                  {'''score''': 0.99_98, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
              ]		, )







       @slow
       @require_torch
       def 						__lowerCamelCase   (				self					:     Tuple      )       ->Tuple:
          lowerCamelCase__       :						Tuple        =							pipeline(
              '''document-question-answering'''		, model='''naver-clova-ix/donut-base-finetuned-docvqa'''		, tokenizer=AutoTokenizer.from_pretrained('''naver-clova-ix/donut-base-finetuned-docvqa'''      )		, feature_extractor='''naver-clova-ix/donut-base-finetuned-docvqa'''		, )

          lowerCamelCase__       :						Optional[int]        =							INVOICE_URL
          lowerCamelCase__       :						Any        =							'''What is the invoice number?'''
          lowerCamelCase__       :						Dict        =							dqa_pipeline(image=A		, question=A		, top_k=2      )
          self.assertEqual(nested_simplify(A		, decimals=4      )		, [{'''answer''': '''us-001'''}]      )







       @require_tf
       @unittest.skip('''Document question answering not implemented in TF'''      )
       def 						__lowerCamelCase   (				self					:     Any      )       ->Optional[Any]:
          pass










import random
import timeit
from functools import wraps
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_tf_auto import TF_MODEL_MAPPING, TF_MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_tf_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_tf_available():
       import tensorflow as tf
       from tensorflow.python.framework.errors_impl import ResourceExhaustedError

       from .benchmark_args_tf import TensorFlowBenchmarkArguments

if is_pyanvml_available():
       import pyanvml.pyanvml as nvml

_A :   Any  			=  logging.get_logger(__name__)



def 	_a						(   UpperCAmelCase      ,  UpperCAmelCase					)    ->    List[str]:







   """simple docstring"""
   def run_func(UpperCAmelCase					):
      @wraps(UpperCAmelCase					)
      def run_in_eager_mode(*UpperCAmelCase      ,  **UpperCAmelCase					):
         return func(*UpperCAmelCase      ,  **UpperCAmelCase					)

      @wraps(UpperCAmelCase					)
      @tf.function(experimental_compile=UpperCAmelCase					)
      def run_in_graph_mode(*UpperCAmelCase      ,  **UpperCAmelCase					):
         return func(*UpperCAmelCase      ,  **UpperCAmelCase					)

      if do_eager_mode is True:
         if use_xla is not False:
            raise ValueError(
                '''Cannot run model in XLA, if `args.eager_mode` is set to `True`. Please set `args.eager_mode=False`.'''					)
         return run_in_eager_mode
      else:
         return run_in_graph_mode

   return run_func



def 	_a						(   UpperCAmelCase      ,  UpperCAmelCase      ,  UpperCAmelCase					)    ->    ["tf.Tensor"]:







   """simple docstring"""
   lowerCamelCase__       :						List[Any]        =							random.Random()
   lowerCamelCase__       :						str        =							[rng.randint(0      ,  vocab_size - 1					) for i in range(batch_size * sequence_length					)]
   return tf.constant(UpperCAmelCase      ,  shape=(batch_size, sequence_length)      ,  dtype=tf.intaa					)




class        __SCREAMING_SNAKE_CASE      (			lowerCAmelCase_ ):
       _UpperCAmelCase				:  TensorFlowBenchmarkArguments
       _UpperCAmelCase				:  PretrainedConfig
       _UpperCAmelCase				:  str			   =   "TensorFlow"







       @property
       def 						__lowerCamelCase   (				self					:     int      )       ->Optional[int]:
          return tf.__version__







       def 						__lowerCamelCase   (				self					:     Optional[int]		, A					:     str		, A					:     int		, A					:     int      )       ->float:
          # initialize GPU on separate process
          lowerCamelCase__       :						Dict        =							self.args.strategy
          if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''      )
          lowerCamelCase__       :						int        =							self._prepare_inference_func(A		, A		, A      )
          return self._measure_speed(_inference      )







       def 						__lowerCamelCase   (				self					:     str		, A					:     str		, A					:     int		, A					:     int      )       ->float:
          lowerCamelCase__       :						Optional[int]        =							self.args.strategy
          if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''      )
          lowerCamelCase__       :						List[Any]        =							self._prepare_train_func(A		, A		, A      )
          return self._measure_speed(_train      )







       def 						__lowerCamelCase   (				self					:     int		, A					:     str		, A					:     int		, A					:     int      )       ->[Memory, Optional[MemorySummary]]:
          # initialize GPU on separate process
          if self.args.is_gpu:
             tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]		, A      )
          lowerCamelCase__       :						int        =							self.args.strategy
          if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''      )
          lowerCamelCase__       :						str        =							self._prepare_inference_func(A		, A		, A      )
          return self._measure_memory(_inference      )







       def 						__lowerCamelCase   (				self					:     List[str]		, A					:     str		, A					:     int		, A					:     int      )       ->[Memory, Optional[MemorySummary]]:
          if self.args.is_gpu:
             tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]		, A      )
          lowerCamelCase__       :						List[Any]        =							self.args.strategy
          if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''      )

          lowerCamelCase__       :						str        =							self._prepare_train_func(A		, A		, A      )
          return self._measure_memory(_train      )







       def 						__lowerCamelCase   (				self					:     Dict		, A					:     str		, A					:     int		, A					:     int      )       ->Callable[[], None]:
          lowerCamelCase__       :						Tuple        =							self.config_dict[model_name]

          if self.args.fpaa:
             raise NotImplementedError('''Mixed precision is currently not supported.'''      )

          lowerCamelCase__       :						Tuple        =							(
              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:
                lowerCamelCase__       :						Any        =							'''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
                lowerCamelCase__       :						List[Any]        =							__import__('''transformers'''		, fromlist=[model_class]      )
                lowerCamelCase__       :						int        =							getattr(A		, A      )
                lowerCamelCase__       :						int        =							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:
             lowerCamelCase__       :						Union[str, Any]        =							TF_MODEL_MAPPING[config.__class__](A      )

          # encoder-decoder has vocab size saved differently
          lowerCamelCase__       :						Tuple        =							config.vocab_size if hasattr(A		, '''vocab_size'''      ) else config.encoder.vocab_size
          lowerCamelCase__       :						Optional[Any]        =							random_input_ids(A		, A		, A      )

          @run_with_tf_optimizations(self.args.eager_mode		, self.args.use_xla      )
          def encoder_decoder_forward():
             return model(A		, decoder_input_ids=A		, training=A      )

          @run_with_tf_optimizations(self.args.eager_mode		, self.args.use_xla      )
          def encoder_forward():
             return model(A		, training=A      )

          lowerCamelCase__       :						int        =							encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

          return _inference







       def 						__lowerCamelCase   (				self					:     List[str]		, A					:     str		, A					:     int		, A					:     int      )       ->Callable[[], None]:
          lowerCamelCase__       :						Tuple        =							self.config_dict[model_name]

          if self.args.eager_mode is not False:
             raise ValueError('''Training cannot be done in eager mode. Please make sure that `args.eager_mode = False`.'''      )

          if self.args.fpaa:
             raise NotImplementedError('''Mixed precision is currently not supported.'''      )

          lowerCamelCase__       :						Optional[int]        =							(
              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:
                lowerCamelCase__       :						Any        =							'''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
                lowerCamelCase__       :						List[str]        =							__import__('''transformers'''		, fromlist=[model_class]      )
                lowerCamelCase__       :						Optional[int]        =							getattr(A		, A      )
                lowerCamelCase__       :						Optional[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:
             lowerCamelCase__       :						List[str]        =							TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](A      )

          # encoder-decoder has vocab size saved differently
          lowerCamelCase__       :						Optional[int]        =							config.vocab_size if hasattr(A		, '''vocab_size'''      ) else config.encoder.vocab_size
          lowerCamelCase__       :						Dict        =							random_input_ids(A		, A		, A      )

          @run_with_tf_optimizations(self.args.eager_mode		, self.args.use_xla      )
          def encoder_decoder_train():
             lowerCamelCase__       :						int        =							model(A		, decoder_input_ids=A		, labels=A		, training=A      )[0]
             lowerCamelCase__       :						List[Any]        =							tf.gradients(A		, model.trainable_variables      )
             return gradients

          @run_with_tf_optimizations(self.args.eager_mode		, self.args.use_xla      )
          def encoder_train():
             lowerCamelCase__       :						Optional[int]        =							model(A		, labels=A		, training=A      )[0]
             lowerCamelCase__       :						List[str]        =							tf.gradients(A		, model.trainable_variables      )
             return gradients

          lowerCamelCase__       :						Tuple        =							encoder_decoder_train if config.is_encoder_decoder else encoder_train

          return _train







       def 						__lowerCamelCase   (				self					:     Tuple		, A					:     Any      )       ->float:
          with self.args.strategy.scope():
             try:
                if self.args.is_tpu or self.args.use_xla:
                   # run additional 10 times to stabilize compilation for tpu
                   logger.info('''Do inference on TPU. 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
                lowerCamelCase__       :						Optional[Any]        =							timeit.repeat(
                    A		, repeat=self.args.repeat		, number=1_0		, )

                return min(A      ) / 10.0
             except ResourceExhaustedError as e:
                self.print_fn(F"Doesn't fit on GPU. {e}"      )







       def 						__lowerCamelCase   (				self					:     List[Any]		, A					:     Callable[[], None]      )       ->[Memory, MemorySummary]:
          logger.info(
              '''Note that TensorFlow allocates more memory than '''
              '''it might need to speed up computation. '''
              '''The memory reported here corresponds to the memory '''
              '''reported by `nvidia-smi`, which can vary depending '''
              '''on total available memory on the GPU that is used.'''      )
          with self.args.strategy.scope():
             try:
                if self.args.trace_memory_line_by_line:
                   if not self.args.eager_mode:
                      raise ValueError(
                          '''`args.eager_mode` is set to `False`. Make sure to run model in eager mode to measure memory'''
                          ''' consumption line by line.'''      )
                   lowerCamelCase__       :						Union[str, Any]        =							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 `args.memory=False`'''      )
                elif self.args.is_gpu:
                   # 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__       :						Any        =							nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx      )
                      lowerCamelCase__       :						Optional[int]        =							nvml.nvmlDeviceGetMemoryInfo(A      )
                      lowerCamelCase__       :						List[Any]        =							meminfo.used
                      lowerCamelCase__       :						Union[str, Any]        =							Memory(A      )
                      # shutdown nvml
                      nvml.nvmlShutdown()
                else:
                   # cpu
                   if self.args.trace_memory_line_by_line:
                      logger.info(
                          '''When enabling line by line tracing, the max peak memory for CPU is inaccurate in'''
                          ''' TensorFlow.'''      )
                      lowerCamelCase__       :						Tuple        =							None
                   else:
                      lowerCamelCase__       :						Dict        =							measure_peak_memory_cpu(A      )
                      lowerCamelCase__       :						Optional[Any]        =							Memory(A      ) if isinstance(A		, A      ) else memory_bytes
                if self.args.trace_memory_line_by_line:
                   lowerCamelCase__       :						Union[str, Any]        =							stop_memory_tracing(A      )
                   if memory is None:
                      lowerCamelCase__       :						Dict        =							summary.total
                else:
                   lowerCamelCase__       :						Optional[int]        =							None

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



