


"""simple docstring"""




from __future__ import annotations

import unittest

from transformers import DistilBertConfig, 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 tensorflow as tf

		from transformers.models.distilbert.modeling_tf_distilbert import (
		    TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFDistilBertForMaskedLM,
		    TFDistilBertForMultipleChoice,
		    TFDistilBertForQuestionAnswering,
		    TFDistilBertForSequenceClassification,
		    TFDistilBertForTokenClassification,
		    TFDistilBertModel,
		)





class  _a							:




	def __init__( self						:		List[str]		, _SCREAMING_SNAKE_CASE						:		List[str]		, )->	Optional[Any]:
			lowerCAmelCase__       :      str   					= parent
			lowerCAmelCase__       :      Optional[int]   					= 13
			lowerCAmelCase__       :      Any   					= 7
			lowerCAmelCase__       :      List[Any]   					= True
			lowerCAmelCase__       :      Optional[int]   					= True
			lowerCAmelCase__       :      Optional[Any]   					= False
			lowerCAmelCase__       :      int   					= True
			lowerCAmelCase__       :      List[Any]   					= 99
			lowerCAmelCase__       :      Tuple   					= 32
			lowerCAmelCase__       :      Optional[int]   					= 2
			lowerCAmelCase__       :      Dict   					= 4
			lowerCAmelCase__       :      int   					= 37
			lowerCAmelCase__       :      Union[str, Any]   					= '''gelu'''
			lowerCAmelCase__       :      Optional[Any]   					= 0.1
			lowerCAmelCase__       :      int   					= 0.1
			lowerCAmelCase__       :      Dict   					= 512
			lowerCAmelCase__       :      Optional[Any]   					= 16
			lowerCAmelCase__       :      List[str]   					= 2
			lowerCAmelCase__       :      Any   					= 0.02
			lowerCAmelCase__       :      List[str]   					= 3
			lowerCAmelCase__       :      Dict   					= 4
			lowerCAmelCase__       :      Optional[Any]   					= None




	def 							UpperCAmelCase__( self						:		Tuple		)->	Dict:
			lowerCAmelCase__       :      int   					= ids_tensor([self.batch_size, self.seq_length]		, self.vocab_size		)

			lowerCAmelCase__       :      int   					= None
			if self.use_input_mask:
					lowerCAmelCase__       :      Dict   					= random_attention_mask([self.batch_size, self.seq_length]		)

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

			lowerCAmelCase__       :      int   					= DistilBertConfig(
			    vocab_size=self.vocab_size		, dim=self.hidden_size		, n_layers=self.num_hidden_layers		, n_heads=self.num_attention_heads		, hidden_dim=self.intermediate_size		, hidden_act=self.hidden_act		, dropout=self.hidden_dropout_prob		, attention_dropout=self.attention_probs_dropout_prob		, max_position_embeddings=self.max_position_embeddings		, initializer_range=self.initializer_range		, )

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




	def 							UpperCAmelCase__( self						:		Union[str, Any]		, _SCREAMING_SNAKE_CASE						:		Dict		, _SCREAMING_SNAKE_CASE						:		List[str]		, _SCREAMING_SNAKE_CASE						:		str		, _SCREAMING_SNAKE_CASE						:		Union[str, Any]		, _SCREAMING_SNAKE_CASE						:		Tuple		, _SCREAMING_SNAKE_CASE						:		List[str]		)->	Optional[Any]:
			lowerCAmelCase__       :      int   					= TFDistilBertModel(config=_SCREAMING_SNAKE_CASE		)
			lowerCAmelCase__       :      Union[str, Any]   					= {'''input_ids''': input_ids, '''attention_mask''': input_mask}

			lowerCAmelCase__       :      List[str]   					= model(_SCREAMING_SNAKE_CASE		)

			lowerCAmelCase__       :      Tuple   					= [input_ids, input_mask]

			lowerCAmelCase__       :      Union[str, Any]   					= model(_SCREAMING_SNAKE_CASE		)

			self.parent.assertEqual(result.last_hidden_state.shape		, (self.batch_size, self.seq_length, self.hidden_size)		)




	def 							UpperCAmelCase__( self						:		int		, _SCREAMING_SNAKE_CASE						:		Any		, _SCREAMING_SNAKE_CASE						:		Any		, _SCREAMING_SNAKE_CASE						:		Dict		, _SCREAMING_SNAKE_CASE						:		Tuple		, _SCREAMING_SNAKE_CASE						:		Optional[int]		, _SCREAMING_SNAKE_CASE						:		List[str]		)->	List[Any]:
			lowerCAmelCase__       :      Tuple   					= TFDistilBertForMaskedLM(config=_SCREAMING_SNAKE_CASE		)
			lowerCAmelCase__       :      str   					= {'''input_ids''': input_ids, '''attention_mask''': input_mask}
			lowerCAmelCase__       :      List[str]   					= model(_SCREAMING_SNAKE_CASE		)
			self.parent.assertEqual(result.logits.shape		, (self.batch_size, self.seq_length, self.vocab_size)		)




	def 							UpperCAmelCase__( self						:		Any		, _SCREAMING_SNAKE_CASE						:		Optional[int]		, _SCREAMING_SNAKE_CASE						:		Optional[Any]		, _SCREAMING_SNAKE_CASE						:		Optional[Any]		, _SCREAMING_SNAKE_CASE						:		int		, _SCREAMING_SNAKE_CASE						:		Any		, _SCREAMING_SNAKE_CASE						:		int		)->	int:
			lowerCAmelCase__       :      str   					= TFDistilBertForQuestionAnswering(config=_SCREAMING_SNAKE_CASE		)
			lowerCAmelCase__       :      int   					= {
			    '''input_ids''': input_ids,
			    '''attention_mask''': input_mask,
			}
			lowerCAmelCase__       :      int   					= model(_SCREAMING_SNAKE_CASE		)
			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)		)




	def 							UpperCAmelCase__( self						:		List[str]		, _SCREAMING_SNAKE_CASE						:		Optional[int]		, _SCREAMING_SNAKE_CASE						:		List[Any]		, _SCREAMING_SNAKE_CASE						:		Tuple		, _SCREAMING_SNAKE_CASE						:		Tuple		, _SCREAMING_SNAKE_CASE						:		Optional[int]		, _SCREAMING_SNAKE_CASE						:		Tuple		)->	List[str]:
			lowerCAmelCase__       :      int   					= self.num_labels
			lowerCAmelCase__       :      str   					= TFDistilBertForSequenceClassification(_SCREAMING_SNAKE_CASE		)
			lowerCAmelCase__       :      Optional[int]   					= {'''input_ids''': input_ids, '''attention_mask''': input_mask}
			lowerCAmelCase__       :      int   					= model(_SCREAMING_SNAKE_CASE		)
			self.parent.assertEqual(result.logits.shape		, (self.batch_size, self.num_labels)		)




	def 							UpperCAmelCase__( self						:		Optional[Any]		, _SCREAMING_SNAKE_CASE						:		Dict		, _SCREAMING_SNAKE_CASE						:		Dict		, _SCREAMING_SNAKE_CASE						:		Tuple		, _SCREAMING_SNAKE_CASE						:		List[Any]		, _SCREAMING_SNAKE_CASE						:		Union[str, Any]		, _SCREAMING_SNAKE_CASE						:		Dict		)->	List[str]:
			lowerCAmelCase__       :      Any   					= self.num_choices
			lowerCAmelCase__       :      Any   					= TFDistilBertForMultipleChoice(_SCREAMING_SNAKE_CASE		)
			lowerCAmelCase__       :      Union[str, Any]   					= tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE		, 1		)		, (1, self.num_choices, 1)		)
			lowerCAmelCase__       :      Any   					= tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE		, 1		)		, (1, self.num_choices, 1)		)
			lowerCAmelCase__       :      Optional[int]   					= {
			    '''input_ids''': multiple_choice_inputs_ids,
			    '''attention_mask''': multiple_choice_input_mask,
			}
			lowerCAmelCase__       :      int   					= model(_SCREAMING_SNAKE_CASE		)
			self.parent.assertEqual(result.logits.shape		, (self.batch_size, self.num_choices)		)




	def 							UpperCAmelCase__( self						:		Optional[Any]		, _SCREAMING_SNAKE_CASE						:		Any		, _SCREAMING_SNAKE_CASE						:		Optional[int]		, _SCREAMING_SNAKE_CASE						:		List[Any]		, _SCREAMING_SNAKE_CASE						:		Tuple		, _SCREAMING_SNAKE_CASE						:		int		, _SCREAMING_SNAKE_CASE						:		Optional[int]		)->	str:
			lowerCAmelCase__       :      Tuple   					= self.num_labels
			lowerCAmelCase__       :      Optional[int]   					= TFDistilBertForTokenClassification(_SCREAMING_SNAKE_CASE		)
			lowerCAmelCase__       :      str   					= {'''input_ids''': input_ids, '''attention_mask''': input_mask}
			lowerCAmelCase__       :      Tuple   					= model(_SCREAMING_SNAKE_CASE		)
			self.parent.assertEqual(result.logits.shape		, (self.batch_size, self.seq_length, self.num_labels)		)




	def 							UpperCAmelCase__( self						:		Optional[int]		)->	Any:
			lowerCAmelCase__       :      List[Any]   					= self.prepare_config_and_inputs()
			(lowerCAmelCase__)       :      Tuple   					= config_and_inputs
			lowerCAmelCase__       :      int   					= {'''input_ids''': input_ids, '''attention_mask''': input_mask}
			return config, inputs_dict





@require_tf
class  _a							(    _lowercase					,						_lowercase					,						unittest.TestCase):
	_a  :  Union[str, Any]				   =							(
	    (
	        TFDistilBertModel,
	        TFDistilBertForMaskedLM,
	        TFDistilBertForQuestionAnswering,
	        TFDistilBertForSequenceClassification,
	        TFDistilBertForTokenClassification,
	        TFDistilBertForMultipleChoice,
	    )
	    if is_tf_available()
	    else None
	)
	_a  :  Tuple				   =							(
	    {
	        '''feature-extraction''': TFDistilBertModel,
	        '''fill-mask''': TFDistilBertForMaskedLM,
	        '''question-answering''': TFDistilBertForQuestionAnswering,
	        '''text-classification''': TFDistilBertForSequenceClassification,
	        '''token-classification''': TFDistilBertForTokenClassification,
	        '''zero-shot''': TFDistilBertForSequenceClassification,
	    }
	    if is_tf_available()
	    else {}
	)
	_a  :  Tuple				   =							False
	_a  :  Tuple				   =							False




	def 							UpperCAmelCase__( self						:		List[Any]		)->	str:
			lowerCAmelCase__       :      Optional[int]   					= TFDistilBertModelTester(self		)
			lowerCAmelCase__       :      int   					= ConfigTester(self		, config_class=_SCREAMING_SNAKE_CASE		, dim=37		)




	def 							UpperCAmelCase__( self						:		int		)->	List[str]:
			self.config_tester.run_common_tests()




	def 							UpperCAmelCase__( self						:		str		)->	Union[str, Any]:
			lowerCAmelCase__       :      str   					= self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_distilbert_model(*_SCREAMING_SNAKE_CASE		)




	def 							UpperCAmelCase__( self						:		Optional[Any]		)->	Optional[int]:
			lowerCAmelCase__       :      Optional[Any]   					= self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_distilbert_for_masked_lm(*_SCREAMING_SNAKE_CASE		)




	def 							UpperCAmelCase__( self						:		Optional[int]		)->	Union[str, Any]:
			lowerCAmelCase__       :      Dict   					= self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_distilbert_for_question_answering(*_SCREAMING_SNAKE_CASE		)




	def 							UpperCAmelCase__( self						:		Optional[int]		)->	Optional[Any]:
			lowerCAmelCase__       :      List[Any]   					= self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_distilbert_for_sequence_classification(*_SCREAMING_SNAKE_CASE		)




	def 							UpperCAmelCase__( self						:		Any		)->	Dict:
			lowerCAmelCase__       :      List[Any]   					= self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_distilbert_for_multiple_choice(*_SCREAMING_SNAKE_CASE		)




	def 							UpperCAmelCase__( self						:		str		)->	str:
			lowerCAmelCase__       :      List[str]   					= self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_distilbert_for_token_classification(*_SCREAMING_SNAKE_CASE		)




	@slow
	def 							UpperCAmelCase__( self						:		List[Any]		)->	Optional[int]:
			for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]		):
					lowerCAmelCase__       :      str   					= TFDistilBertModel.from_pretrained(_SCREAMING_SNAKE_CASE		)
					self.assertIsNotNone(_SCREAMING_SNAKE_CASE		)





@require_tf
class  _a							(    unittest.TestCase):




	@slow
	def 							UpperCAmelCase__( self						:		Tuple		)->	Dict:
			lowerCAmelCase__       :      Any   					= TFDistilBertModel.from_pretrained('''distilbert-base-uncased'''		)
			lowerCAmelCase__       :      Tuple   					= tf.constant([[0, 1, 2, 3, 4, 5]]		)
			lowerCAmelCase__       :      str   					= model(_SCREAMING_SNAKE_CASE		)[0]

			lowerCAmelCase__       :      Tuple   					= [1, 6, 768]
			self.assertEqual(output.shape		, _SCREAMING_SNAKE_CASE		)

			lowerCAmelCase__       :      int   					= tf.constant(
			    [
			        [
			            [0.1926_1885, -0.1373_2955, 0.411_9799],
			            [0.2215_0156, -0.0742_2661, 0.3903_7204],
			            [0.2275_6018, -0.089_6414, 0.370_1467],
			        ]
			    ]		)
			tf.debugging.assert_near(output[:, :3, :3]		, _SCREAMING_SNAKE_CASE		, atol=1E-4		)








class  _a							:




	def __init__( self						:		str		, _SCREAMING_SNAKE_CASE						:		str		, _SCREAMING_SNAKE_CASE						:		Optional[int]		, _SCREAMING_SNAKE_CASE						:		List[Any]		)->	Tuple:
			lowerCAmelCase__       :      List[Any]   					= None
			lowerCAmelCase__       :      int   					= None
			lowerCAmelCase__       :      Union[str, Any]   					= graph

			self._normalize_graph(_SCREAMING_SNAKE_CASE		, _SCREAMING_SNAKE_CASE		)
			lowerCAmelCase__       :      Tuple   					= len(_SCREAMING_SNAKE_CASE		)
			lowerCAmelCase__       :      Optional[int]   					= None




	def 							UpperCAmelCase__( self						:		Union[str, Any]		, _SCREAMING_SNAKE_CASE						:		Optional[Any]		, _SCREAMING_SNAKE_CASE						:		Union[str, Any]		)->	Union[str, Any]:
			if sources is int:
					lowerCAmelCase__       :      List[Any]   					= [sources]
			if sinks is int:
					lowerCAmelCase__       :      Optional[Any]   					= [sinks]

			if len(_SCREAMING_SNAKE_CASE		) == 0 or len(_SCREAMING_SNAKE_CASE		) == 0:
					return

			lowerCAmelCase__       :      Union[str, Any]   					= sources[0]
			lowerCAmelCase__       :      Dict   					= sinks[0]

			# make fake vertex if there are more
			# than one source or sink
			if len(_SCREAMING_SNAKE_CASE		) > 1 or len(_SCREAMING_SNAKE_CASE		) > 1:
					lowerCAmelCase__       :      List[Any]   					= 0
					for i in sources:
							max_input_flow += sum(self.graph[i]		)

					lowerCAmelCase__       :      Any   					= len(self.graph		) + 1
					for room in self.graph:
							room.insert(0		, 0		)
					self.graph.insert(0		, [0] * size		)
					for i in sources:
							lowerCAmelCase__       :      Optional[Any]   					= max_input_flow
					lowerCAmelCase__       :      List[str]   					= 0

					lowerCAmelCase__       :      List[Any]   					= len(self.graph		) + 1
					for room in self.graph:
							room.append(0		)
					self.graph.append([0] * size		)
					for i in sinks:
							lowerCAmelCase__       :      str   					= max_input_flow
					lowerCAmelCase__       :      Tuple   					= size - 1




	def 							UpperCAmelCase__( self						:		Union[str, Any]		)->	Union[str, Any]:
			if self.maximum_flow_algorithm is None:
					raise Exception('''You need to set maximum flow algorithm before.'''		)
			if self.source_index is None or self.sink_index is None:
					return 0

			self.maximum_flow_algorithm.execute()
			return self.maximum_flow_algorithm.getMaximumFlow()




	def 							UpperCAmelCase__( self						:		str		, _SCREAMING_SNAKE_CASE						:		List[Any]		)->	int:
			lowerCAmelCase__       :      int   					= algorithm(self		)





class  _a							:




	def __init__( self						:		Optional[Any]		, _SCREAMING_SNAKE_CASE						:		Tuple		)->	Union[str, Any]:
			lowerCAmelCase__       :      Tuple   					= flow_network
			lowerCAmelCase__       :      Dict   					= flow_network.verticesCount
			lowerCAmelCase__       :      Optional[Any]   					= flow_network.sourceIndex
			lowerCAmelCase__       :      Optional[Any]   					= flow_network.sinkIndex
			# it's just a reference, so you shouldn't change
			# it in your algorithms, use deep copy before doing that
			lowerCAmelCase__       :      str   					= flow_network.graph
			lowerCAmelCase__       :      Optional[int]   					= False




	def 							UpperCAmelCase__( self						:		List[str]		)->	Dict:
			if not self.executed:
					self._algorithm()
					lowerCAmelCase__       :      Any   					= True




	def 							UpperCAmelCase__( self						:		Optional[Any]		)->	int:
			pass





class  _a							(    _lowercase):




	def __init__( self						:		Any		, _SCREAMING_SNAKE_CASE						:		List[Any]		)->	Union[str, Any]:
			super().__init__(_SCREAMING_SNAKE_CASE		)
			# use this to save your result
			lowerCAmelCase__       :      Dict   					= -1




	def 							UpperCAmelCase__( self						:		Any		)->	Optional[Any]:
			if not self.executed:
					raise Exception('''You should execute algorithm before using its result!'''		)

			return self.maximum_flow





class  _a							(    _lowercase):




	def __init__( self						:		Any		, _SCREAMING_SNAKE_CASE						:		List[str]		)->	List[str]:
			super().__init__(_SCREAMING_SNAKE_CASE		)

			lowerCAmelCase__       :      Any   					= [[0] * self.verticies_count for i in range(self.verticies_count		)]

			lowerCAmelCase__       :      Optional[Any]   					= [0] * self.verticies_count
			lowerCAmelCase__       :      str   					= [0] * self.verticies_count




	def 							UpperCAmelCase__( self						:		Any		)->	List[Any]:
			lowerCAmelCase__       :      Optional[Any]   					= self.verticies_count

			# push some substance to graph
			for nextvertex_index, bandwidth in enumerate(self.graph[self.source_index]		):
					self.preflow[self.source_index][nextvertex_index] += bandwidth
					self.preflow[nextvertex_index][self.source_index] -= bandwidth
					self.excesses[nextvertex_index] += bandwidth

			# Relabel-to-front selection rule
			lowerCAmelCase__       :      Any   					= [
			    i
			    for i in range(self.verticies_count		)
			    if i != self.source_index and i != self.sink_index
			]

			# move through list
			lowerCAmelCase__       :      str   					= 0
			while i < len(_SCREAMING_SNAKE_CASE		):
					lowerCAmelCase__       :      Union[str, Any]   					= vertices_list[i]
					lowerCAmelCase__       :      Any   					= self.heights[vertex_index]
					self.process_vertex(_SCREAMING_SNAKE_CASE		)
					if self.heights[vertex_index] > previous_height:
							# if it was relabeled, swap elements
							# and start from 0 index
							vertices_list.insert(0		, vertices_list.pop(_SCREAMING_SNAKE_CASE		)		)
							lowerCAmelCase__       :      Optional[Any]   					= 0
					else:
							i += 1

			lowerCAmelCase__       :      Optional[Any]   					= sum(self.preflow[self.source_index]		)




	def 							UpperCAmelCase__( self						:		List[str]		, _SCREAMING_SNAKE_CASE						:		Optional[Any]		)->	Optional[int]:
			while self.excesses[vertex_index] > 0:
					for neighbour_index in range(self.verticies_count		):
							# if it's neighbour and current vertex is higher
							if (
							    self.graph[vertex_index][neighbour_index]
							    - self.preflow[vertex_index][neighbour_index]
							    > 0
							    and self.heights[vertex_index] > self.heights[neighbour_index]
							):
									self.push(_SCREAMING_SNAKE_CASE		, _SCREAMING_SNAKE_CASE		)

					self.relabel(_SCREAMING_SNAKE_CASE		)




	def 							UpperCAmelCase__( self						:		int		, _SCREAMING_SNAKE_CASE						:		Optional[Any]		, _SCREAMING_SNAKE_CASE						:		Union[str, Any]		)->	Union[str, Any]:
			lowerCAmelCase__       :      Union[str, Any]   					= min(
			    self.excesses[from_index]		, self.graph[from_index][to_index] - self.preflow[from_index][to_index]		, )
			self.preflow[from_index][to_index] += preflow_delta
			self.preflow[to_index][from_index] -= preflow_delta
			self.excesses[from_index] -= preflow_delta
			self.excesses[to_index] += preflow_delta




	def 							UpperCAmelCase__( self						:		Optional[int]		, _SCREAMING_SNAKE_CASE						:		Any		)->	Optional[int]:
			lowerCAmelCase__       :      Optional[int]   					= None
			for to_index in range(self.verticies_count		):
					if (
					    self.graph[vertex_index][to_index]
					    - self.preflow[vertex_index][to_index]
					    > 0
					) and (min_height is None or self.heights[to_index] < min_height):
							lowerCAmelCase__       :      List[Any]   					= self.heights[to_index]

			if min_height is not None:
					lowerCAmelCase__       :      Optional[Any]   					= min_height + 1


if __name__ == "__main__":
		lowerCamelCase						=		[0]
		lowerCamelCase						=		[3]
		# graph = [
		#     [0, 0, 4, 6, 0, 0],
		#     [0, 0, 5, 2, 0, 0],
		#     [0, 0, 0, 0, 4, 4],
		#     [0, 0, 0, 0, 6, 6],
		#     [0, 0, 0, 0, 0, 0],
		#     [0, 0, 0, 0, 0, 0],
		# ]
		lowerCamelCase						=		[[0, 7, 0, 0], [0, 0, 6, 0], [0, 0, 0, 8], [9, 0, 0, 0]]

		# prepare our network
		lowerCamelCase						=		FlowNetwork(graph, entrances, exits)
		# set algorithm
		flow_network.set_maximum_flow_algorithm(PushRelabelExecutor)
		# and calculate
		lowerCamelCase						=		flow_network.find_maximum_flow()

		print(f'''maximum flow is {maximum_flow}''')




