

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 				SCREAMING_SNAKE_CASE__ (__lowerCAmelCase		,	__lowerCAmelCase		,	unittest.TestCase					):
    __lowerCamelCase						:	Any														=					(
        (
            TFMobileBertModel,
            TFMobileBertForMaskedLM,
            TFMobileBertForNextSentencePrediction,
            TFMobileBertForPreTraining,
            TFMobileBertForQuestionAnswering,
            TFMobileBertForSequenceClassification,
            TFMobileBertForTokenClassification,
            TFMobileBertForMultipleChoice,
        )
        if is_tf_available()
        else ()
    )
    __lowerCamelCase						:	Dict														=					(
        {
            '''feature-extraction''': TFMobileBertModel,
            '''fill-mask''': TFMobileBertForMaskedLM,
            '''question-answering''': TFMobileBertForQuestionAnswering,
            '''text-classification''': TFMobileBertForSequenceClassification,
            '''token-classification''': TFMobileBertForTokenClassification,
            '''zero-shot''': TFMobileBertForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )
    __lowerCamelCase						:	int														=					False
    __lowerCamelCase						:	List[str]														=					False





    def 				snake_case_    (				self						,       a						,       a						,       a=False):
         lowercase__   :       Union[str, Any]				=	super()._prepare_for_class(lowerCAmelCase_						,       lowerCAmelCase_						,       return_labels=lowerCAmelCase_)

         if return_labels:
              if model_class in get_values(lowerCAmelCase_):
                   lowercase__   :       Union[str, Any]				=	tf.zeros(self.model_tester.batch_size						,       dtype=tf.intaa)

         return inputs_dict


    class 				SCREAMING_SNAKE_CASE__ (__lowerCAmelCase					):





        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=512						,       a=16						,       a=2						,       a=0.02						,       a=3						,       a=4						,       a=None						,       ):
             lowercase__   :       Tuple				=	parent
             lowercase__   :       int				=	batch_size
             lowercase__   :       Union[str, Any]				=	seq_length
             lowercase__   :       List[Any]				=	is_training
             lowercase__   :       str				=	use_input_mask
             lowercase__   :       str				=	use_token_type_ids
             lowercase__   :       List[str]				=	use_labels
             lowercase__   :       List[Any]				=	vocab_size
             lowercase__   :       Dict				=	hidden_size
             lowercase__   :       Any				=	num_hidden_layers
             lowercase__   :       Optional[int]				=	num_attention_heads
             lowercase__   :       Optional[Any]				=	intermediate_size
             lowercase__   :       Optional[int]				=	hidden_act
             lowercase__   :       List[Any]				=	hidden_dropout_prob
             lowercase__   :       Optional[Any]				=	attention_probs_dropout_prob
             lowercase__   :       Union[str, Any]				=	max_position_embeddings
             lowercase__   :       Any				=	type_vocab_size
             lowercase__   :       Tuple				=	type_sequence_label_size
             lowercase__   :       Union[str, Any]				=	initializer_range
             lowercase__   :       int				=	num_labels
             lowercase__   :       List[str]				=	num_choices
             lowercase__   :       int				=	scope
             lowercase__   :       List[Any]				=	embedding_size





        def 				snake_case_    (				self):
             lowercase__   :       Optional[int]				=	ids_tensor([self.batch_size, self.seq_length]						,       self.vocab_size)

             lowercase__   :       Union[str, Any]				=	None
             if self.use_input_mask:
                  lowercase__   :       Optional[Any]				=	random_attention_mask([self.batch_size, self.seq_length])

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

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

             lowercase__   :       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 				snake_case_    (				self						,       a						,       a						,       a						,       a						,       a						,       a						,       a):
             lowercase__   :       List[Any]				=	TFMobileBertModel(config=lowerCAmelCase_)
             lowercase__   :       List[str]				=	{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
             lowercase__   :       Union[str, Any]				=	model(lowerCAmelCase_)

             lowercase__   :       Optional[Any]				=	[input_ids, input_mask]
             lowercase__   :       str				=	model(lowerCAmelCase_)

             lowercase__   :       str				=	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 				snake_case_    (				self						,       a						,       a						,       a						,       a						,       a						,       a						,       a):
             lowercase__   :       Any				=	TFMobileBertForMaskedLM(config=lowerCAmelCase_)
             lowercase__   :       str				=	{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
             lowercase__   :       Tuple				=	model(lowerCAmelCase_)
             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):
             lowercase__   :       Optional[Any]				=	TFMobileBertForNextSentencePrediction(config=lowerCAmelCase_)
             lowercase__   :       Any				=	{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
             lowercase__   :       Tuple				=	model(lowerCAmelCase_)
             self.parent.assertEqual(result.logits.shape						,       (self.batch_size, 2))





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





        def 				snake_case_    (				self						,       a						,       a						,       a						,       a						,       a						,       a						,       a):
             lowercase__   :       Tuple				=	self.num_choices
             lowercase__   :       Dict				=	TFMobileBertForMultipleChoice(config=lowerCAmelCase_)
             lowercase__   :       List[str]				=	tf.tile(tf.expand_dims(lowerCAmelCase_						,       1)						,       (1, self.num_choices, 1))
             lowercase__   :       Optional[Any]				=	tf.tile(tf.expand_dims(lowerCAmelCase_						,       1)						,       (1, self.num_choices, 1))
             lowercase__   :       str				=	tf.tile(tf.expand_dims(lowerCAmelCase_						,       1)						,       (1, self.num_choices, 1))
             lowercase__   :       str				=	{
                 'input_ids': multiple_choice_inputs_ids,
                 'attention_mask': multiple_choice_input_mask,
                 'token_type_ids': multiple_choice_token_type_ids,
             }
             lowercase__   :       Any				=	model(lowerCAmelCase_)
             self.parent.assertEqual(result.logits.shape						,       (self.batch_size, self.num_choices))





        def 				snake_case_    (				self						,       a						,       a						,       a						,       a						,       a						,       a						,       a):
             lowercase__   :       Any				=	self.num_labels
             lowercase__   :       Tuple				=	TFMobileBertForTokenClassification(config=lowerCAmelCase_)
             lowercase__   :       List[Any]				=	{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
             lowercase__   :       Tuple				=	model(lowerCAmelCase_)
             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):
             lowercase__   :       List[Any]				=	TFMobileBertForQuestionAnswering(config=lowerCAmelCase_)
             lowercase__   :       str				=	{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
             lowercase__   :       List[str]				=	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 				snake_case_    (				self):
             lowercase__   :       Dict				=	self.prepare_config_and_inputs()
             (
                 (
                 lowercase__  
             )				,			(
                 lowercase__  
             )				,			(
                 lowercase__  
             )				,			(
                 lowercase__  
             )				,			(
                 lowercase__  
             )				,			(
                 lowercase__  
             )				,			(
                 lowercase__  
             )				,			  
             )   :       Any				=	config_and_inputs
             lowercase__   :       Dict				=	{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
             return config, inputs_dict





    def 				snake_case_    (				self):
         lowercase__   :       Optional[int]				=	TFMobileBertModelTest.TFMobileBertModelTester(self)
         lowercase__   :       Dict				=	ConfigTester(self						,       config_class=lowerCAmelCase_						,       hidden_size=37)





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





    def 				snake_case_    (				self):
         lowercase__   :       Optional[Any]				=	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_model(*lowerCAmelCase_)





    def 				snake_case_    (				self):
         lowercase__   :       str				=	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_masked_lm(*lowerCAmelCase_)





    def 				snake_case_    (				self):
         lowercase__   :       Union[str, Any]				=	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_multiple_choice(*lowerCAmelCase_)





    def 				snake_case_    (				self):
         lowercase__   :       Union[str, Any]				=	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*lowerCAmelCase_)





    def 				snake_case_    (				self):
         lowercase__   :       Tuple				=	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_pretraining(*lowerCAmelCase_)





    def 				snake_case_    (				self):
         lowercase__   :       Union[str, Any]				=	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_question_answering(*lowerCAmelCase_)





    def 				snake_case_    (				self):
         lowercase__   :       List[Any]				=	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_sequence_classification(*lowerCAmelCase_)





    def 				snake_case_    (				self):
         lowercase__   :       List[str]				=	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_token_classification(*lowerCAmelCase_)





    @slow
    def 				snake_case_    (				self):
         # for model_name in TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         for model_name in ["google/mobilebert-uncased"]:
              lowercase__   :       Any				=	TFMobileBertModel.from_pretrained(lowerCAmelCase_)
              self.assertIsNotNone(lowerCAmelCase_)



@require_tf
class 				SCREAMING_SNAKE_CASE__ (unittest.TestCase					):





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

         lowercase__   :       Optional[int]				=	[1, 6, 3_0522]
         self.assertEqual(output.shape						,       lowerCAmelCase_)

         lowercase__   :       Tuple				=	tf.constant(
             [
                 [
                     [-4.5_919_547, -9.248_295, -9.645_256],
                     [-6.7_306_175, -6.440_284, -6.6_052_837],
                     [-7.2_743_506, -6.7_847_915, -6.024_673],
                 ]
             ])
         tf.debugging.assert_near(output[:, :3, :3]						,       lowerCAmelCase_						,       atol=1e-4)




import math
import sys

def      snake_case__      (   SCREAMING_SNAKE_CASE_		:   int      ):



     '''simple docstring'''





     if number != int(SCREAMING_SNAKE_CASE_      ):
          raise ValueError('the value of input must be a natural number'      )
     if number < 0:
          raise ValueError('the value of input must not be a negative number'      )
     if number == 0:
          return 1
     lowercase__   :       Tuple				=	[-1] * (number + 1)
     lowercase__   :       Tuple				=	0
     for i in range(1     ,					number + 1      ):
          lowercase__   :       Tuple				=	sys.maxsize
          lowercase__   :       str				=	int(math.sqrt(SCREAMING_SNAKE_CASE_      )      )
          for j in range(1     ,					root + 1      ):
               lowercase__   :       List[Any]				=	1 + answers[i - (j**2)]
               lowercase__   :       str				=	min(SCREAMING_SNAKE_CASE_     ,					SCREAMING_SNAKE_CASE_      )
          lowercase__   :       List[str]				=	answer
     return answers[number]


if __name__ == "__main__":
       import doctest

       doctest.testmod()

