

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							:  Tuple            =						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   __lowercase    (			UpperCAmelCase_     ):




			"""simple docstring"""







			def __init__(			self  :					Tuple       ,						*lowerCAmelCase__  :					Dict       ,						lowerCAmelCase__  :					Optional[int]=None       ,						lowerCAmelCase__  :					Optional[int]=None       ,						lowerCAmelCase__  :					Union[str, Any]=None       ,						**lowerCAmelCase__  :					Tuple):
									super().__init__(*lowerCAmelCase__       ,						**lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			Dict  		=	eval_examples
									SCREAMING_SNAKE_CASE_:			str  		=	post_process_function
									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	quant_trainer_args
									SCREAMING_SNAKE_CASE_:			str  		=	128  # default number of calibration samples



			def 	_SCREAMING_SNAKE_CASE      (			self  :					List[Any]       ,						lowerCAmelCase__  :					str=None):
									if calib_dataset is None and self.calib_dataset is None:
															raise ValueError("Trainer: calibration requires an calib_dataset.")
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	calib_dataset if calib_dataset is not None else self.calib_dataset

									SCREAMING_SNAKE_CASE_:			int  		=	self._remove_unused_columns(lowerCAmelCase__       ,						description="Calibration")

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



			def 	_SCREAMING_SNAKE_CASE      (			self  :					str       ,						lowerCAmelCase__  :					Optional[int]=None):
									SCREAMING_SNAKE_CASE_:			Tuple  		=	self.train_dataset if calib_dataset is None else calib_dataset
									SCREAMING_SNAKE_CASE_:			int  		=	self.get_calib_dataloader(lowerCAmelCase__)

									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	self.model
									quant_trainer.configure_model(lowerCAmelCase__       ,						self.quant_trainer_args       ,						calib=lowerCAmelCase__)
									model.eval()
									quant_trainer.enable_calibration(lowerCAmelCase__)

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

									quant_trainer.finish_calibration(lowerCAmelCase__       ,						self.quant_trainer_args)
									SCREAMING_SNAKE_CASE_:			Dict  		=	model



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

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

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

															self.log(lowerCAmelCase__)
									else:
															SCREAMING_SNAKE_CASE_:			List[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())

									SCREAMING_SNAKE_CASE_:			Optional[int]  		=	self.callback_handler.on_evaluate(self.args       ,						self.state       ,						self.control       ,						lowerCAmelCase__)
									return metrics



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Union[str, Any]       ,						lowerCAmelCase__  :					Dict       ,						lowerCAmelCase__  :					Tuple       ,						lowerCAmelCase__  :					Union[str, Any]=None       ,						lowerCAmelCase__  :					str = "test"):
									SCREAMING_SNAKE_CASE_:			str  		=	self.get_test_dataloader(lowerCAmelCase__)

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

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

									SCREAMING_SNAKE_CASE_:			Any  		=	self.post_process_function(lowerCAmelCase__       ,						lowerCAmelCase__       ,						output.predictions       ,						"predict")
									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	self.compute_metrics(lowerCAmelCase__)

									# Prefix all keys with metric_key_prefix + '_'
									for key in list(metrics.keys()):
															if not key.startswith(F"{metric_key_prefix}_"):
																					SCREAMING_SNAKE_CASE_:			Tuple  		=	metrics.pop(lowerCAmelCase__)

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




			def 	_SCREAMING_SNAKE_CASE      (			self  :					Optional[int]       ,						lowerCAmelCase__  :					Union[str, Any]="./"):
									SCREAMING_SNAKE_CASE_:			Tuple  		=	self.eval_dataset
									SCREAMING_SNAKE_CASE_:			List[str]  		=	self.get_eval_dataloader(lowerCAmelCase__)

									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	next(iter(lowerCAmelCase__))

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

									# convert to tuple
									SCREAMING_SNAKE_CASE_:			int  		=	tuple(v.to(lowerCAmelCase__) for k, v in batch.items())

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

									SCREAMING_SNAKE_CASE_:			Optional[int]  		=	True

									SCREAMING_SNAKE_CASE_:			str  		=	self.model.to(lowerCAmelCase__)

									model.eval()
									model.float()

									SCREAMING_SNAKE_CASE_:			Tuple  		=	model.module if hasattr(lowerCAmelCase__       ,						"module") else model
									quant_trainer.configure_model(lowerCAmelCase__       ,						self.quant_trainer_args)

									SCREAMING_SNAKE_CASE_:			Tuple  		=	os.path.join(lowerCAmelCase__       ,						"model.onnx")
									logger.info(F"exporting model to {output_model_file}")

									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	{0: "batch_size", 1: "seq_len"}

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









from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowerCAmelCase							:  Optional[int]            =						{
    """configuration_nllb_moe""": [
        """NLLB_MOE_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """NllbMoeConfig""",
    ]
}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						lowerCAmelCase							:  Optional[int]            =						[
						    """NLLB_MOE_PRETRAINED_MODEL_ARCHIVE_LIST""",
						    """NllbMoeForConditionalGeneration""",
						    """NllbMoeModel""",
						    """NllbMoePreTrainedModel""",
						    """NllbMoeTop2Router""",
						    """NllbMoeSparseMLP""",
						]


if TYPE_CHECKING:
						from .configuration_nllb_moe import (
						    NLLB_MOE_PRETRAINED_CONFIG_ARCHIVE_MAP,
						    NllbMoeConfig,
						)

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_nllb_moe import (
												    NLLB_MOE_PRETRAINED_MODEL_ARCHIVE_LIST,
												    NllbMoeForConditionalGeneration,
												    NllbMoeModel,
												    NllbMoePreTrainedModel,
												    NllbMoeSparseMLP,
												    NllbMoeTopaRouter,
												)


else:
						import sys

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






