


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


_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 							A_      (       lowerCAmelCase__ ):


  def __init__(			self:						Tuple							,*__lowerCAmelCase:						Optional[int]							,__lowerCAmelCase:						List[Any]=None							,__lowerCAmelCase:						int=None							,__lowerCAmelCase:						Dict=None							,**__lowerCAmelCase:						Dict    ):







     '''simple docstring'''







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


  def 							_lowercase					(			self:						Union[str, Any]							,__lowerCAmelCase:						Optional[Any]=None    ):







     '''simple docstring'''







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

     _lowerCamelCase   :				Optional[int]							=     self._remove_unused_columns(_lowerCamelCase							,description="Calibration"    )

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


  def 							_lowercase					(			self:						Optional[int]							,__lowerCAmelCase:						Optional[int]=None    ):







     '''simple docstring'''







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

     _lowerCamelCase   :				List[str]							=     self.model
     quant_trainer.configure_model(_lowerCamelCase							,self.quant_trainer_args							,calib=_lowerCamelCase    )
     model.eval()
     quant_trainer.enable_calibration(_lowerCamelCase    )

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

     quant_trainer.finish_calibration(_lowerCamelCase							,self.quant_trainer_args    )
     _lowerCamelCase   :				Tuple							=     model


  def 							_lowercase					(			self:						List[str]							,__lowerCAmelCase:						Optional[int]=None							,__lowerCAmelCase:						int=None							,__lowerCAmelCase:						int=None							,__lowerCAmelCase:						List[str] = "eval"    ):







     '''simple docstring'''







     _lowerCamelCase   :				Tuple							=     self.eval_dataset if eval_dataset is None else eval_dataset
     _lowerCamelCase   :				List[Any]							=     self.get_eval_dataloader(_lowerCamelCase    )
     _lowerCamelCase   :				Tuple							=     self.eval_examples if eval_examples is None else eval_examples

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

     if self.post_process_function is not None and self.compute_metrics is not None:
        _lowerCamelCase   :				int							=     self.post_process_function(_lowerCamelCase							,_lowerCamelCase							,output.predictions    )
        _lowerCamelCase   :				List[Any]							=     self.compute_metrics(_lowerCamelCase    )

        # Prefix all keys with metric_key_prefix + '_'
        for key in list(metrics.keys()    ):
           if not key.startswith(F"""{metric_key_prefix}_"""    ):
              _lowerCamelCase   :				Optional[int]							=     metrics.pop(_lowerCamelCase    )

        self.log(_lowerCamelCase    )
     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[int]							=     self.callback_handler.on_evaluate(self.args							,self.state							,self.control							,_lowerCamelCase    )
     return metrics


  def 							_lowercase					(			self:						Optional[int]							,__lowerCAmelCase:						Optional[int]							,__lowerCAmelCase:						str							,__lowerCAmelCase:						Dict=None							,__lowerCAmelCase:						int = "test"    ):







     '''simple docstring'''







     _lowerCamelCase   :				Union[str, Any]							=     self.get_test_dataloader(_lowerCamelCase    )

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

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

     _lowerCamelCase   :				Dict							=     self.post_process_function(_lowerCamelCase							,_lowerCamelCase							,output.predictions							,"predict"    )
     _lowerCamelCase   :				Any							=     self.compute_metrics(_lowerCamelCase    )

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

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






  def 							_lowercase					(			self:						int							,__lowerCAmelCase:						Tuple="./"    ):







     '''simple docstring'''







     _lowerCamelCase   :				Tuple							=     self.eval_dataset
     _lowerCamelCase   :				int							=     self.get_eval_dataloader(_lowerCamelCase    )

     _lowerCamelCase   :				Union[str, Any]							=     next(iter(_lowerCamelCase    )    )

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

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

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

     _lowerCamelCase   :				Optional[Any]							=     True

     _lowerCamelCase   :				Tuple							=     self.model.to(_lowerCamelCase    )

     model.eval()
     model.float()

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

     _lowerCamelCase   :				List[str]							=     os.path.join(_lowerCamelCase							,"model.onnx"    )
     logger.info(F"""exporting model to {output_model_file}"""    )

     _lowerCamelCase   :				Optional[Any]							=     {0: """batch_size""", 1: """seq_len"""}

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


"""simple docstring"""
from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_A							=  {"""configuration_sew""": ["""SEW_PRETRAINED_CONFIG_ARCHIVE_MAP""", """SEWConfig"""]}

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    _A							=  [
        """SEW_PRETRAINED_MODEL_ARCHIVE_LIST""",
        """SEWForCTC""",
        """SEWForSequenceClassification""",
        """SEWModel""",
        """SEWPreTrainedModel""",
    ]

if TYPE_CHECKING:
    from .configuration_sew import SEW_PRETRAINED_CONFIG_ARCHIVE_MAP, SEWConfig

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_sew import (
            SEW_PRETRAINED_MODEL_ARCHIVE_LIST,
            SEWForCTC,
            SEWForSequenceClassification,
            SEWModel,
            SEWPreTrainedModel,
        )


else:
    import sys

    _A							=  _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

