





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


    '''simple docstring'''







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



    def  lowercase_					(       self   ,				A_   ,				A_   ,				A_=False     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			super()._prepare_for_class(A_   ,				A_   ,				return_labels=A_     )

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

        return inputs_dict



    class 					__snake_case	(  __SCREAMING_SNAKE_CASE    ):


        '''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.02   ,				A_=3   ,				A_=4   ,				A_=None   ,				):





            '''simple docstring'''

            SCREAMING_SNAKE_CASE__											=			parent
            SCREAMING_SNAKE_CASE__											=			batch_size
            SCREAMING_SNAKE_CASE__											=			seq_length
            SCREAMING_SNAKE_CASE__											=			is_training
            SCREAMING_SNAKE_CASE__											=			use_input_mask
            SCREAMING_SNAKE_CASE__											=			use_token_type_ids
            SCREAMING_SNAKE_CASE__											=			use_labels
            SCREAMING_SNAKE_CASE__											=			vocab_size
            SCREAMING_SNAKE_CASE__											=			hidden_size
            SCREAMING_SNAKE_CASE__											=			num_hidden_layers
            SCREAMING_SNAKE_CASE__											=			num_attention_heads
            SCREAMING_SNAKE_CASE__											=			intermediate_size
            SCREAMING_SNAKE_CASE__											=			hidden_act
            SCREAMING_SNAKE_CASE__											=			hidden_dropout_prob
            SCREAMING_SNAKE_CASE__											=			attention_probs_dropout_prob
            SCREAMING_SNAKE_CASE__											=			max_position_embeddings
            SCREAMING_SNAKE_CASE__											=			type_vocab_size
            SCREAMING_SNAKE_CASE__											=			type_sequence_label_size
            SCREAMING_SNAKE_CASE__											=			initializer_range
            SCREAMING_SNAKE_CASE__											=			num_labels
            SCREAMING_SNAKE_CASE__											=			num_choices
            SCREAMING_SNAKE_CASE__											=			scope
            SCREAMING_SNAKE_CASE__											=			embedding_size



        def  lowercase_					(       self     ):





            '''simple docstring'''

            SCREAMING_SNAKE_CASE__											=			ids_tensor([self.batch_size, self.seq_length]   ,				self.vocab_size     )

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

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

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

            SCREAMING_SNAKE_CASE__											=			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  lowercase_					(       self   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_     ):





            '''simple docstring'''

            SCREAMING_SNAKE_CASE__											=			TFMobileBertModel(config=A_     )
            SCREAMING_SNAKE_CASE__											=			{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
            SCREAMING_SNAKE_CASE__											=			model(A_     )

            SCREAMING_SNAKE_CASE__											=			[input_ids, input_mask]
            SCREAMING_SNAKE_CASE__											=			model(A_     )

            SCREAMING_SNAKE_CASE__											=			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  lowercase_					(       self   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_     ):





            '''simple docstring'''

            SCREAMING_SNAKE_CASE__											=			TFMobileBertForMaskedLM(config=A_     )
            SCREAMING_SNAKE_CASE__											=			{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
            SCREAMING_SNAKE_CASE__											=			model(A_     )
            self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.seq_length, self.vocab_size)     )



        def  lowercase_					(       self   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_     ):





            '''simple docstring'''

            SCREAMING_SNAKE_CASE__											=			TFMobileBertForNextSentencePrediction(config=A_     )
            SCREAMING_SNAKE_CASE__											=			{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
            SCREAMING_SNAKE_CASE__											=			model(A_     )
            self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, 2)     )



        def  lowercase_					(       self   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_     ):





            '''simple docstring'''

            SCREAMING_SNAKE_CASE__											=			TFMobileBertForPreTraining(config=A_     )
            SCREAMING_SNAKE_CASE__											=			{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
            SCREAMING_SNAKE_CASE__											=			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  lowercase_					(       self   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_     ):





            '''simple docstring'''

            SCREAMING_SNAKE_CASE__											=			self.num_labels
            SCREAMING_SNAKE_CASE__											=			TFMobileBertForSequenceClassification(config=A_     )
            SCREAMING_SNAKE_CASE__											=			{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
            SCREAMING_SNAKE_CASE__											=			model(A_     )
            self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.num_labels)     )



        def  lowercase_					(       self   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_     ):





            '''simple docstring'''

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



        def  lowercase_					(       self   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_     ):





            '''simple docstring'''

            SCREAMING_SNAKE_CASE__											=			self.num_labels
            SCREAMING_SNAKE_CASE__											=			TFMobileBertForTokenClassification(config=A_     )
            SCREAMING_SNAKE_CASE__											=			{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
            SCREAMING_SNAKE_CASE__											=			model(A_     )
            self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.seq_length, self.num_labels)     )



        def  lowercase_					(       self   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_     ):





            '''simple docstring'''

            SCREAMING_SNAKE_CASE__											=			TFMobileBertForQuestionAnswering(config=A_     )
            SCREAMING_SNAKE_CASE__											=			{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
            SCREAMING_SNAKE_CASE__											=			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  lowercase_					(       self     ):





            '''simple docstring'''

            SCREAMING_SNAKE_CASE__											=			self.prepare_config_and_inputs()
            (
                (
                SCREAMING_SNAKE_CASE__
            )				,		(
                SCREAMING_SNAKE_CASE__
            )				,		(
                SCREAMING_SNAKE_CASE__
            )				,		(
                SCREAMING_SNAKE_CASE__
            )				,		(
                SCREAMING_SNAKE_CASE__
            )				,		(
                SCREAMING_SNAKE_CASE__
            )				,		(
                SCREAMING_SNAKE_CASE__
            )				,		
            )											=			config_and_inputs
            SCREAMING_SNAKE_CASE__											=			{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
            return config, inputs_dict





    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			TFMobileBertModelTest.TFMobileBertModelTester(self     )
        SCREAMING_SNAKE_CASE__											=			ConfigTester(self   ,				config_class=A_   ,				hidden_size=37     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        self.config_tester.run_common_tests()



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_model(*A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_masked_lm(*A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_multiple_choice(*A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_pretraining(*A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_question_answering(*A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_sequence_classification(*A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_token_classification(*A_     )




    @slow
    def  lowercase_					(       self     ):





        '''simple docstring'''

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





@require_tf
class 					__snake_case	(  unittest.TestCase    ):


    '''simple docstring'''







    @slow
    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			TFMobileBertForPreTraining.from_pretrained('''google/mobilebert-uncased'''     )
        SCREAMING_SNAKE_CASE__											=			tf.constant([[0, 1, 2, 3, 4, 5]]     )
        SCREAMING_SNAKE_CASE__											=			model(A_     )[0]

        SCREAMING_SNAKE_CASE__											=			[1, 6, 3_05_22]
        self.assertEqual(output.shape   ,				A_     )

        SCREAMING_SNAKE_CASE__											=			tf.constant(
            [
                [
                    [-4.5919547, -9.248295, -9.645256],
                    [-6.7306175, -6.440284, -6.6052837],
                    [-7.2743506, -6.7847915, -6.024673],
                ]
            ]     )
        tf.debugging.assert_near(output[:, :3, :3]   ,				A_   ,				atol=1E-4     )




import os
from distutils.util import strtobool
def      UpperCAmelCase__		(UpperCamelCase_				,UpperCamelCase_							):



     """simple docstring"""






     for e in env_keys:
          snake_case											=							int(os.environ.get(UpperCamelCase_				,-1							)							)
          if val >= 0:
               return val
     return default
def      UpperCAmelCase__		(UpperCamelCase_				,UpperCamelCase_=False							):



     """simple docstring"""






     snake_case											=							os.environ.get(UpperCamelCase_				,str(UpperCamelCase_							)							)
     return strtobool(UpperCamelCase_							) == 1  # As its name indicates `strtobool` actually returns an int...

def      UpperCAmelCase__		(UpperCamelCase_				,UpperCamelCase_="no"							):



     """simple docstring"""






     snake_case											=							os.environ.get(UpperCamelCase_				,str(UpperCamelCase_							)							)
     return value




