



"""simple docstring"""
import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTJConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, tooslow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
    import jax
    import jax.numpy as jnp

    from transformers.modeling_flax_pytorch_utils import (
        convert_pytorch_state_dict_to_flax,
        load_flax_weights_in_pytorch_model,
    )
    from transformers.models.gptj.modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel

if is_torch_available():
    import torch






class      _UpperCAmelCase				:
    def __init__(    self						,    snake_case_						,    snake_case_=14						,    snake_case_=7						,    snake_case_=True						,    snake_case_=True						,    snake_case_=False						,    snake_case_=True						,    snake_case_=99						,    snake_case_=32						,    snake_case_=4						,    snake_case_=4						,    snake_case_=4						,    snake_case_=37						,    snake_case_="gelu"						,    snake_case_=0.1						,    snake_case_=0.1						,    snake_case_=5_12						,    snake_case_=0.02						,    ):
      _snake_case     :						int						=   parent
      _snake_case     :						str						=   batch_size
      _snake_case     :						Optional[Any]						=   seq_length
      _snake_case     :						Dict						=   is_training
      _snake_case     :						Optional[Any]						=   use_input_mask
      _snake_case     :						Optional[int]						=   use_token_type_ids
      _snake_case     :						Union[str, Any]						=   use_labels
      _snake_case     :						List[str]						=   vocab_size
      _snake_case     :						Any						=   hidden_size
      _snake_case     :						str						=   rotary_dim
      _snake_case     :						Tuple						=   num_hidden_layers
      _snake_case     :						Optional[Any]						=   num_attention_heads
      _snake_case     :						Tuple						=   intermediate_size
      _snake_case     :						Dict						=   hidden_act
      _snake_case     :						Union[str, Any]						=   hidden_dropout_prob
      _snake_case     :						Union[str, Any]						=   attention_probs_dropout_prob
      _snake_case     :						Optional[Any]						=   max_position_embeddings
      _snake_case     :						Dict						=   initializer_range
      _snake_case     :						Union[str, Any]						=   None
      _snake_case     :						Union[str, Any]						=   vocab_size - 1
      _snake_case     :						int						=   vocab_size - 1
      _snake_case     :						Tuple						=   vocab_size - 1
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						Tuple						=   ids_tensor([self.batch_size, self.seq_length]						,    self.vocab_size     )

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

      _snake_case     :						int						=   GPTJConfig(
          vocab_size=self.vocab_size						,    n_embd=self.hidden_size						,    n_layer=self.num_hidden_layers						,    n_head=self.num_attention_heads						,    n_positions=self.max_position_embeddings						,    use_cache=__A						,    bos_token_id=self.bos_token_id						,    eos_token_id=self.eos_token_id						,    pad_token_id=self.pad_token_id						,    rotary_dim=self.rotary_dim						,    )

      return (config, input_ids, input_mask)
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						Any						=   self.prepare_config_and_inputs()
      _snake_case			,					_snake_case			,					_snake_case     :						List[str]						=   config_and_inputs
      _snake_case     :						List[str]						=   {"input_ids": input_ids, "attention_mask": attention_mask}
      return config, inputs_dict
    def 					lowerCamelCase__	(    self						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_     ):
      _snake_case     :						Union[str, Any]						=   20
      _snake_case     :						List[Any]						=   model_class_name(__A     )

      _snake_case     :						int						=   model.init_cache(input_ids.shape[0]						,    __A     )
      _snake_case     :						Union[str, Any]						=   jnp.ones((input_ids.shape[0], max_decoder_length)						,    dtype="i4"     )

      _snake_case     :						str						=   jnp.broadcast_to(
          jnp.arange(input_ids.shape[-1] - 1     )[None, :]						,    (input_ids.shape[0], input_ids.shape[-1] - 1)     )
      _snake_case     :						List[str]						=   model(
          input_ids[:, :-1]						,    attention_mask=__A						,    past_key_values=__A						,    position_ids=__A						,    )

      _snake_case     :						List[str]						=   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						,    dtype="i4"     )
      _snake_case     :						Union[str, Any]						=   model(
          input_ids[:, -1:]						,    attention_mask=__A						,    past_key_values=outputs_cache.past_key_values						,    position_ids=__A						,    )

      _snake_case     :						Any						=   model(__A     )

      _snake_case     :						Tuple						=   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])     )     )
      self.parent.assertTrue(diff < 1E-3						,    msg=F'Max diff is {diff}'     )


    def 					lowerCamelCase__	(    self						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_     ):
      _snake_case     :						str						=   20
      _snake_case     :						Optional[int]						=   model_class_name(__A     )

      _snake_case     :						Any						=   jnp.concatenate(
          [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])     )]						,    axis=-1						,    )

      _snake_case     :						List[str]						=   model.init_cache(input_ids.shape[0]						,    __A     )
      _snake_case     :						Optional[Any]						=   jnp.broadcast_to(
          jnp.arange(input_ids.shape[-1] - 1     )[None, :]						,    (input_ids.shape[0], input_ids.shape[-1] - 1)     )

      _snake_case     :						Dict						=   model(
          input_ids[:, :-1]						,    attention_mask=__A						,    past_key_values=__A						,    position_ids=__A						,    )
      _snake_case     :						str						=   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						,    dtype="i4"     )
      _snake_case     :						int						=   model(
          input_ids[:, -1:]						,    past_key_values=outputs_cache.past_key_values						,    attention_mask=__A						,    position_ids=__A						,    )

      _snake_case     :						Tuple						=   model(__A						,    attention_mask=__A     )

      _snake_case     :						List[Any]						=   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])     )     )
      self.parent.assertTrue(diff < 1E-3						,    msg=F'Max diff is {diff}'     )






@require_flax
class      _UpperCAmelCase				(					UpperCamelCase__		,					UpperCamelCase__		,					unittest.TestCase):
    __lowercase							: List[Any]  =					(FlaxGPTJModel, FlaxGPTJForCausalLM) if is_flax_available() else ()
    __lowercase							: List[Any]  =					(FlaxGPTJForCausalLM,) if is_flax_available() else ()
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						Any						=   FlaxGPTJModelTester(self     )
    def 					lowerCamelCase__	(    self     ):



      for model_class_name in self.all_model_classes:
        _snake_case			,					_snake_case			,					_snake_case     :						List[Any]						=   self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_use_cache_forward(__A						,    __A						,    __A						,    __A     )
    def 					lowerCamelCase__	(    self     ):



      for model_class_name in self.all_model_classes:
        _snake_case			,					_snake_case			,					_snake_case     :						Optional[Any]						=   self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_use_cache_forward_with_attn_mask(
            __A						,    __A						,    __A						,    __A     )
    @tooslow
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						Union[str, Any]						=   GPTaTokenizer.from_pretrained("gpt2"						,    pad_token="<|endoftext|>"						,    padding_side="left"     )
      _snake_case     :						Tuple						=   tokenizer(["Hello this is a long string", "Hey"]						,    return_tensors="np"						,    padding=__A						,    truncation=__A     )

      _snake_case     :						int						=   FlaxGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"     )
      _snake_case     :						int						=   False
      _snake_case     :						List[Any]						=   model.config.eos_token_id

      _snake_case     :						List[str]						=   jax.jit(model.generate     )

      _snake_case     :						Dict						=   jit_generate(
          inputs["input_ids"]						,    attention_mask=inputs["attention_mask"]						,    pad_token_id=tokenizer.pad_token_id     ).sequences

      _snake_case     :						Union[str, Any]						=   tokenizer.batch_decode(__A						,    skip_special_tokens=__A     )

      _snake_case     :						str						=   [
          "Hello this is a long string of text.\n\nI'm trying to get the text of the",
          "Hey, I'm a little late to the party. I'm going to",
      ]

      self.assertListEqual(__A						,    __A     )
    @is_pt_flax_cross_test
    def 					lowerCamelCase__	(    self     ):
      _snake_case			,					_snake_case     :						Optional[int]						=   self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        with self.subTest(model_class.__name__     ):
          # prepare inputs
          _snake_case     :						Dict						=   self._prepare_for_class(__A						,    __A     )
          _snake_case     :						int						=   {k: torch.tensor(v.tolist()     ) for k, v in prepared_inputs_dict.items()}

          # load corresponding PyTorch class
          _snake_case     :						List[Any]						=   model_class.__name__[4:]  # Skip the "Flax" at the beginning
          _snake_case     :						Union[str, Any]						=   getattr(__A						,    __A     )

          _snake_case			,					_snake_case     :						Any						=   pt_inputs["input_ids"].shape
          _snake_case     :						Optional[int]						=   np.random.randint(0						,    seq_length - 1						,    size=(batch_size,)     )
          for batch_idx, start_index in enumerate(__A     ):
            _snake_case     :						List[Any]						=   0
            _snake_case     :						List[str]						=   1
            _snake_case     :						Tuple						=   0
            _snake_case     :						List[str]						=   1
          _snake_case     :						Optional[Any]						=   pt_model_class(__A     ).eval()
          _snake_case     :						Optional[int]						=   model_class(__A						,    dtype=jnp.floataa     )

          _snake_case     :						List[str]						=   convert_pytorch_state_dict_to_flax(pt_model.state_dict()						,    __A     )
          _snake_case     :						Optional[int]						=   fx_state

          with torch.no_grad():
            _snake_case     :						Dict						=   pt_model(**__A     ).to_tuple()

          _snake_case     :						Dict						=   fx_model(**__A     ).to_tuple()
          self.assertEqual(len(__A     )						,    len(__A     )						,    "Output lengths differ between Flax and PyTorch"     )
          for fx_output, pt_output in zip(__A						,    __A     ):
            self.assert_almost_equals(fx_output[:, -1]						,    pt_output[:, -1].numpy()						,    4E-2     )

          with tempfile.TemporaryDirectory() as tmpdirname:
            pt_model.save_pretrained(__A     )
            _snake_case     :						str						=   model_class.from_pretrained(__A						,    from_pt=__A     )

          _snake_case     :						Optional[int]						=   fx_model_loaded(**__A     ).to_tuple()
          self.assertEqual(
              len(__A     )						,    len(__A     )						,    "Output lengths differ between Flax and PyTorch"     )
          for fx_output_loaded, pt_output in zip(__A						,    __A     ):
            self.assert_almost_equals(fx_output_loaded[:, -1]						,    pt_output[:, -1].numpy()						,    4E-2     )
    @is_pt_flax_cross_test
    def 					lowerCamelCase__	(    self     ):
      _snake_case			,					_snake_case     :						Dict						=   self.model_tester.prepare_config_and_inputs_for_common()
      for model_class in self.all_model_classes:
        with self.subTest(model_class.__name__     ):
          # prepare inputs
          _snake_case     :						Dict						=   self._prepare_for_class(__A						,    __A     )
          _snake_case     :						Any						=   {k: torch.tensor(v.tolist()     ) for k, v in prepared_inputs_dict.items()}

          # load corresponding PyTorch class
          _snake_case     :						Union[str, Any]						=   model_class.__name__[4:]  # Skip the "Flax" at the beginning
          _snake_case     :						str						=   getattr(__A						,    __A     )

          _snake_case     :						Any						=   pt_model_class(__A     ).eval()
          _snake_case     :						List[Any]						=   model_class(__A						,    dtype=jnp.floataa     )

          _snake_case     :						int						=   load_flax_weights_in_pytorch_model(__A						,    fx_model.params     )
          _snake_case			,					_snake_case     :						Dict						=   pt_inputs["input_ids"].shape
          _snake_case     :						List[str]						=   np.random.randint(0						,    seq_length - 1						,    size=(batch_size,)     )
          for batch_idx, start_index in enumerate(__A     ):
            _snake_case     :						Any						=   0
            _snake_case     :						List[str]						=   1
            _snake_case     :						List[str]						=   0
            _snake_case     :						Tuple						=   1

          # make sure weights are tied in PyTorch
          pt_model.tie_weights()

          with torch.no_grad():
            _snake_case     :						str						=   pt_model(**__A     ).to_tuple()

          _snake_case     :						str						=   fx_model(**__A     ).to_tuple()
          self.assertEqual(len(__A     )						,    len(__A     )						,    "Output lengths differ between Flax and PyTorch"     )
          for fx_output, pt_output in zip(__A						,    __A     ):
            self.assert_almost_equals(fx_output[:, -1]						,    pt_output[:, -1].numpy()						,    4E-2     )

          with tempfile.TemporaryDirectory() as tmpdirname:
            fx_model.save_pretrained(__A     )
            _snake_case     :						str						=   pt_model_class.from_pretrained(__A						,    from_flax=__A     )

          with torch.no_grad():
            _snake_case     :						List[Any]						=   pt_model_loaded(**__A     ).to_tuple()

          self.assertEqual(
              len(__A     )						,    len(__A     )						,    "Output lengths differ between Flax and PyTorch"     )
          for fx_output, pt_output in zip(__A						,    __A     ):
            self.assert_almost_equals(fx_output[:, -1]						,    pt_output[:, -1].numpy()						,    4E-2     )


    @tooslow
    def 					lowerCamelCase__	(    self     ):



      for model_class_name in self.all_model_classes:
        _snake_case     :						Dict						=   model_class_name.from_pretrained("EleutherAI/gpt-j-6B"     )
        _snake_case     :						Optional[Any]						=   model(np.ones((1, 1)     )     )
        self.assertIsNotNone(__A     )



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

