import numpy as np

from transformers import BatchFeature
from transformers.testing_utils import require_tf, require_torch

from .test_feature_extraction_common import FeatureExtractionSavingTestMixin
class     a_			(						a__						):






 '''simple docstring'''






 UpperCAmelCase_     			= None
 UpperCAmelCase_     			= None




 @property
 def 		__snake_case						(      self :  int):







        '''simple docstring'''



        return self.feat_extract_tester.prepare_feat_extract_dict()




 def 		__snake_case						(      self :  Any):







        '''simple docstring'''



        lowerCAmelCase__					     = self.feature_extraction_class(**self.feat_extract_dict)
        self.assertTrue(hasattr(_lowerCamelCase      ,   'feature_size'))
        self.assertTrue(hasattr(_lowerCamelCase      ,   'sampling_rate'))
        self.assertTrue(hasattr(_lowerCamelCase      ,   'padding_value'))




 def 		__snake_case						(      self :  Dict):







        '''simple docstring'''



        lowerCAmelCase__					     = self.feat_extract_tester.prepare_inputs_for_common()
        lowerCAmelCase__					     = self.feature_extraction_class(**self.feat_extract_dict)
        lowerCAmelCase__					     = feat_extract.model_input_names[0]

        lowerCAmelCase__					     = BatchFeature({input_name: speech_inputs})

        self.assertTrue(all(len(_lowerCamelCase) == len(_lowerCamelCase) for x, y in zip(_lowerCamelCase      ,   processed_features[input_name])))

        lowerCAmelCase__					     = self.feat_extract_tester.prepare_inputs_for_common(equal_length=_lowerCamelCase)
        lowerCAmelCase__					     = BatchFeature({input_name: speech_inputs}      ,   tensor_type='np')

        lowerCAmelCase__					     = processed_features[input_name]

        if len(batch_features_input.shape) < 3:
               lowerCAmelCase__					     = batch_features_input[:, :, None]

        self.assertTrue(
            batch_features_input.shape
            == (self.feat_extract_tester.batch_size, len(speech_inputs[0]), self.feat_extract_tester.feature_size))




 @require_torch
 def 		__snake_case						(      self :  Any):







        '''simple docstring'''



        lowerCAmelCase__					     = self.feat_extract_tester.prepare_inputs_for_common(equal_length=_lowerCamelCase)
        lowerCAmelCase__					     = self.feature_extraction_class(**self.feat_extract_dict)
        lowerCAmelCase__					     = feat_extract.model_input_names[0]

        lowerCAmelCase__					     = BatchFeature({input_name: speech_inputs}      ,   tensor_type='pt')

        lowerCAmelCase__					     = processed_features[input_name]

        if len(batch_features_input.shape) < 3:
               lowerCAmelCase__					     = batch_features_input[:, :, None]

        self.assertTrue(
            batch_features_input.shape
            == (self.feat_extract_tester.batch_size, len(speech_inputs[0]), self.feat_extract_tester.feature_size))




 @require_tf
 def 		__snake_case						(      self :  List[Any]):







        '''simple docstring'''



        lowerCAmelCase__					     = self.feat_extract_tester.prepare_inputs_for_common(equal_length=_lowerCamelCase)
        lowerCAmelCase__					     = self.feature_extraction_class(**self.feat_extract_dict)
        lowerCAmelCase__					     = feat_extract.model_input_names[0]

        lowerCAmelCase__					     = BatchFeature({input_name: speech_inputs}      ,   tensor_type='tf')

        lowerCAmelCase__					     = processed_features[input_name]

        if len(batch_features_input.shape) < 3:
               lowerCAmelCase__					     = batch_features_input[:, :, None]

        self.assertTrue(
            batch_features_input.shape
            == (self.feat_extract_tester.batch_size, len(speech_inputs[0]), self.feat_extract_tester.feature_size))




 def 		__snake_case						(      self :  Union[str, Any]      ,   lowercase__ :  Dict=False):







        '''simple docstring'''



        def _inputs_have_equal_length(lowercase__ :  str):
               lowerCAmelCase__					     = len(input[0])
               for input_slice in input[1:]:
                      if len(_lowerCamelCase) != length:
                             return False
               return True

        def _inputs_are_equal(lowercase__ :  str      ,   lowercase__ :  Tuple):
               if len(_lowerCamelCase) != len(_lowerCamelCase):
                      return False

               for input_slice_a, input_slice_a in zip(_lowerCamelCase      ,   _lowerCamelCase):
                      if not np.allclose(np.asarray(_lowerCamelCase)      ,   np.asarray(_lowerCamelCase)      ,   atol=1e-3):
                             return False
               return True

        lowerCAmelCase__					     = self.feature_extraction_class(**self.feat_extract_dict)
        lowerCAmelCase__					     = self.feat_extract_tester.prepare_inputs_for_common(numpify=_lowerCamelCase)
        lowerCAmelCase__					     = feat_extract.model_input_names[0]

        lowerCAmelCase__					     = BatchFeature({input_name: speech_inputs})

        lowerCAmelCase__					     = self.feat_extract_tester.seq_length_diff
        lowerCAmelCase__					     = self.feat_extract_tester.max_seq_length + pad_diff
        lowerCAmelCase__					     = self.feat_extract_tester.min_seq_length
        lowerCAmelCase__					     = self.feat_extract_tester.batch_size
        lowerCAmelCase__					     = self.feat_extract_tester.feature_size

        # test padding for List[int] + numpy
        lowerCAmelCase__					     = feat_extract.pad(_lowerCamelCase      ,   padding=_lowerCamelCase)
        lowerCAmelCase__					     = input_a[input_name]

        lowerCAmelCase__					     = feat_extract.pad(_lowerCamelCase      ,   padding='longest')
        lowerCAmelCase__					     = input_a[input_name]

        lowerCAmelCase__					     = feat_extract.pad(_lowerCamelCase      ,   padding='max_length'      ,   max_length=len(speech_inputs[-1]))
        lowerCAmelCase__					     = input_a[input_name]

        lowerCAmelCase__					     = feat_extract.pad(_lowerCamelCase      ,   padding='longest'      ,   return_tensors='np')
        lowerCAmelCase__					     = input_a[input_name]

        # max_length parameter has to be provided when setting `padding="max_length"`
        with self.assertRaises(_lowerCamelCase):
               feat_extract.pad(_lowerCamelCase      ,   padding='max_length')[input_name]

        lowerCAmelCase__					     = feat_extract.pad(
            _lowerCamelCase      ,   padding='max_length'      ,   max_length=_lowerCamelCase      ,   return_tensors='np')
        lowerCAmelCase__					     = input_a[input_name]

        self.assertFalse(_inputs_have_equal_length(_lowerCamelCase))
        self.assertTrue(_inputs_have_equal_length(_lowerCamelCase))
        self.assertTrue(_inputs_have_equal_length(_lowerCamelCase))
        self.assertTrue(_inputs_are_equal(_lowerCamelCase      ,   _lowerCamelCase))
        self.assertTrue(len(input_a[0]) == pad_min_length)
        self.assertTrue(len(input_a[1]) == pad_min_length + pad_diff)
        self.assertTrue(input_a.shape[:2] == (batch_size, len(input_a[0])))
        self.assertTrue(input_a.shape[:2] == (batch_size, pad_max_length))

        if feature_size > 1:
               self.assertTrue(input_a.shape[2] == input_a.shape[2] == feature_size)

        # test padding for `pad_to_multiple_of` for List[int] + numpy
        lowerCAmelCase__					     = feat_extract.pad(_lowerCamelCase      ,   pad_to_multiple_of=10)
        lowerCAmelCase__					     = input_a[input_name]

        lowerCAmelCase__					     = feat_extract.pad(_lowerCamelCase      ,   padding='longest'      ,   pad_to_multiple_of=10)
        lowerCAmelCase__					     = input_a[input_name]

        lowerCAmelCase__					     = feat_extract.pad(
            _lowerCamelCase      ,   padding='max_length'      ,   pad_to_multiple_of=10      ,   max_length=_lowerCamelCase)
        lowerCAmelCase__					     = input_a[input_name]

        lowerCAmelCase__					     = feat_extract.pad(
            _lowerCamelCase      ,   padding='max_length'      ,   pad_to_multiple_of=10      ,   max_length=_lowerCamelCase      ,   return_tensors='np'      ,   )
        lowerCAmelCase__					     = input_a[input_name]

        self.assertTrue(all(len(_lowerCamelCase) % 10 == 0 for x in input_a))
        self.assertTrue(_inputs_are_equal(_lowerCamelCase      ,   _lowerCamelCase))

        lowerCAmelCase__					     = pad_max_length if pad_max_length % 10 == 0 else (pad_max_length // 10 + 1) * 10
        self.assertTrue(all(len(_lowerCamelCase) == expected_mult_pad_length for x in input_a))
        self.assertEqual(input_a.shape[:2]      ,   (batch_size, expected_mult_pad_length))

        if feature_size > 1:
               self.assertTrue(input_a.shape[2] == feature_size)

        # Check padding value is correct
        lowerCAmelCase__					     = (np.ones(self.feat_extract_tester.feature_size) * feat_extract.padding_value).sum()
        self.assertTrue(
            abs(np.asarray(input_a[0])[pad_min_length:].sum() - padding_vector_sum * (pad_max_length - pad_min_length))
            < 1e-3)
        self.assertTrue(
            abs(
                np.asarray(input_a[1])[pad_min_length + pad_diff :].sum()
                - padding_vector_sum * (pad_max_length - pad_min_length - pad_diff))
            < 1e-3)
        self.assertTrue(
            abs(
                np.asarray(input_a[2])[pad_min_length + 2 * pad_diff :].sum()
                - padding_vector_sum * (pad_max_length - pad_min_length - 2 * pad_diff))
            < 1e-3)
        self.assertTrue(
            abs(input_a[0, pad_min_length:].sum() - padding_vector_sum * (pad_max_length - pad_min_length)) < 1e-3)
        self.assertTrue(
            abs(input_a[0, pad_min_length:].sum() - padding_vector_sum * (expected_mult_pad_length - pad_min_length))
            < 1e-3)




 def 		__snake_case						(      self :  Any      ,   lowercase__ :  str=False):







        '''simple docstring'''



        def _inputs_have_equal_length(lowercase__ :  str):
               lowerCAmelCase__					     = len(input[0])
               for input_slice in input[1:]:
                      if len(_lowerCamelCase) != length:
                             return False
               return True

        def _inputs_are_equal(lowercase__ :  Tuple      ,   lowercase__ :  int):
               if len(_lowerCamelCase) != len(_lowerCamelCase):
                      return False

               for input_slice_a, input_slice_a in zip(_lowerCamelCase      ,   _lowerCamelCase):
                      if not np.allclose(np.asarray(_lowerCamelCase)      ,   np.asarray(_lowerCamelCase)      ,   atol=1e-3):
                             return False
               return True

        lowerCAmelCase__					     = self.feature_extraction_class(**self.feat_extract_dict)
        lowerCAmelCase__					     = self.feat_extract_tester.prepare_inputs_for_common(numpify=_lowerCamelCase)
        lowerCAmelCase__					     = feat_extract.model_input_names[0]

        lowerCAmelCase__					     = BatchFeature({input_name: speech_inputs})

        # truncate to smallest
        lowerCAmelCase__					     = feat_extract.pad(
            _lowerCamelCase      ,   padding='max_length'      ,   max_length=len(speech_inputs[0])      ,   truncation=_lowerCamelCase)
        lowerCAmelCase__					     = input_a[input_name]

        lowerCAmelCase__					     = feat_extract.pad(_lowerCamelCase      ,   padding='max_length'      ,   max_length=len(speech_inputs[0]))
        lowerCAmelCase__					     = input_a[input_name]

        self.assertTrue(_inputs_have_equal_length(_lowerCamelCase))
        self.assertFalse(_inputs_have_equal_length(_lowerCamelCase))

        # truncate to smallest with np
        lowerCAmelCase__					     = feat_extract.pad(
            _lowerCamelCase      ,   padding='max_length'      ,   max_length=len(speech_inputs[0])      ,   return_tensors='np'      ,   truncation=_lowerCamelCase      ,   )
        lowerCAmelCase__					     = input_a[input_name]

        lowerCAmelCase__					     = feat_extract.pad(
            _lowerCamelCase      ,   padding='max_length'      ,   max_length=len(speech_inputs[0])      ,   return_tensors='np')
        lowerCAmelCase__					     = input_a[input_name]

        self.assertTrue(_inputs_have_equal_length(_lowerCamelCase))
        self.assertTrue(input_a.shape[1] == len(speech_inputs[0]))

        # since truncation forces padding to be smaller than longest input
        # function can't return `np.ndarray`, but has to return list
        self.assertFalse(_inputs_have_equal_length(_lowerCamelCase))

        # truncate to middle
        lowerCAmelCase__					     = feat_extract.pad(
            _lowerCamelCase      ,   padding='max_length'      ,   max_length=len(speech_inputs[1])      ,   truncation=_lowerCamelCase      ,   return_tensors='np'      ,   )
        lowerCAmelCase__					     = input_a[input_name]

        lowerCAmelCase__					     = feat_extract.pad(
            _lowerCamelCase      ,   padding='max_length'      ,   max_length=len(speech_inputs[1])      ,   truncation=_lowerCamelCase)
        lowerCAmelCase__					     = input_a[input_name]

        lowerCAmelCase__					     = feat_extract.pad(
            _lowerCamelCase      ,   padding='max_length'      ,   max_length=len(speech_inputs[1])      ,   return_tensors='np')
        lowerCAmelCase__					     = input_a[input_name]

        self.assertTrue(input_a.shape[1] == len(speech_inputs[1]))
        self.assertTrue(_inputs_have_equal_length(_lowerCamelCase))
        self.assertTrue(_inputs_have_equal_length(_lowerCamelCase))
        self.assertTrue(_inputs_are_equal(_lowerCamelCase      ,   _lowerCamelCase))

        # since truncation forces padding to be smaller than longest input
        # function can't return `np.ndarray`, but has to return list
        self.assertFalse(_inputs_have_equal_length(_lowerCamelCase))
        self.assertTrue(len(input_a[-1]) == len(speech_inputs[-1]))

        # padding has to be max_length when setting `truncation=True`
        with self.assertRaises(_lowerCamelCase):
               feat_extract.pad(_lowerCamelCase      ,   truncation=_lowerCamelCase)[input_name]

        # padding has to be max_length when setting `truncation=True`
        with self.assertRaises(_lowerCamelCase):
               feat_extract.pad(_lowerCamelCase      ,   padding='longest'      ,   truncation=_lowerCamelCase)[input_name]

        # padding has to be max_length when setting `truncation=True`
        with self.assertRaises(_lowerCamelCase):
               feat_extract.pad(_lowerCamelCase      ,   padding='longest'      ,   truncation=_lowerCamelCase)[input_name]

        # max_length parameter has to be provided when setting `truncation=True` and padding="max_length"
        with self.assertRaises(_lowerCamelCase):
               feat_extract.pad(_lowerCamelCase      ,   padding='max_length'      ,   truncation=_lowerCamelCase)[input_name]

        # test truncation for `pad_to_multiple_of` for List[int] + numpy
        lowerCAmelCase__					     = 12
        lowerCAmelCase__					     = feat_extract.pad(
            _lowerCamelCase      ,   padding='max_length'      ,   max_length=len(speech_inputs[0])      ,   pad_to_multiple_of=_lowerCamelCase      ,   truncation=_lowerCamelCase      ,   )
        lowerCAmelCase__					     = input_a[input_name]

        lowerCAmelCase__					     = feat_extract.pad(
            _lowerCamelCase      ,   padding='max_length'      ,   max_length=len(speech_inputs[0])      ,   pad_to_multiple_of=_lowerCamelCase      ,   )
        lowerCAmelCase__					     = input_a[input_name]

        # retrieve expected_length as multiple of pad_to_multiple_of
        lowerCAmelCase__					     = len(speech_inputs[0])
        if expected_length % pad_to_multiple_of != 0:
               lowerCAmelCase__					     = ((len(speech_inputs[0]) // pad_to_multiple_of) + 1) * pad_to_multiple_of

        self.assertTrue(len(input_a[0]) == expected_length)
        self.assertTrue(_inputs_have_equal_length(_lowerCamelCase))
        self.assertFalse(_inputs_have_equal_length(_lowerCamelCase))




 def 		__snake_case						(      self :  Optional[int]):







        '''simple docstring'''



        self._check_padding(numpify=_lowerCamelCase)




 def 		__snake_case						(      self :  str):







        '''simple docstring'''



        self._check_padding(numpify=_lowerCamelCase)




 def 		__snake_case						(      self :  str):







        '''simple docstring'''



        self._check_truncation(numpify=_lowerCamelCase)




 def 		__snake_case						(      self :  Tuple):







        '''simple docstring'''



        self._check_truncation(numpify=_lowerCamelCase)




 @require_torch
 def 		__snake_case						(      self :  List[Any]):







        '''simple docstring'''



        lowerCAmelCase__					     = self.feature_extraction_class(**self.feat_extract_dict)
        lowerCAmelCase__					     = self.feat_extract_tester.prepare_inputs_for_common()
        lowerCAmelCase__					     = feat_extract.model_input_names[0]

        lowerCAmelCase__					     = BatchFeature({input_name: speech_inputs})

        lowerCAmelCase__					     = feat_extract.pad(_lowerCamelCase      ,   padding='longest'      ,   return_tensors='np')[input_name]
        lowerCAmelCase__					     = feat_extract.pad(_lowerCamelCase      ,   padding='longest'      ,   return_tensors='pt')[input_name]

        self.assertTrue(abs(input_np.astype(np.floataa).sum() - input_pt.numpy().astype(np.floataa).sum()) < 1e-2)




 @require_tf
 def 		__snake_case						(      self :  Optional[int]):







        '''simple docstring'''



        lowerCAmelCase__					     = self.feature_extraction_class(**self.feat_extract_dict)
        lowerCAmelCase__					     = self.feat_extract_tester.prepare_inputs_for_common()
        lowerCAmelCase__					     = feat_extract.model_input_names[0]

        lowerCAmelCase__					     = BatchFeature({input_name: speech_inputs})

        lowerCAmelCase__					     = feat_extract.pad(_lowerCamelCase      ,   padding='longest'      ,   return_tensors='np')[input_name]
        lowerCAmelCase__					     = feat_extract.pad(_lowerCamelCase      ,   padding='longest'      ,   return_tensors='tf')[input_name]

        self.assertTrue(abs(input_np.astype(np.floataa).sum() - input_tf.numpy().astype(np.floataa).sum()) < 1e-2)




 def 		__snake_case						(      self :  int):







        '''simple docstring'''



        lowerCAmelCase__					     = self.feat_extract_dict
        lowerCAmelCase__					     = True
        lowerCAmelCase__					     = self.feature_extraction_class(**_lowerCamelCase)
        lowerCAmelCase__					     = self.feat_extract_tester.prepare_inputs_for_common()
        lowerCAmelCase__					     = [len(_lowerCamelCase) for x in speech_inputs]
        lowerCAmelCase__					     = feat_extract.model_input_names[0]

        lowerCAmelCase__					     = BatchFeature({input_name: speech_inputs})

        lowerCAmelCase__					     = feat_extract.pad(_lowerCamelCase      ,   padding='longest'      ,   return_tensors='np')
        self.assertIn('attention_mask'      ,   _lowerCamelCase)
        self.assertListEqual(list(processed.attention_mask.shape)      ,   list(processed[input_name].shape[:2]))
        self.assertListEqual(processed.attention_mask.sum(-1).tolist()      ,   _lowerCamelCase)




 def 		__snake_case						(      self :  Dict):







        '''simple docstring'''



        lowerCAmelCase__					     = self.feat_extract_dict
        lowerCAmelCase__					     = True
        lowerCAmelCase__					     = self.feature_extraction_class(**_lowerCamelCase)
        lowerCAmelCase__					     = self.feat_extract_tester.prepare_inputs_for_common()
        lowerCAmelCase__					     = [len(_lowerCamelCase) for x in speech_inputs]
        lowerCAmelCase__					     = feat_extract.model_input_names[0]

        lowerCAmelCase__					     = BatchFeature({input_name: speech_inputs})
        lowerCAmelCase__					     = min(_lowerCamelCase)

        lowerCAmelCase__					     = feat_extract.pad(
            _lowerCamelCase      ,   padding='max_length'      ,   max_length=_lowerCamelCase      ,   truncation=_lowerCamelCase      ,   return_tensors='np')
        self.assertIn('attention_mask'      ,   _lowerCamelCase)
        self.assertListEqual(
            list(processed_pad.attention_mask.shape)      ,   [processed_pad[input_name].shape[0], max_length])
        self.assertListEqual(
            processed_pad.attention_mask[:, :max_length].sum(-1).tolist()      ,   [max_length for x in speech_inputs])


import logging
import os

import quant_trainer
import torch
from torch.utils.data import DataLoader

from transformers import Trainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput


lowerCAmelCase__     						=			logging.getLogger(__name__)

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






 '''simple docstring'''






 def __init__(      self :  int      ,   *lowercase__ :  Dict      ,   lowercase__ :  List[Any]=None      ,   lowercase__ :  Optional[Any]=None      ,   lowercase__ :  Dict=None      ,   **lowercase__ :  Optional[int]):







        '''simple docstring'''



        super().__init__(*lowercase__      ,   **lowercase__)
        lowerCAmelCase__					     = eval_examples
        lowerCAmelCase__					     = post_process_function
        lowerCAmelCase__					     = quant_trainer_args
        lowerCAmelCase__					     = 128  # default number of calibration samples




 def 		__snake_case						(      self :  Tuple      ,   lowercase__ :  Any=None):







        '''simple docstring'''



        if calib_dataset is None and self.calib_dataset is None:
               raise ValueError('Trainer: calibration requires an calib_dataset.')
        lowerCAmelCase__					     = calib_dataset if calib_dataset is not None else self.calib_dataset

        lowerCAmelCase__					     = self._remove_unused_columns(lowercase__      ,   description='Calibration')

        return DataLoader(
            lowercase__      ,   batch_size=self.args.eval_batch_size      ,   collate_fn=self.data_collator      ,   drop_last=self.args.dataloader_drop_last      ,   num_workers=self.args.dataloader_num_workers      ,   pin_memory=self.args.dataloader_pin_memory      ,   shuffle=lowercase__      ,   )




 def 		__snake_case						(      self :  List[Any]      ,   lowercase__ :  Union[str, Any]=None):







        '''simple docstring'''



        lowerCAmelCase__					     = self.train_dataset if calib_dataset is None else calib_dataset
        lowerCAmelCase__					     = self.get_calib_dataloader(lowercase__)

        lowerCAmelCase__					     = self.model
        quant_trainer.configure_model(lowercase__      ,   self.quant_trainer_args      ,   calib=lowercase__)
        model.eval()
        quant_trainer.enable_calibration(lowercase__)

        logger.info('***** Running calibration *****')
        logger.info(F"""  Num examples = {self.calib_num}""")
        logger.info(F"""  Batch size = {calib_dataloader.batch_size}""")

        for step, inputs in enumerate(lowercase__):
               # Prediction step
               lowerCAmelCase__     , lowerCAmelCase__     , lowerCAmelCase__					     = self.prediction_step(lowercase__      ,   lowercase__      ,   prediction_loss_only=lowercase__)
               if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
                      break

        quant_trainer.finish_calibration(lowercase__      ,   self.quant_trainer_args)
        lowerCAmelCase__					     = model




 def 		__snake_case						(      self :  Optional[Any]      ,   lowercase__ :  List[Any]=None      ,   lowercase__ :  Optional[Any]=None      ,   lowercase__ :  List[Any]=None      ,   lowercase__ :  str = "eval"):







        '''simple docstring'''



        lowerCAmelCase__					     = self.eval_dataset if eval_dataset is None else eval_dataset
        lowerCAmelCase__					     = self.get_eval_dataloader(lowercase__)
        lowerCAmelCase__					     = self.eval_examples if eval_examples is None else eval_examples

        # Temporarily disable metric computation, we will do it in the loop here.
        lowerCAmelCase__					     = self.compute_metrics
        lowerCAmelCase__					     = None
        lowerCAmelCase__					     = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        try:
               lowerCAmelCase__					     = eval_loop(
                   lowercase__      ,   description='Evaluation'      ,   prediction_loss_only=True if compute_metrics is None else None      ,   ignore_keys=lowercase__      ,   )
        finally:
               lowerCAmelCase__					     = compute_metrics

        if self.post_process_function is not None and self.compute_metrics is not None:
               lowerCAmelCase__					     = self.post_process_function(lowercase__      ,   lowercase__      ,   output.predictions)
               lowerCAmelCase__					     = self.compute_metrics(lowercase__)

               # Prefix all keys with metric_key_prefix + '_'
               for key in list(metrics.keys()):
                      if not key.startswith(F"""{metric_key_prefix}_"""):
                             lowerCAmelCase__					     = metrics.pop(lowercase__)

               self.log(lowercase__)
        else:
               lowerCAmelCase__					     = {}

        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())

        lowerCAmelCase__					     = self.callback_handler.on_evaluate(self.args      ,   self.state      ,   self.control      ,   lowercase__)
        return metrics




 def 		__snake_case						(      self :  Optional[int]      ,   lowercase__ :  str      ,   lowercase__ :  Any      ,   lowercase__ :  List[str]=None      ,   lowercase__ :  str = "test"):







        '''simple docstring'''



        lowerCAmelCase__					     = self.get_test_dataloader(lowercase__)

        # Temporarily disable metric computation, we will do it in the loop here.
        lowerCAmelCase__					     = self.compute_metrics
        lowerCAmelCase__					     = None
        lowerCAmelCase__					     = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        try:
               lowerCAmelCase__					     = eval_loop(
                   lowercase__      ,   description='Prediction'      ,   prediction_loss_only=True if compute_metrics is None else None      ,   ignore_keys=lowercase__      ,   )
        finally:
               lowerCAmelCase__					     = compute_metrics

        if self.post_process_function is None or self.compute_metrics is None:
               return output

        lowerCAmelCase__					     = self.post_process_function(lowercase__      ,   lowercase__      ,   output.predictions      ,   'predict')
        lowerCAmelCase__					     = self.compute_metrics(lowercase__)

        # Prefix all keys with metric_key_prefix + '_'
        for key in list(metrics.keys()):
               if not key.startswith(F"""{metric_key_prefix}_"""):
                      lowerCAmelCase__					     = metrics.pop(lowercase__)

        return PredictionOutput(predictions=predictions.predictions      ,   label_ids=predictions.label_ids      ,   metrics=lowercase__)




 def 		__snake_case						(      self :  List[str]      ,   lowercase__ :  List[str]="./"):







        '''simple docstring'''



        lowerCAmelCase__					     = self.eval_dataset
        lowerCAmelCase__					     = self.get_eval_dataloader(lowercase__)

        lowerCAmelCase__					     = next(iter(lowercase__))

        # saving device - to make it consistent
        lowerCAmelCase__					     = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        # convert to tuple
        lowerCAmelCase__					     = tuple(v.to(lowercase__) for k, v in batch.items())

        logger.info('Converting model to be onnx compatible')
        from pytorch_quantization.nn import TensorQuantizer

        lowerCAmelCase__					     = True

        lowerCAmelCase__					     = self.model.to(lowercase__)

        model.eval()
        model.float()

        lowerCAmelCase__					     = model.module if hasattr(lowercase__      ,   'module') else model
        quant_trainer.configure_model(lowercase__      ,   self.quant_trainer_args)

        lowerCAmelCase__					     = os.path.join(lowercase__      ,   'model.onnx')
        logger.info(F"""exporting model to {output_model_file}""")

        lowerCAmelCase__					     = {0: 'batch_size', 1: 'seq_len'}

        torch.onnx.export(
            lowercase__      ,   lowercase__      ,   lowercase__      ,   export_params=lowercase__      ,   opset_version=13      ,   do_constant_folding=lowercase__      ,   input_names=['input_ids', 'attention_mask', 'token_type_ids']      ,   output_names=['output_start_logits', 'output_end_logits']      ,   dynamic_axes={
                'input_ids': axes,
                'attention_mask': axes,
                'token_type_ids': axes,
                'output_start_logits': axes,
                'output_end_logits': axes,
            }      ,   verbose=lowercase__      ,   )
        logger.info('onnx export finished')
