


"""simple docstring"""







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 				__A    :






      """simple docstring"""




      def __init__(		self			,   __A			,   __A=13			,   __A=7			,   __A=True			,   __A=True			,   __A=True			,   __A=True			,   __A=99			,   __A=64			,   __A=32			,   __A=5			,   __A=4			,   __A=37			,   __A="gelu"			,   __A=0.1			,   __A=0.1			,   __A=512			,   __A=16			,   __A=2			,   __A=0.02			,   __A=3			,   __A=4			,   __A=None			,   ) ->					List[str]:
          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					=embedding_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


      def 				SCREAMING_SNAKE_CASE							(		self						) ->					Optional[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					=self.get_config()

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


      def 				SCREAMING_SNAKE_CASE							(		self						) ->					Optional[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=__A			,   initializer_range=self.initializer_range			,   )


      def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A			,   __A			,   __A			,   __A			,   __A			,   __A						) ->					Dict:
          a					=MobileBertModel(config=__A						)
          model.to(__A						)
          model.eval()
          a					=model(__A			,   attention_mask=__A			,   token_type_ids=__A						)
          a					=model(__A			,   token_type_ids=__A						)
          a					=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 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A			,   __A			,   __A			,   __A			,   __A			,   __A						) ->					Union[str, Any]:
          a					=MobileBertForMaskedLM(config=__A						)
          model.to(__A						)
          model.eval()
          a					=model(__A			,   attention_mask=__A			,   token_type_ids=__A			,   labels=__A						)
          self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.seq_length, self.vocab_size)						)


      def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A			,   __A			,   __A			,   __A			,   __A			,   __A						) ->					str:
          a					=MobileBertForNextSentencePrediction(config=__A						)
          model.to(__A						)
          model.eval()
          a					=model(
              __A			,   attention_mask=__A			,   token_type_ids=__A			,   labels=__A			,   )
          self.parent.assertEqual(result.logits.shape			,   (self.batch_size, 2)						)


      def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A			,   __A			,   __A			,   __A			,   __A			,   __A						) ->					str:
          a					=MobileBertForPreTraining(config=__A						)
          model.to(__A						)
          model.eval()
          a					=model(
              __A			,   attention_mask=__A			,   token_type_ids=__A			,   labels=__A			,   next_sentence_label=__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 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A			,   __A			,   __A			,   __A			,   __A			,   __A						) ->					Optional[Any]:
          a					=MobileBertForQuestionAnswering(config=__A						)
          model.to(__A						)
          model.eval()
          a					=model(
              __A			,   attention_mask=__A			,   token_type_ids=__A			,   start_positions=__A			,   end_positions=__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 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A			,   __A			,   __A			,   __A			,   __A			,   __A						) ->					Any:
          a					=self.num_labels
          a					=MobileBertForSequenceClassification(__A						)
          model.to(__A						)
          model.eval()
          a					=model(__A			,   attention_mask=__A			,   token_type_ids=__A			,   labels=__A						)
          self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.num_labels)						)


      def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A			,   __A			,   __A			,   __A			,   __A			,   __A						) ->					int:
          a					=self.num_labels
          a					=MobileBertForTokenClassification(config=__A						)
          model.to(__A						)
          model.eval()
          a					=model(__A			,   attention_mask=__A			,   token_type_ids=__A			,   labels=__A						)
          self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.seq_length, self.num_labels)						)


      def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A			,   __A			,   __A			,   __A			,   __A			,   __A						) ->					Any:
          a					=self.num_choices
          a					=MobileBertForMultipleChoice(config=__A						)
          model.to(__A						)
          model.eval()
          a					=input_ids.unsqueeze(1						).expand(-1			,   self.num_choices			,   -1						).contiguous()
          a					=token_type_ids.unsqueeze(1						).expand(-1			,   self.num_choices			,   -1						).contiguous()
          a					=input_mask.unsqueeze(1						).expand(-1			,   self.num_choices			,   -1						).contiguous()
          a					=model(
              __A			,   attention_mask=__A			,   token_type_ids=__A			,   labels=__A			,   )
          self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.num_choices)						)





      def 				SCREAMING_SNAKE_CASE							(		self						) ->					Dict:
          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






@require_torch
class 				__A    (	_SCREAMING_SNAKE_CASE,   _SCREAMING_SNAKE_CASE,   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 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A			,   __A=False						) ->					List[str]:
          a					=super()._prepare_for_class(__A			,   __A			,   return_labels=__A						)

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


      def 				SCREAMING_SNAKE_CASE							(		self						) ->					Optional[int]:
          a					=MobileBertModelTester(self						)
          a					=ConfigTester(self			,   config_class=__A			,   hidden_size=37						)


      def 				SCREAMING_SNAKE_CASE							(		self						) ->					List[Any]:
          self.config_tester.run_common_tests()


      def 				SCREAMING_SNAKE_CASE							(		self						) ->					Dict:
          a					=self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_mobilebert_model(*__A						)


      def 				SCREAMING_SNAKE_CASE							(		self						) ->					List[Any]:
          a					=self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_mobilebert_for_masked_lm(*__A						)


      def 				SCREAMING_SNAKE_CASE							(		self						) ->					Any:
          a					=self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_mobilebert_for_multiple_choice(*__A						)


      def 				SCREAMING_SNAKE_CASE							(		self						) ->					Optional[Any]:
          a					=self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*__A						)


      def 				SCREAMING_SNAKE_CASE							(		self						) ->					Any:
          a					=self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_mobilebert_for_pretraining(*__A						)


      def 				SCREAMING_SNAKE_CASE							(		self						) ->					Optional[Any]:
          a					=self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_mobilebert_for_question_answering(*__A						)


      def 				SCREAMING_SNAKE_CASE							(		self						) ->					Any:
          a					=self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_mobilebert_for_sequence_classification(*__A						)





      def 				SCREAMING_SNAKE_CASE							(		self						) ->					str:
          a					=self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_mobilebert_for_token_classification(*__A						)

def _A						(						lowercase		):







    """simple docstring"""
    return torch.tensor(
        lowercase	,   dtype=torch.long	,   device=lowercase	,   )


lowerCamelCase_       :       Optional[int]      						=  1E-3






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






      """simple docstring"""




      @slow
      def 				SCREAMING_SNAKE_CASE							(		self						) ->					List[Any]:
          a					=MobileBertModel.from_pretrained('''google/mobilebert-uncased'''						).to(__A						)
          a					=_long_tensor([[101, 7110, 1005, 1056, 2023, 1_1333, 1_7413, 1029, 102]]						)
          with torch.no_grad():
              a					=model(__A						)[0]
          a					=torch.Size((1, 9, 512)						)
          self.assertEqual(output.shape			,   __A						)
          a					=torch.tensor(
              [
                  [
                      [-2.4_7_3_6_5_2_6E0_7, 8.2_6_9_1_6_5_6E0_4, 1.6_5_2_1_8_3_8E0_5],
                      [-5.7_5_4_1_7_0_4E-0_1, 3.9_0_5_6_0_2_2E0_0, 4.4_0_1_1_5_0_7E0_0],
                      [2.6_0_4_7_3_5_9E0_0, 1.5_6_7_7_6_5_2E0_0, -1.7_3_2_4_1_8_8E-0_1],
                  ]
              ]			,   device=__A			,   )

          # 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
          a					=torch.all((expected_slice / output[..., :3, :3]) >= 1 - TOLERANCE						)
          a					=torch.all((expected_slice / output[..., :3, :3]) <= 1 + TOLERANCE						)

          self.assertTrue(lower_bound and upper_bound						)



"""simple docstring"""







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


lowerCamelCase_       :       int      						=  {"""configuration_mbart""": ["""MBART_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MBartConfig""", """MBartOnnxConfig"""]}

try:
       if not is_sentencepiece_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCamelCase_       :       Dict      						=  ["""MBartTokenizer"""]

try:
       if not is_tokenizers_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCamelCase_       :       str      						=  ["""MBartTokenizerFast"""]

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCamelCase_       :       List[str]      						=  [
           """MBART_PRETRAINED_MODEL_ARCHIVE_LIST""",
           """MBartForCausalLM""",
           """MBartForConditionalGeneration""",
           """MBartForQuestionAnswering""",
           """MBartForSequenceClassification""",
           """MBartModel""",
           """MBartPreTrainedModel""",
       ]

try:
       if not is_tf_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCamelCase_       :       List[Any]      						=  [
           """TFMBartForConditionalGeneration""",
           """TFMBartModel""",
           """TFMBartPreTrainedModel""",
       ]

try:
       if not is_flax_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCamelCase_       :       str      						=  [
           """FlaxMBartForConditionalGeneration""",
           """FlaxMBartForQuestionAnswering""",
           """FlaxMBartForSequenceClassification""",
           """FlaxMBartModel""",
           """FlaxMBartPreTrainedModel""",
       ]


if TYPE_CHECKING:
       from .configuration_mbart import MBART_PRETRAINED_CONFIG_ARCHIVE_MAP, MBartConfig, MBartOnnxConfig

       try:
              if not is_sentencepiece_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .tokenization_mbart import MBartTokenizer

       try:
              if not is_tokenizers_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .tokenization_mbart_fast import MBartTokenizerFast

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_mbart import (
                  MBART_PRETRAINED_MODEL_ARCHIVE_LIST,
                  MBartForCausalLM,
                  MBartForConditionalGeneration,
                  MBartForQuestionAnswering,
                  MBartForSequenceClassification,
                  MBartModel,
                  MBartPreTrainedModel,
              )

       try:
              if not is_tf_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_tf_mbart import TFMBartForConditionalGeneration, TFMBartModel, TFMBartPreTrainedModel

       try:
              if not is_flax_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_flax_mbart import (
                  FlaxMBartForConditionalGeneration,
                  FlaxMBartForQuestionAnswering,
                  FlaxMBartForSequenceClassification,
                  FlaxMBartModel,
                  FlaxMBartPreTrainedModel,
              )

else:
       import sys

       lowerCamelCase_       :       List[Any]      						=  _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)