





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_						       =	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__   (				lowerCAmelCase_       ):



     '''simple docstring'''

     def __init__(							self						:   Optional[Any]						,*lowerCamelCase__						:   int						,lowerCamelCase__						:   str=None						,lowerCamelCase__						:   List[Any]=None						,lowerCamelCase__						:   Tuple=None						,**lowerCamelCase__						:   str	)					->		int:




           '''simple docstring'''


           super().__init__(*lowerCamelCase__						,**lowerCamelCase__	)
           SCREAMING_SNAKE_CASE		=						eval_examples
           SCREAMING_SNAKE_CASE		=						post_process_function
           SCREAMING_SNAKE_CASE		=						quant_trainer_args
           SCREAMING_SNAKE_CASE		=						128  # default number of calibration samples






     def     SCREAMING_SNAKE_CASE__  (							self						:   Optional[Any]						,lowerCamelCase__						:   Optional[Any]=None	)					->		int:




           '''simple docstring'''


           if calib_dataset is None and self.calib_dataset is None:
                 raise ValueError("""Trainer: calibration requires an calib_dataset."""	)
           SCREAMING_SNAKE_CASE		=						calib_dataset if calib_dataset is not None else self.calib_dataset

           SCREAMING_SNAKE_CASE		=						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     SCREAMING_SNAKE_CASE__  (							self						:   List[Any]						,lowerCamelCase__						:   List[str]=None	)					->		Any:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.train_dataset if calib_dataset is None else calib_dataset
           SCREAMING_SNAKE_CASE		=						self.get_calib_dataloader(lowerCamelCase__	)

           SCREAMING_SNAKE_CASE		=						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
                 SCREAMING_SNAKE_CASE, SCREAMING_SNAKE_CASE, SCREAMING_SNAKE_CASE		=						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	)
           SCREAMING_SNAKE_CASE		=						model






     def     SCREAMING_SNAKE_CASE__  (							self						:   int						,lowerCamelCase__						:   List[Any]=None						,lowerCamelCase__						:   Union[str, Any]=None						,lowerCamelCase__						:   Any=None						,lowerCamelCase__						:   str = "eval"	)					->		Dict:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.eval_dataset if eval_dataset is None else eval_dataset
           SCREAMING_SNAKE_CASE		=						self.get_eval_dataloader(lowerCamelCase__	)
           SCREAMING_SNAKE_CASE		=						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		=						self.compute_metrics
           SCREAMING_SNAKE_CASE		=						None
           SCREAMING_SNAKE_CASE		=						self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
           try:
                 SCREAMING_SNAKE_CASE		=						eval_loop(
                     lowerCamelCase__						,description="""Evaluation"""						,prediction_loss_only=True if compute_metrics is None else None						,ignore_keys=lowerCamelCase__						,)
           finally:
                 SCREAMING_SNAKE_CASE		=						compute_metrics

           if self.post_process_function is not None and self.compute_metrics is not None:
                 SCREAMING_SNAKE_CASE		=						self.post_process_function(lowerCamelCase__						,lowerCamelCase__						,output.predictions	)
                 SCREAMING_SNAKE_CASE		=						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}_"""	):
                             SCREAMING_SNAKE_CASE		=						metrics.pop(lowerCamelCase__	)

                 self.log(lowerCamelCase__	)
           else:
                 SCREAMING_SNAKE_CASE		=						{}

           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		=						self.callback_handler.on_evaluate(self.args						,self.state						,self.control						,lowerCamelCase__	)
           return metrics






     def     SCREAMING_SNAKE_CASE__  (							self						:   Optional[Any]						,lowerCamelCase__						:   Optional[int]						,lowerCamelCase__						:   Optional[int]						,lowerCamelCase__						:   Optional[Any]=None						,lowerCamelCase__						:   str = "test"	)					->		Any:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.get_test_dataloader(lowerCamelCase__	)

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

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

           SCREAMING_SNAKE_CASE		=						self.post_process_function(lowerCamelCase__						,lowerCamelCase__						,output.predictions						,"""predict"""	)
           SCREAMING_SNAKE_CASE		=						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}_"""	):
                       SCREAMING_SNAKE_CASE		=						metrics.pop(lowerCamelCase__	)

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






     def     SCREAMING_SNAKE_CASE__  (							self						:   Union[str, Any]						,lowerCamelCase__						:   Dict="./"	)					->		int:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.eval_dataset
           SCREAMING_SNAKE_CASE		=						self.get_eval_dataloader(lowerCamelCase__	)

           SCREAMING_SNAKE_CASE		=						next(iter(lowerCamelCase__	)	)

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

           # convert to tuple
           SCREAMING_SNAKE_CASE		=						tuple(v.to(lowerCamelCase__	) for k, v in batch.items()	)

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

           SCREAMING_SNAKE_CASE		=						True

           SCREAMING_SNAKE_CASE		=						self.model.to(lowerCamelCase__	)

           model.eval()
           model.float()

           SCREAMING_SNAKE_CASE		=						model.module if hasattr(lowerCamelCase__						,"""module"""	) else model
           quant_trainer.configure_model(lowerCamelCase__						,self.quant_trainer_args	)

           SCREAMING_SNAKE_CASE		=						os.path.join(lowerCamelCase__						,"""model.onnx"""	)
           logger.info(F"""exporting model to {output_model_file}"""	)

           SCREAMING_SNAKE_CASE		=						{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"""	)









from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


SCREAMING_SNAKE_CASE_						       =	{
    """configuration_mvp""": ["""MVP_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MvpConfig""", """MvpOnnxConfig"""],
    """tokenization_mvp""": ["""MvpTokenizer"""],
}

try:
      if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE_						       =	["""MvpTokenizerFast"""]

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE_						       =	[
          """MVP_PRETRAINED_MODEL_ARCHIVE_LIST""",
          """MvpForCausalLM""",
          """MvpForConditionalGeneration""",
          """MvpForQuestionAnswering""",
          """MvpForSequenceClassification""",
          """MvpModel""",
          """MvpPreTrainedModel""",
      ]

if TYPE_CHECKING:
      from .configuration_mvp import MVP_PRETRAINED_CONFIG_ARCHIVE_MAP, MvpConfig, MvpOnnxConfig
      from .tokenization_mvp import MvpTokenizer

      try:
            if not is_tokenizers_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_mvp_fast import MvpTokenizerFast

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_mvp import (
                MVP_PRETRAINED_MODEL_ARCHIVE_LIST,
                MvpForCausalLM,
                MvpForConditionalGeneration,
                MvpForQuestionAnswering,
                MvpForSequenceClassification,
                MvpModel,
                MvpPreTrainedModel,
            )

else:
      import sys

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


