


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


a__       :							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):

	"""simple docstring"""





	def __init__(				self    :       Tuple    ,    *UpperCAmelCase__    :       Optional[int]    ,    UpperCAmelCase__    :       str=None    ,    UpperCAmelCase__    :       Dict=None    ,    UpperCAmelCase__    :       Dict=None    ,    **UpperCAmelCase__    :       Tuple     ) ->			int:
						super().__init__(*UpperCAmelCase__    ,    **UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       eval_examples
						__SCREAMING_SNAKE_CASE		     =       post_process_function
						__SCREAMING_SNAKE_CASE		     =       quant_trainer_args
						__SCREAMING_SNAKE_CASE		     =       1_2_8  # default number of calibration samples




	def 						UpperCAmelCase_   (				self    :       Dict    ,    UpperCAmelCase__    :       Tuple=None     ) ->			Optional[int]:
						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(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 						UpperCAmelCase_   (				self    :       Dict    ,    UpperCAmelCase__    :       Tuple=None     ) ->			Optional[int]:
						__SCREAMING_SNAKE_CASE		     =       self.train_dataset if calib_dataset is None else calib_dataset
						__SCREAMING_SNAKE_CASE		     =       self.get_calib_dataloader(UpperCAmelCase__     )

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




	def 						UpperCAmelCase_   (				self    :       Optional[int]    ,    UpperCAmelCase__    :       Tuple=None    ,    UpperCAmelCase__    :       Any=None    ,    UpperCAmelCase__    :       List[Any]=None    ,    UpperCAmelCase__    :       str = "eval"     ) ->			Optional[int]:
						__SCREAMING_SNAKE_CASE		     =       self.eval_dataset if eval_dataset is None else eval_dataset
						__SCREAMING_SNAKE_CASE		     =       self.get_eval_dataloader(UpperCAmelCase__     )
						__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(
											    UpperCAmelCase__    ,    description="Evaluation"    ,    prediction_loss_only=True if compute_metrics is None else None    ,    ignore_keys=UpperCAmelCase__    ,    )
						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(UpperCAmelCase__    ,    UpperCAmelCase__    ,    output.predictions     )
											__SCREAMING_SNAKE_CASE		     =       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}_"""     ):
																					__SCREAMING_SNAKE_CASE		     =       metrics.pop(UpperCAmelCase__     )

											self.log(UpperCAmelCase__     )
						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    ,    UpperCAmelCase__     )
						return metrics




	def 						UpperCAmelCase_   (				self    :       Optional[Any]    ,    UpperCAmelCase__    :       Optional[Any]    ,    UpperCAmelCase__    :       Dict    ,    UpperCAmelCase__    :       Optional[int]=None    ,    UpperCAmelCase__    :       str = "test"     ) ->			Tuple:
						__SCREAMING_SNAKE_CASE		     =       self.get_test_dataloader(UpperCAmelCase__     )

						# 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(
											    UpperCAmelCase__    ,    description="Prediction"    ,    prediction_loss_only=True if compute_metrics is None else None    ,    ignore_keys=UpperCAmelCase__    ,    )
						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(UpperCAmelCase__    ,    UpperCAmelCase__    ,    output.predictions    ,    "predict"     )
						__SCREAMING_SNAKE_CASE		     =       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}_"""     ):
																__SCREAMING_SNAKE_CASE		     =       metrics.pop(UpperCAmelCase__     )

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




	def 						UpperCAmelCase_   (				self    :       int    ,    UpperCAmelCase__    :       int="./"     ) ->			Optional[int]:
						__SCREAMING_SNAKE_CASE		     =       self.eval_dataset
						__SCREAMING_SNAKE_CASE		     =       self.get_eval_dataloader(UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       next(iter(UpperCAmelCase__     )     )

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

						model.eval()
						model.float()

						__SCREAMING_SNAKE_CASE		     =       model.module if hasattr(UpperCAmelCase__    ,    "module"     ) else model
						quant_trainer.configure_model(UpperCAmelCase__    ,    self.quant_trainer_args     )

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

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

						torch.onnx.export(
						    UpperCAmelCase__    ,    UpperCAmelCase__    ,    UpperCAmelCase__    ,    export_params=UpperCAmelCase__    ,    opset_version=1_3    ,    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"     )




"""simple docstring"""







import warnings

from ...utils import logging
from .image_processing_videomae import VideoMAEImageProcessor


a__       :							Dict        =  logging.get_logger(__name__)




class 				UpperCamelCase_			(				UpperCamelCase):

	"""simple docstring"""





	def __init__(				self    :       str    ,    *UpperCAmelCase__    :       List[Any]    ,    **UpperCAmelCase__    :       Tuple     ) ->			None:
						warnings.warn(
						    "The class VideoMAEFeatureExtractor is deprecated and will be removed in version 5 of Transformers."
						    " Please use VideoMAEImageProcessor instead."    ,    UpperCAmelCase__    ,    )
						super().__init__(*UpperCAmelCase__    ,    **UpperCAmelCase__     )
