




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

    '''simple docstring'''






    UpperCAmelCase__				    = (
        (
            TFMobileBertModel,
            TFMobileBertForMaskedLM,
            TFMobileBertForNextSentencePrediction,
            TFMobileBertForPreTraining,
            TFMobileBertForQuestionAnswering,
            TFMobileBertForSequenceClassification,
            TFMobileBertForTokenClassification,
            TFMobileBertForMultipleChoice,
        )
        if is_tf_available()
        else ()
    )
    UpperCAmelCase__				    = (
        {
            '''feature-extraction''': TFMobileBertModel,
            '''fill-mask''': TFMobileBertForMaskedLM,
            '''question-answering''': TFMobileBertForQuestionAnswering,
            '''text-classification''': TFMobileBertForSequenceClassification,
            '''token-classification''': TFMobileBertForTokenClassification,
            '''zero-shot''': TFMobileBertForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )
    UpperCAmelCase__				    = False
    UpperCAmelCase__				    = False


    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[Any]   ,					UpperCAmelCase__ :      Dict   ,					UpperCAmelCase__ :      Tuple   ,					UpperCAmelCase__ :      str=False)		->Optional[Any]:





           '''simple docstring'''

           A__ 					=       super()._prepare_for_class(UpperCAmelCase__   ,					UpperCAmelCase__   ,					return_labels=UpperCAmelCase__)

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

           return inputs_dict


    class 							UpperCamelCase_       (					UpperCAmelCase__	):

        '''simple docstring'''






        def __init__(     self :      List[Any]   ,					UpperCAmelCase__ :      str   ,					UpperCAmelCase__ :      int=13   ,					UpperCAmelCase__ :      str=7   ,					UpperCAmelCase__ :      Optional[Any]=True   ,					UpperCAmelCase__ :      Optional[int]=True   ,					UpperCAmelCase__ :      Union[str, Any]=True   ,					UpperCAmelCase__ :      Optional[int]=True   ,					UpperCAmelCase__ :      str=99   ,					UpperCAmelCase__ :      List[str]=32   ,					UpperCAmelCase__ :      Optional[int]=32   ,					UpperCAmelCase__ :      Any=2   ,					UpperCAmelCase__ :      List[str]=4   ,					UpperCAmelCase__ :      Optional[Any]=37   ,					UpperCAmelCase__ :      Optional[int]="gelu"   ,					UpperCAmelCase__ :      Any=0.1   ,					UpperCAmelCase__ :      Optional[Any]=0.1   ,					UpperCAmelCase__ :      List[Any]=512   ,					UpperCAmelCase__ :      Tuple=16   ,					UpperCAmelCase__ :      Any=2   ,					UpperCAmelCase__ :      Dict=0.02   ,					UpperCAmelCase__ :      int=3   ,					UpperCAmelCase__ :      List[str]=4   ,					UpperCAmelCase__ :      Tuple=None   ,					)		->Any:





               '''simple docstring'''

               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 			SCREAMING_SNAKE_CASE				(     self :      Optional[Any])		->Tuple:





               '''simple docstring'''

               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 			SCREAMING_SNAKE_CASE				(     self :      Union[str, Any]   ,					UpperCAmelCase__ :      List[str]   ,					UpperCAmelCase__ :      str   ,					UpperCAmelCase__ :      str   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      Union[str, Any]   ,					UpperCAmelCase__ :      Optional[Any])		->Any:





               '''simple docstring'''

               A__ 					=       TFMobileBertModel(config=UpperCAmelCase__)
               A__ 					=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
               A__ 					=       model(UpperCAmelCase__)

               A__ 					=       [input_ids, input_mask]
               A__ 					=       model(UpperCAmelCase__)

               A__ 					=       model(UpperCAmelCase__)

               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 			SCREAMING_SNAKE_CASE				(     self :      List[str]   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      str   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      Optional[int]   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      Tuple)		->Optional[Any]:





               '''simple docstring'''

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


        def 			SCREAMING_SNAKE_CASE				(     self :      Dict   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      Optional[Any]   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      Optional[Any])		->int:





               '''simple docstring'''

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


        def 			SCREAMING_SNAKE_CASE				(     self :      Dict   ,					UpperCAmelCase__ :      Dict   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      Union[str, Any]   ,					UpperCAmelCase__ :      Dict   ,					UpperCAmelCase__ :      Dict   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      int)		->List[Any]:





               '''simple docstring'''

               A__ 					=       TFMobileBertForPreTraining(config=UpperCAmelCase__)
               A__ 					=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
               A__ 					=       model(UpperCAmelCase__)
               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 			SCREAMING_SNAKE_CASE				(     self :      Tuple   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      Union[str, Any]   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      List[str]   ,					UpperCAmelCase__ :      Tuple)		->Dict:





               '''simple docstring'''

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


        def 			SCREAMING_SNAKE_CASE				(     self :      Any   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      str   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      Optional[Any]   ,					UpperCAmelCase__ :      str   ,					UpperCAmelCase__ :      int)		->Dict:





               '''simple docstring'''

               A__ 					=       self.num_choices
               A__ 					=       TFMobileBertForMultipleChoice(config=UpperCAmelCase__)
               A__ 					=       tf.tile(tf.expand_dims(UpperCAmelCase__   ,					1)   ,					(1, self.num_choices, 1))
               A__ 					=       tf.tile(tf.expand_dims(UpperCAmelCase__   ,					1)   ,					(1, self.num_choices, 1))
               A__ 					=       tf.tile(tf.expand_dims(UpperCAmelCase__   ,					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(UpperCAmelCase__)
               self.parent.assertEqual(result.logits.shape   ,					(self.batch_size, self.num_choices))


        def 			SCREAMING_SNAKE_CASE				(     self :      Dict   ,					UpperCAmelCase__ :      Union[str, Any]   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      Optional[int]   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      Dict   ,					UpperCAmelCase__ :      Optional[int])		->int:





               '''simple docstring'''

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


        def 			SCREAMING_SNAKE_CASE				(     self :      List[Any]   ,					UpperCAmelCase__ :      Optional[Any]   ,					UpperCAmelCase__ :      Dict   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      Tuple   ,					UpperCAmelCase__ :      Union[str, Any]   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      Union[str, Any])		->Union[str, Any]:





               '''simple docstring'''

               A__ 					=       TFMobileBertForQuestionAnswering(config=UpperCAmelCase__)
               A__ 					=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
               A__ 					=       model(UpperCAmelCase__)
               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 			SCREAMING_SNAKE_CASE				(     self :      Any)		->str:





               '''simple docstring'''

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


    def 			SCREAMING_SNAKE_CASE				(     self :      Any)		->Union[str, Any]:





           '''simple docstring'''

           A__ 					=       TFMobileBertModelTest.TFMobileBertModelTester(self)
           A__ 					=       ConfigTester(self   ,					config_class=UpperCAmelCase__   ,					hidden_size=37)


    def 			SCREAMING_SNAKE_CASE				(     self :      Tuple)		->Optional[int]:





           '''simple docstring'''

           self.config_tester.run_common_tests()


    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[Any])		->Dict:





           '''simple docstring'''

           A__ 					=       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_model(*UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      List[str])		->Tuple:





           '''simple docstring'''

           A__ 					=       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_masked_lm(*UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      Union[str, Any])		->Dict:





           '''simple docstring'''

           A__ 					=       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_multiple_choice(*UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      Tuple)		->Dict:





           '''simple docstring'''

           A__ 					=       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      List[Any])		->Union[str, Any]:





           '''simple docstring'''

           A__ 					=       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_pretraining(*UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      int)		->Optional[int]:





           '''simple docstring'''

           A__ 					=       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_question_answering(*UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      Tuple)		->Optional[int]:





           '''simple docstring'''

           A__ 					=       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_sequence_classification(*UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      Tuple)		->List[str]:





           '''simple docstring'''

           A__ 					=       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_token_classification(*UpperCAmelCase__)





    @slow
    def 			SCREAMING_SNAKE_CASE				(     self :      str)		->List[Any]:





           '''simple docstring'''

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


@require_tf
class 							UpperCamelCase_       (					unittest.TestCase	):

    '''simple docstring'''






    @slow
    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[Any])		->Any:





           '''simple docstring'''

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

           A__ 					=       [1, 6, 30_522]
           self.assertEqual(output.shape   ,					UpperCAmelCase__)

           A__ 					=       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]   ,					UpperCAmelCase__   ,					atol=1e-4)









import unittest

import torch

from diffusers import VQModel
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()

class 							__SCREAMING_SNAKE_CASE       (  _a ,		_a ,		unittest.TestCase							):
							snake_case  :   List[str]      =							VQModel
							snake_case  :   Dict      =							"""sample"""






							@property
							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase=(32, 32)	):
														UpperCamelCase__					  =      4
														UpperCamelCase__					  =      3

														UpperCamelCase__					  =      floats_tensor((batch_size, num_channels) + sizes	).to(__lowerCAmelCase	)

														return {"sample": image}






							@property
							def 					_lowerCamelCase		(    self	):
														return (3, 32, 32)






							@property
							def 					_lowerCamelCase		(    self	):
														return (3, 32, 32)






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      {
														    """block_out_channels""": [32, 64],
														    """in_channels""": 3,
														    """out_channels""": 3,
														    """down_block_types""": ["""DownEncoderBlock2D""", """DownEncoderBlock2D"""],
														    """up_block_types""": ["""UpDecoderBlock2D""", """UpDecoderBlock2D"""],
														    """latent_channels""": 3,
														}
														UpperCamelCase__					  =      self.dummy_input
														return init_dict, inputs_dict






							def 					_lowerCamelCase		(    self	):
														pass






							def 					_lowerCamelCase		(    self	):
														pass






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__  ,  UpperCamelCase__					  =      VQModel.from_pretrained("""fusing/vqgan-dummy"""	,			output_loading_info=__lowerCAmelCase	)
														self.assertIsNotNone(__lowerCAmelCase	)
														self.assertEqual(len(loading_info["""missing_keys"""]	)	,			0	)

														model.to(__lowerCAmelCase	)
														UpperCamelCase__					  =      model(**self.dummy_input	)

														assert image is not None, "Make sure output is not None"







							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      VQModel.from_pretrained("""fusing/vqgan-dummy"""	)
														model.to(__lowerCAmelCase	).eval()

														torch.manual_seed(0	)
														if torch.cuda.is_available():
																					torch.cuda.manual_seed_all(0	)

														UpperCamelCase__					  =      torch.randn(1	,			model.config.in_channels	,			model.config.sample_size	,			model.config.sample_size	)
														UpperCamelCase__					  =      image.to(__lowerCAmelCase	)
														with torch.no_grad():
																					UpperCamelCase__					  =      model(__lowerCAmelCase	).sample

														UpperCamelCase__					  =      output[0, -1, -3:, -3:].flatten().cpu()
														# fmt: off
														UpperCamelCase__					  =      torch.tensor([-0.0153, -0.4044, -0.1880, -0.5161, -0.2418, -0.4072, -0.1612, -0.0633, -0.0143]	)
														# fmt: on
														self.assertTrue(torch.allclose(__lowerCAmelCase	,			__lowerCAmelCase	,			atol=1E-3	)	)




