



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


a_						:     List[str]	     =		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__  (						UpperCAmelCase_):
     """simple docstring"""



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






           '''simple docstring'''



           super().__init__(*__a				,			**__a    )
           lowerCamelCase								=		eval_examples
           lowerCamelCase								=		post_process_function
           lowerCamelCase								=		quant_trainer_args
           lowerCamelCase								=		1_28  # default number of calibration samples







     def 						_a      (self				,			__a=None    ):






           '''simple docstring'''



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

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






           '''simple docstring'''



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

           lowerCamelCase								=		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								=		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								=		model







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






           '''simple docstring'''



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

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

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

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

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







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






           '''simple docstring'''



           lowerCamelCase								=		self.get_test_dataloader(__a    )

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

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

           lowerCamelCase								=		self.post_process_function(__a				,			__a				,			output.predictions				,			"predict"    )
           lowerCamelCase								=		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								=		metrics.pop(__a    )

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







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






           '''simple docstring'''



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

           lowerCamelCase								=		next(iter(__a    )    )

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

           # convert to tuple
           lowerCamelCase								=		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								=		True

           lowerCamelCase								=		self.model.to(__a    )

           model.eval()
           model.float()

           lowerCamelCase								=		model.module if hasattr(__a				,			"module"    ) else model
           quant_trainer.configure_model(__a				,			self.quant_trainer_args    )

           lowerCamelCase								=		os.path.join(__a				,			"model.onnx"    )
           logger.info(F"""exporting model to {output_model_file}"""    )

           lowerCamelCase								=		{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"    )




from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
     import numpy
     import tensorflow as tf

     from transformers import (
         TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
         TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
         TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
         BertConfig,
         DPRConfig,
         TFDPRContextEncoder,
         TFDPRQuestionEncoder,
         TFDPRReader,
     )





class    lowerCamelCase__  :
     """simple docstring"""



     def __init__(self				,			__a				,			__a=13				,			__a=7				,			__a=True				,			__a=True				,			__a=True				,			__a=True				,			__a=99				,			__a=32				,			__a=2				,			__a=4				,			__a=37				,			__a="gelu"				,			__a=0.1				,			__a=0.1				,			__a=5_12				,			__a=16				,			__a=2				,			__a=0.02				,			__a=3				,			__a=4				,			__a=None				,			__a=0				,			):






           '''simple docstring'''



           lowerCamelCase								=		parent
           lowerCamelCase								=		batch_size
           lowerCamelCase								=		seq_length
           lowerCamelCase								=		is_training
           lowerCamelCase								=		use_input_mask
           lowerCamelCase								=		use_token_type_ids
           lowerCamelCase								=		use_labels
           lowerCamelCase								=		vocab_size
           lowerCamelCase								=		hidden_size
           lowerCamelCase								=		num_hidden_layers
           lowerCamelCase								=		num_attention_heads
           lowerCamelCase								=		intermediate_size
           lowerCamelCase								=		hidden_act
           lowerCamelCase								=		hidden_dropout_prob
           lowerCamelCase								=		attention_probs_dropout_prob
           lowerCamelCase								=		max_position_embeddings
           lowerCamelCase								=		type_vocab_size
           lowerCamelCase								=		type_sequence_label_size
           lowerCamelCase								=		initializer_range
           lowerCamelCase								=		num_labels
           lowerCamelCase								=		num_choices
           lowerCamelCase								=		scope
           lowerCamelCase								=		projection_dim







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		ids_tensor([self.batch_size, self.seq_length]				,			self.vocab_size    )

           lowerCamelCase								=		None
           if self.use_input_mask:
                 # follow test_modeling_tf_ctrl.py
                 lowerCamelCase								=		random_attention_mask([self.batch_size, self.seq_length]    )

           lowerCamelCase								=		None
           if self.use_token_type_ids:
                 lowerCamelCase								=		ids_tensor([self.batch_size, self.seq_length]				,			self.type_vocab_size    )

           lowerCamelCase								=		None
           lowerCamelCase								=		None
           lowerCamelCase								=		None
           if self.use_labels:
                 lowerCamelCase								=		ids_tensor([self.batch_size]				,			self.type_sequence_label_size    )
                 lowerCamelCase								=		ids_tensor([self.batch_size, self.seq_length]				,			self.num_labels    )
                 lowerCamelCase								=		ids_tensor([self.batch_size]				,			self.num_choices    )

           lowerCamelCase								=		BertConfig(
               vocab_size=self.vocab_size				,			hidden_size=self.hidden_size				,			num_hidden_layers=self.num_hidden_layers				,			num_attention_heads=self.num_attention_heads				,			intermediate_size=self.intermediate_size				,			hidden_act=self.hidden_act				,			hidden_dropout_prob=self.hidden_dropout_prob				,			attention_probs_dropout_prob=self.attention_probs_dropout_prob				,			max_position_embeddings=self.max_position_embeddings				,			type_vocab_size=self.type_vocab_size				,			is_decoder=__a				,			initializer_range=self.initializer_range				,			)
           lowerCamelCase								=		DPRConfig(projection_dim=self.projection_dim				,			**config.to_dict()    )

           return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels







     def 						_a      (self				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a    ):






           '''simple docstring'''



           lowerCamelCase								=		TFDPRContextEncoder(config=__a    )
           lowerCamelCase								=		model(__a				,			attention_mask=__a				,			token_type_ids=__a    )
           lowerCamelCase								=		model(__a				,			token_type_ids=__a    )
           lowerCamelCase								=		model(__a    )
           self.parent.assertEqual(result.pooler_output.shape				,			(self.batch_size, self.projection_dim or self.hidden_size)    )







     def 						_a      (self				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a    ):






           '''simple docstring'''



           lowerCamelCase								=		TFDPRQuestionEncoder(config=__a    )
           lowerCamelCase								=		model(__a				,			attention_mask=__a				,			token_type_ids=__a    )
           lowerCamelCase								=		model(__a				,			token_type_ids=__a    )
           lowerCamelCase								=		model(__a    )
           self.parent.assertEqual(result.pooler_output.shape				,			(self.batch_size, self.projection_dim or self.hidden_size)    )







     def 						_a      (self				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a    ):






           '''simple docstring'''



           lowerCamelCase								=		TFDPRReader(config=__a    )
           lowerCamelCase								=		model(__a				,			attention_mask=__a    )

           self.parent.assertEqual(result.start_logits.shape				,			(self.batch_size, self.seq_length)    )
           self.parent.assertEqual(result.end_logits.shape				,			(self.batch_size, self.seq_length)    )
           self.parent.assertEqual(result.relevance_logits.shape				,			(self.batch_size,)    )







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.prepare_config_and_inputs()
           (
               (
               lowerCamelCase
           )			,							(
               lowerCamelCase
           )			,							(
               lowerCamelCase
           )			,							(
               lowerCamelCase
           )			,							(
               lowerCamelCase
           )			,							(
               lowerCamelCase
           )			,							(
               lowerCamelCase
           )			,							
           )								=		config_and_inputs
           lowerCamelCase								=		{"input_ids": input_ids}
           return config, inputs_dict





@require_tf
class    lowerCamelCase__  (						UpperCAmelCase_       ,    UpperCAmelCase_       ,    unittest.TestCase):
     """simple docstring"""



     _A       =			(
         (
             TFDPRContextEncoder,
             TFDPRQuestionEncoder,
             TFDPRReader,
         )
         if is_tf_available()
         else ()
     )
     _A       =			{'feature-extraction': TFDPRQuestionEncoder} if is_tf_available() else {}

     _A       =			False
     _A       =			False
     _A       =			False
     _A       =			False
     _A       =			False







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		TFDPRModelTester(self    )
           lowerCamelCase								=		ConfigTester(self				,			config_class=__a				,			hidden_size=37    )







     def 						_a      (self    ):






           '''simple docstring'''



           self.config_tester.run_common_tests()







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_dpr_context_encoder(*__a    )







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_dpr_question_encoder(*__a    )







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_dpr_reader(*__a    )







     @slow
     def 						_a      (self    ):






           '''simple docstring'''



           for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 lowerCamelCase								=		TFDPRContextEncoder.from_pretrained(__a    )
                 self.assertIsNotNone(__a    )

           for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 lowerCamelCase								=		TFDPRContextEncoder.from_pretrained(__a    )
                 self.assertIsNotNone(__a    )

           for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 lowerCamelCase								=		TFDPRQuestionEncoder.from_pretrained(__a    )
                 self.assertIsNotNone(__a    )

           for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 lowerCamelCase								=		TFDPRReader.from_pretrained(__a    )
                 self.assertIsNotNone(__a    )





@require_tf
class    lowerCamelCase__  (						unittest.TestCase):
     """simple docstring"""



     @slow
     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		TFDPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base"    )

           lowerCamelCase								=		tf.constant(
               [[1_01, 75_92, 10_10, 20_03, 20_26, 38_99, 1_01_40, 10_29, 1_02]]    )  # [CLS] hello, is my dog cute? [SEP]
           lowerCamelCase								=		model(__a    )[0]  # embedding shape = (1, 768)
           # compare the actual values for a slice.
           lowerCamelCase								=		tf.constant(
               [
                   [
                       0.0323_6253,
                       0.1275_3335,
                       0.1681_8509,
                       0.0027_9786,
                       0.389_6933,
                       0.2426_4945,
                       0.217_8971,
                       -0.0233_5227,
                       -0.0848_1959,
                       -0.1432_4117,
                   ]
               ]    )
           self.assertTrue(numpy.allclose(output[:, :10].numpy()				,			expected_slice.numpy()				,			atol=1E-4    )    )