


"""simple docstring"""


from __future__ import annotations

import unittest

from transformers import MobileBertConfig, is_tf_available
from transformers.models.auto import get_values
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 import (
	    TF_MODEL_FOR_PRETRAINING_MAPPING,
	    TFMobileBertForMaskedLM,
	    TFMobileBertForMultipleChoice,
	    TFMobileBertForNextSentencePrediction,
	    TFMobileBertForPreTraining,
	    TFMobileBertForQuestionAnswering,
	    TFMobileBertForSequenceClassification,
	    TFMobileBertForTokenClassification,
	    TFMobileBertModel,
	)
@require_tf
class     _UpperCamelCase    (  lowerCAmelCase__			,lowerCAmelCase__			,unittest.TestCase						):
							'''simple docstring'''







							__UpperCAmelCase		:       Optional[int]							 =(
							    (
							        TFMobileBertModel,
							        TFMobileBertForMaskedLM,
							        TFMobileBertForNextSentencePrediction,
							        TFMobileBertForPreTraining,
							        TFMobileBertForQuestionAnswering,
							        TFMobileBertForSequenceClassification,
							        TFMobileBertForTokenClassification,
							        TFMobileBertForMultipleChoice,
							    )
							    if is_tf_available()
							    else ()
							)
							__UpperCAmelCase		:       Optional[int]							 =(
							    {
							        """feature-extraction""": TFMobileBertModel,
							        """fill-mask""": TFMobileBertForMaskedLM,
							        """question-answering""": TFMobileBertForQuestionAnswering,
							        """text-classification""": TFMobileBertForSequenceClassification,
							        """token-classification""": TFMobileBertForTokenClassification,
							        """zero-shot""": TFMobileBertForSequenceClassification,
							    }
							    if is_tf_available()
							    else {}
							)
							__UpperCAmelCase		:       Tuple							 =False
							__UpperCAmelCase		:       Any							 =False

							def      snake_case							(	self       ,  __a       ,  __a       ,  __a=False      ):
										__lowerCAmelCase							       =     super()._prepare_for_class(__a       ,  __a       ,  return_labels=__a      )

										if return_labels:
													if model_class in get_values(__a      ):
																__lowerCAmelCase							       =     tf.zeros(self.model_tester.batch_size       ,  dtype=tf.intaa      )

										return inputs_dict
							class     _UpperCamelCase    (  lowerCAmelCase__						):
														'''simple docstring'''







														def __init__(	self       ,  __a       ,  __a=13       ,  __a=7       ,  __a=True       ,  __a=True       ,  __a=True       ,  __a=True       ,  __a=99       ,  __a=32       ,  __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.0_2       ,  __a=3       ,  __a=4       ,  __a=None       ,  ):
																	__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							       =     embedding_size

														def      snake_case							(	self      ):
																	__lowerCAmelCase							       =     ids_tensor([self.batch_size, self.seq_length]       ,  self.vocab_size      )

																	__lowerCAmelCase							       =     None
																	if self.use_input_mask:
																				__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							       =     MobileBertConfig(
																	    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       ,  initializer_range=self.initializer_range       ,  embedding_size=self.embedding_size       ,  )

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

														def      snake_case							(	self       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a      ):
																	__lowerCAmelCase							       =     TFMobileBertModel(config=__a      )
																	__lowerCAmelCase							       =     {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
																	__lowerCAmelCase							       =     model(__a      )

																	__lowerCAmelCase							       =     [input_ids, input_mask]
																	__lowerCAmelCase							       =     model(__a      )

																	__lowerCAmelCase							       =     model(__a      )

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

														def      snake_case							(	self       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a      ):
																	__lowerCAmelCase							       =     TFMobileBertForMaskedLM(config=__a      )
																	__lowerCAmelCase							       =     {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
																	__lowerCAmelCase							       =     model(__a      )
																	self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.seq_length, self.vocab_size)      )

														def      snake_case							(	self       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a      ):
																	__lowerCAmelCase							       =     TFMobileBertForNextSentencePrediction(config=__a      )
																	__lowerCAmelCase							       =     {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
																	__lowerCAmelCase							       =     model(__a      )
																	self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, 2)      )

														def      snake_case							(	self       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a      ):
																	__lowerCAmelCase							       =     TFMobileBertForPreTraining(config=__a      )
																	__lowerCAmelCase							       =     {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
																	__lowerCAmelCase							       =     model(__a      )
																	self.parent.assertEqual(
																	    result.prediction_logits.shape       ,  (self.batch_size, self.seq_length, self.vocab_size)      )
																	self.parent.assertEqual(result.seq_relationship_logits.shape       ,  (self.batch_size, 2)      )

														def      snake_case							(	self       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a      ):
																	__lowerCAmelCase							       =     self.num_labels
																	__lowerCAmelCase							       =     TFMobileBertForSequenceClassification(config=__a      )
																	__lowerCAmelCase							       =     {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
																	__lowerCAmelCase							       =     model(__a      )
																	self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.num_labels)      )

														def      snake_case							(	self       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a      ):
																	__lowerCAmelCase							       =     self.num_choices
																	__lowerCAmelCase							       =     TFMobileBertForMultipleChoice(config=__a      )
																	__lowerCAmelCase							       =     tf.tile(tf.expand_dims(__a       ,  1      )       ,  (1, self.num_choices, 1)      )
																	__lowerCAmelCase							       =     tf.tile(tf.expand_dims(__a       ,  1      )       ,  (1, self.num_choices, 1)      )
																	__lowerCAmelCase							       =     tf.tile(tf.expand_dims(__a       ,  1      )       ,  (1, self.num_choices, 1)      )
																	__lowerCAmelCase							       =     {
																	    "input_ids": multiple_choice_inputs_ids,
																	    "attention_mask": multiple_choice_input_mask,
																	    "token_type_ids": multiple_choice_token_type_ids,
																	}
																	__lowerCAmelCase							       =     model(__a      )
																	self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.num_choices)      )

														def      snake_case							(	self       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a      ):
																	__lowerCAmelCase							       =     self.num_labels
																	__lowerCAmelCase							       =     TFMobileBertForTokenClassification(config=__a      )
																	__lowerCAmelCase							       =     {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
																	__lowerCAmelCase							       =     model(__a      )
																	self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.seq_length, self.num_labels)      )

														def      snake_case							(	self       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a      ):
																	__lowerCAmelCase							       =     TFMobileBertForQuestionAnswering(config=__a      )
																	__lowerCAmelCase							       =     {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
																	__lowerCAmelCase							       =     model(__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)      )



														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, "token_type_ids": token_type_ids, "attention_mask": input_mask}
																	return config, inputs_dict

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     TFMobileBertModelTest.TFMobileBertModelTester(self      )
										__lowerCAmelCase							       =     ConfigTester(self       ,  config_class=__a       ,  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_mobilebert_model(*__a      )

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_mobilebert_for_masked_lm(*__a      )

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_mobilebert_for_multiple_choice(*__a      )

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*__a      )

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_mobilebert_for_pretraining(*__a      )

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_mobilebert_for_question_answering(*__a      )

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_mobilebert_for_sequence_classification(*__a      )

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_mobilebert_for_token_classification(*__a      )



							@slow
							def      snake_case							(	self      ):
										# for model_name in TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										for model_name in ["google/mobilebert-uncased"]:
													__lowerCAmelCase							       =     TFMobileBertModel.from_pretrained(__a      )
													self.assertIsNotNone(__a      )
@require_tf
class     _UpperCamelCase    (  unittest.TestCase						):
							'''simple docstring'''







							@slow
							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     TFMobileBertForPreTraining.from_pretrained("google/mobilebert-uncased"      )
										__lowerCAmelCase							       =     tf.constant([[0, 1, 2, 3, 4, 5]]      )
										__lowerCAmelCase							       =     model(__a      )[0]

										__lowerCAmelCase							       =     [1, 6, 3_05_22]
										self.assertEqual(output.shape       ,  __a      )

										__lowerCAmelCase							       =     tf.constant(
										    [
										        [
										            [-4.5_9_1_9_5_4_7, -9.2_4_8_2_9_5, -9.6_4_5_2_5_6],
										            [-6.7_3_0_6_1_7_5, -6.4_4_0_2_8_4, -6.6_0_5_2_8_3_7],
										            [-7.2_7_4_3_5_0_6, -6.7_8_4_7_9_1_5, -6.0_2_4_6_7_3],
										        ]
										    ]      )
										tf.debugging.assert_near(output[:, :3, :3]       ,  __a       ,  atol=1e-4      )








"""simple docstring"""


import heapq as hq
import math
from collections.abc import Iterator
class     _UpperCamelCase    :
							'''simple docstring'''







							def __init__(	self       ,  __a      ):
										__lowerCAmelCase							       =     str(id_      )
										__lowerCAmelCase							       =     None
										__lowerCAmelCase							       =     None
										__lowerCAmelCase							       =     []
										__lowerCAmelCase							       =     {}  # {vertex:distance}

							def __lt__(	self       ,  __a      ):
										return self.key < other.key

							def __repr__(	self      ):
										return self.id

							def      snake_case							(	self       ,  __a      ):
										self.neighbors.append(__a      )



							def      snake_case							(	self       ,  __a       ,  __a      ):
										__lowerCAmelCase							       =     weight


def _lowerCamelCase      (						_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase ):

			'''simple docstring'''




			graph[a - 1].add_neighbor(graph[b - 1] )
			graph[b - 1].add_neighbor(graph[a - 1] )
			# add the edges:
			graph[a - 1].add_edge(graph[b - 1]       ,					_UpperCamelCase )
			graph[b - 1].add_edge(graph[a - 1]       ,					_UpperCamelCase )


def _lowerCamelCase      (						_UpperCamelCase       ,					_UpperCamelCase ):

			'''simple docstring'''




			__lowerCAmelCase							       =     []
			for u in graph:
						__lowerCAmelCase							       =     math.inf
						__lowerCAmelCase							       =     None
			__lowerCAmelCase							       =     0
			__lowerCAmelCase							       =     graph[:]
			while q:
						__lowerCAmelCase							       =     min(_UpperCamelCase )
						q.remove(_UpperCamelCase )
						for v in u.neighbors:
									if (v in q) and (u.edges[v.id] < v.key):
												__lowerCAmelCase							       =     u
												__lowerCAmelCase							       =     u.edges[v.id]
			for i in range(1       ,					len(_UpperCamelCase ) ):
						a.append((int(graph[i].id ) + 1, int(graph[i].pi.id ) + 1) )
			return a


def _lowerCamelCase      (						_UpperCamelCase       ,					_UpperCamelCase ):

			'''simple docstring'''




			for u in graph:
						__lowerCAmelCase							       =     math.inf
						__lowerCAmelCase							       =     None
			__lowerCAmelCase							       =     0

			__lowerCAmelCase							       =     list(_UpperCamelCase )
			hq.heapify(_UpperCamelCase )

			while h:
						__lowerCAmelCase							       =     hq.heappop(_UpperCamelCase )
						for v in u.neighbors:
									if (v in h) and (u.edges[v.id] < v.key):
												__lowerCAmelCase							       =     u
												__lowerCAmelCase							       =     u.edges[v.id]
												hq.heapify(_UpperCamelCase )

			for i in range(1       ,					len(_UpperCamelCase ) ):
						yield (int(graph[i].id ) + 1, int(graph[i].pi.id ) + 1)


def _lowerCamelCase      (						):

			'''simple docstring'''






if __name__ == "__main__":
	import doctest

	doctest.testmod()




