





'''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							:	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	(       UpperCamelCase__				):
 def __init__(				self      :Any						,       *lowercase_      :Optional[Any]						,       lowercase_      :Any=None						,       lowercase_      :Union[str, Any]=None						,       lowercase_      :List[Any]=None						,       **lowercase_      :List[str]							)->							List[Any]:
     super().__init__(*lowercase_						,       **lowercase_							)
     A__				=			eval_examples
     A__				=			post_process_function
     A__				=			quant_trainer_args
     A__				=			1_28  # default number of calibration samples
 def      UpperCAmelCase_							(				self      :Optional[int]						,       lowercase_      :List[str]=None							)->							Union[str, Any]:
     if calib_dataset is None and self.calib_dataset is None:
         raise ValueError("Trainer: calibration requires an calib_dataset."							)
     A__				=			calib_dataset if calib_dataset is not None else self.calib_dataset

     A__				=			self._remove_unused_columns(lowercase_						,       description="Calibration"							)

     return DataLoader(
         lowercase_						,       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=lowercase_						,       )
 def      UpperCAmelCase_							(				self      :List[str]						,       lowercase_      :Any=None							)->							Union[str, Any]:
     A__				=			self.train_dataset if calib_dataset is None else calib_dataset
     A__				=			self.get_calib_dataloader(lowercase_							)

     A__				=			self.model
     quant_trainer.configure_model(lowercase_						,       self.quant_trainer_args						,       calib=lowercase_							)
     model.eval()
     quant_trainer.enable_calibration(lowercase_							)

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

     quant_trainer.finish_calibration(lowercase_						,       self.quant_trainer_args							)
     A__				=			model
 def      UpperCAmelCase_							(				self      :Any						,       lowercase_      :List[str]=None						,       lowercase_      :Dict=None						,       lowercase_      :Any=None						,       lowercase_      :str = "eval"							)->							Union[str, Any]:
     A__				=			self.eval_dataset if eval_dataset is None else eval_dataset
     A__				=			self.get_eval_dataloader(lowercase_							)
     A__				=			self.eval_examples if eval_examples is None else eval_examples

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

     if self.post_process_function is not None and self.compute_metrics is not None:
         A__				=			self.post_process_function(lowercase_						,       lowercase_						,       output.predictions							)
         A__				=			self.compute_metrics(lowercase_							)

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

         self.log(lowercase_							)
     else:
         A__				=			{}

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

     A__				=			self.callback_handler.on_evaluate(self.args						,       self.state						,       self.control						,       lowercase_							)
     return metrics
 def      UpperCAmelCase_							(				self      :List[str]						,       lowercase_      :str						,       lowercase_      :Optional[Any]						,       lowercase_      :str=None						,       lowercase_      :str = "test"							)->							List[Any]:
     A__				=			self.get_test_dataloader(lowercase_							)

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

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

     A__				=			self.post_process_function(lowercase_						,       lowercase_						,       output.predictions						,       "predict"							)
     A__				=			self.compute_metrics(lowercase_							)

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

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



 def      UpperCAmelCase_							(				self      :Any						,       lowercase_      :Dict="./"							)->							int:
     A__				=			self.eval_dataset
     A__				=			self.get_eval_dataloader(lowercase_							)

     A__				=			next(iter(lowercase_							)							)

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

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

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

     A__				=			True

     A__				=			self.model.to(lowercase_							)

     model.eval()
     model.float()

     A__				=			model.module if hasattr(lowercase_						,       "module"							) else model
     quant_trainer.configure_model(lowercase_						,       self.quant_trainer_args							)

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

     A__				=			{0: "batch_size", 1: "seq_len"}

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










'''simple docstring'''

from .integrations import (
    is_optuna_available,
    is_ray_available,
    is_sigopt_available,
    is_wandb_available,
    run_hp_search_optuna,
    run_hp_search_ray,
    run_hp_search_sigopt,
    run_hp_search_wandb,
)
from .trainer_utils import (
    HPSearchBackend,
    default_hp_space_optuna,
    default_hp_space_ray,
    default_hp_space_sigopt,
    default_hp_space_wandb,
)
from .utils import logging


__lowerCAmelCase							:	Optional[Any]			=logging.get_logger(__name__)






class    UpperCAmelCase	:
 __lowercase   		=				42
 __lowercase   		=				None
 @staticmethod
 def      UpperCAmelCase_							(				)->							Dict:
     raise NotImplementedError
 def      UpperCAmelCase_							(				self      :List[Any]						,       lowercase_      :str						,       lowercase_      :int						,       lowercase_      :str						,       **lowercase_      :Dict							)->							str:
     raise NotImplementedError
 def      UpperCAmelCase_							(				self      :Optional[int]						,       lowercase_      :int							)->							Any:
     raise NotImplementedError
 def      UpperCAmelCase_							(				self      :List[Any]							)->							Optional[Any]:
     if not self.is_available():
         raise RuntimeError(
             F"You picked the {self.name} backend, but it is not installed. Run {self.pip_install()}."							)



 @classmethod
 def      UpperCAmelCase_							(				cls      :int							)->							Any:
     return F"`pip install {cls.pip_package or cls.name}`"






class    UpperCAmelCase	(       UpperCamelCase__				):
 __lowercase   		=				"""optuna"""
 @staticmethod
 def      UpperCAmelCase_							(				)->							int:
     return is_optuna_available()
 def      UpperCAmelCase_							(				self      :List[str]						,       lowercase_      :str						,       lowercase_      :int						,       lowercase_      :str						,       **lowercase_      :List[Any]							)->							Tuple:
     return run_hp_search_optuna(lowercase_						,       lowercase_						,       lowercase_						,       **lowercase_							)



 def      UpperCAmelCase_							(				self      :str						,       lowercase_      :Optional[int]							)->							Optional[Any]:
     return default_hp_space_optuna(lowercase_							)






class    UpperCAmelCase	(       UpperCamelCase__				):
 __lowercase   		=				"""ray"""
 __lowercase   		=				"""'ray[tune]'"""
 @staticmethod
 def      UpperCAmelCase_							(				)->							str:
     return is_ray_available()
 def      UpperCAmelCase_							(				self      :int						,       lowercase_      :Dict						,       lowercase_      :int						,       lowercase_      :str						,       **lowercase_      :List[str]							)->							int:
     return run_hp_search_ray(lowercase_						,       lowercase_						,       lowercase_						,       **lowercase_							)



 def      UpperCAmelCase_							(				self      :Optional[int]						,       lowercase_      :Dict							)->							int:
     return default_hp_space_ray(lowercase_							)






class    UpperCAmelCase	(       UpperCamelCase__				):
 __lowercase   		=				"""sigopt"""
 @staticmethod
 def      UpperCAmelCase_							(				)->							Union[str, Any]:
     return is_sigopt_available()
 def      UpperCAmelCase_							(				self      :Any						,       lowercase_      :Union[str, Any]						,       lowercase_      :int						,       lowercase_      :str						,       **lowercase_      :Dict							)->							Dict:
     return run_hp_search_sigopt(lowercase_						,       lowercase_						,       lowercase_						,       **lowercase_							)



 def      UpperCAmelCase_							(				self      :Optional[int]						,       lowercase_      :Optional[int]							)->							List[str]:
     return default_hp_space_sigopt(lowercase_							)






class    UpperCAmelCase	(       UpperCamelCase__				):
 __lowercase   		=				"""wandb"""
 @staticmethod
 def      UpperCAmelCase_							(				)->							List[str]:
     return is_wandb_available()
 def      UpperCAmelCase_							(				self      :Dict						,       lowercase_      :Optional[Any]						,       lowercase_      :int						,       lowercase_      :str						,       **lowercase_      :Dict							)->							List[str]:
     return run_hp_search_wandb(lowercase_						,       lowercase_						,       lowercase_						,       **lowercase_							)



 def      UpperCAmelCase_							(				self      :Union[str, Any]						,       lowercase_      :str							)->							Dict:
     return default_hp_space_wandb(lowercase_							)

__lowerCAmelCase							:	int			={
    HPSearchBackend(backend.name): backend for backend in [OptunaBackend, RayTuneBackend, SigOptBackend, WandbBackend]
}





def 						UpperCamelCase				(							):
    A__				=			[backend for backend in ALL_HYPERPARAMETER_SEARCH_BACKENDS.values() if backend.is_available()]
    if len(_lowerCamelCase       ) > 0:
        A__				=			available_backends[0].name
        if len(_lowerCamelCase       ) > 1:
            logger.info(
                F"{len(_lowerCamelCase       )} hyperparameter search backends available. Using {name} as the default."       )
        return name
    raise RuntimeError(
        "No hyperparameter search backend available.\n"
        + "\n".join(
            F" - To install {backend.name} run {backend.pip_install()}"
            for backend in ALL_HYPERPARAMETER_SEARCH_BACKENDS.values()       )       )



