

import logging

import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.albert.modeling_albert import (
    ALBERT_INPUTS_DOCSTRING,
    ALBERT_START_DOCSTRING,
    AlbertModel,
    AlbertPreTrainedModel,
    AlbertTransformer,
)


_SCREAMING_SNAKE_CASE							 =			logging.getLogger(__name__)


class      lowerCamelCase_ (			_lowerCAmelCase):




 def     a       (			self  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase=None  ,  _UpperCAmelCase=None							):
   if current_layer == 0:
     __UpperCAmelCase	:  Union[str, Any]		=     self.embedding_hidden_mapping_in(_UpperCAmelCase							)
   else:
     __UpperCAmelCase	:  List[Any]		=     hidden_states[0]

   __UpperCAmelCase	:  Tuple		=     int(self.config.num_hidden_layers / self.config.num_hidden_groups							)

   # Index of the hidden group
   __UpperCAmelCase	:  Optional[int]		=     int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)							)

   __UpperCAmelCase	:  Any		=     self.albert_layer_groups[group_idx](
       _UpperCAmelCase  ,  _UpperCAmelCase  ,  head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]  ,  )
   __UpperCAmelCase	:  Optional[int]		=     layer_group_output[0]

   return (hidden_states,)



@add_start_docstrings(
    '''The bare ALBERT Model transformer with PABEE outputting raw hidden-states without any specific head on top.'''	,				_lowerCAmelCase	,				)
class      lowerCamelCase_ (			_lowerCAmelCase):




 def __init__(			self  ,  _UpperCAmelCase							):
   super().__init__(_UpperCAmelCase							)

   __UpperCAmelCase	:  Optional[int]		=     AlbertTransformerWithPabee(_UpperCAmelCase							)

   self.init_weights()
   __UpperCAmelCase	:  int		=     0
   __UpperCAmelCase	:  Optional[int]		=     0
   __UpperCAmelCase	:  Dict		=     0

   __UpperCAmelCase	:  Optional[int]		=     0




 def     a       (			self  ,  _UpperCAmelCase							):
   __UpperCAmelCase	:  int		=     threshold




 def     a       (			self  ,  _UpperCAmelCase							):
   __UpperCAmelCase	:  Union[str, Any]		=     patience




 def     a       (			self							):
   __UpperCAmelCase	:  List[Any]		=     0
   __UpperCAmelCase	:  Tuple		=     0




 def     a       (			self							):
   __UpperCAmelCase	:  List[str]		=     self.inference_layers_num / self.inference_instances_num
   __UpperCAmelCase	:  List[Any]		=     (
       F"*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ="
       F" {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***"
   )
   print(_UpperCAmelCase							)




 @add_start_docstrings_to_model_forward(_UpperCAmelCase							)
 def     a       (			self  ,  _UpperCAmelCase=None  ,  _UpperCAmelCase=None  ,  _UpperCAmelCase=None  ,  _UpperCAmelCase=None  ,  _UpperCAmelCase=None  ,  _UpperCAmelCase=None  ,  _UpperCAmelCase=None  ,  _UpperCAmelCase=None  ,  _UpperCAmelCase=False  ,  ):

   if input_ids is not None and inputs_embeds is not None:
     raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time'							)
   elif input_ids is not None:
     __UpperCAmelCase	:  List[Any]		=     input_ids.size()
   elif inputs_embeds is not None:
     __UpperCAmelCase	:  Optional[Any]		=     inputs_embeds.size()[:-1]
   else:
     raise ValueError('You have to specify either input_ids or inputs_embeds'							)

   __UpperCAmelCase	:  Any		=     input_ids.device if input_ids is not None else inputs_embeds.device

   if attention_mask is None:
     __UpperCAmelCase	:  Any		=     torch.ones(_UpperCAmelCase  ,  device=_UpperCAmelCase							)
   if token_type_ids is None:
     __UpperCAmelCase	:  Dict		=     torch.zeros(_UpperCAmelCase  ,  dtype=torch.long  ,  device=_UpperCAmelCase							)

   __UpperCAmelCase	:  int		=     attention_mask.unsqueeze(1							).unsqueeze(2							)
   __UpperCAmelCase	:  List[Any]		=     extended_attention_mask.to(dtype=self.dtype							)  # fp16 compatibility
   __UpperCAmelCase	:  Dict		=     (1.0 - extended_attention_mask) * -1_0_0_0_0.0
   __UpperCAmelCase	:  Dict		=     self.get_head_mask(_UpperCAmelCase  ,  self.config.num_hidden_layers							)

   __UpperCAmelCase	:  Optional[Any]		=     self.embeddings(
       _UpperCAmelCase  ,  position_ids=_UpperCAmelCase  ,  token_type_ids=_UpperCAmelCase  ,  inputs_embeds=_UpperCAmelCase							)
   __UpperCAmelCase	:  int		=     embedding_output

   if self.training:
     __UpperCAmelCase	:  Dict		=     []
     for i in range(self.config.num_hidden_layers							):
       __UpperCAmelCase	:  Union[str, Any]		=     self.encoder.adaptive_forward(
           _UpperCAmelCase  ,  current_layer=_UpperCAmelCase  ,  attention_mask=_UpperCAmelCase  ,  head_mask=_UpperCAmelCase  ,  )

       __UpperCAmelCase	:  int		=     self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]							)							)
       __UpperCAmelCase	:  Optional[Any]		=     output_layers[i](output_dropout(_UpperCAmelCase							)							)
       res.append(_UpperCAmelCase							)
   elif self.patience == 0:  # Use all layers for inference
     __UpperCAmelCase	:  str		=     self.encoder(_UpperCAmelCase  ,  _UpperCAmelCase  ,  head_mask=_UpperCAmelCase							)
     __UpperCAmelCase	:  List[str]		=     self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]							)							)
     __UpperCAmelCase	:  Tuple		=     [output_layers[self.config.num_hidden_layers - 1](_UpperCAmelCase							)]
   else:
     __UpperCAmelCase	:  Tuple		=     0
     __UpperCAmelCase	:  List[Any]		=     None
     __UpperCAmelCase	:  Any		=     0
     for i in range(self.config.num_hidden_layers							):
       calculated_layer_num += 1
       __UpperCAmelCase	:  Union[str, Any]		=     self.encoder.adaptive_forward(
           _UpperCAmelCase  ,  current_layer=_UpperCAmelCase  ,  attention_mask=_UpperCAmelCase  ,  head_mask=_UpperCAmelCase  ,  )

       __UpperCAmelCase	:  List[str]		=     self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]							)							)
       __UpperCAmelCase	:  Tuple		=     output_layers[i](_UpperCAmelCase							)
       if regression:
         __UpperCAmelCase	:  Dict		=     logits.detach()
         if patient_result is not None:
           __UpperCAmelCase	:  Any		=     patient_result.detach()
         if (patient_result is not None) and torch.abs(patient_result - labels							) < self.regression_threshold:
           patient_counter += 1
         else:
           __UpperCAmelCase	:  Union[str, Any]		=     0
       else:
         __UpperCAmelCase	:  List[str]		=     logits.detach().argmax(dim=1							)
         if patient_result is not None:
           __UpperCAmelCase	:  List[Any]		=     patient_result.detach().argmax(dim=1							)
         if (patient_result is not None) and torch.all(labels.eq(_UpperCAmelCase							)							):
           patient_counter += 1
         else:
           __UpperCAmelCase	:  Optional[int]		=     0

       __UpperCAmelCase	:  Union[str, Any]		=     logits
       if patient_counter == self.patience:
         break
     __UpperCAmelCase	:  List[Any]		=     [patient_result]
     self.inference_layers_num += calculated_layer_num
     self.inference_instances_num += 1

   return res



@add_start_docstrings(
    '''Albert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of
    the pooled output) e.g. for GLUE tasks. '''	,				_lowerCAmelCase	,				)
class      lowerCamelCase_ (			_lowerCAmelCase):




 def __init__(			self  ,  _UpperCAmelCase							):
   super().__init__(_UpperCAmelCase							)
   __UpperCAmelCase	:  str		=     config.num_labels

   __UpperCAmelCase	:  Any		=     AlbertModelWithPabee(_UpperCAmelCase							)
   __UpperCAmelCase	:  Tuple		=     nn.Dropout(config.classifier_dropout_prob							)
   __UpperCAmelCase	:  Optional[int]		=     nn.ModuleList(
       [nn.Linear(config.hidden_size  ,  self.config.num_labels							) for _ in range(config.num_hidden_layers							)]							)

   self.init_weights()




 @add_start_docstrings_to_model_forward(_UpperCAmelCase							)
 def     a       (			self  ,  _UpperCAmelCase=None  ,  _UpperCAmelCase=None  ,  _UpperCAmelCase=None  ,  _UpperCAmelCase=None  ,  _UpperCAmelCase=None  ,  _UpperCAmelCase=None  ,  _UpperCAmelCase=None  ,  ):
   __UpperCAmelCase	:  Tuple		=     self.albert(
       input_ids=_UpperCAmelCase  ,  attention_mask=_UpperCAmelCase  ,  token_type_ids=_UpperCAmelCase  ,  position_ids=_UpperCAmelCase  ,  head_mask=_UpperCAmelCase  ,  inputs_embeds=_UpperCAmelCase  ,  output_dropout=self.dropout  ,  output_layers=self.classifiers  ,  regression=self.num_labels == 1  ,  )

   __UpperCAmelCase	:  List[Any]		=     (logits[-1],)

   if labels is not None:
     __UpperCAmelCase	:  Optional[int]		=     None
     __UpperCAmelCase	:  Dict		=     0
     for ix, logits_item in enumerate(_UpperCAmelCase							):
       if self.num_labels == 1:
         #  We are doing regression
         __UpperCAmelCase	:  List[Any]		=     MSELoss()
         __UpperCAmelCase	:  Any		=     loss_fct(logits_item.view(-1							)  ,  labels.view(-1							)							)
       else:
         __UpperCAmelCase	:  Tuple		=     CrossEntropyLoss()
         __UpperCAmelCase	:  List[Any]		=     loss_fct(logits_item.view(-1  ,  self.num_labels							)  ,  labels.view(-1							)							)
       if total_loss is None:
         __UpperCAmelCase	:  Union[str, Any]		=     loss
       else:
         total_loss += loss * (ix + 1)
       total_weights += ix + 1
     __UpperCAmelCase	:  List[Any]		=     (total_loss / total_weights,) + outputs

   return outputs







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



       """simple docstring"""






       def __init__(		self				,       _lowerCamelCase				,       _lowerCamelCase=13				,       _lowerCamelCase=7				,       _lowerCamelCase=True				,       _lowerCamelCase=True				,       _lowerCamelCase=True				,       _lowerCamelCase=True				,       _lowerCamelCase=99				,       _lowerCamelCase=16				,       _lowerCamelCase=32				,       _lowerCamelCase=2				,       _lowerCamelCase=4				,       _lowerCamelCase=37				,       _lowerCamelCase="gelu"				,       _lowerCamelCase=0.1				,       _lowerCamelCase=0.1				,       _lowerCamelCase=512				,       _lowerCamelCase=16				,       _lowerCamelCase=2				,       _lowerCamelCase=0.0_2				,       _lowerCamelCase=3				,       _lowerCamelCase=4				,       _lowerCamelCase=None				,       )					->List[Any]:
           SCREAMING_SNAKE_CASE				:				int				      = parent
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = 13
           SCREAMING_SNAKE_CASE				:				int				      = 7
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = True
           SCREAMING_SNAKE_CASE				:				List[str]				      = True
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = True
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = True
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = 99
           SCREAMING_SNAKE_CASE				:				str				      = 16
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = 32
           SCREAMING_SNAKE_CASE				:				str				      = 2
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = 4
           SCREAMING_SNAKE_CASE				:				int				      = 37
           SCREAMING_SNAKE_CASE				:				Any				      = '''gelu'''
           SCREAMING_SNAKE_CASE				:				str				      = 0.1
           SCREAMING_SNAKE_CASE				:				Dict				      = 0.1
           SCREAMING_SNAKE_CASE				:				Tuple				      = 512
           SCREAMING_SNAKE_CASE				:				Any				      = 16
           SCREAMING_SNAKE_CASE				:				str				      = 2
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = 0.0_2
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = 3
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = 4
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = None






       def 	__lowerCAmelCase     (		self   )					->List[str]:
           SCREAMING_SNAKE_CASE				:				Tuple				      = ids_tensor([self.batch_size, self.seq_length]				,       self.vocab_size   )

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

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

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

           SCREAMING_SNAKE_CASE				:				List[Any]				      = 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 	__lowerCAmelCase     (		self				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase   )					->List[str]:
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = TFAlbertModel(config=_lowerCamelCase   )
           # inputs = {'input_ids': input_ids,
           #           'attention_mask': input_mask,
           #           'token_type_ids': token_type_ids}
           # sequence_output, pooled_output = model(**inputs)
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = model(_lowerCamelCase   )

           SCREAMING_SNAKE_CASE				:				List[str]				      = [input_ids, input_mask]
           SCREAMING_SNAKE_CASE				:				List[str]				      = model(_lowerCamelCase   )

           SCREAMING_SNAKE_CASE				:				Any				      = model(_lowerCamelCase   )

           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				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase   )					->List[Any]:
           SCREAMING_SNAKE_CASE				:				str				      = self.num_labels
           SCREAMING_SNAKE_CASE				:				List[str]				      = TFAlbertForPreTraining(config=_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Any				      = {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
           SCREAMING_SNAKE_CASE				:				int				      = model(_lowerCamelCase   )
           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 	__lowerCAmelCase     (		self				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase   )					->str:
           SCREAMING_SNAKE_CASE				:				str				      = TFAlbertForMaskedLM(config=_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				int				      = {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = model(_lowerCamelCase   )
           self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.seq_length, self.vocab_size)   )






       def 	__lowerCAmelCase     (		self				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase   )					->int:
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = self.num_labels
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = TFAlbertForSequenceClassification(config=_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
           SCREAMING_SNAKE_CASE				:				Tuple				      = model(_lowerCamelCase   )
           self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.num_labels)   )






       def 	__lowerCAmelCase     (		self				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase   )					->Dict:
           SCREAMING_SNAKE_CASE				:				List[str]				      = TFAlbertForQuestionAnswering(config=_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
           SCREAMING_SNAKE_CASE				:				int				      = model(_lowerCamelCase   )
           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				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase   )					->Any:
           SCREAMING_SNAKE_CASE				:				List[Any]				      = self.num_choices
           SCREAMING_SNAKE_CASE				:				Any				      = TFAlbertForMultipleChoice(config=_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Tuple				      = tf.tile(tf.expand_dims(_lowerCamelCase				,       1   )				,       (1, self.num_choices, 1)   )
           SCREAMING_SNAKE_CASE				:				Tuple				      = tf.tile(tf.expand_dims(_lowerCamelCase				,       1   )				,       (1, self.num_choices, 1)   )
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = tf.tile(tf.expand_dims(_lowerCamelCase				,       1   )				,       (1, self.num_choices, 1)   )
           SCREAMING_SNAKE_CASE				:				Tuple				      = {
               '''input_ids''': multiple_choice_inputs_ids,
               '''attention_mask''': multiple_choice_input_mask,
               '''token_type_ids''': multiple_choice_token_type_ids,
           }
           SCREAMING_SNAKE_CASE				:				str				      = model(_lowerCamelCase   )
           self.parent.assertListEqual(list(result['''logits'''].shape   )				,       [self.batch_size, self.num_choices]   )






       def 	__lowerCAmelCase     (		self				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase   )					->Optional[Any]:
           SCREAMING_SNAKE_CASE				:				str				      = self.num_labels
           SCREAMING_SNAKE_CASE				:				List[Any]				      = TFAlbertForTokenClassification(config=_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = {
               '''input_ids''': input_ids,
               '''attention_mask''': input_mask,
               '''token_type_ids''': token_type_ids,
           }
           SCREAMING_SNAKE_CASE				:				str				      = model(_lowerCamelCase   )
           self.parent.assertListEqual(list(result['''logits'''].shape   )				,       [self.batch_size, self.seq_length, self.num_labels]   )






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






@require_tf
class     a_       (			a__				,    a__				,    unittest.TestCase							):



       """simple docstring"""
       __SCREAMING_SNAKE_CASE       :					int						=  (
           (
               TFAlbertModel,
               TFAlbertForPreTraining,
               TFAlbertForMaskedLM,
               TFAlbertForSequenceClassification,
               TFAlbertForQuestionAnswering,
               TFAlbertForTokenClassification,
               TFAlbertForMultipleChoice,
           )
           if is_tf_available()
           else ()
       )
       __SCREAMING_SNAKE_CASE       :					Optional[int]						=  (
           {
               '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       :					Tuple						=  False
       __SCREAMING_SNAKE_CASE       :					Any						=  False






       def 	__lowerCAmelCase     (		self				,       _lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase=False   )					->int:
           SCREAMING_SNAKE_CASE				:				Tuple				      = super()._prepare_for_class(_lowerCamelCase				,       _lowerCamelCase				,       return_labels=_lowerCamelCase   )

           if return_labels:
               if model_class in get_values(_lowerCamelCase   ):
                   SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = tf.zeros(self.model_tester.batch_size				,       dtype=tf.intaa   )

           return inputs_dict






       def 	__lowerCAmelCase     (		self   )					->Tuple:
           SCREAMING_SNAKE_CASE				:				Dict				      = TFAlbertModelTester(self   )
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = ConfigTester(self				,       config_class=_lowerCamelCase				,       hidden_size=37   )






       def 	__lowerCAmelCase     (		self   )					->Union[str, Any]:
           self.config_tester.run_common_tests()






       def 	__lowerCAmelCase     (		self   )					->str:
           SCREAMING_SNAKE_CASE				:				Tuple				      = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_albert_model(*_lowerCamelCase   )






       def 	__lowerCAmelCase     (		self   )					->str:
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_albert_for_pretraining(*_lowerCamelCase   )






       def 	__lowerCAmelCase     (		self   )					->Any:
           SCREAMING_SNAKE_CASE				:				Dict				      = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_albert_for_masked_lm(*_lowerCamelCase   )






       def 	__lowerCAmelCase     (		self   )					->Optional[Any]:
           SCREAMING_SNAKE_CASE				:				List[str]				      = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_albert_for_multiple_choice(*_lowerCamelCase   )






       def 	__lowerCAmelCase     (		self   )					->Dict:
           SCREAMING_SNAKE_CASE				:				List[Any]				      = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_albert_for_sequence_classification(*_lowerCamelCase   )






       def 	__lowerCAmelCase     (		self   )					->str:
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_albert_for_question_answering(*_lowerCamelCase   )






       @slow
       def 	__lowerCAmelCase     (		self   )					->int:
           for model_name in TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               SCREAMING_SNAKE_CASE				:				Tuple				      = TFAlbertModel.from_pretrained(_lowerCamelCase   )
               self.assertIsNotNone(_lowerCamelCase   )






@require_tf
class     a_       (			unittest.TestCase							):



       """simple docstring"""






       @slow
       def 	__lowerCAmelCase     (		self   )					->List[Any]:
           SCREAMING_SNAKE_CASE				:				Dict				      = TFAlbertForPreTraining.from_pretrained('''albert-base-v2'''   )
           SCREAMING_SNAKE_CASE				:				str				      = tf.constant([[0, 1, 2, 3, 4, 5]]   )
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = model(_lowerCamelCase   )[0]

           SCREAMING_SNAKE_CASE				:				List[str]				      = [1, 6, 3_0000]
           self.assertEqual(output.shape				,       _lowerCamelCase   )

           SCREAMING_SNAKE_CASE				:				Dict				      = tf.constant(
               [
                   [
                       [4.5_9_5_6_6_8, 0.7_4_4_6_2_7_5_4, -1.8_1_8_1_4_7],
                       [4.5_9_5_4_3_4_7, 0.7_4_5_4_1_8_4, -1.8_1_8_8_2_5_8],
                       [4.5_9_5_4_9_0_5, 0.7_4_4_8_2_3_5, -1.8_1_8_2_3_1_6],
                   ]
               ]   )
           tf.debugging.assert_near(output[:, :3, :3]				,       _lowerCamelCase				,       atol=1e-4   )
