





import unittest

from transformers import MobileBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         MODEL_FOR_PRETRAINING_MAPPING,
         MobileBertForMaskedLM,
         MobileBertForMultipleChoice,
         MobileBertForNextSentencePrediction,
         MobileBertForPreTraining,
         MobileBertForQuestionAnswering,
         MobileBertForSequenceClassification,
         MobileBertForTokenClassification,
         MobileBertModel,
     )







class    __snake_case	:





       '''simple docstring'''







       def __init__(  self			:   Dict							,    A			:   Optional[Any]							,    A			:   Optional[Any]=13							,    A			:   List[str]=7							,    A			:   List[str]=True							,    A			:   Optional[int]=True							,    A			:   List[Any]=True							,    A			:   Tuple=True							,    A			:   Optional[int]=99							,    A			:   Any=64							,    A			:   Union[str, Any]=32							,    A			:   Optional[int]=5							,    A			:   Optional[Any]=4							,    A			:   Optional[int]=37							,    A			:   Union[str, Any]="gelu"							,    A			:   Optional[Any]=0.1							,    A			:   List[Any]=0.1							,    A			:   Tuple=512							,    A			:   int=16							,    A			:   List[str]=2							,    A			:   List[str]=0.02							,    A			:   Tuple=3							,    A			:   Optional[int]=4							,    A			:   Dict=None							,    ):
         __snake_case:						Optional[Any]			  =						parent
         __snake_case:						Optional[int]			  =						batch_size
         __snake_case:						Tuple			  =						seq_length
         __snake_case:						List[str]			  =						is_training
         __snake_case:						List[Any]			  =						use_input_mask
         __snake_case:						str			  =						use_token_type_ids
         __snake_case:						int			  =						use_labels
         __snake_case:						str			  =						vocab_size
         __snake_case:						int			  =						hidden_size
         __snake_case:						Optional[Any]			  =						embedding_size
         __snake_case:						List[str]			  =						num_hidden_layers
         __snake_case:						List[str]			  =						num_attention_heads
         __snake_case:						int			  =						intermediate_size
         __snake_case:						Optional[int]			  =						hidden_act
         __snake_case:						Tuple			  =						hidden_dropout_prob
         __snake_case:						Any			  =						attention_probs_dropout_prob
         __snake_case:						List[str]			  =						max_position_embeddings
         __snake_case:						Any			  =						type_vocab_size
         __snake_case:						Optional[int]			  =						type_sequence_label_size
         __snake_case:						Optional[Any]			  =						initializer_range
         __snake_case:						Optional[int]			  =						num_labels
         __snake_case:						Optional[Any]			  =						num_choices
         __snake_case:						Any			  =						scope




       def 							UpperCAmelCase__			(  self			:   int		):
         __snake_case:						Union[str, Any]			  =						ids_tensor([self.batch_size, self.seq_length]							,    self.vocab_size		)

         __snake_case:						Any			  =						None
         if self.use_input_mask:
           __snake_case:						List[str]			  =						random_attention_mask([self.batch_size, self.seq_length]		)

         __snake_case:						Any			  =						None
         if self.use_token_type_ids:
           __snake_case:						Union[str, Any]			  =						ids_tensor([self.batch_size, self.seq_length]							,    self.type_vocab_size		)

         __snake_case:						Dict			  =						None
         __snake_case:						Optional[Any]			  =						None
         __snake_case:						List[str]			  =						None
         if self.use_labels:
           __snake_case:						Dict			  =						ids_tensor([self.batch_size]							,    self.type_sequence_label_size		)
           __snake_case:						Optional[int]			  =						ids_tensor([self.batch_size, self.seq_length]							,    self.num_labels		)
           __snake_case:						str			  =						ids_tensor([self.batch_size]							,    self.num_choices		)

         __snake_case:						Tuple			  =						self.get_config()

         return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels




       def 							UpperCAmelCase__			(  self			:   List[Any]		):
         return 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							,    embedding_size=self.embedding_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							,    is_decoder=__lowerCAmelCase							,    initializer_range=self.initializer_range							,    )




       def 							UpperCAmelCase__			(  self			:   Tuple							,    A			:   Optional[int]							,    A			:   str							,    A			:   Dict							,    A			:   int							,    A			:   int							,    A			:   Any							,    A			:   Optional[Any]		):
         __snake_case:						Dict			  =						MobileBertModel(config=__lowerCAmelCase		)
         model.to(__lowerCAmelCase		)
         model.eval()
         __snake_case:						Dict			  =						model(__lowerCAmelCase							,    attention_mask=__lowerCAmelCase							,    token_type_ids=__lowerCAmelCase		)
         __snake_case:						Optional[int]			  =						model(__lowerCAmelCase							,    token_type_ids=__lowerCAmelCase		)
         __snake_case:						Dict			  =						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 							UpperCAmelCase__			(  self			:   int							,    A			:   Union[str, Any]							,    A			:   Union[str, Any]							,    A			:   int							,    A			:   List[str]							,    A			:   List[str]							,    A			:   Tuple							,    A			:   Optional[int]		):
         __snake_case:						str			  =						MobileBertForMaskedLM(config=__lowerCAmelCase		)
         model.to(__lowerCAmelCase		)
         model.eval()
         __snake_case:						List[str]			  =						model(__lowerCAmelCase							,    attention_mask=__lowerCAmelCase							,    token_type_ids=__lowerCAmelCase							,    labels=__lowerCAmelCase		)
         self.parent.assertEqual(result.logits.shape							,    (self.batch_size, self.seq_length, self.vocab_size)		)




       def 							UpperCAmelCase__			(  self			:   Tuple							,    A			:   List[Any]							,    A			:   Optional[Any]							,    A			:   Any							,    A			:   Optional[Any]							,    A			:   int							,    A			:   List[str]							,    A			:   List[str]		):
         __snake_case:						Tuple			  =						MobileBertForNextSentencePrediction(config=__lowerCAmelCase		)
         model.to(__lowerCAmelCase		)
         model.eval()
         __snake_case:						str			  =						model(
             __lowerCAmelCase							,    attention_mask=__lowerCAmelCase							,    token_type_ids=__lowerCAmelCase							,    labels=__lowerCAmelCase							,    )
         self.parent.assertEqual(result.logits.shape							,    (self.batch_size, 2)		)




       def 							UpperCAmelCase__			(  self			:   Tuple							,    A			:   Any							,    A			:   List[str]							,    A			:   str							,    A			:   str							,    A			:   Tuple							,    A			:   int							,    A			:   int		):
         __snake_case:						Any			  =						MobileBertForPreTraining(config=__lowerCAmelCase		)
         model.to(__lowerCAmelCase		)
         model.eval()
         __snake_case:						str			  =						model(
             __lowerCAmelCase							,    attention_mask=__lowerCAmelCase							,    token_type_ids=__lowerCAmelCase							,    labels=__lowerCAmelCase							,    next_sentence_label=__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 							UpperCAmelCase__			(  self			:   Tuple							,    A			:   Any							,    A			:   List[str]							,    A			:   Dict							,    A			:   int							,    A			:   Optional[int]							,    A			:   Optional[Any]							,    A			:   List[str]		):
         __snake_case:						str			  =						MobileBertForQuestionAnswering(config=__lowerCAmelCase		)
         model.to(__lowerCAmelCase		)
         model.eval()
         __snake_case:						int			  =						model(
             __lowerCAmelCase							,    attention_mask=__lowerCAmelCase							,    token_type_ids=__lowerCAmelCase							,    start_positions=__lowerCAmelCase							,    end_positions=__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 							UpperCAmelCase__			(  self			:   Any							,    A			:   Tuple							,    A			:   Tuple							,    A			:   Any							,    A			:   Optional[int]							,    A			:   int							,    A			:   Dict							,    A			:   Dict		):
         __snake_case:						Optional[Any]			  =						self.num_labels
         __snake_case:						Optional[Any]			  =						MobileBertForSequenceClassification(__lowerCAmelCase		)
         model.to(__lowerCAmelCase		)
         model.eval()
         __snake_case:						int			  =						model(__lowerCAmelCase							,    attention_mask=__lowerCAmelCase							,    token_type_ids=__lowerCAmelCase							,    labels=__lowerCAmelCase		)
         self.parent.assertEqual(result.logits.shape							,    (self.batch_size, self.num_labels)		)




       def 							UpperCAmelCase__			(  self			:   List[str]							,    A			:   List[str]							,    A			:   str							,    A			:   int							,    A			:   str							,    A			:   Union[str, Any]							,    A			:   int							,    A			:   int		):
         __snake_case:						str			  =						self.num_labels
         __snake_case:						List[str]			  =						MobileBertForTokenClassification(config=__lowerCAmelCase		)
         model.to(__lowerCAmelCase		)
         model.eval()
         __snake_case:						Dict			  =						model(__lowerCAmelCase							,    attention_mask=__lowerCAmelCase							,    token_type_ids=__lowerCAmelCase							,    labels=__lowerCAmelCase		)
         self.parent.assertEqual(result.logits.shape							,    (self.batch_size, self.seq_length, self.num_labels)		)




       def 							UpperCAmelCase__			(  self			:   Dict							,    A			:   Optional[int]							,    A			:   List[Any]							,    A			:   List[Any]							,    A			:   Tuple							,    A			:   str							,    A			:   List[str]							,    A			:   str		):
         __snake_case:						Optional[Any]			  =						self.num_choices
         __snake_case:						Any			  =						MobileBertForMultipleChoice(config=__lowerCAmelCase		)
         model.to(__lowerCAmelCase		)
         model.eval()
         __snake_case:						int			  =						input_ids.unsqueeze(1		).expand(-1							,    self.num_choices							,    -1		).contiguous()
         __snake_case:						Optional[int]			  =						token_type_ids.unsqueeze(1		).expand(-1							,    self.num_choices							,    -1		).contiguous()
         __snake_case:						Optional[Any]			  =						input_mask.unsqueeze(1		).expand(-1							,    self.num_choices							,    -1		).contiguous()
         __snake_case:						Optional[Any]			  =						model(
             __lowerCAmelCase							,    attention_mask=__lowerCAmelCase							,    token_type_ids=__lowerCAmelCase							,    labels=__lowerCAmelCase							,    )
         self.parent.assertEqual(result.logits.shape							,    (self.batch_size, self.num_choices)		)






       def 							UpperCAmelCase__			(  self			:   str		):
         __snake_case:						Optional[int]			  =						self.prepare_config_and_inputs()
         (
             (
             __snake_case       
         )  ,       (
             __snake_case       
         )  ,       (
             __snake_case       
         )  ,       (
             __snake_case       
         )  ,       (
             __snake_case       
         )  ,       (
             __snake_case       
         )  ,       (
             __snake_case       
         )  ,              
         ):						str			  =						config_and_inputs
         __snake_case:						List[str]			  =						{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
         return config, inputs_dict







@require_torch
class    __snake_case	(							_a       ,      _a       ,      unittest.TestCase      ):





       '''simple docstring'''







       lowerCAmelCase__						=					(
           (
               MobileBertModel,
               MobileBertForMaskedLM,
               MobileBertForMultipleChoice,
               MobileBertForNextSentencePrediction,
               MobileBertForPreTraining,
               MobileBertForQuestionAnswering,
               MobileBertForSequenceClassification,
               MobileBertForTokenClassification,
           )
           if is_torch_available()
           else ()
       )
       lowerCAmelCase__						=					(
           {
               """feature-extraction""": MobileBertModel,
               """fill-mask""": MobileBertForMaskedLM,
               """question-answering""": MobileBertForQuestionAnswering,
               """text-classification""": MobileBertForSequenceClassification,
               """token-classification""": MobileBertForTokenClassification,
               """zero-shot""": MobileBertForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       lowerCAmelCase__						=					True




       def 							UpperCAmelCase__			(  self			:   Optional[int]							,    A			:   Optional[int]							,    A			:   str							,    A			:   Dict=False		):
         __snake_case:						Any			  =						super()._prepare_for_class(__lowerCAmelCase							,    __lowerCAmelCase							,    return_labels=__lowerCAmelCase		)

         if return_labels:
           if model_class in get_values(__lowerCAmelCase		):
             __snake_case:						Any			  =						torch.zeros(
                 (self.model_tester.batch_size, self.model_tester.seq_length)							,    dtype=torch.long							,    device=__lowerCAmelCase		)
             __snake_case:						Any			  =						torch.zeros(
                 self.model_tester.batch_size							,    dtype=torch.long							,    device=__lowerCAmelCase		)
         return inputs_dict




       def 							UpperCAmelCase__			(  self			:   Tuple		):
         __snake_case:						Optional[int]			  =						MobileBertModelTester(self		)
         __snake_case:						List[str]			  =						ConfigTester(self							,    config_class=__lowerCAmelCase							,    hidden_size=37		)




       def 							UpperCAmelCase__			(  self			:   Dict		):
         self.config_tester.run_common_tests()




       def 							UpperCAmelCase__			(  self			:   List[Any]		):
         __snake_case:						Tuple			  =						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_model(*__lowerCAmelCase		)




       def 							UpperCAmelCase__			(  self			:   str		):
         __snake_case:						Optional[int]			  =						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_masked_lm(*__lowerCAmelCase		)




       def 							UpperCAmelCase__			(  self			:   List[str]		):
         __snake_case:						Union[str, Any]			  =						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_multiple_choice(*__lowerCAmelCase		)




       def 							UpperCAmelCase__			(  self			:   Dict		):
         __snake_case:						List[str]			  =						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*__lowerCAmelCase		)




       def 							UpperCAmelCase__			(  self			:   Optional[Any]		):
         __snake_case:						Optional[Any]			  =						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_pretraining(*__lowerCAmelCase		)




       def 							UpperCAmelCase__			(  self			:   Tuple		):
         __snake_case:						Optional[Any]			  =						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_question_answering(*__lowerCAmelCase		)




       def 							UpperCAmelCase__			(  self			:   Dict		):
         __snake_case:						List[Any]			  =						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_sequence_classification(*__lowerCAmelCase		)






       def 							UpperCAmelCase__			(  self			:   Tuple		):
         __snake_case:						Optional[int]			  =						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_token_classification(*__lowerCAmelCase		)







def       A__							(   SCREAMING_SNAKE_CASE__)       ->   Union[str, Any]:
  return torch.tensor(
      a__ , dtype=torch.long , device=a__ , )


__UpperCAmelCase	:       Union[str, Any]    	=							1e-3







@require_torch
@require_sentencepiece
@require_tokenizers
class    __snake_case	(							unittest.TestCase      ):





       '''simple docstring'''







       @slow
       def 							UpperCAmelCase__			(  self			:   Any		):
         __snake_case:						Optional[int]			  =						MobileBertModel.from_pretrained("""google/mobilebert-uncased"""		).to(__lowerCAmelCase		)
         __snake_case:						Tuple			  =						_long_tensor([[101, 7_110, 1_005, 1_056, 2_023, 11_333, 17_413, 1_029, 102]]		)
         with torch.no_grad():
           __snake_case:						Any			  =						model(__lowerCAmelCase		)[0]
         __snake_case:						Tuple			  =						torch.Size((1, 9, 512)		)
         self.assertEqual(output.shape							,    __lowerCAmelCase		)
         __snake_case:						Optional[Any]			  =						torch.tensor(
             [
                 [
                     [-2.4_73_65_26E07, 8.2_69_16_56E04, 1.6_52_18_38E05],
                     [-5.7_54_17_04E-01, 3.9_05_60_22E00, 4.4_01_15_07E00],
                     [2.6_04_73_59E00, 1.5_67_76_52E00, -1.7_32_41_88E-01],
                 ]
             ]							,    device=__lowerCAmelCase							,    )

         # MobileBERT results range from 10e0 to 10e8. Even a 0.0000001% difference with a value of 10e8 results in a
         # ~1 difference, it's therefore not a good idea to measure using addition.
         # Here, we instead divide the expected result with the result in order to obtain ~1. We then check that the
         # result is held between bounds: 1 - TOLERANCE < expected_result / result < 1 + TOLERANCE
         __snake_case:						Union[str, Any]			  =						torch.all((expected_slice / output[..., :3, :3]) >= 1 - TOLERANCE		)
         __snake_case:						Any			  =						torch.all((expected_slice / output[..., :3, :3]) <= 1 + TOLERANCE		)

         self.assertTrue(lower_bound and upper_bound		)













from ...utils import (
    OptionalDependencyNotAvailable,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)


try:
      if not (is_transformers_available() and is_torch_available()):
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      from ...utils.dummy_torch_and_transformers_objects import (
          ImageTextPipelineOutput,
          UniDiffuserPipeline,
      )
else:
      from .modeling_text_decoder import UniDiffuserTextDecoder
      from .modeling_uvit import UniDiffuserModel, UTransformeraDModel
      from .pipeline_unidiffuser import ImageTextPipelineOutput, UniDiffuserPipeline






