



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       _SCREAMING_SNAKE_CASE	(UpperCamelCase ):





     def __init__(       self :							int     ,  *UpperCamelCase :							Dict     ,  UpperCamelCase :							List[Any]=None     ,  UpperCamelCase :							Optional[Any]=None     ,  UpperCamelCase :							List[str]=None     ,  **UpperCamelCase :							Dict				)->Dict:
      super().__init__(*UpperCamelCase     ,  **UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    str       =	eval_examples
      __SCREAMING_SNAKE_CASE       :    Any       =	post_process_function
      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	quant_trainer_args
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	1_2_8  # default number of calibration samples





     def 	__snake_case    (       self :							Dict     ,  UpperCamelCase :							List[str]=None				)->Union[str, Any]:
      if calib_dataset is None and self.calib_dataset is None:
       raise ValueError("Trainer: calibration requires an calib_dataset."				)
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	calib_dataset if calib_dataset is not None else self.calib_dataset

      __SCREAMING_SNAKE_CASE       :    List[Any]       =	self._remove_unused_columns(UpperCamelCase     ,  description="Calibration"				)

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





     def 	__snake_case    (       self :							int     ,  UpperCamelCase :							Tuple=None				)->Tuple:
      __SCREAMING_SNAKE_CASE       :    List[str]       =	self.train_dataset if calib_dataset is None else calib_dataset
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	self.get_calib_dataloader(UpperCamelCase				)

      __SCREAMING_SNAKE_CASE       :    Any       =	self.model
      quant_trainer.configure_model(UpperCamelCase     ,  self.quant_trainer_args     ,  calib=UpperCamelCase				)
      model.eval()
      quant_trainer.enable_calibration(UpperCamelCase				)

      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(UpperCamelCase				):
       # Prediction step
       __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       :    Optional[Any]       =	self.prediction_step(UpperCamelCase     ,  UpperCamelCase     ,  prediction_loss_only=UpperCamelCase				)
       if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
        break

      quant_trainer.finish_calibration(UpperCamelCase     ,  self.quant_trainer_args				)
      __SCREAMING_SNAKE_CASE       :    Dict       =	model





     def 	__snake_case    (       self :							List[str]     ,  UpperCamelCase :							Tuple=None     ,  UpperCamelCase :							List[Any]=None     ,  UpperCamelCase :							Optional[int]=None     ,  UpperCamelCase :							str = "eval"				)->Optional[int]:
      __SCREAMING_SNAKE_CASE       :    Any       =	self.eval_dataset if eval_dataset is None else eval_dataset
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	self.get_eval_dataloader(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	self.eval_examples if eval_examples is None else eval_examples

      # Temporarily disable metric computation, we will do it in the loop here.
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	self.compute_metrics
      __SCREAMING_SNAKE_CASE       :    str       =	None
      __SCREAMING_SNAKE_CASE       :    str       =	self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
      try:
       __SCREAMING_SNAKE_CASE       :    int       =	eval_loop(
           UpperCamelCase     ,  description="Evaluation"     ,  prediction_loss_only=True if compute_metrics is None else None     ,  ignore_keys=UpperCamelCase     ,  )
      finally:
       __SCREAMING_SNAKE_CASE       :    int       =	compute_metrics

      if self.post_process_function is not None and self.compute_metrics is not None:
       __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	self.post_process_function(UpperCamelCase     ,  UpperCamelCase     ,  output.predictions				)
       __SCREAMING_SNAKE_CASE       :    Any       =	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}_"""				):
         __SCREAMING_SNAKE_CASE       :    Optional[int]       =	metrics.pop(UpperCamelCase				)

       self.log(UpperCamelCase				)
      else:
       __SCREAMING_SNAKE_CASE       :    Optional[int]       =	{}

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

      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	self.callback_handler.on_evaluate(self.args     ,  self.state     ,  self.control     ,  UpperCamelCase				)
      return metrics





     def 	__snake_case    (       self :							Dict     ,  UpperCamelCase :							List[Any]     ,  UpperCamelCase :							Optional[Any]     ,  UpperCamelCase :							Any=None     ,  UpperCamelCase :							str = "test"				)->List[Any]:
      __SCREAMING_SNAKE_CASE       :    Tuple       =	self.get_test_dataloader(UpperCamelCase				)

      # Temporarily disable metric computation, we will do it in the loop here.
      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	self.compute_metrics
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	None
      __SCREAMING_SNAKE_CASE       :    int       =	self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
      try:
       __SCREAMING_SNAKE_CASE       :    Tuple       =	eval_loop(
           UpperCamelCase     ,  description="Prediction"     ,  prediction_loss_only=True if compute_metrics is None else None     ,  ignore_keys=UpperCamelCase     ,  )
      finally:
       __SCREAMING_SNAKE_CASE       :    Tuple       =	compute_metrics

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

      __SCREAMING_SNAKE_CASE       :    Any       =	self.post_process_function(UpperCamelCase     ,  UpperCamelCase     ,  output.predictions     ,  "predict"				)
      __SCREAMING_SNAKE_CASE       :    Tuple       =	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}_"""				):
        __SCREAMING_SNAKE_CASE       :    Optional[int]       =	metrics.pop(UpperCamelCase				)

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






     def 	__snake_case    (       self :							str     ,  UpperCamelCase :							List[str]="./"				)->Optional[int]:
      __SCREAMING_SNAKE_CASE       :    int       =	self.eval_dataset
      __SCREAMING_SNAKE_CASE       :    Tuple       =	self.get_eval_dataloader(UpperCamelCase				)

      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	next(iter(UpperCamelCase				)				)

      # saving device - to make it consistent
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	torch.device("cuda" if torch.cuda.is_available() else "cpu"				)

      # convert to tuple
      __SCREAMING_SNAKE_CASE       :    Any       =	tuple(v.to(UpperCamelCase				) for k, v in batch.items()				)

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

      __SCREAMING_SNAKE_CASE       :    Dict       =	True

      __SCREAMING_SNAKE_CASE       :    Any       =	self.model.to(UpperCamelCase				)

      model.eval()
      model.float()

      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	model.module if hasattr(UpperCamelCase     ,  "module"				) else model
      quant_trainer.configure_model(UpperCamelCase     ,  self.quant_trainer_args				)

      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	os.path.join(UpperCamelCase     ,  "model.onnx"				)
      logger.info(F"""exporting model to {output_model_file}"""				)

      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	{0: "batch_size", 1: "seq_len"}

      torch.onnx.export(
          UpperCamelCase     ,  UpperCamelCase     ,  UpperCamelCase     ,  export_params=UpperCamelCase     ,  opset_version=1_3     ,  do_constant_folding=UpperCamelCase     ,  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=UpperCamelCase     ,  )
      logger.info("onnx export finished"				)











from diffusers.utils.testing_utils import require_onnxruntime

@require_onnxruntime
class       _SCREAMING_SNAKE_CASE	:
     pass






