






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   lowerCamelCase      (_SCREAMING_SNAKE_CASE ,  _SCREAMING_SNAKE_CASE ,  unittest.TestCase			):







		'''simple docstring'''



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



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase=False   )	->							Dict:
							UpperCAmelCase_    :      Union[str, Any]            =						super()._prepare_for_class(_lowerCAmelCase				,		_lowerCAmelCase				,		return_labels=_lowerCAmelCase   )

							if return_labels:
												if model_class in get_values(_lowerCAmelCase   ):
																	UpperCAmelCase_    :      int            =						tf.zeros(self.model_tester.batch_size				,		dtype=tf.intaa   )

							return inputs_dict



		class   lowerCamelCase      (_SCREAMING_SNAKE_CASE			):







				'''simple docstring'''



				def __init__(							self				,		_UpperCamelCase				,		_UpperCamelCase=1_3				,		_UpperCamelCase=7				,		_UpperCamelCase=True				,		_UpperCamelCase=True				,		_UpperCamelCase=True				,		_UpperCamelCase=True				,		_UpperCamelCase=9_9				,		_UpperCamelCase=3_2				,		_UpperCamelCase=3_2				,		_UpperCamelCase=2				,		_UpperCamelCase=4				,		_UpperCamelCase=3_7				,		_UpperCamelCase="gelu"				,		_UpperCamelCase=0.1				,		_UpperCamelCase=0.1				,		_UpperCamelCase=5_1_2				,		_UpperCamelCase=1_6				,		_UpperCamelCase=2				,		_UpperCamelCase=0.02				,		_UpperCamelCase=3				,		_UpperCamelCase=4				,		_UpperCamelCase=None				,		)	->							Optional[Any]:
									UpperCAmelCase_    :      List[Any]            =						parent
									UpperCAmelCase_    :      Tuple            =						batch_size
									UpperCAmelCase_    :      int            =						seq_length
									UpperCAmelCase_    :      Tuple            =						is_training
									UpperCAmelCase_    :      Tuple            =						use_input_mask
									UpperCAmelCase_    :      Union[str, Any]            =						use_token_type_ids
									UpperCAmelCase_    :      List[Any]            =						use_labels
									UpperCAmelCase_    :      List[Any]            =						vocab_size
									UpperCAmelCase_    :      Optional[Any]            =						hidden_size
									UpperCAmelCase_    :      Optional[int]            =						num_hidden_layers
									UpperCAmelCase_    :      List[Any]            =						num_attention_heads
									UpperCAmelCase_    :      Union[str, Any]            =						intermediate_size
									UpperCAmelCase_    :      Optional[int]            =						hidden_act
									UpperCAmelCase_    :      List[Any]            =						hidden_dropout_prob
									UpperCAmelCase_    :      List[Any]            =						attention_probs_dropout_prob
									UpperCAmelCase_    :      Optional[Any]            =						max_position_embeddings
									UpperCAmelCase_    :      Dict            =						type_vocab_size
									UpperCAmelCase_    :      Optional[int]            =						type_sequence_label_size
									UpperCAmelCase_    :      Union[str, Any]            =						initializer_range
									UpperCAmelCase_    :      List[str]            =						num_labels
									UpperCAmelCase_    :      str            =						num_choices
									UpperCAmelCase_    :      Tuple            =						scope
									UpperCAmelCase_    :      List[str]            =						embedding_size



				def 					__UpperCAmelCase			(							self   )	->							List[str]:
									UpperCAmelCase_    :      Any            =						ids_tensor([self.batch_size, self.seq_length]				,		self.vocab_size   )

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

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

									UpperCAmelCase_    :      Tuple            =						None
									UpperCAmelCase_    :      Tuple            =						None
									UpperCAmelCase_    :      int            =						None
									if self.use_labels:
														UpperCAmelCase_    :      Tuple            =						ids_tensor([self.batch_size]				,		self.type_sequence_label_size   )
														UpperCAmelCase_    :      Optional[int]            =						ids_tensor([self.batch_size, self.seq_length]				,		self.num_labels   )
														UpperCAmelCase_    :      Optional[Any]            =						ids_tensor([self.batch_size]				,		self.num_choices   )

									UpperCAmelCase_    :      Union[str, Any]            =						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 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							List[str]:
									UpperCAmelCase_    :      List[str]            =						TFMobileBertModel(config=_lowerCAmelCase   )
									UpperCAmelCase_    :      Any            =						{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
									UpperCAmelCase_    :      List[Any]            =						model(_lowerCAmelCase   )

									UpperCAmelCase_    :      Tuple            =						[input_ids, input_mask]
									UpperCAmelCase_    :      Optional[Any]            =						model(_lowerCAmelCase   )

									UpperCAmelCase_    :      List[Any]            =						model(_lowerCAmelCase   )

									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 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							Optional[int]:
									UpperCAmelCase_    :      Tuple            =						TFMobileBertForMaskedLM(config=_lowerCAmelCase   )
									UpperCAmelCase_    :      str            =						{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
									UpperCAmelCase_    :      Optional[Any]            =						model(_lowerCAmelCase   )
									self.parent.assertEqual(result.logits.shape				,		(self.batch_size, self.seq_length, self.vocab_size)   )



				def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							Tuple:
									UpperCAmelCase_    :      Union[str, Any]            =						TFMobileBertForNextSentencePrediction(config=_lowerCAmelCase   )
									UpperCAmelCase_    :      List[str]            =						{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
									UpperCAmelCase_    :      List[Any]            =						model(_lowerCAmelCase   )
									self.parent.assertEqual(result.logits.shape				,		(self.batch_size, 2)   )



				def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							List[Any]:
									UpperCAmelCase_    :      Optional[Any]            =						TFMobileBertForPreTraining(config=_lowerCAmelCase   )
									UpperCAmelCase_    :      Tuple            =						{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
									UpperCAmelCase_    :      Any            =						model(_lowerCAmelCase   )
									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 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							int:
									UpperCAmelCase_    :      Dict            =						self.num_labels
									UpperCAmelCase_    :      List[Any]            =						TFMobileBertForSequenceClassification(config=_lowerCAmelCase   )
									UpperCAmelCase_    :      Optional[int]            =						{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
									UpperCAmelCase_    :      List[str]            =						model(_lowerCAmelCase   )
									self.parent.assertEqual(result.logits.shape				,		(self.batch_size, self.num_labels)   )



				def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							List[str]:
									UpperCAmelCase_    :      Dict            =						self.num_choices
									UpperCAmelCase_    :      Union[str, Any]            =						TFMobileBertForMultipleChoice(config=_lowerCAmelCase   )
									UpperCAmelCase_    :      Any            =						tf.tile(tf.expand_dims(_lowerCAmelCase				,		1   )				,		(1, self.num_choices, 1)   )
									UpperCAmelCase_    :      Optional[int]            =						tf.tile(tf.expand_dims(_lowerCAmelCase				,		1   )				,		(1, self.num_choices, 1)   )
									UpperCAmelCase_    :      str            =						tf.tile(tf.expand_dims(_lowerCAmelCase				,		1   )				,		(1, self.num_choices, 1)   )
									UpperCAmelCase_    :      Dict            =						{
									    'input_ids': multiple_choice_inputs_ids,
									    'attention_mask': multiple_choice_input_mask,
									    'token_type_ids': multiple_choice_token_type_ids,
									}
									UpperCAmelCase_    :      Dict            =						model(_lowerCAmelCase   )
									self.parent.assertEqual(result.logits.shape				,		(self.batch_size, self.num_choices)   )



				def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							Union[str, Any]:
									UpperCAmelCase_    :      Any            =						self.num_labels
									UpperCAmelCase_    :      Union[str, Any]            =						TFMobileBertForTokenClassification(config=_lowerCAmelCase   )
									UpperCAmelCase_    :      Union[str, Any]            =						{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
									UpperCAmelCase_    :      List[str]            =						model(_lowerCAmelCase   )
									self.parent.assertEqual(result.logits.shape				,		(self.batch_size, self.seq_length, self.num_labels)   )



				def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							List[str]:
									UpperCAmelCase_    :      int            =						TFMobileBertForQuestionAnswering(config=_lowerCAmelCase   )
									UpperCAmelCase_    :      Optional[Any]            =						{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
									UpperCAmelCase_    :      Any            =						model(_lowerCAmelCase   )
									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   )	->							str:
									UpperCAmelCase_    :      str            =						self.prepare_config_and_inputs()
									(
									    (
									    UpperCAmelCase_ 
									)       ,  (
									    UpperCAmelCase_ 
									)       ,  (
									    UpperCAmelCase_ 
									)       ,  (
									    UpperCAmelCase_ 
									)       ,  (
									    UpperCAmelCase_ 
									)       ,  (
									    UpperCAmelCase_ 
									)       ,  (
									    UpperCAmelCase_ 
									)       ,   
									)    :      str            =						config_and_inputs
									UpperCAmelCase_    :      Optional[int]            =						{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
									return config, inputs_dict





		def 					__UpperCAmelCase			(							self   )	->							Union[str, Any]:
							UpperCAmelCase_    :      Any            =						TFMobileBertModelTest.TFMobileBertModelTester(self   )
							UpperCAmelCase_    :      str            =						ConfigTester(self				,		config_class=_lowerCAmelCase				,		hidden_size=3_7   )



		def 					__UpperCAmelCase			(							self   )	->							Union[str, Any]:
							self.config_tester.run_common_tests()



		def 					__UpperCAmelCase			(							self   )	->							Union[str, Any]:
							UpperCAmelCase_    :      List[str]            =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mobilebert_model(*_lowerCAmelCase   )



		def 					__UpperCAmelCase			(							self   )	->							int:
							UpperCAmelCase_    :      Any            =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mobilebert_for_masked_lm(*_lowerCAmelCase   )



		def 					__UpperCAmelCase			(							self   )	->							List[str]:
							UpperCAmelCase_    :      Any            =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mobilebert_for_multiple_choice(*_lowerCAmelCase   )



		def 					__UpperCAmelCase			(							self   )	->							Tuple:
							UpperCAmelCase_    :      Tuple            =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*_lowerCAmelCase   )



		def 					__UpperCAmelCase			(							self   )	->							List[Any]:
							UpperCAmelCase_    :      Optional[Any]            =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mobilebert_for_pretraining(*_lowerCAmelCase   )



		def 					__UpperCAmelCase			(							self   )	->							List[str]:
							UpperCAmelCase_    :      Tuple            =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mobilebert_for_question_answering(*_lowerCAmelCase   )



		def 					__UpperCAmelCase			(							self   )	->							int:
							UpperCAmelCase_    :      int            =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mobilebert_for_sequence_classification(*_lowerCAmelCase   )



		def 					__UpperCAmelCase			(							self   )	->							List[str]:
							UpperCAmelCase_    :      List[str]            =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mobilebert_for_token_classification(*_lowerCAmelCase   )



		@slow
		def 					__UpperCAmelCase			(							self   )	->							Dict:
							# for model_name in TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							for model_name in ["google/mobilebert-uncased"]:
												UpperCAmelCase_    :      Optional[Any]            =						TFMobileBertModel.from_pretrained(_lowerCAmelCase   )
												self.assertIsNotNone(_lowerCAmelCase   )





@require_tf
class   lowerCamelCase      (unittest.TestCase			):







		'''simple docstring'''



		@slow
		def 					__UpperCAmelCase			(							self   )	->							str:
							UpperCAmelCase_    :      Optional[Any]            =						TFMobileBertForPreTraining.from_pretrained('google/mobilebert-uncased'   )
							UpperCAmelCase_    :      str            =						tf.constant([[0, 1, 2, 3, 4, 5]]   )
							UpperCAmelCase_    :      List[Any]            =						model(_lowerCAmelCase   )[0]

							UpperCAmelCase_    :      Optional[int]            =						[1, 6, 3_0_5_2_2]
							self.assertEqual(output.shape				,		_lowerCAmelCase   )

							UpperCAmelCase_    :      Optional[Any]            =						tf.constant(
							    [
							        [
							            [-4.5_91_95_47, -9.24_82_95, -9.64_52_56],
							            [-6.7_30_61_75, -6.44_02_84, -6.6_05_28_37],
							            [-7.2_74_35_06, -6.7_84_79_15, -6.02_46_73],
							        ]
							    ]   )
							tf.debugging.assert_near(output[:, :3, :3]				,		_lowerCAmelCase				,		atol=1E-4   )












from math import sqrt

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	int							)					-> bool:

      if 1 < number < 4:
            # 2 and 3 are primes
            return True
      elif number < 2 or number % 2 == 0 or number % 3 == 0:
            # Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
            return False

      # All primes number are in format of 6k +/- 1
      for i in range(5			,			int(sqrt(__UpperCAmelCase							) + 1							)			,			6							):
            if number % i == 0 or number % (i + 2) == 0:
                  return False
      return True

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	int = 1_00_01							)					-> int:
      SCREAMING_SNAKE_CASE_							=			0
      SCREAMING_SNAKE_CASE_							=			1
      while count != nth and number < 3:
            number += 1
            if is_prime(__UpperCAmelCase							):
                  count += 1
      while count != nth:
            number += 2
            if is_prime(__UpperCAmelCase							):
                  count += 1
      return number


if __name__ == "__main__":
   print(f'''{solution() = }''')