"""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  snake_case__      (       UpperCamelCase_,						UpperCamelCase_,						unittest.TestCase       ):
 _snake_case       :					Optional[int]							  =						(
     (
         TFMobileBertModel,
         TFMobileBertForMaskedLM,
         TFMobileBertForNextSentencePrediction,
         TFMobileBertForPreTraining,
         TFMobileBertForQuestionAnswering,
         TFMobileBertForSequenceClassification,
         TFMobileBertForTokenClassification,
         TFMobileBertForMultipleChoice,
     )
     if is_tf_available()
     else ()
 )
 _snake_case       :					Dict							  =						(
     {
         """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       :					Any							  =						False
 _snake_case       :					str							  =						False


 def 						a__   ( self    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase=False						):
        __a       =    super()._prepare_for_class(__SCREAMING_SNAKE_CASE    ,		__SCREAMING_SNAKE_CASE    ,		return_labels=__SCREAMING_SNAKE_CASE						)

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

        return inputs_dict






 class  snake_case__      (       UpperCamelCase_       ):


  def __init__( self    ,		lowerCamelCase    ,		lowerCamelCase=13    ,		lowerCamelCase=7    ,		lowerCamelCase=True    ,		lowerCamelCase=True    ,		lowerCamelCase=True    ,		lowerCamelCase=True    ,		lowerCamelCase=99    ,		lowerCamelCase=32    ,		lowerCamelCase=32    ,		lowerCamelCase=2    ,		lowerCamelCase=4    ,		lowerCamelCase=37    ,		lowerCamelCase="gelu"    ,		lowerCamelCase=0.1    ,		lowerCamelCase=0.1    ,		lowerCamelCase=512    ,		lowerCamelCase=16    ,		lowerCamelCase=2    ,		lowerCamelCase=0.02    ,		lowerCamelCase=3    ,		lowerCamelCase=4    ,		lowerCamelCase=None    ,		):
         __a       =    parent
         __a       =    batch_size
         __a       =    seq_length
         __a       =    is_training
         __a       =    use_input_mask
         __a       =    use_token_type_ids
         __a       =    use_labels
         __a       =    vocab_size
         __a       =    hidden_size
         __a       =    num_hidden_layers
         __a       =    num_attention_heads
         __a       =    intermediate_size
         __a       =    hidden_act
         __a       =    hidden_dropout_prob
         __a       =    attention_probs_dropout_prob
         __a       =    max_position_embeddings
         __a       =    type_vocab_size
         __a       =    type_sequence_label_size
         __a       =    initializer_range
         __a       =    num_labels
         __a       =    num_choices
         __a       =    scope
         __a       =    embedding_size


  def 						a__   ( self						):
         __a       =    ids_tensor([self.batch_size, self.seq_length]    ,		self.vocab_size						)

         __a       =    None
         if self.use_input_mask:
                __a       =    random_attention_mask([self.batch_size, self.seq_length]						)

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

         __a       =    None
         __a       =    None
         __a       =    None
         if self.use_labels:
                __a       =    ids_tensor([self.batch_size]    ,		self.type_sequence_label_size						)
                __a       =    ids_tensor([self.batch_size, self.seq_length]    ,		self.num_labels						)
                __a       =    ids_tensor([self.batch_size]    ,		self.num_choices						)

         __a       =    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 						a__   ( self    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase						):
         __a       =    TFMobileBertModel(config=__SCREAMING_SNAKE_CASE						)
         __a       =    {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
         __a       =    model(__SCREAMING_SNAKE_CASE						)

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

         __a       =    model(__SCREAMING_SNAKE_CASE						)

         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 						a__   ( self    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase						):
         __a       =    TFMobileBertForMaskedLM(config=__SCREAMING_SNAKE_CASE						)
         __a       =    {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
         __a       =    model(__SCREAMING_SNAKE_CASE						)
         self.parent.assertEqual(result.logits.shape    ,		(self.batch_size, self.seq_length, self.vocab_size)						)


  def 						a__   ( self    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase						):
         __a       =    TFMobileBertForNextSentencePrediction(config=__SCREAMING_SNAKE_CASE						)
         __a       =    {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
         __a       =    model(__SCREAMING_SNAKE_CASE						)
         self.parent.assertEqual(result.logits.shape    ,		(self.batch_size, 2)						)


  def 						a__   ( self    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase						):
         __a       =    TFMobileBertForPreTraining(config=__SCREAMING_SNAKE_CASE						)
         __a       =    {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
         __a       =    model(__SCREAMING_SNAKE_CASE						)
         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 						a__   ( self    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase						):
         __a       =    self.num_labels
         __a       =    TFMobileBertForSequenceClassification(config=__SCREAMING_SNAKE_CASE						)
         __a       =    {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
         __a       =    model(__SCREAMING_SNAKE_CASE						)
         self.parent.assertEqual(result.logits.shape    ,		(self.batch_size, self.num_labels)						)


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


  def 						a__   ( self    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase						):
         __a       =    self.num_labels
         __a       =    TFMobileBertForTokenClassification(config=__SCREAMING_SNAKE_CASE						)
         __a       =    {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
         __a       =    model(__SCREAMING_SNAKE_CASE						)
         self.parent.assertEqual(result.logits.shape    ,		(self.batch_size, self.seq_length, self.num_labels)						)


  def 						a__   ( self    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase						):
         __a       =    TFMobileBertForQuestionAnswering(config=__SCREAMING_SNAKE_CASE						)
         __a       =    {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
         __a       =    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 						a__   ( self						):
         __a       =    self.prepare_config_and_inputs()
         (
             __a
         )       =    config_and_inputs
         __a       =    {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
         return config, inputs_dict


 def 						a__   ( self						):
        __a       =    TFMobileBertModelTest.TFMobileBertModelTester(self						)
        __a       =    ConfigTester(self    ,		config_class=__SCREAMING_SNAKE_CASE    ,		hidden_size=37						)


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


 def 						a__   ( self						):
        __a       =    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_model(*__SCREAMING_SNAKE_CASE						)


 def 						a__   ( self						):
        __a       =    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_masked_lm(*__SCREAMING_SNAKE_CASE						)


 def 						a__   ( self						):
        __a       =    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_multiple_choice(*__SCREAMING_SNAKE_CASE						)


 def 						a__   ( self						):
        __a       =    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*__SCREAMING_SNAKE_CASE						)


 def 						a__   ( self						):
        __a       =    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_pretraining(*__SCREAMING_SNAKE_CASE						)


 def 						a__   ( self						):
        __a       =    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_question_answering(*__SCREAMING_SNAKE_CASE						)


 def 						a__   ( self						):
        __a       =    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_sequence_classification(*__SCREAMING_SNAKE_CASE						)


 def 						a__   ( self						):
        __a       =    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_token_classification(*__SCREAMING_SNAKE_CASE						)






 @slow
 def 						a__   ( self						):







        for model_name in ["google/mobilebert-uncased"]:
               __a       =    TFMobileBertModel.from_pretrained(__SCREAMING_SNAKE_CASE						)
               self.assertIsNotNone(__SCREAMING_SNAKE_CASE						)






@require_tf
class  snake_case__      (       unittest.TestCase       ):


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

        __a       =    [1, 6, 30522]
        self.assertEqual(output.shape    ,		__SCREAMING_SNAKE_CASE						)

        __a       =    tf.constant(
            [
                [
                    [-4.591_9547, -9.24_8295, -9.64_5256],
                    [-6.730_6175, -6.44_0284, -6.605_2837],
                    [-7.274_3506, -6.784_7915, -6.02_4673],
                ]
            ]						)
        tf.debugging.assert_near(output[:, :3, :3]    ,		__SCREAMING_SNAKE_CASE    ,		atol=1E-4						)


def 				SCREAMING_SNAKE_CASE_   (      UpperCamelCase__			,		UpperCamelCase__ = " "		):
      UpperCamelCase__     : Optional[int]											=					[]

      UpperCamelCase__     : int											=					0
      for index, char in enumerate(UpperCamelCase__		):
            if char == separator:
                  split_words.append(string[last_index:index]		)
                  UpperCamelCase__     : int											=					index + 1
            elif index + 1 == len(UpperCamelCase__		):
                  split_words.append(string[last_index : index + 1]		)
      return split_words


if __name__ == "__main__":
    from doctest import testmod

    testmod()


