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

    if return_labels:
     if model_class in get_values(snake_case_       ):
      __a			:  List[str]		    =  tf.zeros(self.model_tester.batch_size ,    dtype=tf.intaa       )

    return inputs_dict



   class 		UpperCamelCase__  (	__lowercase   ):
      def __init__(self			: Optional[Any] ,    snake_case_			: Optional[int] ,    snake_case_			: int=1_3 ,    snake_case_			: Tuple=7 ,    snake_case_			: Union[str, Any]=True ,    snake_case_			: Union[str, Any]=True ,    snake_case_			: Tuple=True ,    snake_case_			: Union[str, Any]=True ,    snake_case_			: str=9_9 ,    snake_case_			: Dict=3_2 ,    snake_case_			: List[str]=3_2 ,    snake_case_			: Any=2 ,    snake_case_			: Union[str, Any]=4 ,    snake_case_			: Optional[int]=3_7 ,    snake_case_			: Union[str, Any]="gelu" ,    snake_case_			: Any=0.1 ,    snake_case_			: Any=0.1 ,    snake_case_			: Dict=5_1_2 ,    snake_case_			: Optional[int]=1_6 ,    snake_case_			: Union[str, Any]=2 ,    snake_case_			: Dict=0.02 ,    snake_case_			: Optional[Any]=3 ,    snake_case_			: str=4 ,    snake_case_			: Tuple=None ,    ):
       __a			:  Dict		    =  parent
       __a			:  Optional[int]		    =  batch_size
       __a			:  int		    =  seq_length
       __a			:  Optional[int]		    =  is_training
       __a			:  Any		    =  use_input_mask
       __a			:  int		    =  use_token_type_ids
       __a			:  Union[str, Any]		    =  use_labels
       __a			:  int		    =  vocab_size
       __a			:  int		    =  hidden_size
       __a			:  Any		    =  num_hidden_layers
       __a			:  List[Any]		    =  num_attention_heads
       __a			:  Union[str, Any]		    =  intermediate_size
       __a			:  Optional[int]		    =  hidden_act
       __a			:  Optional[int]		    =  hidden_dropout_prob
       __a			:  List[Any]		    =  attention_probs_dropout_prob
       __a			:  Optional[int]		    =  max_position_embeddings
       __a			:  str		    =  type_vocab_size
       __a			:  List[str]		    =  type_sequence_label_size
       __a			:  List[str]		    =  initializer_range
       __a			:  List[Any]		    =  num_labels
       __a			:  List[str]		    =  num_choices
       __a			:  Any		    =  scope
       __a			:  str		    =  embedding_size
      def 				lowerCAmelCase      (self			: Union[str, Any]       ):
       __a			:  Optional[Any]		    =  ids_tensor([self.batch_size, self.seq_length] ,    self.vocab_size       )

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

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

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

       __a			:  Dict		    =  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			: List[str] ,    snake_case_			: List[Any] ,    snake_case_			: Optional[int] ,    snake_case_			: Dict ,    snake_case_			: Union[str, Any] ,    snake_case_			: Union[str, Any] ,    snake_case_			: Optional[int] ,    snake_case_			: Any       ):
       __a			:  Dict		    =  TFMobileBertModel(config=snake_case_       )
       __a			:  int		    =  {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
       __a			:  Tuple		    =  model(snake_case_       )

       __a			:  Any		    =  [input_ids, input_mask]
       __a			:  Tuple		    =  model(snake_case_       )

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





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







   @slow
   def 				lowerCAmelCase      (self			: Optional[int]       ):
    # for model_name in TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
    for model_name in ["google/mobilebert-uncased"]:
     __a			:  Optional[int]		    =  TFMobileBertModel.from_pretrained(snake_case_       )
     self.assertIsNotNone(snake_case_       )





@require_tf
class 		UpperCamelCase__  (	unittest.TestCase   ):
   @slow
   def 				lowerCAmelCase      (self			: Dict       ):
    __a			:  Any		    =  TFMobileBertForPreTraining.from_pretrained('''google/mobilebert-uncased'''       )
    __a			:  int		    =  tf.constant([[0, 1, 2, 3, 4, 5]]       )
    __a			:  str		    =  model(snake_case_       )[0]

    __a			:  Optional[int]		    =  [1, 6, 3_0_5_2_2]
    self.assertEqual(output.shape ,    snake_case_       )

    __a			:  str		    =  tf.constant(
        [
            [
                [-4.591_9547, -9.24_8295, -9.64_5256],
                [-6.730_6175, -6.44_0284, -6.605_2837],
                [-7.274_3506, -6.784_7915, -6.02_4673],
            ]
        ]       )
    tf.debugging.assert_near(output[:, :3, :3] ,    snake_case_ ,    atol=1E-4       )


import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
      from transformers import AutoModelForSeqaSeqLM, AutoTokenizer



@require_torch
@require_sentencepiece
@require_tokenizers
class 		UpperCamelCase__  (	unittest.TestCase   ):
   @slow
   def 				lowerCAmelCase      (self			: Tuple       ):
    __a			:  List[str]		    =  AutoModelForSeqaSeqLM.from_pretrained('''google/mt5-small''' ,    return_dict=snake_case_       ).to(snake_case_       )
    __a			:  List[Any]		    =  AutoTokenizer.from_pretrained('''google/mt5-small'''       )

    __a			:  Optional[int]		    =  tokenizer('''Hello there''' ,    return_tensors='''pt'''       ).input_ids
    __a			:  Dict		    =  tokenizer('''Hi I am''' ,    return_tensors='''pt'''       ).input_ids

    __a			:  Optional[Any]		    =  model(input_ids.to(snake_case_       ) ,    labels=labels.to(snake_case_       )       ).loss
    __a			:  Tuple		    =  -(labels.shape[-1] * loss.item())

    __a			:  Dict		    =  -84.9127
    self.assertTrue(abs(mtf_score - EXPECTED_SCORE       ) < 1E-4       )
