





'''simple docstring'''






# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion
# and https://github.com/hojonathanho/diffusion
import math
from dataclasses import dataclass
from typing import List, Optional, Tuple, Union

import numpy as np
import torch

from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.schedulers.scheduling_utils import SchedulerMixin
from diffusers.utils import BaseOutput, deprecate
@dataclass
# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->DDIM
class 		lowerCamelCase	(   _A							):
	snake_case_   = 42
	snake_case_   = None






def   __A    (				a_    :	str   ,a_    :	Tuple=0.9_9_9   ,a_    :	List[str]="cosine"   ,):
		if alpha_transform_type == "cosine":

				def alpha_bar_fn(a_    :	List[str]						):
						return math.cos((t + 0.0_0_8) / 1.0_0_8 * math.pi / 2						) ** 2

		elif alpha_transform_type == "exp":

				def alpha_bar_fn(a_    :	Any						):
						return math.exp(t * -1_2.0						)

		else:
				raise ValueError(f'''Unsupported alpha_tranform_type: {alpha_transform_type}'''						)

		lowerCAmelCase  :     int       						= []
		for i in range(a_						):
				lowerCAmelCase  :     int       						= i / num_diffusion_timesteps
				lowerCAmelCase  :     Tuple       						= (i + 1) / num_diffusion_timesteps
				betas.append(min(1 - alpha_bar_fn(a_						) / alpha_bar_fn(a_						)   ,a_						)						)
		return torch.tensor(a_   ,dtype=torch.floataa						)

class 		lowerCamelCase	(   _A   , _A							):
	snake_case_   = 1



	@register_to_config
	def __init__( self		,   a_ = 1_000		,   a_ = 0.0001		,   a_ = 0.02		,   a_ = "linear"		,   a_ = None		,   a_ = True		,   a_ = True		,   a_ = 0		,   a_ = "epsilon"		,   a_ = 1.0		,   **a_		,   ):
			if kwargs.get("set_alpha_to_one"		,   a_						) is not None:
					lowerCAmelCase  :     List[str]       						= (
					    "The `set_alpha_to_one` argument is deprecated. Please use `set_alpha_to_zero` instead."
					)
					deprecate("set_alpha_to_one"		,   "1.0.0"		,   a_		,   standard_warn=a_						)
					lowerCAmelCase  :     str       						= kwargs["set_alpha_to_one"]
			if trained_betas is not None:
					lowerCAmelCase  :     Optional[Any]       						= torch.tensor(a_		,   dtype=torch.floataa						)
			elif beta_schedule == "linear":
					lowerCAmelCase  :     List[Any]       						= torch.linspace(a_		,   a_		,   a_		,   dtype=torch.floataa						)
			elif beta_schedule == "scaled_linear":
					# this schedule is very specific to the latent diffusion model.
					lowerCAmelCase  :     Optional[int]       						= (
					    torch.linspace(beta_start**0.5		,   beta_end**0.5		,   a_		,   dtype=torch.floataa						) ** 2
					)
			elif beta_schedule == "squaredcos_cap_v2":
					# Glide cosine schedule
					lowerCAmelCase  :     Dict       						= betas_for_alpha_bar(a_						)
			else:
					raise NotImplementedError(F'''{beta_schedule} does is not implemented for {self.__class__}'''						)

			lowerCAmelCase  :     Union[str, Any]       						= 1.0 - self.betas
			lowerCAmelCase  :     str       						= torch.cumprod(self.alphas		,   dim=0						)

			# At every step in inverted ddim, we are looking into the next alphas_cumprod
			# For the final step, there is no next alphas_cumprod, and the index is out of bounds
			# `set_alpha_to_zero` decides whether we set this parameter simply to zero
			# in this case, self.step() just output the predicted noise
			# or whether we use the final alpha of the "non-previous" one.
			lowerCAmelCase  :     Union[str, Any]       						= torch.tensor(0.0						) if set_alpha_to_zero else self.alphas_cumprod[-1]

			# standard deviation of the initial noise distribution
			lowerCAmelCase  :     Any       						= 1.0

			# setable values
			lowerCAmelCase  :     Optional[Any]       						= None
			lowerCAmelCase  :     Optional[Any]       						= torch.from_numpy(np.arange(0		,   a_						).copy().astype(np.intaa						)						)



	def 	_lowerCamelCase     ( self		,   a_		,   a_ = None						):
			return sample



	def 	_lowerCamelCase     ( self		,   a_		,   a_ = None						):

			if num_inference_steps > self.config.num_train_timesteps:
					raise ValueError(
					    F'''`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:'''
					    F''' {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle'''
					    F''' maximal {self.config.num_train_timesteps} timesteps.'''						)

			lowerCAmelCase  :     str       						= num_inference_steps
			lowerCAmelCase  :     Union[str, Any]       						= self.config.num_train_timesteps // self.num_inference_steps
			# creates integer timesteps by multiplying by ratio
			# casting to int to avoid issues when num_inference_step is power of 3
			lowerCAmelCase  :     Dict       						= (np.arange(0		,   a_						) * step_ratio).round().copy().astype(np.intaa						)
			lowerCAmelCase  :     str       						= torch.from_numpy(a_						).to(a_						)
			self.timesteps += self.config.steps_offset



	def 	_lowerCamelCase     ( self		,   a_		,   a_		,   a_		,   a_ = 0.0		,   a_ = False		,   a_ = None		,   a_ = True		,   ):
			# 1. get previous step value (=t+1)
			lowerCAmelCase  :     List[str]       						= timestep + self.config.num_train_timesteps // self.num_inference_steps

			# 2. compute alphas, betas
			# change original implementation to exactly match noise levels for analogous forward process
			lowerCAmelCase  :     Dict       						= self.alphas_cumprod[timestep]
			lowerCAmelCase  :     Optional[Any]       						= (
			    self.alphas_cumprod[prev_timestep]
			    if prev_timestep < self.config.num_train_timesteps
			    else self.final_alpha_cumprod
			)

			lowerCAmelCase  :     Dict       						= 1 - alpha_prod_t

			# 3. compute predicted original sample from predicted noise also called
			# "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
			if self.config.prediction_type == "epsilon":
					lowerCAmelCase  :     List[str]       						= (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
					lowerCAmelCase  :     Optional[Any]       						= model_output
			elif self.config.prediction_type == "sample":
					lowerCAmelCase  :     Union[str, Any]       						= model_output
					lowerCAmelCase  :     Union[str, Any]       						= (sample - alpha_prod_t ** 0.5 * pred_original_sample) / beta_prod_t ** 0.5
			elif self.config.prediction_type == "v_prediction":
					lowerCAmelCase  :     Optional[int]       						= (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
					lowerCAmelCase  :     Any       						= (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample
			else:
					raise ValueError(
					    F'''prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or'''
					    " `v_prediction`"						)

			# 4. Clip or threshold "predicted x_0"
			if self.config.clip_sample:
					lowerCAmelCase  :     Tuple       						= pred_original_sample.clamp(
					    -self.config.clip_sample_range		,   self.config.clip_sample_range						)

			# 5. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
			lowerCAmelCase  :     Tuple       						= (1 - alpha_prod_t_prev) ** 0.5 * pred_epsilon

			# 6. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
			lowerCAmelCase  :     Optional[Any]       						= alpha_prod_t_prev ** 0.5 * pred_original_sample + pred_sample_direction

			if not return_dict:
					return (prev_sample, pred_original_sample)
			return DDIMSchedulerOutput(prev_sample=a_		,   pred_original_sample=a_						)





	def __len__( self						):
			return self.config.num_train_timesteps










'''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									=   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	(   _A							):



	def __init__( self		,   *a_		,   a_=None		,   a_=None		,   a_=None		,   **a_						):
			super().__init__(*a_		,   **a_						)
			lowerCAmelCase  :     Tuple       						= eval_examples
			lowerCAmelCase  :     List[Any]       						= post_process_function
			lowerCAmelCase  :     List[Any]       						= quant_trainer_args
			lowerCAmelCase  :     Dict       						= 128  # default number of calibration samples



	def 	_lowerCamelCase     ( self		,   a_=None						):
			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(a_		,   description="Calibration"						)

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



	def 	_lowerCamelCase     ( self		,   a_=None						):
			lowerCAmelCase  :     str       						= self.train_dataset if calib_dataset is None else calib_dataset
			lowerCAmelCase  :     str       						= self.get_calib_dataloader(a_						)

			lowerCAmelCase  :     Tuple       						= self.model
			quant_trainer.configure_model(a_		,   self.quant_trainer_args		,   calib=a_						)
			model.eval()
			quant_trainer.enable_calibration(a_						)

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

			quant_trainer.finish_calibration(a_		,   self.quant_trainer_args						)
			lowerCAmelCase  :     Union[str, Any]       						= model



	def 	_lowerCamelCase     ( self		,   a_=None		,   a_=None		,   a_=None		,   a_ = "eval"						):
			lowerCAmelCase  :     Optional[Any]       						= self.eval_dataset if eval_dataset is None else eval_dataset
			lowerCAmelCase  :     Optional[Any]       						= self.get_eval_dataloader(a_						)
			lowerCAmelCase  :     Union[str, 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  :     Dict       						= self.compute_metrics
			lowerCAmelCase  :     List[str]       						= None
			lowerCAmelCase  :     str       						= self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
			try:
					lowerCAmelCase  :     Union[str, Any]       						= eval_loop(
					    a_		,   description="Evaluation"		,   prediction_loss_only=True if compute_metrics is None else None		,   ignore_keys=a_		,   )
			finally:
					lowerCAmelCase  :     List[str]       						= compute_metrics

			if self.post_process_function is not None and self.compute_metrics is not None:
					lowerCAmelCase  :     Optional[Any]       						= self.post_process_function(a_		,   a_		,   output.predictions						)
					lowerCAmelCase  :     Optional[Any]       						= self.compute_metrics(a_						)

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

					self.log(a_						)
			else:
					lowerCAmelCase  :     str       						= {}

			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  :     Dict       						= self.callback_handler.on_evaluate(self.args		,   self.state		,   self.control		,   a_						)
			return metrics



	def 	_lowerCamelCase     ( self		,   a_		,   a_		,   a_=None		,   a_ = "test"						):
			lowerCAmelCase  :     int       						= self.get_test_dataloader(a_						)

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

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

			lowerCAmelCase  :     List[str]       						= self.post_process_function(a_		,   a_		,   output.predictions		,   "predict"						)
			lowerCAmelCase  :     Any       						= self.compute_metrics(a_						)

			# Prefix all keys with metric_key_prefix + '_'
			for key in list(metrics.keys()						):
					if not key.startswith(F'''{metric_key_prefix}_'''						):
							lowerCAmelCase  :     Optional[Any]       						= metrics.pop(a_						)

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





	def 	_lowerCamelCase     ( self		,   a_="./"						):
			lowerCAmelCase  :     Optional[Any]       						= self.eval_dataset
			lowerCAmelCase  :     Optional[int]       						= self.get_eval_dataloader(a_						)

			lowerCAmelCase  :     Union[str, Any]       						= next(iter(a_						)						)

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

			# convert to tuple
			lowerCAmelCase  :     str       						= tuple(v.to(a_						) for k, v in batch.items()						)

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

			lowerCAmelCase  :     Any       						= True

			lowerCAmelCase  :     Dict       						= self.model.to(a_						)

			model.eval()
			model.float()

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

			lowerCAmelCase  :     Any       						= os.path.join(a_		,   "model.onnx"						)
			logger.info(F'''exporting model to {output_model_file}'''						)

			lowerCAmelCase  :     List[Any]       						= {0: "batch_size", 1: "seq_len"}

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



