



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


__snake_case				:  Optional[int] =       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__  (    __SCREAMING_SNAKE_CASE				):




	'''simple docstring'''





	def __init__(	self:						List[Any]			,     *_SCREAMING_SNAKE_CASE:						Dict			,     _SCREAMING_SNAKE_CASE:						Dict=None			,     _SCREAMING_SNAKE_CASE:						Tuple=None			,     _SCREAMING_SNAKE_CASE:						List[Any]=None			,     **_SCREAMING_SNAKE_CASE:						Union[str, Any])							->			Any:

				"""simple docstring"""

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

	def 				_SCREAMING_SNAKE_CASE  (	self:						Any			,     _SCREAMING_SNAKE_CASE:						Tuple=None)							->			Optional[Any]:

				"""simple docstring"""

				if calib_dataset is None and self.calib_dataset is None:
							raise ValueError("Trainer: calibration requires an calib_dataset.")
				__lowerCAmelCase :    Union[str, Any]				=			calib_dataset if calib_dataset is not None else self.calib_dataset

				__lowerCAmelCase :    List[str]				=			self._remove_unused_columns(_SCREAMING_SNAKE_CASE			,     description="Calibration")

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

	def 				_SCREAMING_SNAKE_CASE  (	self:						Optional[int]			,     _SCREAMING_SNAKE_CASE:						int=None)							->			List[str]:

				"""simple docstring"""

				__lowerCAmelCase :    Union[str, Any]				=			self.train_dataset if calib_dataset is None else calib_dataset
				__lowerCAmelCase :    List[str]				=			self.get_calib_dataloader(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Dict				=			self.model
				quant_trainer.configure_model(_SCREAMING_SNAKE_CASE			,     self.quant_trainer_args			,     calib=_SCREAMING_SNAKE_CASE)
				model.eval()
				quant_trainer.enable_calibration(_SCREAMING_SNAKE_CASE)

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

				quant_trainer.finish_calibration(_SCREAMING_SNAKE_CASE			,     self.quant_trainer_args)
				__lowerCAmelCase :    Dict				=			model

	def 				_SCREAMING_SNAKE_CASE  (	self:						int			,     _SCREAMING_SNAKE_CASE:						str=None			,     _SCREAMING_SNAKE_CASE:						int=None			,     _SCREAMING_SNAKE_CASE:						str=None			,     _SCREAMING_SNAKE_CASE:						str = "eval")							->			Dict:

				"""simple docstring"""

				__lowerCAmelCase :    List[Any]				=			self.eval_dataset if eval_dataset is None else eval_dataset
				__lowerCAmelCase :    Dict				=			self.get_eval_dataloader(_SCREAMING_SNAKE_CASE)
				__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[int]				=			self.compute_metrics
				__lowerCAmelCase :    Tuple				=			None
				__lowerCAmelCase :    str				=			self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
				try:
							__lowerCAmelCase :    Dict				=			eval_loop(
							    _SCREAMING_SNAKE_CASE			,     description="Evaluation"			,     prediction_loss_only=True if compute_metrics is None else None			,     ignore_keys=_SCREAMING_SNAKE_CASE			,     )
				finally:
							__lowerCAmelCase :    List[Any]				=			compute_metrics

				if self.post_process_function is not None and self.compute_metrics is not None:
							__lowerCAmelCase :    Any				=			self.post_process_function(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE			,     output.predictions)
							__lowerCAmelCase :    Any				=			self.compute_metrics(_SCREAMING_SNAKE_CASE)

							# Prefix all keys with metric_key_prefix + '_'
							for key in list(metrics.keys()):
										if not key.startswith(F"""{metric_key_prefix}_"""):
													__lowerCAmelCase :    Union[str, Any]				=			metrics.pop(_SCREAMING_SNAKE_CASE)

							self.log(_SCREAMING_SNAKE_CASE)
				else:
							__lowerCAmelCase :    Tuple				=			{}

				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 :    List[Any]				=			self.callback_handler.on_evaluate(self.args			,     self.state			,     self.control			,     _SCREAMING_SNAKE_CASE)
				return metrics

	def 				_SCREAMING_SNAKE_CASE  (	self:						List[str]			,     _SCREAMING_SNAKE_CASE:						Dict			,     _SCREAMING_SNAKE_CASE:						Optional[int]			,     _SCREAMING_SNAKE_CASE:						Tuple=None			,     _SCREAMING_SNAKE_CASE:						str = "test")							->			Optional[int]:

				"""simple docstring"""

				__lowerCAmelCase :    int				=			self.get_test_dataloader(_SCREAMING_SNAKE_CASE)

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

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

				__lowerCAmelCase :    Optional[Any]				=			self.post_process_function(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE			,     output.predictions			,     "predict")
				__lowerCAmelCase :    Optional[int]				=			self.compute_metrics(_SCREAMING_SNAKE_CASE)

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

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







	def 				_SCREAMING_SNAKE_CASE  (	self:						str			,     _SCREAMING_SNAKE_CASE:						Dict="./")							->			List[Any]:

				"""simple docstring"""

				__lowerCAmelCase :    Dict				=			self.eval_dataset
				__lowerCAmelCase :    Optional[int]				=			self.get_eval_dataloader(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Optional[Any]				=			next(iter(_SCREAMING_SNAKE_CASE))

				# saving device - to make it consistent
				__lowerCAmelCase :    Union[str, Any]				=			torch.device("cuda" if torch.cuda.is_available() else "cpu")

				# convert to tuple
				__lowerCAmelCase :    Optional[int]				=			tuple(v.to(_SCREAMING_SNAKE_CASE) for k, v in batch.items())

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

				__lowerCAmelCase :    Dict				=			True

				__lowerCAmelCase :    Optional[int]				=			self.model.to(_SCREAMING_SNAKE_CASE)

				model.eval()
				model.float()

				__lowerCAmelCase :    Dict				=			model.module if hasattr(_SCREAMING_SNAKE_CASE			,     "module") else model
				quant_trainer.configure_model(_SCREAMING_SNAKE_CASE			,     self.quant_trainer_args)

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

				__lowerCAmelCase :    Dict				=			{0: "batch_size", 1: "seq_len"}

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




"""simple docstring"""






# NOTE: This file is deprecated and will be removed in a future version.
# It only exists so that temporarely `from diffusers.pipelines import DiffusionPipeline` works

from ...utils import deprecate
from ..controlnet.pipeline_flax_controlnet import FlaxStableDiffusionControlNetPipeline  # noqa: F401


deprecate(
    'stable diffusion controlnet',
    '0.22.0',
    'Importing `FlaxStableDiffusionControlNetPipeline` from diffusers.pipelines.stable_diffusion.flax_pipeline_stable_diffusion_controlnet is deprecated. Please import `from diffusers import FlaxStableDiffusionControlNetPipeline` instead.',
    standard_warn=False,
    stacklevel=3,
)