



"""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,
)


if is_sentencepiece_available():
    from ..ta.tokenization_ta import TaTokenizer
else:
    from ...utils.dummy_sentencepiece_objects import TaTokenizer

_a      :			List[str]					=							TaTokenizer

if is_tokenizers_available():
    from ..ta.tokenization_ta_fast import TaTokenizerFast
else:
    from ...utils.dummy_tokenizers_objects import TaTokenizerFast

_a      :			Optional[Any]					=							TaTokenizerFast

_a      :			Union[str, Any]					=							{"""configuration_mt5""": ["""MT5Config""", """MT5OnnxConfig"""]}

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    _a      :			Optional[Any]					=							[
        """MT5EncoderModel""",
        """MT5ForConditionalGeneration""",
        """MT5ForQuestionAnswering""",
        """MT5Model""",
        """MT5PreTrainedModel""",
        """MT5Stack""",
    ]

try:
    if not is_tf_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    _a      :			Optional[Any]					=							["""TFMT5EncoderModel""", """TFMT5ForConditionalGeneration""", """TFMT5Model"""]

try:
    if not is_flax_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    _a      :			List[str]					=							["""FlaxMT5EncoderModel""", """FlaxMT5ForConditionalGeneration""", """FlaxMT5Model"""]


if TYPE_CHECKING:
    from .configuration_mta import MTaConfig, MTaOnnxConfig

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_mta import (
            MTaEncoderModel,
            MTaForConditionalGeneration,
            MTaForQuestionAnswering,
            MTaModel,
            MTaPreTrainedModel,
            MTaStack,
        )

    try:
        if not is_tf_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_tf_mta import TFMTaEncoderModel, TFMTaForConditionalGeneration, TFMTaModel

    try:
        if not is_flax_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_flax_mta import FlaxMTaEncoderModel, FlaxMTaForConditionalGeneration, FlaxMTaModel

else:
    import sys

    _a      :			List[str]					=							_LazyModule(
        __name__,
        globals()["""__file__"""],
        _import_structure,
        extra_objects={"""MT5Tokenizer""": MTaTokenizer, """MT5TokenizerFast""": MTaTokenizerFast},
        module_spec=__spec__,
    )



"""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      _UpperCAmelCase				(					_snake_case		,					_snake_case		,					unittest.TestCase):
    __lowercase							: Dict  =					(
        (
            TFMobileBertModel,
            TFMobileBertForMaskedLM,
            TFMobileBertForNextSentencePrediction,
            TFMobileBertForPreTraining,
            TFMobileBertForQuestionAnswering,
            TFMobileBertForSequenceClassification,
            TFMobileBertForTokenClassification,
            TFMobileBertForMultipleChoice,
        )
        if is_tf_available()
        else ()
    )
    __lowercase							: Optional[Any]  =					(
        {
            """feature-extraction""": TFMobileBertModel,
            """fill-mask""": TFMobileBertForMaskedLM,
            """question-answering""": TFMobileBertForQuestionAnswering,
            """text-classification""": TFMobileBertForSequenceClassification,
            """token-classification""": TFMobileBertForTokenClassification,
            """zero-shot""": TFMobileBertForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )
    __lowercase							: Union[str, Any]  =					False
    __lowercase							: Optional[int]  =					False
    def 					lowerCamelCase__	(    self						,    snake_case_						,    snake_case_						,    snake_case_=False     ):
      _snake_case     :						Union[str, Any]						=   super()._prepare_for_class(snake_case_						,    snake_case_						,    return_labels=snake_case_     )

      if return_labels:
        if model_class in get_values(snake_case_     ):
          _snake_case     :						Any						=   tf.zeros(self.model_tester.batch_size						,    dtype=tf.intaa     )

      return inputs_dict






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

          _snake_case     :						Optional[Any]						=   None
          if self.use_input_mask:
            _snake_case     :						Optional[Any]						=   random_attention_mask([self.batch_size, self.seq_length]     )

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

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

          _snake_case     :						Tuple						=   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 					lowerCamelCase__	(    self						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_     ):
          _snake_case     :						Dict						=   TFMobileBertModel(config=snake_case_     )
          _snake_case     :						Optional[Any]						=   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
          _snake_case     :						Optional[int]						=   model(snake_case_     )

          _snake_case     :						Union[str, Any]						=   [input_ids, input_mask]
          _snake_case     :						Optional[Any]						=   model(snake_case_     )

          _snake_case     :						Dict						=   model(snake_case_     )

          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 					lowerCamelCase__	(    self						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_     ):
          _snake_case     :						List[Any]						=   TFMobileBertForMaskedLM(config=snake_case_     )
          _snake_case     :						Tuple						=   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
          _snake_case     :						List[str]						=   model(snake_case_     )
          self.parent.assertEqual(result.logits.shape						,    (self.batch_size, self.seq_length, self.vocab_size)     )
        def 					lowerCamelCase__	(    self						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_     ):
          _snake_case     :						Union[str, Any]						=   TFMobileBertForNextSentencePrediction(config=snake_case_     )
          _snake_case     :						Dict						=   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
          _snake_case     :						Tuple						=   model(snake_case_     )
          self.parent.assertEqual(result.logits.shape						,    (self.batch_size, 2)     )
        def 					lowerCamelCase__	(    self						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_     ):
          _snake_case     :						str						=   TFMobileBertForPreTraining(config=snake_case_     )
          _snake_case     :						Union[str, Any]						=   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
          _snake_case     :						List[Any]						=   model(snake_case_     )
          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 					lowerCamelCase__	(    self						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_     ):
          _snake_case     :						str						=   self.num_labels
          _snake_case     :						str						=   TFMobileBertForSequenceClassification(config=snake_case_     )
          _snake_case     :						Union[str, Any]						=   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
          _snake_case     :						Optional[int]						=   model(snake_case_     )
          self.parent.assertEqual(result.logits.shape						,    (self.batch_size, self.num_labels)     )
        def 					lowerCamelCase__	(    self						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_     ):
          _snake_case     :						Any						=   self.num_choices
          _snake_case     :						Tuple						=   TFMobileBertForMultipleChoice(config=snake_case_     )
          _snake_case     :						List[Any]						=   tf.tile(tf.expand_dims(snake_case_						,    1     )						,    (1, self.num_choices, 1)     )
          _snake_case     :						List[str]						=   tf.tile(tf.expand_dims(snake_case_						,    1     )						,    (1, self.num_choices, 1)     )
          _snake_case     :						Tuple						=   tf.tile(tf.expand_dims(snake_case_						,    1     )						,    (1, self.num_choices, 1)     )
          _snake_case     :						int						=   {
              "input_ids": multiple_choice_inputs_ids,
              "attention_mask": multiple_choice_input_mask,
              "token_type_ids": multiple_choice_token_type_ids,
          }
          _snake_case     :						Optional[Any]						=   model(snake_case_     )
          self.parent.assertEqual(result.logits.shape						,    (self.batch_size, self.num_choices)     )
        def 					lowerCamelCase__	(    self						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_     ):
          _snake_case     :						Union[str, Any]						=   self.num_labels
          _snake_case     :						Optional[int]						=   TFMobileBertForTokenClassification(config=snake_case_     )
          _snake_case     :						List[Any]						=   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
          _snake_case     :						List[Any]						=   model(snake_case_     )
          self.parent.assertEqual(result.logits.shape						,    (self.batch_size, self.seq_length, self.num_labels)     )
        def 					lowerCamelCase__	(    self						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_     ):
          _snake_case     :						int						=   TFMobileBertForQuestionAnswering(config=snake_case_     )
          _snake_case     :						str						=   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
          _snake_case     :						Union[str, Any]						=   model(snake_case_     )
          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 					lowerCamelCase__	(    self     ):
          _snake_case     :						Optional[Any]						=   self.prepare_config_and_inputs()
          (
              (
              _snake_case
          )			,					(
              _snake_case
          )			,					(
              _snake_case
          )			,					(
              _snake_case
          )			,					(
              _snake_case
          )			,					(
              _snake_case
          )			,					(
              _snake_case
          )			,					
          )     :						Tuple						=   config_and_inputs
          _snake_case     :						Tuple						=   {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
          return config, inputs_dict


    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						int						=   TFMobileBertModelTest.TFMobileBertModelTester(self     )
      _snake_case     :						Optional[Any]						=   ConfigTester(self						,    config_class=snake_case_						,    hidden_size=37     )
    def 					lowerCamelCase__	(    self     ):
      self.config_tester.run_common_tests()
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						int						=   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_model(*snake_case_     )
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						List[Any]						=   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_masked_lm(*snake_case_     )
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						Optional[Any]						=   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_multiple_choice(*snake_case_     )
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						Tuple						=   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*snake_case_     )
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						str						=   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_pretraining(*snake_case_     )
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						Dict						=   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_question_answering(*snake_case_     )
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						Any						=   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_sequence_classification(*snake_case_     )
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						Any						=   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_token_classification(*snake_case_     )


    @slow
    def 					lowerCamelCase__	(    self     ):
      # for model_name in TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      for model_name in ["google/mobilebert-uncased"]:
        _snake_case     :						str						=   TFMobileBertModel.from_pretrained(snake_case_     )
        self.assertIsNotNone(snake_case_     )






@require_tf
class      _UpperCAmelCase				(					unittest.TestCase):
    @slow
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						Any						=   TFMobileBertForPreTraining.from_pretrained("google/mobilebert-uncased"     )
      _snake_case     :						str						=   tf.constant([[0, 1, 2, 3, 4, 5]]     )
      _snake_case     :						Union[str, Any]						=   model(snake_case_     )[0]

      _snake_case     :						int						=   [1, 6, 3_05_22]
      self.assertEqual(output.shape						,    snake_case_     )

      _snake_case     :						Optional[Any]						=   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]						,    snake_case_						,    atol=1E-4     )

