'''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_	:		Any			=     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 		SCREAMING_SNAKE_CASE     (					SCREAMING_SNAKE_CASE	):







				'''simple docstring'''





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


										'''simple docstring'''




										super().__init__(*lowercase__       ,		**lowercase__					)
										_UpperCamelCase			:			Union[str, Any]   			=						eval_examples
										_UpperCamelCase			:			List[Any]   			=						post_process_function
										_UpperCamelCase			:			Dict   			=						quant_trainer_args
										_UpperCamelCase			:			Dict   			=						128  # default number of calibration samples





				def  snake_case__				(						self					:						Optional[Any]       ,		lowercase__					:						str=None					)							->Dict:


										'''simple docstring'''




										if calib_dataset is None and self.calib_dataset is None:
																raise ValueError("Trainer: calibration requires an calib_dataset."					)
										_UpperCamelCase			:			str   			=						calib_dataset if calib_dataset is not None else self.calib_dataset

										_UpperCamelCase			:			int   			=						self._remove_unused_columns(lowercase__       ,		description="Calibration"					)

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





				def  snake_case__				(						self					:						Dict       ,		lowercase__					:						str=None					)							->Optional[Any]:


										'''simple docstring'''




										_UpperCamelCase			:			Any   			=						self.train_dataset if calib_dataset is None else calib_dataset
										_UpperCamelCase			:			str   			=						self.get_calib_dataloader(lowercase__					)

										_UpperCamelCase			:			Optional[int]   			=						self.model
										quant_trainer.configure_model(lowercase__       ,		self.quant_trainer_args       ,		calib=lowercase__					)
										model.eval()
										quant_trainer.enable_calibration(lowercase__					)

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

										quant_trainer.finish_calibration(lowercase__       ,		self.quant_trainer_args					)
										_UpperCamelCase			:			List[Any]   			=						model





				def  snake_case__				(						self					:						Union[str, Any]       ,		lowercase__					:						Tuple=None       ,		lowercase__					:						Dict=None       ,		lowercase__					:						str=None       ,		lowercase__					:						str = "eval"					)							->Optional[int]:


										'''simple docstring'''




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

										if self.post_process_function is not None and self.compute_metrics is not None:
																_UpperCamelCase			:			str   			=						self.post_process_function(lowercase__       ,		lowercase__       ,		output.predictions					)
																_UpperCamelCase			:			List[str]   			=						self.compute_metrics(lowercase__					)

																# Prefix all keys with metric_key_prefix + '_'
																for key in list(metrics.keys()					):
																						if not key.startswith(f'''{metric_key_prefix}_'''					):
																												_UpperCamelCase			:			Dict   			=						metrics.pop(lowercase__					)

																self.log(lowercase__					)
										else:
																_UpperCamelCase			:			int   			=						{}

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

										_UpperCamelCase			:			List[Any]   			=						self.callback_handler.on_evaluate(self.args       ,		self.state       ,		self.control       ,		lowercase__					)
										return metrics





				def  snake_case__				(						self					:						int       ,		lowercase__					:						str       ,		lowercase__					:						Tuple       ,		lowercase__					:						Dict=None       ,		lowercase__					:						str = "test"					)							->Union[str, Any]:


										'''simple docstring'''




										_UpperCamelCase			:			List[str]   			=						self.get_test_dataloader(lowercase__					)

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

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

										_UpperCamelCase			:			Dict   			=						self.post_process_function(lowercase__       ,		lowercase__       ,		output.predictions       ,		"predict"					)
										_UpperCamelCase			:			Optional[int]   			=						self.compute_metrics(lowercase__					)

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

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







				def  snake_case__				(						self					:						str       ,		lowercase__					:						Optional[Any]="./"					)							->Tuple:


										'''simple docstring'''




										_UpperCamelCase			:			Union[str, Any]   			=						self.eval_dataset
										_UpperCamelCase			:			Dict   			=						self.get_eval_dataloader(lowercase__					)

										_UpperCamelCase			:			List[Any]   			=						next(iter(lowercase__					)					)

										# saving device - to make it consistent
										_UpperCamelCase			:			Optional[int]   			=						torch.device("cuda" if torch.cuda.is_available() else "cpu"					)

										# convert to tuple
										_UpperCamelCase			:			List[Any]   			=						tuple(v.to(lowercase__					) for k, v in batch.items()					)

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

										_UpperCamelCase			:			Any   			=						True

										_UpperCamelCase			:			int   			=						self.model.to(lowercase__					)

										model.eval()
										model.float()

										_UpperCamelCase			:			Optional[int]   			=						model.module if hasattr(lowercase__       ,		"module"					) else model
										quant_trainer.configure_model(lowercase__       ,		self.quant_trainer_args					)

										_UpperCamelCase			:			Any   			=						os.path.join(lowercase__       ,		"model.onnx"					)
										logger.info(f'''exporting model to {output_model_file}'''					)

										_UpperCamelCase			:			str   			=						{0: "batch_size", 1: "seq_len"}

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



'''simple docstring'''





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


lowerCAmelCase_	:		Union[str, Any]			=     {"""configuration_opt""": ["""OPT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """OPTConfig"""]}

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCAmelCase_	:		Tuple			=     [
			    """OPT_PRETRAINED_MODEL_ARCHIVE_LIST""",
			    """OPTForCausalLM""",
			    """OPTModel""",
			    """OPTPreTrainedModel""",
			    """OPTForSequenceClassification""",
			    """OPTForQuestionAnswering""",
			]

try:
			if not is_tf_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCAmelCase_	:		str			=     ["""TFOPTForCausalLM""", """TFOPTModel""", """TFOPTPreTrainedModel"""]

try:
			if not is_flax_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCAmelCase_	:		List[Any]			=     [
			    """FlaxOPTForCausalLM""",
			    """FlaxOPTModel""",
			    """FlaxOPTPreTrainedModel""",
			]


if TYPE_CHECKING:
			from .configuration_opt import OPT_PRETRAINED_CONFIG_ARCHIVE_MAP, OPTConfig

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_opt import (
						    OPT_PRETRAINED_MODEL_ARCHIVE_LIST,
						    OPTForCausalLM,
						    OPTForQuestionAnswering,
						    OPTForSequenceClassification,
						    OPTModel,
						    OPTPreTrainedModel,
						)

			try:
						if not is_tf_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_tf_opt import TFOPTForCausalLM, TFOPTModel, TFOPTPreTrainedModel

			try:
						if not is_flax_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_flax_opt import FlaxOPTForCausalLM, FlaxOPTModel, FlaxOPTPreTrainedModel

else:
			import sys

			lowerCAmelCase_	:		Any			=     _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


