






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


SCREAMING_SNAKE_CASE						:						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       snake_case_       (    _lowerCamelCase       ):



	"""simple docstring"""

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


					"""simple docstring"""

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


	def       _UpperCAmelCase						(      self  ,					__a=None						):


					"""simple docstring"""

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

					A__							 =       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       _UpperCAmelCase						(      self  ,					__a=None						):


					"""simple docstring"""

					A__							 =       self.train_dataset if calib_dataset is None else calib_dataset
					A__							 =       self.get_calib_dataloader(__a						)

					A__							 =       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
									A__ ,			A__ ,			A__							 =       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						)
					A__							 =       model


	def       _UpperCAmelCase						(      self  ,					__a=None  ,					__a=None  ,					__a=None  ,					__a = "eval"						):


					"""simple docstring"""

					A__							 =       self.eval_dataset if eval_dataset is None else eval_dataset
					A__							 =       self.get_eval_dataloader(__a						)
					A__							 =       self.eval_examples if eval_examples is None else eval_examples

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

					if self.post_process_function is not None and self.compute_metrics is not None:
									A__							 =       self.post_process_function(__a  ,					__a  ,					output.predictions						)
									A__							 =       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}_'''						):
																	A__							 =       metrics.pop(__a						)

									self.log(__a						)
					else:
									A__							 =       {}

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

					A__							 =       self.callback_handler.on_evaluate(self.args  ,					self.state  ,					self.control  ,					__a						)
					return metrics


	def       _UpperCAmelCase						(      self  ,					__a  ,					__a  ,					__a=None  ,					__a = "test"						):


					"""simple docstring"""

					A__							 =       self.get_test_dataloader(__a						)

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

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

					A__							 =       self.post_process_function(__a  ,					__a  ,					output.predictions  ,					'predict'						)
					A__							 =       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}_'''						):
													A__							 =       metrics.pop(__a						)

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






	def       _UpperCAmelCase						(      self  ,					__a="./"						):


					"""simple docstring"""

					A__							 =       self.eval_dataset
					A__							 =       self.get_eval_dataloader(__a						)

					A__							 =       next(iter(__a						)						)

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

					# convert to tuple
					A__							 =       tuple(v.to(__a						) for k, v in batch.items()						)

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

					A__							 =       True

					A__							 =       self.model.to(__a						)

					model.eval()
					model.float()

					A__							 =       model.module if hasattr(__a  ,					'module'						) else model
					quant_trainer.configure_model(__a  ,					self.quant_trainer_args						)

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

					A__							 =       {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'						)










"""simple docstring"""






import torch

from transformers import AutoModel





class       snake_case_       (    torch.nn.Module       ):



	"""simple docstring"""

	def __init__(      self  ,					__a="sayef/fsner-bert-base-uncased"						):


					"""simple docstring"""

					super(__a  ,					self						).__init__()

					A__							 =       AutoModel.from_pretrained(__a  ,					return_dict=__a						)
					A__							 =       torch.nn.CosineSimilarity(3  ,					1E-08						)
					A__							 =       torch.nn.Softmax(dim=1						)


	def       _UpperCAmelCase						(      self  ,					**__a						):


					"""simple docstring"""

					return self.bert(**__a						).last_hidden_state


	def       _UpperCAmelCase						(      self  ,					__a						):


					"""simple docstring"""

					return token_embeddings.sum(2  ,					keepdim=__a						)


	def       _UpperCAmelCase						(      self  ,					__a  ,					__a  ,					__a=1						):


					"""simple docstring"""

					return self.softmax(T * self.cos(__a  ,					__a						)						)






	def       _UpperCAmelCase						(      self  ,					__a  ,					__a						):


					"""simple docstring"""

					A__							 =       W_supports['sizes'].tolist()
					A__							 =       W_supports['start_token_id'].item()
					A__							 =       W_supports['end_token_id'].item()

					del W_supports["sizes"]
					del W_supports["start_token_id"]
					del W_supports["end_token_id"]

					A__							 =       self.BERT(**__a						)
					A__							 =       self.BERT(**__a						)

					A__							 =       None
					A__							 =       None

					A__							 =       W_supports['input_ids'] == start_token_id
					A__							 =       W_supports['input_ids'] == end_token_id

					for i, size in enumerate(__a						):
									if i == 0:
													A__							 =       0
									else:
													A__							 =       support_sizes[i - 1]

									A__							 =       S[s : s + size][start_token_masks[s : s + size]]
									A__							 =       S[s : s + size][end_token_masks[s : s + size]]

									A__							 =       torch.matmul(q[i]  ,					s_start.T						).sum(1						).softmax(0						)
									A__							 =       torch.matmul(q[i]  ,					s_end.T						).sum(1						).softmax(0						)

									if p_starts is not None:
													A__							 =       torch.vstack((p_starts, p_start)						)
													A__							 =       torch.vstack((p_ends, p_end)						)
									else:
													A__							 =       p_start
													A__							 =       p_end

					return p_starts, p_ends


