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


_lowercase			:  Dict   							=	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 			lowerCAmelCase__       (  lowerCamelCase_					):






					def __init__(       self    ,							*__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE=None    ,							__SCREAMING_SNAKE_CASE=None    ,							__SCREAMING_SNAKE_CASE=None    ,							**__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											super().__init__(*__SCREAMING_SNAKE_CASE    ,							**__SCREAMING_SNAKE_CASE		)
											lowercase_     :					Optional[int]													=	eval_examples
											lowercase_     :					List[str]													=	post_process_function
											lowercase_     :					Optional[Any]													=	quant_trainer_args
											lowercase_     :					str													=	1_28  # default number of calibration samples






					def      _snake_case   (       self    ,							__SCREAMING_SNAKE_CASE=None		):


											"""simple docstring"""





											if calib_dataset is None and self.calib_dataset is None:
																	raise ValueError('''Trainer: calibration requires an calib_dataset.'''		)
											lowercase_     :					List[Any]													=	calib_dataset if calib_dataset is not None else self.calib_dataset

											lowercase_     :					Union[str, Any]													=	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      _snake_case   (       self    ,							__SCREAMING_SNAKE_CASE=None		):


											"""simple docstring"""





											lowercase_     :					Any													=	self.train_dataset if calib_dataset is None else calib_dataset
											lowercase_     :					Union[str, Any]													=	self.get_calib_dataloader(__SCREAMING_SNAKE_CASE		)

											lowercase_     :					str													=	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
																	lowercase_						,			lowercase_						,			lowercase_     :					List[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		)
											lowercase_     :					Optional[Any]													=	model






					def      _snake_case   (       self    ,							__SCREAMING_SNAKE_CASE=None    ,							__SCREAMING_SNAKE_CASE=None    ,							__SCREAMING_SNAKE_CASE=None    ,							__SCREAMING_SNAKE_CASE = "eval"		):


											"""simple docstring"""





											lowercase_     :					Any													=	self.eval_dataset if eval_dataset is None else eval_dataset
											lowercase_     :					Optional[Any]													=	self.get_eval_dataloader(__SCREAMING_SNAKE_CASE		)
											lowercase_     :					List[str]													=	self.eval_examples if eval_examples is None else eval_examples

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

											if self.post_process_function is not None and self.compute_metrics is not None:
																	lowercase_     :					Any													=	self.post_process_function(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE    ,							output.predictions		)
																	lowercase_     :					Optional[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}_'''		):
																													lowercase_     :					Optional[Any]													=	metrics.pop(__SCREAMING_SNAKE_CASE		)

																	self.log(__SCREAMING_SNAKE_CASE		)
											else:
																	lowercase_     :					Union[str, 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()		)

											lowercase_     :					str													=	self.callback_handler.on_evaluate(self.args    ,							self.state    ,							self.control    ,							__SCREAMING_SNAKE_CASE		)
											return metrics






					def      _snake_case   (       self    ,							__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE=None    ,							__SCREAMING_SNAKE_CASE = "test"		):


											"""simple docstring"""





											lowercase_     :					List[str]													=	self.get_test_dataloader(__SCREAMING_SNAKE_CASE		)

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

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

											lowercase_     :					int													=	self.post_process_function(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE    ,							output.predictions    ,							'''predict'''		)
											lowercase_     :					Optional[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}_'''		):
																							lowercase_     :					int													=	metrics.pop(__SCREAMING_SNAKE_CASE		)

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






					def      _snake_case   (       self    ,							__SCREAMING_SNAKE_CASE="./"		):


											"""simple docstring"""





											lowercase_     :					Optional[int]													=	self.eval_dataset
											lowercase_     :					Dict													=	self.get_eval_dataloader(__SCREAMING_SNAKE_CASE		)

											lowercase_     :					Union[str, Any]													=	next(iter(__SCREAMING_SNAKE_CASE		)		)

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

											# convert to tuple
											lowercase_     :					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

											lowercase_     :					Optional[int]													=	True

											lowercase_     :					Optional[int]													=	self.model.to(__SCREAMING_SNAKE_CASE		)

											model.eval()
											model.float()

											lowercase_     :					str													=	model.module if hasattr(__SCREAMING_SNAKE_CASE    ,							'''module'''		) else model
											quant_trainer.configure_model(__SCREAMING_SNAKE_CASE    ,							self.quant_trainer_args		)

											lowercase_     :					Dict													=	os.path.join(__SCREAMING_SNAKE_CASE    ,							'''model.onnx'''		)
											logger.info(F'''exporting model to {output_model_file}'''		)

											lowercase_     :					int													=	{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'''







import torch

from diffusers import DDPMScheduler

from .test_schedulers import SchedulerCommonTest






class 			lowerCAmelCase__       (  lowerCamelCase_					):
					lowerCAmelCase_					       =					(DDPMScheduler,)






					def      _snake_case   (       self    ,							**__SCREAMING_SNAKE_CASE		):


											"""simple docstring"""





											lowercase_     :					Optional[Any]													=	{
											    '''num_train_timesteps''': 10_00,
											    '''beta_start''': 0.0_001,
											    '''beta_end''': 0.02,
											    '''beta_schedule''': '''linear''',
											    '''variance_type''': '''fixed_small''',
											    '''clip_sample''': True,
											}

											config.update(**__SCREAMING_SNAKE_CASE		)
											return config






					def      _snake_case   (       self		):


											"""simple docstring"""





											for timesteps in [1, 5, 1_00, 10_00]:
																	self.check_over_configs(num_train_timesteps=__SCREAMING_SNAKE_CASE		)






					def      _snake_case   (       self		):


											"""simple docstring"""





											for beta_start, beta_end in zip([0.0_001, 0.001, 0.01, 0.1]    ,							[0.002, 0.02, 0.2, 2]		):
																	self.check_over_configs(beta_start=__SCREAMING_SNAKE_CASE    ,							beta_end=__SCREAMING_SNAKE_CASE		)






					def      _snake_case   (       self		):


											"""simple docstring"""





											for schedule in ["linear", "squaredcos_cap_v2"]:
																	self.check_over_configs(beta_schedule=__SCREAMING_SNAKE_CASE		)






					def      _snake_case   (       self		):


											"""simple docstring"""





											for variance in ["fixed_small", "fixed_large", "other"]:
																	self.check_over_configs(variance_type=__SCREAMING_SNAKE_CASE		)






					def      _snake_case   (       self		):


											"""simple docstring"""





											for clip_sample in [True, False]:
																	self.check_over_configs(clip_sample=__SCREAMING_SNAKE_CASE		)






					def      _snake_case   (       self		):


											"""simple docstring"""





											self.check_over_configs(thresholding=__SCREAMING_SNAKE_CASE		)
											for threshold in [0.5, 1.0, 2.0]:
																	for prediction_type in ["epsilon", "sample", "v_prediction"]:
																							self.check_over_configs(
																							    thresholding=__SCREAMING_SNAKE_CASE    ,							prediction_type=__SCREAMING_SNAKE_CASE    ,							sample_max_value=__SCREAMING_SNAKE_CASE    ,							)






					def      _snake_case   (       self		):


											"""simple docstring"""





											for prediction_type in ["epsilon", "sample", "v_prediction"]:
																	self.check_over_configs(prediction_type=__SCREAMING_SNAKE_CASE		)






					def      _snake_case   (       self		):


											"""simple docstring"""





											for t in [0, 5_00, 9_99]:
																	self.check_over_forward(time_step=__SCREAMING_SNAKE_CASE		)






					def      _snake_case   (       self		):


											"""simple docstring"""





											lowercase_     :					List[str]													=	self.scheduler_classes[0]
											lowercase_     :					int													=	self.get_scheduler_config()
											lowercase_     :					str													=	scheduler_class(**__SCREAMING_SNAKE_CASE		)

											assert torch.sum(torch.abs(scheduler._get_variance(0		) - 0.0		)		) < 1E-5
											assert torch.sum(torch.abs(scheduler._get_variance(4_87		) - 0.00_979		)		) < 1E-5
											assert torch.sum(torch.abs(scheduler._get_variance(9_99		) - 0.02		)		) < 1E-5






					def      _snake_case   (       self		):


											"""simple docstring"""





											lowercase_     :					Union[str, Any]													=	self.scheduler_classes[0]
											lowercase_     :					Any													=	self.get_scheduler_config()
											lowercase_     :					Optional[int]													=	scheduler_class(**__SCREAMING_SNAKE_CASE		)

											lowercase_     :					Any													=	len(__SCREAMING_SNAKE_CASE		)

											lowercase_     :					Optional[int]													=	self.dummy_model()
											lowercase_     :					Any													=	self.dummy_sample_deter
											lowercase_     :					List[str]													=	torch.manual_seed(0		)

											for t in reversed(range(__SCREAMING_SNAKE_CASE		)		):
																	# 1. predict noise residual
																	lowercase_     :					Optional[int]													=	model(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE		)

																	# 2. predict previous mean of sample x_t-1
																	lowercase_     :					Union[str, Any]													=	scheduler.step(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE    ,							generator=__SCREAMING_SNAKE_CASE		).prev_sample

																	# if t > 0:
																	#     noise = self.dummy_sample_deter
																	#     variance = scheduler.get_variance(t) ** (0.5) * noise
																	#
																	# sample = pred_prev_sample + variance
																	lowercase_     :					List[str]													=	pred_prev_sample

											lowercase_     :					str													=	torch.sum(torch.abs(__SCREAMING_SNAKE_CASE		)		)
											lowercase_     :					Tuple													=	torch.mean(torch.abs(__SCREAMING_SNAKE_CASE		)		)

											assert abs(result_sum.item() - 258.9_606		) < 1E-2
											assert abs(result_mean.item() - 0.3_372		) < 1E-3






					def      _snake_case   (       self		):


											"""simple docstring"""





											lowercase_     :					Optional[Any]													=	self.scheduler_classes[0]
											lowercase_     :					Tuple													=	self.get_scheduler_config(prediction_type='''v_prediction'''		)
											lowercase_     :					Any													=	scheduler_class(**__SCREAMING_SNAKE_CASE		)

											lowercase_     :					Optional[int]													=	len(__SCREAMING_SNAKE_CASE		)

											lowercase_     :					List[str]													=	self.dummy_model()
											lowercase_     :					int													=	self.dummy_sample_deter
											lowercase_     :					str													=	torch.manual_seed(0		)

											for t in reversed(range(__SCREAMING_SNAKE_CASE		)		):
																	# 1. predict noise residual
																	lowercase_     :					Optional[Any]													=	model(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE		)

																	# 2. predict previous mean of sample x_t-1
																	lowercase_     :					Optional[Any]													=	scheduler.step(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE    ,							generator=__SCREAMING_SNAKE_CASE		).prev_sample

																	# if t > 0:
																	#     noise = self.dummy_sample_deter
																	#     variance = scheduler.get_variance(t) ** (0.5) * noise
																	#
																	# sample = pred_prev_sample + variance
																	lowercase_     :					int													=	pred_prev_sample

											lowercase_     :					str													=	torch.sum(torch.abs(__SCREAMING_SNAKE_CASE		)		)
											lowercase_     :					Tuple													=	torch.mean(torch.abs(__SCREAMING_SNAKE_CASE		)		)

											assert abs(result_sum.item() - 202.0_296		) < 1E-2
											assert abs(result_mean.item() - 0.2_631		) < 1E-3






					def      _snake_case   (       self		):


											"""simple docstring"""





											lowercase_     :					Optional[int]													=	self.scheduler_classes[0]
											lowercase_     :					int													=	self.get_scheduler_config()
											lowercase_     :					int													=	scheduler_class(**__SCREAMING_SNAKE_CASE		)

											lowercase_     :					int													=	[1_00, 87, 50, 1, 0]

											scheduler.set_timesteps(timesteps=__SCREAMING_SNAKE_CASE		)

											lowercase_     :					List[str]													=	scheduler.timesteps

											for i, timestep in enumerate(__SCREAMING_SNAKE_CASE		):
																	if i == len(__SCREAMING_SNAKE_CASE		) - 1:
																							lowercase_     :					str													=	-1
																	else:
																							lowercase_     :					Any													=	timesteps[i + 1]

																	lowercase_     :					Optional[Any]													=	scheduler.previous_timestep(__SCREAMING_SNAKE_CASE		)
																	lowercase_     :					int													=	prev_t.item()

																	self.assertEqual(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE		)






					def      _snake_case   (       self		):


											"""simple docstring"""





											lowercase_     :					Tuple													=	self.scheduler_classes[0]
											lowercase_     :					List[Any]													=	self.get_scheduler_config()
											lowercase_     :					str													=	scheduler_class(**__SCREAMING_SNAKE_CASE		)

											lowercase_     :					Optional[int]													=	[1_00, 87, 50, 51, 0]

											with self.assertRaises(__SCREAMING_SNAKE_CASE    ,							msg='''`custom_timesteps` must be in descending order.'''		):
																	scheduler.set_timesteps(timesteps=__SCREAMING_SNAKE_CASE		)






					def      _snake_case   (       self		):


											"""simple docstring"""





											lowercase_     :					Optional[Any]													=	self.scheduler_classes[0]
											lowercase_     :					str													=	self.get_scheduler_config()
											lowercase_     :					Tuple													=	scheduler_class(**__SCREAMING_SNAKE_CASE		)

											lowercase_     :					Tuple													=	[1_00, 87, 50, 1, 0]
											lowercase_     :					List[str]													=	len(__SCREAMING_SNAKE_CASE		)

											with self.assertRaises(__SCREAMING_SNAKE_CASE    ,							msg='''Can only pass one of `num_inference_steps` or `custom_timesteps`.'''		):
																	scheduler.set_timesteps(num_inference_steps=__SCREAMING_SNAKE_CASE    ,							timesteps=__SCREAMING_SNAKE_CASE		)






					def      _snake_case   (       self		):


											"""simple docstring"""





											lowercase_     :					List[Any]													=	self.scheduler_classes[0]
											lowercase_     :					Optional[int]													=	self.get_scheduler_config()
											lowercase_     :					Any													=	scheduler_class(**__SCREAMING_SNAKE_CASE		)

											lowercase_     :					str													=	[scheduler.config.num_train_timesteps]

											with self.assertRaises(
											    __SCREAMING_SNAKE_CASE    ,							msg='''`timesteps` must start before `self.config.train_timesteps`: {scheduler.config.num_train_timesteps}}'''    ,							):
																	scheduler.set_timesteps(timesteps=__SCREAMING_SNAKE_CASE		)



