




import inspect
import unittest
from typing import List

import numpy as np

from transformers import EfficientFormerConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
  import tensorflow as tf

  from transformers import (
      TFEfficientFormerForImageClassification,
      TFEfficientFormerForImageClassificationWithTeacher,
      TFEfficientFormerModel,
  )
  from transformers.models.efficientformer.modeling_tf_efficientformer import (
      TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
  )


if is_vision_available():
  from PIL import Image

  from transformers import EfficientFormerImageProcessor


class 							UpperCamelCase_       :

    '''simple docstring'''






    def __init__(     self :      List[Any]   ,					UpperCAmelCase__ :      Optional[int]   ,					UpperCAmelCase__ :      int = 13   ,					UpperCAmelCase__ :      int = 64   ,					UpperCAmelCase__ :      int = 2   ,					UpperCAmelCase__ :      int = 3   ,					UpperCAmelCase__ :      int = 3   ,					UpperCAmelCase__ :      bool = True   ,					UpperCAmelCase__ :      bool = True   ,					UpperCAmelCase__ :      int = 128   ,					UpperCAmelCase__ :      int=[16, 32, 64, 128]   ,					UpperCAmelCase__ :      int = 7   ,					UpperCAmelCase__ :      int = 4   ,					UpperCAmelCase__ :      int = 37   ,					UpperCAmelCase__ :      str = "gelu"   ,					UpperCAmelCase__ :      float = 0.1   ,					UpperCAmelCase__ :      float = 0.1   ,					UpperCAmelCase__ :      int = 10   ,					UpperCAmelCase__ :      float = 0.02   ,					UpperCAmelCase__ :      int = 2   ,					UpperCAmelCase__ :      int = 1   ,					UpperCAmelCase__ :      int = 128   ,					UpperCAmelCase__ :      List[int] = [2, 2, 2, 2]   ,					UpperCAmelCase__ :      int = 2   ,					UpperCAmelCase__ :      int = 2   ,					)		->List[str]:





           '''simple docstring'''

           A__ 					=       parent
           A__ 					=       batch_size
           A__ 					=       image_size
           A__ 					=       patch_size
           A__ 					=       num_channels
           A__ 					=       is_training
           A__ 					=       use_labels
           A__ 					=       hidden_size
           A__ 					=       num_hidden_layers
           A__ 					=       num_attention_heads
           A__ 					=       intermediate_size
           A__ 					=       hidden_act
           A__ 					=       hidden_dropout_prob
           A__ 					=       attention_probs_dropout_prob
           A__ 					=       type_sequence_label_size
           A__ 					=       initializer_range
           A__ 					=       encoder_stride
           A__ 					=       num_attention_outputs
           A__ 					=       embed_dim
           A__ 					=       embed_dim + 1
           A__ 					=       resolution
           A__ 					=       depths
           A__ 					=       hidden_sizes
           A__ 					=       dim
           A__ 					=       mlp_expansion_ratio


    def 			SCREAMING_SNAKE_CASE				(     self :      List[str])		->str:





           '''simple docstring'''

           A__ 					=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

           A__ 					=       None
           if self.use_labels:
                  A__ 					=       ids_tensor([self.batch_size]   ,					self.type_sequence_label_size)

           A__ 					=       self.get_config()

           return config, pixel_values, labels


    def 			SCREAMING_SNAKE_CASE				(     self :      str)		->Tuple:





           '''simple docstring'''

           return EfficientFormerConfig(
               image_size=self.image_size   ,					patch_size=self.patch_size   ,					num_channels=self.num_channels   ,					hidden_size=self.hidden_size   ,					num_hidden_layers=self.num_hidden_layers   ,					num_attention_heads=self.num_attention_heads   ,					intermediate_size=self.intermediate_size   ,					hidden_act=self.hidden_act   ,					hidden_dropout_prob=self.hidden_dropout_prob   ,					attention_probs_dropout_prob=self.attention_probs_dropout_prob   ,					is_decoder=UpperCAmelCase__   ,					initializer_range=self.initializer_range   ,					encoder_stride=self.encoder_stride   ,					resolution=self.resolution   ,					depths=self.depths   ,					hidden_sizes=self.hidden_sizes   ,					dim=self.dim   ,					mlp_expansion_ratio=self.mlp_expansion_ratio   ,					)


    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[int]   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      str)		->Any:





           '''simple docstring'''

           A__ 					=       TFEfficientFormerModel(config=UpperCAmelCase__)
           A__ 					=       model(UpperCAmelCase__   ,					training=UpperCAmelCase__)
           self.parent.assertEqual(result.last_hidden_state.shape   ,					(self.batch_size, self.seq_length, self.hidden_size))


    def 			SCREAMING_SNAKE_CASE				(     self :      Tuple   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      Optional[int])		->Optional[Any]:





           '''simple docstring'''

           A__ 					=       self.type_sequence_label_size
           A__ 					=       TFEfficientFormerForImageClassification(UpperCAmelCase__)
           A__ 					=       model(UpperCAmelCase__   ,					labels=UpperCAmelCase__   ,					training=UpperCAmelCase__)
           self.parent.assertEqual(result.logits.shape   ,					(self.batch_size, self.type_sequence_label_size))

           # test greyscale images
           A__ 					=       1
           A__ 					=       TFEfficientFormerForImageClassification(UpperCAmelCase__)

           A__ 					=       floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
           A__ 					=       model(UpperCAmelCase__   ,					labels=UpperCAmelCase__)
           self.parent.assertEqual(result.logits.shape   ,					(self.batch_size, self.type_sequence_label_size))





    def 			SCREAMING_SNAKE_CASE				(     self :      Dict)		->List[str]:





           '''simple docstring'''

           A__ 					=       self.prepare_config_and_inputs()
           A__					,				A__					,				A__ 					=       config_and_inputs
           A__ 					=       {'''pixel_values''': pixel_values}
           return config, inputs_dict


@require_tf
class 							UpperCamelCase_       (					UpperCAmelCase__						,   UpperCAmelCase__						,   unittest.TestCase	):

    '''simple docstring'''






    UpperCAmelCase__				    = (
        (
            TFEfficientFormerModel,
            TFEfficientFormerForImageClassificationWithTeacher,
            TFEfficientFormerForImageClassification,
        )
        if is_tf_available()
        else ()
    )
    UpperCAmelCase__				    = (
        {
            '''feature-extraction''': TFEfficientFormerModel,
            '''image-classification''': (
                TFEfficientFormerForImageClassification,
                TFEfficientFormerForImageClassificationWithTeacher,
            ),
        }
        if is_tf_available()
        else {}
    )

    UpperCAmelCase__				    = False

    UpperCAmelCase__				    = False
    UpperCAmelCase__				    = False
    UpperCAmelCase__				    = False
    UpperCAmelCase__				    = False


    def 			SCREAMING_SNAKE_CASE				(     self :      List[Any])		->List[Any]:





           '''simple docstring'''

           A__ 					=       TFEfficientFormerModelTester(self)
           A__ 					=       ConfigTester(
               self   ,					config_class=UpperCAmelCase__   ,					has_text_modality=UpperCAmelCase__   ,					hidden_size=37)


    def 			SCREAMING_SNAKE_CASE				(     self :      List[str])		->int:





           '''simple docstring'''

           self.config_tester.run_common_tests()


    @unittest.skip(reason='''EfficientFormer does not use inputs_embeds''')
    def 			SCREAMING_SNAKE_CASE				(     self :      Tuple)		->Optional[Any]:





           '''simple docstring'''

           pass


    @unittest.skip(reason='''EfficientFormer does not support input and output embeddings''')
    def 			SCREAMING_SNAKE_CASE				(     self :      str)		->int:





           '''simple docstring'''

           pass


    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[Any])		->Union[str, Any]:





           '''simple docstring'''

           A__					,				A__ 					=       self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  A__ 					=       model_class(UpperCAmelCase__)
                  A__ 					=       inspect.signature(model.call)
                  # signature.parameters is an OrderedDict => so arg_names order is deterministic
                  A__ 					=       [*signature.parameters.keys()]

                  A__ 					=       ['''pixel_values''']
                  self.assertListEqual(arg_names[:1]   ,					UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      List[str])		->Optional[Any]:





           '''simple docstring'''

           def check_hidden_states_output(UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      Union[str, Any]   ,					UpperCAmelCase__ :      Dict):
                  A__ 					=       model_class(UpperCAmelCase__)

                  A__ 					=       model(**self._prepare_for_class(UpperCAmelCase__   ,					UpperCAmelCase__)   ,					training=UpperCAmelCase__)
                  A__ 					=       outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

                  A__ 					=       getattr(
                      self.model_tester   ,					'''expected_num_hidden_layers'''   ,					self.model_tester.num_hidden_layers + 1)
                  self.assertEqual(len(UpperCAmelCase__)   ,					UpperCAmelCase__)

                  if hasattr(self.model_tester   ,					'''encoder_seq_length'''):
                         A__ 					=       self.model_tester.encoder_seq_length
                         if hasattr(self.model_tester   ,					'''chunk_length''') and self.model_tester.chunk_length > 1:
                                A__ 					=       seq_length * self.model_tester.chunk_length
                  else:
                         A__ 					=       self.model_tester.seq_length

                  self.assertListEqual(
                      list(hidden_states[-1].shape[-2:])   ,					[seq_length, self.model_tester.hidden_size]   ,					)

                  if config.is_encoder_decoder:
                         A__ 					=       outputs.decoder_hidden_states

                         self.asseretIsInstance(UpperCAmelCase__   ,					(list, tuple))
                         self.assertEqual(len(UpperCAmelCase__)   ,					UpperCAmelCase__)
                         A__ 					=       getattr(self.model_tester   ,					'''seq_length'''   ,					UpperCAmelCase__)
                         A__ 					=       getattr(self.model_tester   ,					'''decoder_seq_length'''   ,					UpperCAmelCase__)

                         self.assertListEqual(
                             list(hidden_states[-1].shape[-2:])   ,					[decoder_seq_length, self.model_tester.hidden_size]   ,					)

           A__					,				A__ 					=       self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  A__ 					=       True
                  check_hidden_states_output(UpperCAmelCase__   ,					UpperCAmelCase__   ,					UpperCAmelCase__)

                  # check that output_hidden_states also work using config
                  del inputs_dict["output_hidden_states"]
                  A__ 					=       True

                  check_hidden_states_output(UpperCAmelCase__   ,					UpperCAmelCase__   ,					UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      List[str]   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      Tuple   ,					UpperCAmelCase__ :      Any=False)		->Tuple:





           '''simple docstring'''

           A__ 					=       super()._prepare_for_class(UpperCAmelCase__   ,					UpperCAmelCase__   ,					return_labels=UpperCAmelCase__)

           if return_labels:
                  if model_class.__name__ == "TFEfficientFormerForImageClassificationWithTeacher":
                         del inputs_dict["labels"]

           return inputs_dict


    def 			SCREAMING_SNAKE_CASE				(     self :      List[str])		->str:





           '''simple docstring'''

           A__ 					=       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCAmelCase__)


    @unittest.skip(reason='''EfficientFormer does not implement masked image modeling yet''')
    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[Any])		->Any:





           '''simple docstring'''

           A__ 					=       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_image_modeling(*UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[Any])		->Optional[int]:





           '''simple docstring'''

           A__ 					=       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase__)


    @slow
    def 			SCREAMING_SNAKE_CASE				(     self :      Tuple)		->Optional[int]:





           '''simple docstring'''

           for model_name in TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  A__ 					=       TFEfficientFormerModel.from_pretrained(UpperCAmelCase__)
                  self.assertIsNotNone(UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      int)		->Optional[Any]:





           '''simple docstring'''

           A__					,				A__ 					=       self.model_tester.prepare_config_and_inputs_for_common()
           A__ 					=       True

           A__ 					=       getattr(self.model_tester   ,					'''seq_length'''   ,					UpperCAmelCase__)
           A__ 					=       getattr(self.model_tester   ,					'''encoder_seq_length'''   ,					UpperCAmelCase__)
           A__ 					=       getattr(self.model_tester   ,					'''key_length'''   ,					UpperCAmelCase__)
           A__ 					=       getattr(self.model_tester   ,					'''chunk_length'''   ,					UpperCAmelCase__)

           if chunk_length is not None and hasattr(self.model_tester   ,					'''num_hashes'''):
                  A__ 					=       encoder_seq_length * self.model_tester.num_hashes

           for model_class in self.all_model_classes:
                  A__ 					=       True
                  A__ 					=       False
                  A__ 					=       True
                  A__ 					=       model_class(UpperCAmelCase__)

                  A__ 					=       model(**self._prepare_for_class(UpperCAmelCase__   ,					UpperCAmelCase__)   ,					training=UpperCAmelCase__)
                  A__ 					=       outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
                  self.assertEqual(len(UpperCAmelCase__)   ,					self.model_tester.num_attention_outputs)

                  # check that output_attentions also work using config
                  del inputs_dict["output_attentions"]
                  A__ 					=       True
                  A__ 					=       model_class(UpperCAmelCase__)
                  A__ 					=       model(**self._prepare_for_class(UpperCAmelCase__   ,					UpperCAmelCase__)   ,					training=UpperCAmelCase__)

                  A__ 					=       outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
                  self.assertEqual(len(UpperCAmelCase__)   ,					self.model_tester.num_attention_outputs)

                  if chunk_length is not None:
                         self.assertListEqual(
                             list(attentions[0].shape[-4:])   ,					[self.model_tester.num_attention_heads, encoder_seq_length, chunk_length, encoder_key_length]   ,					)
                  else:
                         self.assertListEqual(
                             list(attentions[0].shape[-3:])   ,					[self.model_tester.num_attention_heads, encoder_seq_length, encoder_key_length]   ,					)





    def 			SCREAMING_SNAKE_CASE				(     self :      Dict)		->str:





           '''simple docstring'''

           A__					,				A__ 					=       self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  # Prepare our model
                  A__ 					=       model_class(UpperCAmelCase__)
                  # These are maximally general inputs for the model, with multiple None dimensions
                  # Hopefully this will catch any conditionals that fail for flexible shapes
                  A__ 					=       {
                      key: tf.keras.Input(shape=val.shape[1:]   ,					dtype=val.dtype   ,					name=UpperCAmelCase__)
                      for key, val in model.input_signature.items()
                      if key in model.dummy_inputs
                  }
                  A__ 					=       model(UpperCAmelCase__)
                  self.assertTrue(outputs_dict is not None)







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



       """simple docstring"""

       A__ 					=       Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''		)
       return image


@require_tf
@require_vision
class 							UpperCamelCase_       (					unittest.TestCase	):

    '''simple docstring'''






    @cached_property
    def 			SCREAMING_SNAKE_CASE				(     self :      List[Any])		->Optional[Any]:





           '''simple docstring'''

           return (
               EfficientFormerImageProcessor.from_pretrained('''snap-research/efficientformer-l1-300''')
               if is_vision_available()
               else None
           )


    @slow
    def 			SCREAMING_SNAKE_CASE				(     self :      str)		->Union[str, Any]:





           '''simple docstring'''

           A__ 					=       TFEfficientFormerForImageClassification.from_pretrained('''snap-research/efficientformer-l1-300''')
           A__ 					=       self.default_image_processor
           A__ 					=       prepare_img()
           A__ 					=       image_processor(images=UpperCAmelCase__   ,					return_tensors='''tf''')
           # forward pass
           A__ 					=       model(**UpperCAmelCase__   ,					training=UpperCAmelCase__)
           # verify the logits
           A__ 					=       tf.TensorShape((1, 1_000))
           self.assertEqual(outputs.logits.shape   ,					UpperCAmelCase__)
           A__ 					=       tf.constant([-0.0555, 0.4825, -0.0852])
           self.assertTrue(np.allclose(outputs.logits[0, :3]   ,					UpperCAmelCase__   ,					atol=1e-4))





    @slow
    def 			SCREAMING_SNAKE_CASE				(     self :      Tuple)		->List[str]:





           '''simple docstring'''

           A__ 					=       TFEfficientFormerForImageClassificationWithTeacher.from_pretrained(
               '''snap-research/efficientformer-l1-300''')
           A__ 					=       self.default_image_processor
           A__ 					=       prepare_img()
           A__ 					=       image_processor(images=UpperCAmelCase__   ,					return_tensors='''tf''')
           # forward pass
           A__ 					=       model(**UpperCAmelCase__   ,					training=UpperCAmelCase__)
           # verify the logits
           A__ 					=       tf.TensorShape((1, 1_000))
           self.assertEqual(outputs.logits.shape   ,					UpperCAmelCase__)
           A__ 					=       tf.constant([-0.1312, 0.4353, -1.0499])
           self.assertTrue(np.allclose(outputs.logits[0, :3]   ,					UpperCAmelCase__   ,					atol=1e-4))










import math
import time
from typing import Dict, List, Optional

from torch.utils.data import Dataset

from transformers import SeqaSeqTrainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput, speed_metrics


if is_torch_tpu_available(check_device=False):
			import torch_xla.core.xla_model as xm
			import torch_xla.debug.metrics as met




class 							snake_case_       (    __UpperCamelCase      ):

						"""simple docstring"""



						def __init__(self:					str ,					*__UpperCAmelCase:					Optional[int] ,					__UpperCAmelCase:					Optional[int]=None ,					__UpperCAmelCase:					Dict=None ,					**__UpperCAmelCase:					Tuple   )    ->     Dict:




											'''simple docstring'''



											super().__init__(*__UpperCAmelCase ,					**__UpperCAmelCase   )
											__a					:						str      =					eval_examples
											__a					:						List[Any]      =					post_process_function



						def        UpperCAmelCase__					(self:					Optional[int] ,					__UpperCAmelCase:					Optional[Dataset] = None ,					__UpperCAmelCase:					Optional[int]=None ,					__UpperCAmelCase:					Optional[List[str]] = None ,					__UpperCAmelCase:					str = "eval" ,					**__UpperCAmelCase:					Any ,					)    ->     Dict[str, float]:




											'''simple docstring'''



											__a					:						Optional[int]      =					gen_kwargs.copy()
											__a					:						Union[str, Any]      =					(
											    gen_kwargs["max_length"] if gen_kwargs.get("max_length"   ) is not None else self.args.generation_max_length
											)
											__a					:						Dict      =					(
											    gen_kwargs["num_beams"] if gen_kwargs.get("num_beams"   ) is not None else self.args.generation_num_beams
											)
											__a					:						Optional[Any]      =					gen_kwargs

											__a					:						Union[str, Any]      =					self.eval_dataset if eval_dataset is None else eval_dataset
											__a					:						Optional[int]      =					self.get_eval_dataloader(__UpperCAmelCase   )
											__a					:						Tuple      =					self.eval_examples if eval_examples is None else eval_examples

											# Temporarily disable metric computation, we will do it in the loop here.
											__a					:						int      =					self.compute_metrics
											__a					:						Dict      =					None
											__a					:						List[Any]      =					time.time()
											__a					:						str      =					self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
											try:
																__a					:						str      =					eval_loop(
																    __UpperCAmelCase ,					description="Evaluation" ,					prediction_loss_only=True if compute_metrics is None else None ,					ignore_keys=__UpperCAmelCase ,					metric_key_prefix=__UpperCAmelCase ,					)
											finally:
																__a					:						Union[str, Any]      =					compute_metrics
											__a					:						str      =					self.args.eval_batch_size * self.args.world_size
											if f'{metric_key_prefix}_jit_compilation_time' in output.metrics:
																start_time += output.metrics[f'{metric_key_prefix}_jit_compilation_time']
											output.metrics.update(
											    speed_metrics(
											        __UpperCAmelCase ,					__UpperCAmelCase ,					num_samples=output.num_samples ,					num_steps=math.ceil(output.num_samples / total_batch_size   ) ,					)   )

											if self.post_process_function is not None and self.compute_metrics is not None and self.args.should_save:
																# Only the main node write the results by default
																__a					:						Union[str, Any]      =					self.post_process_function(__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase   )
																__a					:						Dict      =					self.compute_metrics(__UpperCAmelCase   )

																# Prefix all keys with metric_key_prefix + '_'
																for key in list(metrics.keys()   ):
																					if not key.startswith(f'{metric_key_prefix}_'   ):
																										__a					:						int      =					metrics.pop(__UpperCAmelCase   )

																metrics.update(output.metrics   )
											else:
																__a					:						Optional[int]      =					output.metrics

											if self.args.should_log:
																# Only the main node log the results by default
																self.log(__UpperCAmelCase   )

											if self.args.tpu_metrics_debug or self.args.debug:
																# tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
																xm.master_print(met.metrics_report()   )

											__a					:						Tuple      =					self.callback_handler.on_evaluate(self.args ,					self.state ,					self.control ,					__UpperCAmelCase   )
											return metrics



						def        UpperCAmelCase__					(self:					Dict ,					__UpperCAmelCase:					List[str] ,					__UpperCAmelCase:					List[str] ,					__UpperCAmelCase:					Union[str, Any]=None ,					__UpperCAmelCase:					str = "test" ,					**__UpperCAmelCase:					List[Any]   )    ->     Union[str, Any]:




											'''simple docstring'''



											__a					:						Any      =					gen_kwargs.copy()

											__a					:						Tuple      =					self.get_test_dataloader(__UpperCAmelCase   )

											# Temporarily disable metric computation, we will do it in the loop here.
											__a					:						Optional[Any]      =					self.compute_metrics
											__a					:						int      =					None
											__a					:						List[Any]      =					time.time()
											__a					:						Dict      =					self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
											try:
																__a					:						Optional[Any]      =					eval_loop(
																    __UpperCAmelCase ,					description="Prediction" ,					prediction_loss_only=True if compute_metrics is None else None ,					ignore_keys=__UpperCAmelCase ,					metric_key_prefix=__UpperCAmelCase ,					)
											finally:
																__a					:						List[Any]      =					compute_metrics

											__a					:						Optional[Any]      =					self.args.eval_batch_size * self.args.world_size
											if f'{metric_key_prefix}_jit_compilation_time' in output.metrics:
																start_time += output.metrics[f'{metric_key_prefix}_jit_compilation_time']
											output.metrics.update(
											    speed_metrics(
											        __UpperCAmelCase ,					__UpperCAmelCase ,					num_samples=output.num_samples ,					num_steps=math.ceil(output.num_samples / total_batch_size   ) ,					)   )
											if self.post_process_function is None or self.compute_metrics is None:
																return output

											__a					:						Optional[int]      =					self.post_process_function(__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					"predict"   )
											__a					:						Dict      =					self.compute_metrics(__UpperCAmelCase   )

											# Prefix all keys with metric_key_prefix + '_'
											for key in list(metrics.keys()   ):
																if not key.startswith(f'{metric_key_prefix}_'   ):
																					__a					:						List[str]      =					metrics.pop(__UpperCAmelCase   )
											metrics.update(output.metrics   )
											return PredictionOutput(predictions=predictions.predictions ,					label_ids=predictions.label_ids ,					metrics=__UpperCAmelCase   )





