




'''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  __lowercase    (	__magic_name__							,				__magic_name__							,				unittest.TestCase			):
   _a    						=   (
       (
           TFMobileBertModel,
           TFMobileBertForMaskedLM,
           TFMobileBertForNextSentencePrediction,
           TFMobileBertForPreTraining,
           TFMobileBertForQuestionAnswering,
           TFMobileBertForSequenceClassification,
           TFMobileBertForTokenClassification,
           TFMobileBertForMultipleChoice,
       )
       if is_tf_available()
       else ()
   )
   _a    						=   (
       {
           """feature-extraction""": TFMobileBertModel,
           """fill-mask""": TFMobileBertForMaskedLM,
           """question-answering""": TFMobileBertForQuestionAnswering,
           """text-classification""": TFMobileBertForSequenceClassification,
           """token-classification""": TFMobileBertForTokenClassification,
           """zero-shot""": TFMobileBertForSequenceClassification,
       }
       if is_tf_available()
       else {}
   )
   _a    						=   False
   _a    						=   False







   def 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase=False   )					->     List[Any]:
     __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  __lowercase    (	__magic_name__			):







      def __init__(						self   ,				UpperCamelCase   ,				UpperCamelCase=13   ,				UpperCamelCase=7   ,				UpperCamelCase=True   ,				UpperCamelCase=True   ,				UpperCamelCase=True   ,				UpperCamelCase=True   ,				UpperCamelCase=99   ,				UpperCamelCase=32   ,				UpperCamelCase=32   ,				UpperCamelCase=2   ,				UpperCamelCase=4   ,				UpperCamelCase=37   ,				UpperCamelCase="gelu"   ,				UpperCamelCase=0.1   ,				UpperCamelCase=0.1   ,				UpperCamelCase=512   ,				UpperCamelCase=16   ,				UpperCamelCase=2   ,				UpperCamelCase=0.02   ,				UpperCamelCase=3   ,				UpperCamelCase=4   ,				UpperCamelCase=None   ,				)					->     Optional[Any]:
        __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 		UpperCamelCase__      (						self   )					->     int:
        __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 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   )					->     Tuple:
        __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 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   )					->     Optional[int]:
        __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 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   )					->     str:
        __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 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   )					->     int:
        __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 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   )					->     Optional[Any]:
        __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 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   )					->     Tuple:
        __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 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   )					->     str:
        __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 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   )					->     Tuple:
        __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 		UpperCamelCase__      (						self   )					->     List[Any]:
        __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 		UpperCamelCase__      (						self   )					->     List[str]:
     __a								=       TFMobileBertModelTest.TFMobileBertModelTester(self   )
     __a								=       ConfigTester(self   ,				config_class=UpperCamelCase   ,				hidden_size=37   )







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







   def 		UpperCamelCase__      (						self   )					->     Optional[int]:
     __a								=       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_mobilebert_model(*UpperCamelCase   )







   def 		UpperCamelCase__      (						self   )					->     str:
     __a								=       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_mobilebert_for_masked_lm(*UpperCamelCase   )







   def 		UpperCamelCase__      (						self   )					->     int:
     __a								=       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_mobilebert_for_multiple_choice(*UpperCamelCase   )







   def 		UpperCamelCase__      (						self   )					->     int:
     __a								=       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*UpperCamelCase   )







   def 		UpperCamelCase__      (						self   )					->     Dict:
     __a								=       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_mobilebert_for_pretraining(*UpperCamelCase   )







   def 		UpperCamelCase__      (						self   )					->     int:
     __a								=       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_mobilebert_for_question_answering(*UpperCamelCase   )







   def 		UpperCamelCase__      (						self   )					->     Optional[Any]:
     __a								=       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_mobilebert_for_sequence_classification(*UpperCamelCase   )







   def 		UpperCamelCase__      (						self   )					->     List[Any]:
     __a								=       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_mobilebert_for_token_classification(*UpperCamelCase   )







   @slow
   def 		UpperCamelCase__      (						self   )					->     Union[str, Any]:
     # for model_name in TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     for model_name in ["google/mobilebert-uncased"]:
       __a								=       TFMobileBertModel.from_pretrained(UpperCamelCase   )
       self.assertIsNotNone(UpperCamelCase   )





@require_tf
class  __lowercase    (	unittest.TestCase			):







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

     __a								=       [1, 6, 3_0522]
     self.assertEqual(output.shape   ,				UpperCamelCase   )

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









'''simple docstring'''


import gc
import inspect
import unittest

import torch
from parameterized import parameterized

from diffusers import PriorTransformer
from diffusers.utils import floats_tensor, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin


enable_full_determinism()





class  __lowercase    (	__magic_name__							,				unittest.TestCase			):
   _a    						=   PriorTransformer
   _a    						=   """hidden_states"""







   @property
   def 		UpperCamelCase__      (						self   )					->     int:
     __a								=       4
     __a								=       8
     __a								=       7

     __a								=       floats_tensor((batch_size, embedding_dim)   ).to(UpperCamelCase   )

     __a								=       floats_tensor((batch_size, embedding_dim)   ).to(UpperCamelCase   )
     __a								=       floats_tensor((batch_size, num_embeddings, embedding_dim)   ).to(UpperCamelCase   )

     return {
         "hidden_states": hidden_states,
         "timestep": 2,
         "proj_embedding": proj_embedding,
         "encoder_hidden_states": encoder_hidden_states,
     }







   def 		UpperCamelCase__      (						self   ,				UpperCamelCase=0   )					->     Union[str, Any]:
     torch.manual_seed(UpperCamelCase   )
     __a								=       4
     __a								=       8
     __a								=       7

     __a								=       torch.randn((batch_size, embedding_dim)   ).to(UpperCamelCase   )

     __a								=       torch.randn((batch_size, embedding_dim)   ).to(UpperCamelCase   )
     __a								=       torch.randn((batch_size, num_embeddings, embedding_dim)   ).to(UpperCamelCase   )

     return {
         "hidden_states": hidden_states,
         "timestep": 2,
         "proj_embedding": proj_embedding,
         "encoder_hidden_states": encoder_hidden_states,
     }







   @property
   def 		UpperCamelCase__      (						self   )					->     List[str]:
     return (4, 8)







   @property
   def 		UpperCamelCase__      (						self   )					->     str:
     return (4, 8)







   def 		UpperCamelCase__      (						self   )					->     Dict:
     __a								=       {
         'num_attention_heads': 2,
         'attention_head_dim': 4,
         'num_layers': 2,
         'embedding_dim': 8,
         'num_embeddings': 7,
         'additional_embeddings': 4,
     }
     __a								=       self.dummy_input
     return init_dict, inputs_dict







   def 		UpperCamelCase__      (						self   )					->     List[Any]:
     __a						,							__a								=       PriorTransformer.from_pretrained(
         'hf-internal-testing/prior-dummy'   ,				output_loading_info=UpperCamelCase   )
     self.assertIsNotNone(UpperCamelCase   )
     self.assertEqual(len(loading_info['missing_keys']   )   ,				0   )

     model.to(UpperCamelCase   )
     __a								=       model(**self.dummy_input   )[0]

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







   def 		UpperCamelCase__      (						self   )					->     Optional[Any]:
     __a						,							__a								=       self.prepare_init_args_and_inputs_for_common()

     __a								=       self.model_class(**UpperCamelCase   )
     __a								=       inspect.signature(model.forward   )
     # signature.parameters is an OrderedDict => so arg_names order is deterministic
     __a								=       [*signature.parameters.keys()]

     __a								=       ['hidden_states', 'timestep']
     self.assertListEqual(arg_names[:2]   ,				UpperCamelCase   )







   def 		UpperCamelCase__      (						self   )					->     str:
     __a								=       PriorTransformer.from_pretrained('hf-internal-testing/prior-dummy'   )
     __a								=       model.to(UpperCamelCase   )

     if hasattr(UpperCamelCase   ,				'set_default_attn_processor'   ):
       model.set_default_attn_processor()

     __a								=       self.get_dummy_seed_input()

     with torch.no_grad():
       __a								=       model(**UpperCamelCase   )[0]

     __a								=       output[0, :5].flatten().cpu()
     print(UpperCamelCase   )

     # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
     # the expected output slices are not the same for CPU and GPU.
     __a								=       torch.tensor([-1.3_436, -0.2_870, 0.7_538, 0.4_368, -0.0_239]   )
     self.assertTrue(torch_all_close(UpperCamelCase   ,				UpperCamelCase   ,				rtol=1e-2   )   )





@slow
class  __lowercase    (	unittest.TestCase			):







   def 		UpperCamelCase__      (						self   ,				UpperCamelCase=1   ,				UpperCamelCase=768   ,				UpperCamelCase=77   ,				UpperCamelCase=0   )					->     List[str]:
     torch.manual_seed(UpperCamelCase   )
     __a								=       batch_size
     __a								=       embedding_dim
     __a								=       num_embeddings

     __a								=       torch.randn((batch_size, embedding_dim)   ).to(UpperCamelCase   )

     __a								=       torch.randn((batch_size, embedding_dim)   ).to(UpperCamelCase   )
     __a								=       torch.randn((batch_size, num_embeddings, embedding_dim)   ).to(UpperCamelCase   )

     return {
         "hidden_states": hidden_states,
         "timestep": 2,
         "proj_embedding": proj_embedding,
         "encoder_hidden_states": encoder_hidden_states,
     }







   def 		UpperCamelCase__      (						self   )					->     List[str]:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()







   @parameterized.expand(
       [
           # fmt: off
           [13, [-0.5_861,  0.1_283, -0.0_931,  0.0_882,  0.4_476,  0.1_329, -0.0_498,  0.0_640]],
           [37, [-0.4_913,  0.0_110, -0.0_483,  0.0_541,  0.4_954, -0.0_170,  0.0_354,  0.1_651]],
           # fmt: on
       ]   )
   def 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase   )					->     Optional[Any]:
     __a								=       PriorTransformer.from_pretrained('kandinsky-community/kandinsky-2-1-prior'   ,				subfolder='prior'   )
     model.to(UpperCamelCase   )
     __a								=       self.get_dummy_seed_input(seed=UpperCamelCase   )

     with torch.no_grad():
       __a								=       model(**UpperCamelCase   )[0]

     assert list(sample.shape   ) == [1, 768]

     __a								=       sample[0, :8].flatten().cpu()
     print(UpperCamelCase   )
     __a								=       torch.tensor(UpperCamelCase   )

     assert torch_all_close(UpperCamelCase   ,				UpperCamelCase   ,				atol=1e-3   )



