




"""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


A_      			=       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       lowercase(  __a				):







 '''simple docstring'''





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




   '''simple docstring'''



   super().__init__(*a_,   **a_	)
   _snake_case							:	int			=						eval_examples
   _snake_case							:	Optional[int]			=						post_process_function
   _snake_case							:	Dict			=						quant_trainer_args
   _snake_case							:	Optional[int]			=						128  # default number of calibration samples


 def 				UpperCamelCase_     (       self:	Tuple,   a_:	Any=None	):




   '''simple docstring'''



   if calib_dataset is None and self.calib_dataset is None:
     raise ValueError("""Trainer: calibration requires an calib_dataset."""	)
   _snake_case							:	List[str]			=						calib_dataset if calib_dataset is not None else self.calib_dataset

   _snake_case							:	List[Any]			=						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 				UpperCamelCase_     (       self:	List[str],   a_:	Any=None	):




   '''simple docstring'''



   _snake_case							:	List[Any]			=						self.train_dataset if calib_dataset is None else calib_dataset
   _snake_case							:	Optional[int]			=						self.get_calib_dataloader(a_	)

   _snake_case							:	int			=						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
     _snake_case	,      _snake_case	,      _snake_case							:	str			=						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	)
   _snake_case							:	str			=						model


 def 				UpperCamelCase_     (       self:	List[Any],   a_:	Optional[int]=None,   a_:	str=None,   a_:	Optional[int]=None,   a_:	str = "eval"	):




   '''simple docstring'''



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

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

   if self.post_process_function is not None and self.compute_metrics is not None:
     _snake_case							:	Any			=						self.post_process_function(a_,   a_,   output.predictions	)
     _snake_case							:	Union[str, Any]			=						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}_"	):
         _snake_case							:	Dict			=						metrics.pop(a_	)

     self.log(a_	)
   else:
     _snake_case							:	Union[str, Any]			=						{}

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

   _snake_case							:	List[Any]			=						self.callback_handler.on_evaluate(self.args,   self.state,   self.control,   a_	)
   return metrics


 def 				UpperCamelCase_     (       self:	Optional[Any],   a_:	int,   a_:	int,   a_:	Optional[Any]=None,   a_:	str = "test"	):




   '''simple docstring'''



   _snake_case							:	Dict			=						self.get_test_dataloader(a_	)

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

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

   _snake_case							:	List[Any]			=						self.post_process_function(a_,   a_,   output.predictions,   """predict"""	)
   _snake_case							:	int			=						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}_"	):
       _snake_case							:	List[str]			=						metrics.pop(a_	)

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


 def 				UpperCamelCase_     (       self:	Optional[int],   a_:	int="./"	):




   '''simple docstring'''



   _snake_case							:	Optional[int]			=						self.eval_dataset
   _snake_case							:	Union[str, Any]			=						self.get_eval_dataloader(a_	)

   _snake_case							:	str			=						next(iter(a_	)	)

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

   # convert to tuple
   _snake_case							:	List[Any]			=						tuple(v.to(a_	) for k, v in batch.items()	)

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

   _snake_case							:	Dict			=						True

   _snake_case							:	str			=						self.model.to(a_	)

   model.eval()
   model.float()

   _snake_case							:	List[Any]			=						model.module if hasattr(a_,   """module"""	) else model
   quant_trainer.configure_model(a_,   self.quant_trainer_args	)

   _snake_case							:	Dict			=						os.path.join(a_,   """model.onnx"""	)
   logger.info(f"exporting model to {output_model_file}"	)

   _snake_case							:	List[Any]			=						{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"""


from abc import ABC, abstractmethod
from argparse import ArgumentParser






class       lowercase(  __a				):







 '''simple docstring'''





 @staticmethod
 @abstractmethod
 def 				UpperCamelCase_     (       a_:	ArgumentParser	):




   '''simple docstring'''



   raise NotImplementedError()


 @abstractmethod
 def 				UpperCamelCase_     (       self:	List[str]	):




   '''simple docstring'''



   raise NotImplementedError()


