





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				:
 def __init__( self    : Tuple     ,	UpperCAmelCase    : Union[str, Any]     ,	UpperCAmelCase    : str=13     ,	UpperCAmelCase    : Tuple=7     ,	UpperCAmelCase    : str=True     ,	UpperCAmelCase    : List[Any]=True     ,	UpperCAmelCase    : List[Any]=True     ,	UpperCAmelCase    : Union[str, Any]=True     ,	UpperCAmelCase    : int=99     ,	UpperCAmelCase    : Tuple=32     ,	UpperCAmelCase    : Optional[int]=2     ,	UpperCAmelCase    : Union[str, Any]=4     ,	UpperCAmelCase    : Optional[Any]=37     ,	UpperCAmelCase    : Tuple="gelu"     ,	UpperCAmelCase    : str=0.1     ,	UpperCAmelCase    : Optional[int]=0.1     ,	UpperCAmelCase    : List[Any]=512     ,	UpperCAmelCase    : List[Any]=16     ,	UpperCAmelCase    : Optional[Any]=2     ,	UpperCAmelCase    : List[str]=0.0_2     ,	UpperCAmelCase    : Tuple=3     ,	UpperCAmelCase    : Tuple=4     ,	UpperCAmelCase    : Tuple=None     ,	UpperCAmelCase    : Dict=0     ,	)     ->      List[str]:
     lowerCamelCase__       : Optional[int]   					=						parent
     lowerCamelCase__       : List[str]   					=						batch_size
     lowerCamelCase__       : Any   					=						seq_length
     lowerCamelCase__       : str   					=						is_training
     lowerCamelCase__       : Any   					=						use_input_mask
     lowerCamelCase__       : Optional[int]   					=						use_token_type_ids
     lowerCamelCase__       : str   					=						use_labels
     lowerCamelCase__       : Optional[Any]   					=						vocab_size
     lowerCamelCase__       : Any   					=						hidden_size
     lowerCamelCase__       : Any   					=						num_hidden_layers
     lowerCamelCase__       : Union[str, Any]   					=						num_attention_heads
     lowerCamelCase__       : str   					=						intermediate_size
     lowerCamelCase__       : List[str]   					=						hidden_act
     lowerCamelCase__       : List[str]   					=						hidden_dropout_prob
     lowerCamelCase__       : str   					=						attention_probs_dropout_prob
     lowerCamelCase__       : Optional[int]   					=						max_position_embeddings
     lowerCamelCase__       : Any   					=						type_vocab_size
     lowerCamelCase__       : List[Any]   					=						type_sequence_label_size
     lowerCamelCase__       : Optional[Any]   					=						initializer_range
     lowerCamelCase__       : int   					=						num_labels
     lowerCamelCase__       : Dict   					=						num_choices
     lowerCamelCase__       : str   					=						scope
     lowerCamelCase__       : Optional[Any]   					=						projection_dim
 def 	A_      ( self    : Optional[Any]      )     ->      Any:
     lowerCamelCase__       : Optional[Any]   					=						ids_tensor([self.batch_size, self.seq_length]     ,	self.vocab_size      )

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

     lowerCamelCase__       : List[Any]   					=						None
     if self.use_token_type_ids:
         lowerCamelCase__       : List[str]   					=						ids_tensor([self.batch_size, self.seq_length]     ,	self.type_vocab_size      )

     lowerCamelCase__       : List[str]   					=						None
     lowerCamelCase__       : Union[str, Any]   					=						None
     lowerCamelCase__       : Union[str, Any]   					=						None
     if self.use_labels:
         lowerCamelCase__       : List[Any]   					=						ids_tensor([self.batch_size]     ,	self.type_sequence_label_size      )
         lowerCamelCase__       : Any   					=						ids_tensor([self.batch_size, self.seq_length]     ,	self.num_labels      )
         lowerCamelCase__       : Dict   					=						ids_tensor([self.batch_size]     ,	self.num_choices      )

     lowerCamelCase__       : str   					=						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=UpperCAmelCase     ,	initializer_range=self.initializer_range     ,	)
     lowerCamelCase__       : Dict   					=						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    : Any     ,	UpperCAmelCase    : Any     ,	UpperCAmelCase    : str     ,	UpperCAmelCase    : Optional[Any]     ,	UpperCAmelCase    : Optional[Any]     ,	UpperCAmelCase    : Tuple     ,	UpperCAmelCase    : List[str]     ,	UpperCAmelCase    : Dict      )     ->      Optional[Any]:
     lowerCamelCase__       : Optional[Any]   					=						TFDPRContextEncoder(config=UpperCAmelCase      )
     lowerCamelCase__       : List[Any]   					=						model(UpperCAmelCase     ,	attention_mask=UpperCAmelCase     ,	token_type_ids=UpperCAmelCase      )
     lowerCamelCase__       : Optional[int]   					=						model(UpperCAmelCase     ,	token_type_ids=UpperCAmelCase      )
     lowerCamelCase__       : Union[str, Any]   					=						model(UpperCAmelCase      )
     self.parent.assertEqual(result.pooler_output.shape     ,	(self.batch_size, self.projection_dim or self.hidden_size)      )
 def 	A_      ( self    : Optional[Any]     ,	UpperCAmelCase    : Tuple     ,	UpperCAmelCase    : List[Any]     ,	UpperCAmelCase    : List[Any]     ,	UpperCAmelCase    : Any     ,	UpperCAmelCase    : List[Any]     ,	UpperCAmelCase    : str     ,	UpperCAmelCase    : Tuple      )     ->      Union[str, Any]:
     lowerCamelCase__       : Optional[int]   					=						TFDPRQuestionEncoder(config=UpperCAmelCase      )
     lowerCamelCase__       : Optional[int]   					=						model(UpperCAmelCase     ,	attention_mask=UpperCAmelCase     ,	token_type_ids=UpperCAmelCase      )
     lowerCamelCase__       : Any   					=						model(UpperCAmelCase     ,	token_type_ids=UpperCAmelCase      )
     lowerCamelCase__       : Union[str, Any]   					=						model(UpperCAmelCase      )
     self.parent.assertEqual(result.pooler_output.shape     ,	(self.batch_size, self.projection_dim or self.hidden_size)      )
 def 	A_      ( self    : Any     ,	UpperCAmelCase    : List[Any]     ,	UpperCAmelCase    : Tuple     ,	UpperCAmelCase    : Dict     ,	UpperCAmelCase    : Optional[Any]     ,	UpperCAmelCase    : Union[str, Any]     ,	UpperCAmelCase    : List[str]     ,	UpperCAmelCase    : List[Any]      )     ->      Optional[Any]:
     lowerCamelCase__       : List[Any]   					=						TFDPRReader(config=UpperCAmelCase      )
     lowerCamelCase__       : str   					=						model(UpperCAmelCase     ,	attention_mask=UpperCAmelCase      )

     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    : Any      )     ->      Union[str, Any]:
     lowerCamelCase__       : List[str]   					=						self.prepare_config_and_inputs()
     (
         (
         lowerCamelCase__
     )					,				(
         lowerCamelCase__
     )					,				(
         lowerCamelCase__
     )					,				(
         lowerCamelCase__
     )					,				(
         lowerCamelCase__
     )					,				(
         lowerCamelCase__
     )					,				(
         lowerCamelCase__
     )					,				
     )       : Optional[Any]   					=						config_and_inputs
     lowerCamelCase__       : int   					=						{'input_ids': input_ids}
     return config, inputs_dict


@require_tf
class   lowerCAmelCase				(   __UpperCamelCase,    __UpperCamelCase,    unittest.TestCase     ):
 UpperCAmelCase__						=						(
     (
         TFDPRContextEncoder,
         TFDPRQuestionEncoder,
         TFDPRReader,
     )
     if is_tf_available()
     else ()
 )
 UpperCAmelCase__						=						{"""feature-extraction""": TFDPRQuestionEncoder} if is_tf_available() else {}

 UpperCAmelCase__						=						False
 UpperCAmelCase__						=						False
 UpperCAmelCase__						=						False
 UpperCAmelCase__						=						False
 UpperCAmelCase__						=						False
 def 	A_      ( self    : Optional[Any]      )     ->      Union[str, Any]:
     lowerCamelCase__       : int   					=						TFDPRModelTester(self      )
     lowerCamelCase__       : List[Any]   					=						ConfigTester(self     ,	config_class=UpperCAmelCase     ,	hidden_size=37      )
 def 	A_      ( self    : Optional[Any]      )     ->      Optional[int]:
     self.config_tester.run_common_tests()
 def 	A_      ( self    : Tuple      )     ->      Optional[int]:
     lowerCamelCase__       : Any   					=						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_dpr_context_encoder(*UpperCAmelCase      )
 def 	A_      ( self    : Optional[Any]      )     ->      Any:
     lowerCamelCase__       : Optional[Any]   					=						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_dpr_question_encoder(*UpperCAmelCase      )
 def 	A_      ( self    : Tuple      )     ->      Optional[int]:
     lowerCamelCase__       : int   					=						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_dpr_reader(*UpperCAmelCase      )







 @slow
 def 	A_      ( self    : List[Any]      )     ->      Tuple:
     for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         lowerCamelCase__       : Optional[int]   					=						TFDPRContextEncoder.from_pretrained(UpperCAmelCase      )
         self.assertIsNotNone(UpperCAmelCase      )

     for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         lowerCamelCase__       : str   					=						TFDPRContextEncoder.from_pretrained(UpperCAmelCase      )
         self.assertIsNotNone(UpperCAmelCase      )

     for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         lowerCamelCase__       : int   					=						TFDPRQuestionEncoder.from_pretrained(UpperCAmelCase      )
         self.assertIsNotNone(UpperCAmelCase      )

     for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         lowerCamelCase__       : List[Any]   					=						TFDPRReader.from_pretrained(UpperCAmelCase      )
         self.assertIsNotNone(UpperCAmelCase      )


@require_tf
class   lowerCAmelCase				(   unittest.TestCase     ):
 @slow
 def 	A_      ( self    : Dict      )     ->      str:
     lowerCamelCase__       : int   					=						TFDPRQuestionEncoder.from_pretrained('facebook/dpr-question_encoder-single-nq-base'      )

     lowerCamelCase__       : Optional[int]   					=						tf.constant(
         [[101, 7592, 1010, 2003, 2026, 3899, 10140, 1029, 102]]      )  # [CLS] hello, is my dog cute? [SEP]
     lowerCamelCase__       : Union[str, Any]   					=						model(UpperCAmelCase      )[0]  # embedding shape = (1, 768)
     # compare the actual values for a slice.
     lowerCamelCase__       : List[Any]   					=						tf.constant(
         [
             [
                 0.0_3_2_3_6_2_5_3,
                 0.1_2_7_5_3_3_3_5,
                 0.1_6_8_1_8_5_0_9,
                 0.0_0_2_7_9_7_8_6,
                 0.3_8_9_6_9_3_3,
                 0.2_4_2_6_4_9_4_5,
                 0.2_1_7_8_9_7_1,
                 -0.0_2_3_3_5_2_2_7,
                 -0.0_8_4_8_1_9_5_9,
                 -0.1_4_3_2_4_1_1_7,
             ]
         ]      )
     self.assertTrue(numpy.allclose(output[:, :10].numpy()     ,	expected_slice.numpy()     ,	atol=1e-4      )      )












import gc
import threading
import time

import psutil
import torch

class   lowerCAmelCase				:
 def __init__( self    : str      )     ->      Union[str, Any]:
     lowerCamelCase__       : Optional[Any]   					=						psutil.Process()
     lowerCamelCase__       : Union[str, Any]   					=						False
 def 	A_      ( self    : Optional[int]      )     ->      int:
     lowerCamelCase__       : Optional[Any]   					=						-1

     while True:
         lowerCamelCase__       : Dict   					=						max(self.process.memory_info().rss     ,	self.cpu_memory_peak      )

         # can't sleep or will not catch the peak right (this comment is here on purpose)
         if not self.peak_monitoring:
             break
 def 	A_      ( self    : Tuple      )     ->      Dict:
     lowerCamelCase__       : List[Any]   					=						True
     lowerCamelCase__       : List[str]   					=						threading.Thread(target=self.peak_monitor      )
     lowerCamelCase__       : Union[str, Any]   					=						True
     self.thread.start()







 def 	A_      ( self    : str      )     ->      Dict:
     lowerCamelCase__       : int   					=						False
     self.thread.join()
     return self.cpu_memory_peak


_UpperCAmelCase							:			Dict        =							PeakCPUMemory()


def 			SCREAMING_SNAKE_CASE (			)     ->	Union[str, Any]:
    # Time
    lowerCamelCase__       : List[Any]   					=						{'time': time.time()}

    gc.collect()
    torch.cuda.empty_cache()

    # CPU mem
    lowerCamelCase__       : List[str]   					=						psutil.Process().memory_info().rss
    cpu_peak_tracker.start()

    # GPU mem
    for i in range(torch.cuda.device_count()					):
        lowerCamelCase__       : Union[str, Any]   					=						torch.cuda.memory_allocated(_UpperCAmelCase					)
    torch.cuda.reset_peak_memory_stats()

    return measures


def 			SCREAMING_SNAKE_CASE (			_UpperCAmelCase					)     ->	Tuple:
    # Time
    lowerCamelCase__       : Optional[int]   					=						{'time': time.time() - start_measures['time']}

    gc.collect()
    torch.cuda.empty_cache()

    # CPU mem
    lowerCamelCase__       : Dict   					=						(psutil.Process().memory_info().rss - start_measures['cpu']) / 2**20
    lowerCamelCase__       : int   					=						(cpu_peak_tracker.stop() - start_measures['cpu']) / 2**20

    # GPU mem
    for i in range(torch.cuda.device_count()					):
        lowerCamelCase__       : List[str]   					=						(torch.cuda.memory_allocated(_UpperCAmelCase					) - start_measures[str(_UpperCAmelCase					)]) / 2**20
        lowerCamelCase__       : Optional[Any]   					=						(torch.cuda.max_memory_allocated(_UpperCAmelCase					) - start_measures[str(_UpperCAmelCase					)]) / 2**20

    return measures


def 			SCREAMING_SNAKE_CASE (			_UpperCAmelCase		,      _UpperCAmelCase					)     ->	Dict:
    print(F"""{description}:"""					)
    print(F"""- Time: {measures["time"]:.2f}s"""					)
    for i in range(torch.cuda.device_count()					):
        print(F"""- GPU {i} allocated: {measures[str(_UpperCAmelCase					)]:.2f}MiB"""					)
        lowerCamelCase__       : List[str]   					=						measures[F"""{i}-peak"""]
        print(F"""- GPU {i} peak: {peak:.2f}MiB"""					)
    print(F"""- CPU RAM allocated: {measures["cpu"]:.2f}MiB"""					)
    print(F"""- CPU RAM peak: {measures["cpu-peak"]:.2f}MiB"""					)





