




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 							__UpperCamelCase			:





				'''simple docstring'''



				def __init__( self       , UpperCAmelCase_       , UpperCAmelCase_=13       , UpperCAmelCase_=7       , UpperCAmelCase_=True       , UpperCAmelCase_=True       , UpperCAmelCase_=True       , UpperCAmelCase_=True       , UpperCAmelCase_=99       , UpperCAmelCase_=32       , UpperCAmelCase_=2       , UpperCAmelCase_=4       , UpperCAmelCase_=37       , UpperCAmelCase_="gelu"       , UpperCAmelCase_=0.1       , UpperCAmelCase_=0.1       , UpperCAmelCase_=5_12       , UpperCAmelCase_=16       , UpperCAmelCase_=2       , UpperCAmelCase_=0.02       , UpperCAmelCase_=3       , UpperCAmelCase_=4       , UpperCAmelCase_=None       , UpperCAmelCase_=0       , ):
											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       __snake_case		( self						):
											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=_lowercase       , 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       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						):
											lowerCAmelCase   =		TFDPRContextEncoder(config=_lowercase						)
											lowerCAmelCase   =		model(_lowercase       , attention_mask=_lowercase       , token_type_ids=_lowercase						)
											lowerCAmelCase   =		model(_lowercase       , token_type_ids=_lowercase						)
											lowerCAmelCase   =		model(_lowercase						)
											self.parent.assertEqual(result.pooler_output.shape       , (self.batch_size, self.projection_dim or self.hidden_size)						)



				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						):
											lowerCAmelCase   =		TFDPRQuestionEncoder(config=_lowercase						)
											lowerCAmelCase   =		model(_lowercase       , attention_mask=_lowercase       , token_type_ids=_lowercase						)
											lowerCAmelCase   =		model(_lowercase       , token_type_ids=_lowercase						)
											lowerCAmelCase   =		model(_lowercase						)
											self.parent.assertEqual(result.pooler_output.shape       , (self.batch_size, self.projection_dim or self.hidden_size)						)



				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						):
											lowerCAmelCase   =		TFDPRReader(config=_lowercase						)
											lowerCAmelCase   =		model(_lowercase       , attention_mask=_lowercase						)

											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       __snake_case		( self						):
											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 							__UpperCamelCase			( UpperCAmelCase_     ,					UpperCAmelCase_     ,					unittest.TestCase  ):





				'''simple docstring'''



				__a							: Tuple										=(
				    (
				        TFDPRContextEncoder,
				        TFDPRQuestionEncoder,
				        TFDPRReader,
				    )
				    if is_tf_available()
				    else ()
				)
				__a							: Optional[int]										={"""feature-extraction""": TFDPRQuestionEncoder} if is_tf_available() else {}

				__a							: Optional[int]										=False
				__a							: str										=False
				__a							: str										=False
				__a							: Any										=False
				__a							: str										=False



				def       __snake_case		( self						):
											lowerCAmelCase   =		TFDPRModelTester(self						)
											lowerCAmelCase   =		ConfigTester(self       , config_class=_lowercase       , hidden_size=37						)



				def       __snake_case		( self						):
											self.config_tester.run_common_tests()



				def       __snake_case		( self						):
											lowerCAmelCase   =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_dpr_context_encoder(*_lowercase						)



				def       __snake_case		( self						):
											lowerCAmelCase   =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_dpr_question_encoder(*_lowercase						)



				def       __snake_case		( self						):
											lowerCAmelCase   =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_dpr_reader(*_lowercase						)






				@slow
				def       __snake_case		( self						):
											for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		lowerCAmelCase   =		TFDPRContextEncoder.from_pretrained(_lowercase						)
																		self.assertIsNotNone(_lowercase						)

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

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

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




@require_tf
class 							__UpperCamelCase			( unittest.TestCase  ):





				'''simple docstring'''



				@slow
				def       __snake_case		( self						):
											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(_lowercase						)[0]  # embedding shape = (1, 768)
											# compare the actual values for a slice.
											lowerCAmelCase   =		tf.constant(
											    [
											        [
											            0.03236253,
											            0.12753335,
											            0.16818509,
											            0.00279786,
											            0.3896933,
											            0.24264945,
											            0.2178971,
											            -0.02335227,
											            -0.08481959,
											            -0.14324117,
											        ]
											    ]						)
											self.assertTrue(numpy.allclose(output[:, :10].numpy()       , expected_slice.numpy()       , atol=1E-4						)						)









from __future__ import annotations

from typing import Generic, TypeVar

UpperCAmelCase_						=TypeVar("""T""")


class 							__UpperCamelCase			( Generic[T]  ):





				'''simple docstring'''



				def __init__( self       , UpperCAmelCase_						):
											lowerCAmelCase   =		data
											lowerCAmelCase   =		self
											lowerCAmelCase   =		0




class 							__UpperCamelCase			( Generic[T]  ):





				'''simple docstring'''



				def __init__( self						):
											# map from node name to the node object
											lowerCAmelCase   =		{}



				def       __snake_case		( self       , UpperCAmelCase_						):
											# create a new set with x as its member
											lowerCAmelCase   =		DisjointSetTreeNode(UpperCAmelCase_						)



				def       __snake_case		( self       , UpperCAmelCase_						):
											# find the set x belongs to (with path-compression)
											lowerCAmelCase   =		self.map[data]
											if elem_ref != elem_ref.parent:
																		lowerCAmelCase   =		self.find_set(elem_ref.parent.data						)
											return elem_ref.parent



				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_						):
											# helper function for union operation
											if nodea.rank > nodea.rank:
																		lowerCAmelCase   =		nodea
											else:
																		lowerCAmelCase   =		nodea
																		if nodea.rank == nodea.rank:
																									nodea.rank += 1






				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_						):
											# merge 2 disjoint sets
											self.link(self.find_set(UpperCAmelCase_						)       , self.find_set(UpperCAmelCase_						)						)




class 							__UpperCamelCase			( Generic[T]  ):





				'''simple docstring'''



				def __init__( self						):
											# connections: map from the node to the neighbouring nodes (with weights)
											lowerCAmelCase   =		{}



				def       __snake_case		( self       , UpperCAmelCase_						):
											# add a node ONLY if its not present in the graph
											if node not in self.connections:
																		lowerCAmelCase   =		{}



				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						):
											# add an edge with the given weight
											self.add_node(UpperCAmelCase_						)
											self.add_node(UpperCAmelCase_						)
											lowerCAmelCase   =		weight
											lowerCAmelCase   =		weight






				def       __snake_case		( self						):
											lowerCAmelCase   =		[]
											lowerCAmelCase   =		set()
											for start in self.connections:
																		for end in self.connections[start]:
																									if (start, end) not in seen:
																																seen.add((end, start)						)
																																edges.append((start, end, self.connections[start][end])						)
											edges.sort(key=lambda UpperCAmelCase_						: x[2]						)

											# creating the disjoint set
											lowerCAmelCase   =		DisjointSetTree[T]()
											for node in self.connections:
																		disjoint_set.make_set(UpperCAmelCase_						)

											# MST generation
											lowerCAmelCase   =		0
											lowerCAmelCase   =		0
											lowerCAmelCase   =		GraphUndirectedWeighted[T]()
											while num_edges < len(self.connections						) - 1:
																		lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   =		edges[index]
																		index += 1
																		lowerCAmelCase   =		disjoint_set.find_set(UpperCAmelCase_						)
																		lowerCAmelCase   =		disjoint_set.find_set(UpperCAmelCase_						)
																		if parent_u != parent_v:
																									num_edges += 1
																									graph.add_edge(UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						)
																									disjoint_set.union(UpperCAmelCase_       , UpperCAmelCase_						)
											return graph



