


from __future__ import annotations

import unittest

from transformers import AlbertConfig, 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
  from transformers.models.albert.modeling_tf_albert import (
      TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFAlbertForMaskedLM,
      TFAlbertForMultipleChoice,
      TFAlbertForPreTraining,
      TFAlbertForQuestionAnswering,
      TFAlbertForSequenceClassification,
      TFAlbertForTokenClassification,
      TFAlbertModel,
  )







class    lowerCamelCase__     :





     def __init__(self	,  UpperCAmelCase	,  UpperCAmelCase=1_3	,  UpperCAmelCase=7	,  UpperCAmelCase=True	,  UpperCAmelCase=True	,  UpperCAmelCase=True	,  UpperCAmelCase=True	,  UpperCAmelCase=9_9	,  UpperCAmelCase=1_6	,  UpperCAmelCase=3_2	,  UpperCAmelCase=2	,  UpperCAmelCase=4	,  UpperCAmelCase=3_7	,  UpperCAmelCase="gelu"	,  UpperCAmelCase=0.1	,  UpperCAmelCase=0.1	,  UpperCAmelCase=5_1_2	,  UpperCAmelCase=1_6	,  UpperCAmelCase=2	,  UpperCAmelCase=0.02	,  UpperCAmelCase=3	,  UpperCAmelCase=4	,  UpperCAmelCase=None	,  )  ->     str:
      _lowercase							   =parent
      _lowercase							   =1_3
      _lowercase							   =7
      _lowercase							   =True
      _lowercase							   =True
      _lowercase							   =True
      _lowercase							   =True
      _lowercase							   =9_9
      _lowercase							   =1_6
      _lowercase							   =3_2
      _lowercase							   =2
      _lowercase							   =4
      _lowercase							   =3_7
      _lowercase							   ='''gelu'''
      _lowercase							   =0.1
      _lowercase							   =0.1
      _lowercase							   =5_1_2
      _lowercase							   =1_6
      _lowercase							   =2
      _lowercase							   =0.02
      _lowercase							   =3
      _lowercase							   =4
      _lowercase							   =None





     def   __A      (self )  ->     Dict:
      _lowercase							   =ids_tensor([self.batch_size, self.seq_length]	,  self.vocab_size )

      _lowercase							   =None
      if self.use_input_mask:
       _lowercase							   =random_attention_mask([self.batch_size, self.seq_length] )

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

      _lowercase							   =None
      _lowercase							   =None
      _lowercase							   =None
      if self.use_labels:
       _lowercase							   =ids_tensor([self.batch_size]	,  self.type_sequence_label_size )
       _lowercase							   =ids_tensor([self.batch_size, self.seq_length]	,  self.num_labels )
       _lowercase							   =ids_tensor([self.batch_size]	,  self.num_choices )

      _lowercase							   =AlbertConfig(
          vocab_size=self.vocab_size	,  hidden_size=self.hidden_size	,  embedding_size=self.embedding_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	,  )

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





     def   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase )  ->     Tuple:
      _lowercase							   =TFAlbertModel(config=UpperCAmelCase )
      # inputs = {'input_ids': input_ids,
      #           'attention_mask': input_mask,
      #           'token_type_ids': token_type_ids}
      # sequence_output, pooled_output = model(**inputs)
      _lowercase							   ={'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
      _lowercase							   =model(UpperCAmelCase )

      _lowercase							   =[input_ids, input_mask]
      _lowercase							   =model(UpperCAmelCase )

      _lowercase							   =model(UpperCAmelCase )

      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   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase )  ->     int:
      _lowercase							   =self.num_labels
      _lowercase							   =TFAlbertForPreTraining(config=UpperCAmelCase )
      _lowercase							   ={'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
      _lowercase							   =model(UpperCAmelCase )
      self.parent.assertEqual(result.prediction_logits.shape	,  (self.batch_size, self.seq_length, self.vocab_size) )
      self.parent.assertEqual(result.sop_logits.shape	,  (self.batch_size, self.num_labels) )





     def   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase )  ->     str:
      _lowercase							   =TFAlbertForMaskedLM(config=UpperCAmelCase )
      _lowercase							   ={'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
      _lowercase							   =model(UpperCAmelCase )
      self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.seq_length, self.vocab_size) )





     def   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase )  ->     str:
      _lowercase							   =self.num_labels
      _lowercase							   =TFAlbertForSequenceClassification(config=UpperCAmelCase )
      _lowercase							   ={'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
      _lowercase							   =model(UpperCAmelCase )
      self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.num_labels) )





     def   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase )  ->     Any:
      _lowercase							   =TFAlbertForQuestionAnswering(config=UpperCAmelCase )
      _lowercase							   ={'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
      _lowercase							   =model(UpperCAmelCase )
      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   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase )  ->     List[Any]:
      _lowercase							   =self.num_choices
      _lowercase							   =TFAlbertForMultipleChoice(config=UpperCAmelCase )
      _lowercase							   =tf.tile(tf.expand_dims(UpperCAmelCase	,  1 )	,  (1, self.num_choices, 1) )
      _lowercase							   =tf.tile(tf.expand_dims(UpperCAmelCase	,  1 )	,  (1, self.num_choices, 1) )
      _lowercase							   =tf.tile(tf.expand_dims(UpperCAmelCase	,  1 )	,  (1, self.num_choices, 1) )
      _lowercase							   ={
          '''input_ids''': multiple_choice_inputs_ids,
          '''attention_mask''': multiple_choice_input_mask,
          '''token_type_ids''': multiple_choice_token_type_ids,
      }
      _lowercase							   =model(UpperCAmelCase )
      self.parent.assertListEqual(list(result['''logits'''].shape )	,  [self.batch_size, self.num_choices] )





     def   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase )  ->     List[str]:
      _lowercase							   =self.num_labels
      _lowercase							   =TFAlbertForTokenClassification(config=UpperCAmelCase )
      _lowercase							   ={
          '''input_ids''': input_ids,
          '''attention_mask''': input_mask,
          '''token_type_ids''': token_type_ids,
      }
      _lowercase							   =model(UpperCAmelCase )
      self.parent.assertListEqual(list(result['''logits'''].shape )	,  [self.batch_size, self.seq_length, self.num_labels] )





     def   __A      (self )  ->     List[Any]:
      _lowercase							   =self.prepare_config_and_inputs()
      (
          (
          _lowercase
      )   ,					(
          _lowercase
      )   ,					(
          _lowercase
      )   ,					(
          _lowercase
      )   ,					(
          _lowercase
      )   ,					(
          _lowercase
      )   ,					(
          _lowercase
      )   ,					
      )							   =config_and_inputs
      _lowercase							   ={'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
      return config, inputs_dict







@require_tf
class    lowerCamelCase__     ( lowerCAmelCase     ,							lowerCAmelCase     ,							unittest.TestCase):
     SCREAMING_SNAKE_CASE__       =    (
         (
             TFAlbertModel,
             TFAlbertForPreTraining,
             TFAlbertForMaskedLM,
             TFAlbertForSequenceClassification,
             TFAlbertForQuestionAnswering,
             TFAlbertForTokenClassification,
             TFAlbertForMultipleChoice,
         )
         if is_tf_available()
         else ()
     )
     SCREAMING_SNAKE_CASE__       =    (
         {
             '''feature-extraction''': TFAlbertModel,
             '''fill-mask''': TFAlbertForMaskedLM,
             '''question-answering''': TFAlbertForQuestionAnswering,
             '''text-classification''': TFAlbertForSequenceClassification,
             '''token-classification''': TFAlbertForTokenClassification,
             '''zero-shot''': TFAlbertForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )
     SCREAMING_SNAKE_CASE__       =    False
     SCREAMING_SNAKE_CASE__       =    False





     def   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase=False )  ->     Any:
      _lowercase							   =super()._prepare_for_class(UpperCAmelCase	,  UpperCAmelCase	,  return_labels=UpperCAmelCase )

      if return_labels:
       if model_class in get_values(UpperCAmelCase ):
        _lowercase							   =tf.zeros(self.model_tester.batch_size	,  dtype=tf.intaa )

      return inputs_dict





     def   __A      (self )  ->     Optional[Any]:
      _lowercase							   =TFAlbertModelTester(self )
      _lowercase							   =ConfigTester(self	,  config_class=UpperCAmelCase	,  hidden_size=3_7 )





     def   __A      (self )  ->     Any:
      self.config_tester.run_common_tests()





     def   __A      (self )  ->     Optional[Any]:
      _lowercase							   =self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_albert_model(*UpperCAmelCase )





     def   __A      (self )  ->     Tuple:
      _lowercase							   =self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_albert_for_pretraining(*UpperCAmelCase )





     def   __A      (self )  ->     str:
      _lowercase							   =self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_albert_for_masked_lm(*UpperCAmelCase )





     def   __A      (self )  ->     str:
      _lowercase							   =self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_albert_for_multiple_choice(*UpperCAmelCase )





     def   __A      (self )  ->     List[Any]:
      _lowercase							   =self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_albert_for_sequence_classification(*UpperCAmelCase )





     def   __A      (self )  ->     Optional[int]:
      _lowercase							   =self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_albert_for_question_answering(*UpperCAmelCase )





     @slow
     def   __A      (self )  ->     Any:
      for model_name in TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       _lowercase							   =TFAlbertModel.from_pretrained(UpperCAmelCase )
       self.assertIsNotNone(UpperCAmelCase )







@require_tf
class    lowerCamelCase__     ( unittest.TestCase):





     @slow
     def   __A      (self )  ->     str:
      _lowercase							   =TFAlbertForPreTraining.from_pretrained('''albert-base-v2''' )
      _lowercase							   =tf.constant([[0, 1, 2, 3, 4, 5]] )
      _lowercase							   =model(UpperCAmelCase )[0]

      _lowercase							   =[1, 6, 3_0_0_0_0]
      self.assertEqual(output.shape	,  UpperCAmelCase )

      _lowercase							   =tf.constant(
          [
              [
                  [4.59_5668, 0.7446_2754, -1.81_8147],
                  [4.595_4347, 0.745_4184, -1.818_8258],
                  [4.595_4905, 0.744_8235, -1.818_2316],
              ]
          ] )
      tf.debugging.assert_near(output[:, :3, :3]	,  UpperCAmelCase	,  atol=1e-4 )





from ..utils import DummyObject, requires_backends







class    lowerCamelCase__     ( metaclass=lowerCAmelCase):
     SCREAMING_SNAKE_CASE__       =    ['''tensorflow_text''']





     def __init__(self	,  *UpperCAmelCase	,  **UpperCAmelCase )  ->     Any:
      requires_backends(self	,  ['''tensorflow_text'''] )

