


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


SCREAMING_SNAKE_CASE       :		List[str]         =				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      UpperCamelCase					(		lowercase__     ):


   '''simple docstring'''







   def __init__(     self   ,							*UpperCamelCase_   ,							UpperCamelCase_=None   ,							UpperCamelCase_=None   ,							UpperCamelCase_=None   ,							**UpperCamelCase_		):
        super().__init__(*UpperCamelCase_   ,							**UpperCamelCase_		)
        lowercase_       :Union[str, Any]     =       eval_examples
        lowercase_       :Dict     =       post_process_function
        lowercase_       :List[Any]     =       quant_trainer_args
        lowercase_       :Dict     =       128  # default number of calibration samples
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_=None		):
        if calib_dataset is None and self.calib_dataset is None:
             raise ValueError('''Trainer: calibration requires an calib_dataset.'''		)
        lowercase_       :Optional[Any]     =       calib_dataset if calib_dataset is not None else self.calib_dataset

        lowercase_       :Union[str, 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 		UpperCamelCase  (     self   ,							UpperCamelCase_=None		):
        lowercase_       :Dict     =       self.train_dataset if calib_dataset is None else calib_dataset
        lowercase_       :Any     =       self.get_calib_dataloader(UpperCamelCase_		)

        lowercase_       :List[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
             lowercase_    ,      lowercase_    ,      lowercase_       :str     =       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		)
        lowercase_       :List[Any]     =       model
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_=None   ,							UpperCamelCase_=None   ,							UpperCamelCase_=None   ,							UpperCamelCase_ = "eval"		):
        lowercase_       :List[str]     =       self.eval_dataset if eval_dataset is None else eval_dataset
        lowercase_       :Optional[int]     =       self.get_eval_dataloader(UpperCamelCase_		)
        lowercase_       :List[str]     =       self.eval_examples if eval_examples is None else eval_examples

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

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

             self.log(UpperCamelCase_		)
        else:
             lowercase_       :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()		)

        lowercase_       :Any     =       self.callback_handler.on_evaluate(self.args   ,							self.state   ,							self.control   ,							UpperCamelCase_		)
        return metrics
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_=None   ,							UpperCamelCase_ = "test"		):
        lowercase_       :List[str]     =       self.get_test_dataloader(UpperCamelCase_		)

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

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

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

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




   def 		UpperCamelCase  (     self   ,							UpperCamelCase_="./"		):
        lowercase_       :Dict     =       self.eval_dataset
        lowercase_       :str     =       self.get_eval_dataloader(UpperCamelCase_		)

        lowercase_       :List[Any]     =       next(iter(UpperCamelCase_		)		)

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

        # convert to tuple
        lowercase_       :str     =       tuple(v.to(UpperCamelCase_		) for k, v in batch.items()		)

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

        lowercase_       :Tuple     =       True

        lowercase_       :List[Any]     =       self.model.to(UpperCamelCase_		)

        model.eval()
        model.float()

        lowercase_       :Dict     =       model.module if hasattr(UpperCamelCase_   ,							'''module'''		) else model
        quant_trainer.configure_model(UpperCamelCase_   ,							self.quant_trainer_args		)

        lowercase_       :List[Any]     =       os.path.join(UpperCamelCase_   ,							'''model.onnx'''		)
        logger.info(f"exporting model to {output_model_file}"		)

        lowercase_       :Dict     =       {0: '''batch_size''', 1: '''seq_len'''}

        torch.onnx.export(
            UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							export_params=UpperCamelCase_   ,							opset_version=13   ,							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'''		)





def     UpperCamelCase    (      _a       ,		_a		)					->							int:
     '''simple docstring'''





     while a != 0:
          lowercase_    ,      lowercase_       :Union[str, Any]     =       b % a, a
     return b




def     UpperCamelCase    (      _a       ,		_a		)					->							int:
     '''simple docstring'''





     if gcd(_a       ,		_a		) != 1:
          lowercase_       :Union[str, Any]     =       f"mod inverse of {a!r} and {m!r} does not exist"
          raise ValueError(_a		)
     lowercase_    ,      lowercase_    ,      lowercase_       :Any     =       1, 0, a
     lowercase_    ,      lowercase_    ,      lowercase_       :List[str]     =       0, 1, m
     while va != 0:
          lowercase_       :Tuple     =       ua // va
          lowercase_    ,      lowercase_    ,      lowercase_    ,      lowercase_    ,      lowercase_    ,      lowercase_       :Optional[int]     =       (ua - q * va), (ua - q * va), (ua - q * va), va, va, va
     return ua % m

