'''simple docstring'''




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


__UpperCAmelCase												=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__				(       UpperCAmelCase__			):






     def __init__(			self		:		Any  ,    *a		:		Tuple  ,    a		:		Dict=None  ,    a		:		Dict=None  ,    a		:		str=None  ,    **a		:		Optional[Any]   ):





         """simple docstring"""



         super().__init__(*a  ,    **a   )
         __lowerCamelCase          =  eval_examples
         __lowerCamelCase          =  post_process_function
         __lowerCamelCase          =  quant_trainer_args
         __lowerCamelCase          =  1_28  # default number of calibration samples






     def        SCREAMING_SNAKE_CASE__       (			self		:		Any  ,    a		:		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(a  ,    description='''Calibration'''   )

         return DataLoader(
             a  ,    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=a  ,    )






     def        SCREAMING_SNAKE_CASE__       (			self		:		List[Any]  ,    a		:		Dict=None   ):





         """simple docstring"""



         __lowerCamelCase          =  self.train_dataset if calib_dataset is None else calib_dataset
         __lowerCamelCase          =  self.get_calib_dataloader(a   )

         __lowerCamelCase          =  self.model
         quant_trainer.configure_model(a  ,    self.quant_trainer_args  ,    calib=a   )
         model.eval()
         quant_trainer.enable_calibration(a   )

         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(a   ):
             # Prediction step
             __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase          =  self.prediction_step(a  ,    a  ,    prediction_loss_only=a   )
             if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
                 break

         quant_trainer.finish_calibration(a  ,    self.quant_trainer_args   )
         __lowerCamelCase          =  model






     def        SCREAMING_SNAKE_CASE__       (			self		:		List[str]  ,    a		:		Any=None  ,    a		:		List[str]=None  ,    a		:		Tuple=None  ,    a		:		str = "eval"   ):





         """simple docstring"""



         __lowerCamelCase          =  self.eval_dataset if eval_dataset is None else eval_dataset
         __lowerCamelCase          =  self.get_eval_dataloader(a   )
         __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(
                 a  ,    description='''Evaluation'''  ,    prediction_loss_only=True if compute_metrics is None else None  ,    ignore_keys=a  ,    )
         finally:
             __lowerCamelCase          =  compute_metrics

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

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

             self.log(a   )
         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  ,    a   )
         return metrics






     def        SCREAMING_SNAKE_CASE__       (			self		:		Union[str, Any]  ,    a		:		List[str]  ,    a		:		List[str]  ,    a		:		Any=None  ,    a		:		str = "test"   ):





         """simple docstring"""



         __lowerCamelCase          =  self.get_test_dataloader(a   )

         # 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(
                 a  ,    description='''Prediction'''  ,    prediction_loss_only=True if compute_metrics is None else None  ,    ignore_keys=a  ,    )
         finally:
             __lowerCamelCase          =  compute_metrics

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

         __lowerCamelCase          =  self.post_process_function(a  ,    a  ,    output.predictions  ,    '''predict'''   )
         __lowerCamelCase          =  self.compute_metrics(a   )

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

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






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[int]  ,    a		:		Any="./"   ):





         """simple docstring"""



         __lowerCamelCase          =  self.eval_dataset
         __lowerCamelCase          =  self.get_eval_dataloader(a   )

         __lowerCamelCase          =  next(iter(a   )   )

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

         # convert to tuple
         __lowerCamelCase          =  tuple(v.to(a   ) 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(a   )

         model.eval()
         model.float()

         __lowerCamelCase          =  model.module if hasattr(a  ,    '''module'''   ) else model
         quant_trainer.configure_model(a  ,    self.quant_trainer_args   )

         __lowerCamelCase          =  os.path.join(a  ,    '''model.onnx'''   )
         logger.info(f"""exporting model to {output_model_file}"""   )

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

         torch.onnx.export(
             a  ,    a  ,    a  ,    export_params=a  ,    opset_version=13  ,    do_constant_folding=a  ,    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=a  ,    )
         logger.info('''onnx export finished'''   )

'''simple docstring'''




import time
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, torch_device

from ..test_modeling_common import ids_tensor


if is_torch_available():
  import torch

  from transformers.generation import (
      MaxLengthCriteria,
      MaxNewTokensCriteria,
      MaxTimeCriteria,
      StoppingCriteriaList,
      validate_stopping_criteria,
  )




@require_torch
class    a__				(       unittest.TestCase			):






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[int]  ,    a		:		int   ):





         """simple docstring"""



         __lowerCamelCase          =  3
         __lowerCamelCase          =  2_50

         __lowerCamelCase          =  ids_tensor((batch_size, length)  ,    a   )
         __lowerCamelCase          =  torch.ones((batch_size, length)  ,    device=a  ,    dtype=torch.float   ) / length
         return input_ids, scores






     def        SCREAMING_SNAKE_CASE__       (			self		:		List[Any]   ):





         """simple docstring"""



         __lowerCamelCase     ,  __lowerCamelCase          =  self._get_tensors(5   )

         __lowerCamelCase          =  StoppingCriteriaList(
             [
                 MaxLengthCriteria(max_length=10   ),
                 MaxTimeCriteria(max_time=0.1   ),
             ]   )

         self.assertFalse(criteria(a  ,    a   )   )

         __lowerCamelCase     ,  __lowerCamelCase          =  self._get_tensors(9   )
         self.assertFalse(criteria(a  ,    a   )   )

         __lowerCamelCase     ,  __lowerCamelCase          =  self._get_tensors(10   )
         self.assertTrue(criteria(a  ,    a   )   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		Union[str, Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  MaxLengthCriteria(max_length=10   )

         __lowerCamelCase     ,  __lowerCamelCase          =  self._get_tensors(5   )
         self.assertFalse(criteria(a  ,    a   )   )

         __lowerCamelCase     ,  __lowerCamelCase          =  self._get_tensors(9   )
         self.assertFalse(criteria(a  ,    a   )   )

         __lowerCamelCase     ,  __lowerCamelCase          =  self._get_tensors(10   )
         self.assertTrue(criteria(a  ,    a   )   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		List[Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  MaxNewTokensCriteria(start_length=5  ,    max_new_tokens=5   )

         __lowerCamelCase     ,  __lowerCamelCase          =  self._get_tensors(5   )
         self.assertFalse(criteria(a  ,    a   )   )

         __lowerCamelCase     ,  __lowerCamelCase          =  self._get_tensors(9   )
         self.assertFalse(criteria(a  ,    a   )   )

         __lowerCamelCase     ,  __lowerCamelCase          =  self._get_tensors(10   )
         self.assertTrue(criteria(a  ,    a   )   )

         __lowerCamelCase          =  StoppingCriteriaList([criteria]   )
         self.assertEqual(criteria_list.max_length  ,    10   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		int   ):





         """simple docstring"""



         __lowerCamelCase     ,  __lowerCamelCase          =  self._get_tensors(5   )

         __lowerCamelCase          =  MaxTimeCriteria(max_time=0.1   )
         self.assertFalse(criteria(a  ,    a   )   )

         __lowerCamelCase          =  MaxTimeCriteria(max_time=0.1  ,    initial_timestamp=time.time() - 0.2   )
         self.assertTrue(criteria(a  ,    a   )   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		List[Any]   ):





         """simple docstring"""



         validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10   )]   )  ,    10   )

         with self.assertWarns(a   ):
             validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10   )]   )  ,    11   )

         __lowerCamelCase          =  validate_stopping_criteria(StoppingCriteriaList()  ,    11   )

         self.assertEqual(len(a   )  ,    1   )
