


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


_lowercase       		= 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		( __UpperCamelCase ):





		"""simple docstring"""







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




								'''simple docstring'''

								super().__init__(*lowerCAmelCase__						,**lowerCAmelCase__      )
								lowerCAmelCase_	:     Any          = eval_examples
								lowerCAmelCase_	:     Dict          = post_process_function
								lowerCAmelCase_	:     Any          = quant_trainer_args
								lowerCAmelCase_	:     str          = 1_28  # default number of calibration samples



		def  UpperCAmelCase_							(							self    :   Optional[Any]						,lowerCAmelCase__    :   Optional[int]=None      )      ->							Dict:




								'''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_	:     Union[str, Any]          = 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  UpperCAmelCase_							(							self    :   Union[str, Any]						,lowerCAmelCase__    :   List[Any]=None      )      ->							int:




								'''simple docstring'''

								lowerCAmelCase_	:     str          = self.train_dataset if calib_dataset is None else calib_dataset
								lowerCAmelCase_	:     List[Any]          = 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_    ,					lowerCAmelCase_    ,					lowerCAmelCase_	:     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_	:     str          = model



		def  UpperCAmelCase_							(							self    :   Optional[Any]						,lowerCAmelCase__    :   Union[str, Any]=None						,lowerCAmelCase__    :   Optional[Any]=None						,lowerCAmelCase__    :   Union[str, Any]=None						,lowerCAmelCase__    :   Dict = "eval"      )      ->							Any:




								'''simple docstring'''

								lowerCAmelCase_	:     Optional[Any]          = self.eval_dataset if eval_dataset is None else eval_dataset
								lowerCAmelCase_	:     Optional[int]          = self.get_eval_dataloader(lowerCAmelCase__      )
								lowerCAmelCase_	:     Any          = self.eval_examples if eval_examples is None else eval_examples

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

								if self.post_process_function is not None and self.compute_metrics is not None:
														lowerCAmelCase_	:     str          = self.post_process_function(lowerCAmelCase__						,lowerCAmelCase__						,output.predictions      )
														lowerCAmelCase_	:     Optional[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_	:     str          = 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_	:     int          = self.callback_handler.on_evaluate(self.args						,self.state						,self.control						,lowerCAmelCase__      )
								return metrics



		def  UpperCAmelCase_							(							self    :   List[str]						,lowerCAmelCase__    :   int						,lowerCAmelCase__    :   Optional[Any]						,lowerCAmelCase__    :   Any=None						,lowerCAmelCase__    :   Tuple = "test"      )      ->							str:




								'''simple docstring'''

								lowerCAmelCase_	:     Optional[Any]          = self.get_test_dataloader(lowerCAmelCase__      )

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

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

								lowerCAmelCase_	:     Any          = self.post_process_function(lowerCAmelCase__						,lowerCAmelCase__						,output.predictions						,"predict"      )
								lowerCAmelCase_	:     Optional[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_	:     Tuple          = metrics.pop(lowerCAmelCase__      )

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



		def  UpperCAmelCase_							(							self    :   Any						,lowerCAmelCase__    :   List[Any]="./"      )      ->							Tuple:




								'''simple docstring'''

								lowerCAmelCase_	:     List[str]          = self.eval_dataset
								lowerCAmelCase_	:     int          = self.get_eval_dataloader(lowerCAmelCase__      )

								lowerCAmelCase_	:     List[Any]          = next(iter(lowerCAmelCase__      )      )

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

								# convert to tuple
								lowerCAmelCase_	:     Optional[int]          = 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_	:     Any          = True

								lowerCAmelCase_	:     List[str]          = self.model.to(lowerCAmelCase__      )

								model.eval()
								model.float()

								lowerCAmelCase_	:     Union[str, Any]          = model.module if hasattr(lowerCAmelCase__						,"module"      ) else model
								quant_trainer.configure_model(lowerCAmelCase__						,self.quant_trainer_args      )

								lowerCAmelCase_	:     Tuple          = os.path.join(lowerCAmelCase__						,"model.onnx"      )
								logger.info(f'''exporting model to {output_model_file}'''      )

								lowerCAmelCase_	:     Dict          = {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"      )




import math





def 						SCREAMING_SNAKE_CASE__     (   __lowerCAmelCase     ,		__lowerCAmelCase				):

    if initial_intensity < 0:
        raise ValueError("The value of intensity cannot be negative"				)
        # handling of negative values of initial intensity
    if angle < 0 or angle > 360:
        raise ValueError("In Malus Law, the angle is in the range 0-360 degrees"				)
        # handling of values out of allowed range
    return initial_intensity * (math.cos(math.radians(__lowerCAmelCase				)				) ** 2)


if __name__ == "__main__":
     import doctest

     doctest.testmod(name='''malus_law''')



