


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        _a							(			UpperCamelCase__			):

							def __init__(							self:	str    ,						*UpperCamelCase_:	Optional[int]    ,						UpperCamelCase_:	Optional[Any]=None    ,						UpperCamelCase_:	Optional[int]=None    ,						UpperCamelCase_:	Tuple=None    ,						**UpperCamelCase_:	Any							)			->   Union[str, Any]:





									"""simple docstring"""







									super().__init__(*UpperCamelCase_    ,						**UpperCamelCase_							)
									lowercase__            =      eval_examples
									lowercase__            =      post_process_function
									lowercase__            =      quant_trainer_args
									lowercase__            =      128  # default number of calibration samples

							def       lowerCamelCase_ (							self:	Union[str, Any]    ,						UpperCamelCase_:	Optional[Any]=None							)			->   Any:





									"""simple docstring"""







									if calib_dataset is None and self.calib_dataset is None:
											raise ValueError('''Trainer: calibration requires an calib_dataset.'''							)
									lowercase__            =      calib_dataset if calib_dataset is not None else self.calib_dataset

									lowercase__            =      self._remove_unused_columns(UpperCamelCase_    ,						description='''Calibration'''							)

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

							def       lowerCamelCase_ (							self:	Union[str, Any]    ,						UpperCamelCase_:	str=None							)			->   Dict:





									"""simple docstring"""







									lowercase__            =      self.train_dataset if calib_dataset is None else calib_dataset
									lowercase__            =      self.get_calib_dataloader(UpperCamelCase_							)

									lowercase__            =      self.model
									quant_trainer.configure_model(UpperCamelCase_    ,						self.quant_trainer_args    ,						calib=UpperCamelCase_							)
									model.eval()
									quant_trainer.enable_calibration(UpperCamelCase_							)

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

									quant_trainer.finish_calibration(UpperCamelCase_    ,						self.quant_trainer_args							)
									lowercase__            =      model

							def       lowerCamelCase_ (							self:	Union[str, Any]    ,						UpperCamelCase_:	Union[str, Any]=None    ,						UpperCamelCase_:	List[Any]=None    ,						UpperCamelCase_:	List[str]=None    ,						UpperCamelCase_:	str = "eval"							)			->   Tuple:





									"""simple docstring"""







									lowercase__            =      self.eval_dataset if eval_dataset is None else eval_dataset
									lowercase__            =      self.get_eval_dataloader(UpperCamelCase_							)
									lowercase__            =      self.eval_examples if eval_examples is None else eval_examples

									# Temporarily disable metric computation, we will do it in the loop here.
									lowercase__            =      self.compute_metrics
									lowercase__            =      None
									lowercase__            =      self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
									try:
											lowercase__            =      eval_loop(
											    UpperCamelCase_    ,						description='''Evaluation'''    ,						prediction_loss_only=True if compute_metrics is None else None    ,						ignore_keys=UpperCamelCase_    ,						)
									finally:
											lowercase__            =      compute_metrics

									if self.post_process_function is not None and self.compute_metrics is not None:
											lowercase__            =      self.post_process_function(UpperCamelCase_    ,						UpperCamelCase_    ,						output.predictions							)
											lowercase__            =      self.compute_metrics(UpperCamelCase_							)

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

											self.log(UpperCamelCase_							)
									else:
											lowercase__            =      {}

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

							def       lowerCamelCase_ (							self:	int    ,						UpperCamelCase_:	Optional[Any]    ,						UpperCamelCase_:	Tuple    ,						UpperCamelCase_:	List[Any]=None    ,						UpperCamelCase_:	str = "test"							)			->   Optional[int]:





									"""simple docstring"""







									lowercase__            =      self.get_test_dataloader(UpperCamelCase_							)

									# Temporarily disable metric computation, we will do it in the loop here.
									lowercase__            =      self.compute_metrics
									lowercase__            =      None
									lowercase__            =      self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
									try:
											lowercase__            =      eval_loop(
											    UpperCamelCase_    ,						description='''Prediction'''    ,						prediction_loss_only=True if compute_metrics is None else None    ,						ignore_keys=UpperCamelCase_    ,						)
									finally:
											lowercase__            =      compute_metrics

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

									lowercase__            =      self.post_process_function(UpperCamelCase_    ,						UpperCamelCase_    ,						output.predictions    ,						'''predict'''							)
									lowercase__            =      self.compute_metrics(UpperCamelCase_							)

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

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


							def       lowerCamelCase_ (							self:	int    ,						UpperCamelCase_:	Tuple="./"							)			->   str:





									"""simple docstring"""







									lowercase__            =      self.eval_dataset
									lowercase__            =      self.get_eval_dataloader(UpperCamelCase_							)

									lowercase__            =      next(iter(UpperCamelCase_							)							)

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

									# convert to tuple
									lowercase__            =      tuple(v.to(UpperCamelCase_							) for k, v in batch.items()							)

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

									lowercase__            =      True

									lowercase__            =      self.model.to(UpperCamelCase_							)

									model.eval()
									model.float()

									lowercase__            =      model.module if hasattr(UpperCamelCase_    ,						'''module'''							) else model
									quant_trainer.configure_model(UpperCamelCase_    ,						self.quant_trainer_args							)

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

									lowercase__            =      {0: '''batch_size''', 1: '''seq_len'''}

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






'''simple docstring'''



import re
import string

import numpy as np

import datasets


__lowercase	:						Tuple       =	'\nReturns the rate at which the input predicted strings exactly match their references, ignoring any strings input as part of the regexes_to_ignore list.\n'

__lowercase	:						List[str]       =	'\nArgs:\n    predictions: List of predicted texts.\n    references: List of reference texts.\n    regexes_to_ignore: List, defaults to None. Regex expressions of characters to\n        ignore when calculating the exact matches. Note: these regexes are removed\n        from the input data before the changes based on the options below (e.g. ignore_case,\n        ignore_punctuation, ignore_numbers) are applied.\n    ignore_case: Boolean, defaults to False. If true, turns everything\n        to lowercase so that capitalization differences are ignored.\n    ignore_punctuation: Boolean, defaults to False. If true, removes all punctuation before\n        comparing predictions and references.\n    ignore_numbers: Boolean, defaults to False. If true, removes all punctuation before\n        comparing predictions and references.\nReturns:\n    exact_match: Dictionary containing exact_match rate. Possible values are between 0.0 and 100.0, inclusive.\nExamples:\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds)\n    >>> print(round(results["exact_match"], 1))\n    25.0\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=["the ", "yell"], ignore_case=True, ignore_punctuation=True)\n    >>> print(round(results["exact_match"], 1))\n    50.0\n\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=["the ", "yell", "YELL"], ignore_case=True, ignore_punctuation=True)\n    >>> print(round(results["exact_match"], 1))\n    75.0\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=["the ", "yell", "YELL"], ignore_case=True, ignore_punctuation=True, ignore_numbers=True)\n    >>> print(round(results["exact_match"], 1))\n    100.0\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["The cat sat on the mat.", "Theaters are great.", "It\'s like comparing oranges and apples."]\n    >>> preds = ["The cat sat on the mat?", "Theaters are great.", "It\'s like comparing apples and oranges."]\n    >>> results = exact_match.compute(references=refs, predictions=preds)\n    >>> print(round(results["exact_match"], 1))\n    33.3\n\n'

__lowercase	:						Any       =	'\n'





@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION		,      _KWARGS_DESCRIPTION					)
class 							__UpperCamelCase				(		datasets.Metric					):




	def       __UpperCAmelCase							(   self							):

						'''simple docstring'''





						return datasets.MetricInfo(
						    description=_DESCRIPTION						,				citation=_CITATION						,				inputs_description=_KWARGS_DESCRIPTION						,				features=datasets.Features(
						        {
						            'predictions': datasets.Value('string'						,				id='sequence'							),
						            'references': datasets.Value('string'						,				id='sequence'							),
						        }							)						,				reference_urls=[]						,				)





	def       __UpperCAmelCase							(   self						,				__a						,				__a						,				__a=None						,				__a=False						,				__a=False						,				__a=False						,				):

						'''simple docstring'''





						if regexes_to_ignore is not None:
											for s in regexes_to_ignore:
																__a			:       Tuple					    =  np.array([re.sub(__a						,				''						,				__a							) for x in predictions]							)
																__a			:       List[Any]					    =  np.array([re.sub(__a						,				''						,				__a							) for x in references]							)
						else:
											__a			:       int					    =  np.asarray(__a							)
											__a			:       str					    =  np.asarray(__a							)

						if ignore_case:
											__a			:       Dict					    =  np.char.lower(__a							)
											__a			:       List[str]					    =  np.char.lower(__a							)

						if ignore_punctuation:
											__a			:       Dict					    =  string.punctuation.maketrans(''						,				''						,				string.punctuation							)
											__a			:       Tuple					    =  np.char.translate(__a						,				table=__a							)
											__a			:       Dict					    =  np.char.translate(__a						,				table=__a							)

						if ignore_numbers:
											__a			:       Optional[int]					    =  string.digits.maketrans(''						,				''						,				string.digits							)
											__a			:       Tuple					    =  np.char.translate(__a						,				table=__a							)
											__a			:       Optional[int]					    =  np.char.translate(__a						,				table=__a							)

						__a			:       Any					    =  predictions == references

						return {"exact_match": np.mean(__a							) * 100}



