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			:							Union[str, Any]     =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 	__snake_case	( __lowerCAmelCase		):





      '''simple docstring'''






      def __init__(					self     :				Any    ,  *_UpperCamelCase     :				Dict    ,  _UpperCamelCase     :				Optional[int]=None    ,  _UpperCamelCase     :				Union[str, Any]=None    ,  _UpperCamelCase     :				Optional[int]=None    ,  **_UpperCamelCase     :				Union[str, Any])				->int:


          """simple docstring"""


          super().__init__(*_UpperCamelCase    ,  **_UpperCamelCase)
          _lowerCamelCase     :		Optional[Any] 						=						eval_examples
          _lowerCamelCase     :		Dict 						=						post_process_function
          _lowerCamelCase     :		int 						=						quant_trainer_args
          _lowerCamelCase     :		Union[str, Any] 						=						128  # default number of calibration samples



      def 	_SCREAMING_SNAKE_CASE    (					self     :				Dict    ,  _UpperCamelCase     :				Optional[int]=None)				->Optional[int]:


          """simple docstring"""


          if calib_dataset is None and self.calib_dataset is None:
              raise ValueError("""Trainer: calibration requires an calib_dataset.""")
          _lowerCamelCase     :		Union[str, Any] 						=						calib_dataset if calib_dataset is not None else self.calib_dataset

          _lowerCamelCase     :		str 						=						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 	_SCREAMING_SNAKE_CASE    (					self     :				int    ,  _UpperCamelCase     :				int=None)				->Optional[int]:


          """simple docstring"""


          _lowerCamelCase     :		Optional[int] 						=						self.train_dataset if calib_dataset is None else calib_dataset
          _lowerCamelCase     :		Optional[int] 						=						self.get_calib_dataloader(_UpperCamelCase)

          _lowerCamelCase     :		List[str] 						=						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
              _lowerCamelCase						,							_lowerCamelCase						,							_lowerCamelCase     :		List[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)
          _lowerCamelCase     :		Any 						=						model



      def 	_SCREAMING_SNAKE_CASE    (					self     :				List[str]    ,  _UpperCamelCase     :				Optional[Any]=None    ,  _UpperCamelCase     :				List[str]=None    ,  _UpperCamelCase     :				Dict=None    ,  _UpperCamelCase     :				str = "eval")				->Union[str, Any]:


          """simple docstring"""


          _lowerCamelCase     :		Optional[Any] 						=						self.eval_dataset if eval_dataset is None else eval_dataset
          _lowerCamelCase     :		Dict 						=						self.get_eval_dataloader(_UpperCamelCase)
          _lowerCamelCase     :		int 						=						self.eval_examples if eval_examples is None else eval_examples

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

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

              self.log(_UpperCamelCase)
          else:
              _lowerCamelCase     :		List[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())

          _lowerCamelCase     :		Optional[Any] 						=						self.callback_handler.on_evaluate(self.args    ,  self.state    ,  self.control    ,  _UpperCamelCase)
          return metrics



      def 	_SCREAMING_SNAKE_CASE    (					self     :				Optional[Any]    ,  _UpperCamelCase     :				List[str]    ,  _UpperCamelCase     :				int    ,  _UpperCamelCase     :				Union[str, Any]=None    ,  _UpperCamelCase     :				str = "test")				->Union[str, Any]:


          """simple docstring"""


          _lowerCamelCase     :		Tuple 						=						self.get_test_dataloader(_UpperCamelCase)

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

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

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

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




      def 	_SCREAMING_SNAKE_CASE    (					self     :				Optional[int]    ,  _UpperCamelCase     :				List[str]="./")				->Union[str, Any]:


          """simple docstring"""


          _lowerCamelCase     :		List[Any] 						=						self.eval_dataset
          _lowerCamelCase     :		List[Any] 						=						self.get_eval_dataloader(_UpperCamelCase)

          _lowerCamelCase     :		Optional[Any] 						=						next(iter(_UpperCamelCase))

          # saving device - to make it consistent
          _lowerCamelCase     :		Dict 						=						torch.device("""cuda""" if torch.cuda.is_available() else """cpu""")

          # convert to tuple
          _lowerCamelCase     :		Optional[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

          _lowerCamelCase     :		Optional[int] 						=						True

          _lowerCamelCase     :		Union[str, Any] 						=						self.model.to(_UpperCamelCase)

          model.eval()
          model.float()

          _lowerCamelCase     :		Any 						=						model.module if hasattr(_UpperCamelCase    ,  """module""") else model
          quant_trainer.configure_model(_UpperCamelCase    ,  self.quant_trainer_args)

          _lowerCamelCase     :		Tuple 						=						os.path.join(_UpperCamelCase    ,  """model.onnx""")
          logger.info(F"""exporting model to {output_model_file}""")

          _lowerCamelCase     :		List[str] 						=						{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""")





import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp







def 						A__  (					):







    '''simple docstring'''
    _lowerCamelCase     :		Optional[int] 						=						ArgumentParser(
        description=(
            """PyTorch TPU distributed training launch helper utility that will spawn up multiple distributed processes"""
        )	)

    # Optional arguments for the launch helper
    parser.add_argument("""--num_cores"""   , type=__A   , default=1   , help="""Number of TPU cores to use (1 or 8)."""	)

    # positional
    parser.add_argument(
        """training_script"""   , type=__A   , help=(
            """The full path to the single TPU training """
            """program/script to be launched in parallel, """
            """followed by all the arguments for the """
            """training script"""
        )   , )

    # rest from the training program
    parser.add_argument("""training_script_args"""   , nargs=__A	)

    return parser.parse_args()







def 						A__  (					):







    '''simple docstring'''
    _lowerCamelCase     :		List[str] 						=						parse_args()

    # Import training_script as a module.
    _lowerCamelCase     :		List[Any] 						=						Path(args.training_script	)
    sys.path.append(str(script_fpath.parent.resolve()	)	)
    _lowerCamelCase     :		Optional[Any] 						=						script_fpath.stem
    _lowerCamelCase     :		Dict 						=						importlib.import_module(__A	)

    # Patch sys.argv
    _lowerCamelCase     :		Union[str, Any] 						=						[args.training_script] + args.training_script_args + ["""--tpu_num_cores""", str(args.num_cores	)]

    xmp.spawn(mod._mp_fn   , args=()   , nprocs=args.num_cores	)


if __name__ == "__main__":
    main()




