


from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFImageClassifierOutputWithNoAttention,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_resnet import ResNetConfig


SCREAMING_SNAKE_CASE__				:    List[Any]         =    logging.get_logger(__name__)

# General docstring
SCREAMING_SNAKE_CASE__				:    Optional[int]         =    'ResNetConfig'

# Base docstring
SCREAMING_SNAKE_CASE__				:    Any         =    'microsoft/resnet-50'
SCREAMING_SNAKE_CASE__				:    str         =    [1, 2_048, 7, 7]

# Image classification docstring
SCREAMING_SNAKE_CASE__				:    Any         =    'microsoft/resnet-50'
SCREAMING_SNAKE_CASE__				:    str         =    'tiger cat'

SCREAMING_SNAKE_CASE__				:    List[Any]         =    [
    'microsoft/resnet-50',
    # See all resnet models at https://huggingface.co/models?filter=resnet
]


class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):


     def __init__(     self       ,				A_       ,				A_ = 3       ,				A_ = 1       ,				A_ = "relu"       ,				**A_						):
       super().__init__(**A_						)
       _UpperCAmelCase		:       Union[str, Any]       =    kernel_size // 2
       _UpperCAmelCase		:       Tuple       =    tf.keras.layers.ConvaD(
           A_       ,				kernel_size=A_       ,				strides=A_       ,				padding="""valid"""       ,				use_bias=A_       ,				name="""convolution"""						)
       # Use same default momentum and epsilon as PyTorch equivalent
       _UpperCAmelCase		:       Optional[int]       =    tf.keras.layers.BatchNormalization(epsilon=1e-5       ,				momentum=0.9       ,				name="""normalization"""						)
       _UpperCAmelCase		:       Tuple       =    ACTaFN[activation] if activation is not None else tf.keras.layers.Activation("""linear"""						)


     def 					__snake_case(     self       ,				A_						):
       # Pad to match that done in the PyTorch Conv2D model
       _UpperCAmelCase		:       Any       =    (self.pad_value, self.pad_value)
       _UpperCAmelCase		:       Tuple       =    tf.pad(A_       ,				[(0, 0), height_pad, width_pad, (0, 0)]						)
       _UpperCAmelCase		:       List[Any]       =    self.conv(A_						)
       return hidden_state




     def 					__snake_case(     self       ,				A_       ,				A_ = False						):
       _UpperCAmelCase		:       Any       =    self.convolution(A_						)
       _UpperCAmelCase		:       Optional[Any]       =    self.normalization(A_       ,				training=A_						)
       _UpperCAmelCase		:       List[str]       =    self.activation(A_						)
       return hidden_state


class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):


     def __init__(     self       ,				A_       ,				**A_						):
       super().__init__(**A_						)
       _UpperCAmelCase		:       str       =    TFResNetConvLayer(
           config.embedding_size       ,				kernel_size=7       ,				stride=2       ,				activation=config.hidden_act       ,				name="""embedder"""       ,				)
       _UpperCAmelCase		:       Any       =    tf.keras.layers.MaxPoolaD(pool_size=3       ,				strides=2       ,				padding="""valid"""       ,				name="""pooler"""						)
       _UpperCAmelCase		:       Dict       =    config.num_channels




     def 					__snake_case(     self       ,				A_       ,				A_ = False						):
       _UpperCAmelCase,_UpperCAmelCase,_UpperCAmelCase,_UpperCAmelCase		:       Optional[int]       =    shape_list(A_						)
       if tf.executing_eagerly() and num_channels != self.num_channels:
         raise ValueError(
             """Make sure that the channel dimension of the pixel values match with the one set in the configuration."""						)
       _UpperCAmelCase		:       Optional[Any]       =    pixel_values
       _UpperCAmelCase		:       List[str]       =    self.embedder(A_						)
       _UpperCAmelCase		:       str       =    tf.pad(A_       ,				[[0, 0], [1, 1], [1, 1], [0, 0]]						)
       _UpperCAmelCase		:       List[Any]       =    self.pooler(A_						)
       return hidden_state


class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):


     def __init__(     self       ,				A_       ,				A_ = 2       ,				**A_						):
       super().__init__(**A_						)
       _UpperCAmelCase		:       Optional[int]       =    tf.keras.layers.ConvaD(
           A_       ,				kernel_size=1       ,				strides=A_       ,				use_bias=A_       ,				name="""convolution"""						)
       # Use same default momentum and epsilon as PyTorch equivalent
       _UpperCAmelCase		:       List[str]       =    tf.keras.layers.BatchNormalization(epsilon=1e-5       ,				momentum=0.9       ,				name="""normalization"""						)




     def 					__snake_case(     self       ,				A_       ,				A_ = False						):
       _UpperCAmelCase		:       Union[str, Any]       =    x
       _UpperCAmelCase		:       Optional[int]       =    self.convolution(A_						)
       _UpperCAmelCase		:       int       =    self.normalization(A_       ,				training=A_						)
       return hidden_state


class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):


     def __init__(     self       ,				A_       ,				A_       ,				A_ = 1       ,				A_ = "relu"       ,				**A_						):
       super().__init__(**A_						)
       _UpperCAmelCase		:       Optional[int]       =    in_channels != out_channels or stride != 1
       _UpperCAmelCase		:       str       =    TFResNetConvLayer(A_       ,				stride=A_       ,				name="""layer.0"""						)
       _UpperCAmelCase		:       Optional[Any]       =    TFResNetConvLayer(A_       ,				activation=A_       ,				name="""layer.1"""						)
       _UpperCAmelCase		:       Union[str, Any]       =    (
           TFResNetShortCut(A_       ,				stride=A_       ,				name="""shortcut"""						)
           if should_apply_shortcut
           else tf.keras.layers.Activation("""linear"""       ,				name="""shortcut"""						)
       )
       _UpperCAmelCase		:       Tuple       =    ACTaFN[activation]




     def 					__snake_case(     self       ,				A_       ,				A_ = False						):
       _UpperCAmelCase		:       Dict       =    hidden_state
       _UpperCAmelCase		:       Union[str, Any]       =    self.conva(A_       ,				training=A_						)
       _UpperCAmelCase		:       Union[str, Any]       =    self.conva(A_       ,				training=A_						)
       _UpperCAmelCase		:       int       =    self.shortcut(A_       ,				training=A_						)
       hidden_state += residual
       _UpperCAmelCase		:       List[str]       =    self.activation(A_						)
       return hidden_state


class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):


     def __init__(     self       ,				A_       ,				A_       ,				A_ = 1       ,				A_ = "relu"       ,				A_ = 4       ,				**A_       ,				):
       super().__init__(**A_						)
       _UpperCAmelCase		:       Any       =    in_channels != out_channels or stride != 1
       _UpperCAmelCase		:       List[str]       =    out_channels // reduction
       _UpperCAmelCase		:       Any       =    TFResNetConvLayer(A_       ,				kernel_size=1       ,				name="""layer.0"""						)
       _UpperCAmelCase		:       List[Any]       =    TFResNetConvLayer(A_       ,				stride=A_       ,				name="""layer.1"""						)
       _UpperCAmelCase		:       Tuple       =    TFResNetConvLayer(A_       ,				kernel_size=1       ,				activation=A_       ,				name="""layer.2"""						)
       _UpperCAmelCase		:       int       =    (
           TFResNetShortCut(A_       ,				stride=A_       ,				name="""shortcut"""						)
           if should_apply_shortcut
           else tf.keras.layers.Activation("""linear"""       ,				name="""shortcut"""						)
       )
       _UpperCAmelCase		:       int       =    ACTaFN[activation]




     def 					__snake_case(     self       ,				A_       ,				A_ = False						):
       _UpperCAmelCase		:       Tuple       =    hidden_state
       _UpperCAmelCase		:       str       =    self.conva(A_       ,				training=A_						)
       _UpperCAmelCase		:       List[Any]       =    self.conva(A_       ,				training=A_						)
       _UpperCAmelCase		:       int       =    self.conva(A_       ,				training=A_						)
       _UpperCAmelCase		:       List[Any]       =    self.shortcut(A_       ,				training=A_						)
       hidden_state += residual
       _UpperCAmelCase		:       str       =    self.activation(A_						)
       return hidden_state


class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):


     def __init__(     self       ,				A_       ,				A_       ,				A_       ,				A_ = 2       ,				A_ = 2       ,				**A_						):
       super().__init__(**A_						)

       _UpperCAmelCase		:       Optional[int]       =    TFResNetBottleNeckLayer if config.layer_type == """bottleneck""" else TFResNetBasicLayer

       _UpperCAmelCase		:       Tuple       =    [layer(A_       ,				A_       ,				stride=A_       ,				activation=config.hidden_act       ,				name="""layers.0"""						)]
       layers += [
           layer(A_       ,				A_       ,				activation=config.hidden_act       ,				name=F'''layers.{i + 1}'''						)
           for i in range(depth - 1						)
       ]
       _UpperCAmelCase		:       Any       =    layers




     def 					__snake_case(     self       ,				A_       ,				A_ = False						):
       for layer in self.stage_layers:
         _UpperCAmelCase		:       Union[str, Any]       =    layer(A_       ,				training=A_						)
       return hidden_state


class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):


     def __init__(     self       ,				A_       ,				**A_						):
       super().__init__(**A_						)
       # based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
       _UpperCAmelCase		:       Dict       =    [
           TFResNetStage(
               A_       ,				config.embedding_size       ,				config.hidden_sizes[0]       ,				stride=2 if config.downsample_in_first_stage else 1       ,				depth=config.depths[0]       ,				name="""stages.0"""       ,				)
       ]
       for i, (in_channels, out_channels, depth) in enumerate(
           zip(config.hidden_sizes       ,				config.hidden_sizes[1:]       ,				config.depths[1:]						)						):
         self.stages.append(TFResNetStage(A_       ,				A_       ,				A_       ,				depth=A_       ,				name=F'''stages.{i + 1}'''						)						)




     def 					__snake_case(     self       ,				A_       ,				A_ = False       ,				A_ = True       ,				A_ = False       ,				):
       _UpperCAmelCase		:       Dict       =    () if output_hidden_states else None

       for stage_module in self.stages:
         if output_hidden_states:
           _UpperCAmelCase		:       int       =    hidden_states + (hidden_state,)

         _UpperCAmelCase		:       List[Any]       =    stage_module(A_       ,				training=A_						)

       if output_hidden_states:
         _UpperCAmelCase		:       Dict       =    hidden_states + (hidden_state,)

       if not return_dict:
         return tuple(v for v in [hidden_state, hidden_states] if v is not None						)

       return TFBaseModelOutputWithNoAttention(last_hidden_state=A_       ,				hidden_states=A_						)


class 		_SCREAMING_SNAKE_CASE    (	A    ):
     __SCREAMING_SNAKE_CASE    				=						ResNetConfig
     __SCREAMING_SNAKE_CASE    				=						'''resnet'''
     __SCREAMING_SNAKE_CASE    				=						'''pixel_values'''


     @property
     def 					__snake_case(     self						):
       return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 2_24, 2_24)       ,				dtype=tf.floataa						)}


SCREAMING_SNAKE_CASE__				:    List[Any]         =    R'\n    This model is a TensorFlow\n    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a\n    regular TensorFlow Module and refer to the TensorFlow documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`ResNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.\n'


SCREAMING_SNAKE_CASE__				:    str         =    R'\n    Args:\n        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConvNextImageProcessor.__call__`] for details.\n\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n'


@keras_serializable
class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):
     __SCREAMING_SNAKE_CASE    				=						ResNetConfig


     def __init__(     self       ,				A_       ,				**A_						):
       super().__init__(**A_						)
       _UpperCAmelCase		:       int       =    config
       _UpperCAmelCase		:       str       =    TFResNetEmbeddings(A_       ,				name="""embedder"""						)
       _UpperCAmelCase		:       int       =    TFResNetEncoder(A_       ,				name="""encoder"""						)
       _UpperCAmelCase		:       List[Any]       =    tf.keras.layers.GlobalAveragePoolingaD(keepdims=A_						)




     @unpack_inputs
     def 					__snake_case(     self       ,				A_       ,				A_ = None       ,				A_ = None       ,				A_ = False       ,				):
       _UpperCAmelCase		:       Optional[Any]       =    (
           output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
       )
       _UpperCAmelCase		:       int       =    return_dict if return_dict is not None else self.config.use_return_dict

       # TF 2.0 image layers can't use NCHW format when running on CPU.
       # We transpose to NHWC format and then transpose back after the full forward pass.
       # (batch_size, num_channels, height, width) -> (batch_size, height, width, num_channels)
       _UpperCAmelCase		:       Any       =    tf.transpose(A_       ,				perm=[0, 2, 3, 1]						)
       _UpperCAmelCase		:       Union[str, Any]       =    self.embedder(A_       ,				training=A_						)

       _UpperCAmelCase		:       Tuple       =    self.encoder(
           A_       ,				output_hidden_states=A_       ,				return_dict=A_       ,				training=A_						)

       _UpperCAmelCase		:       Any       =    encoder_outputs[0]

       _UpperCAmelCase		:       Any       =    self.pooler(A_						)

       # Transpose all the outputs to the NCHW format
       # (batch_size, height, width, num_channels) -> (batch_size, num_channels, height, width)
       _UpperCAmelCase		:       Any       =    tf.transpose(A_       ,				(0, 3, 1, 2)						)
       _UpperCAmelCase		:       Dict       =    tf.transpose(A_       ,				(0, 3, 1, 2)						)
       _UpperCAmelCase		:       int       =    ()
       for hidden_state in encoder_outputs[1:]:
         _UpperCAmelCase		:       Tuple       =    hidden_states + tuple(tf.transpose(A_       ,				(0, 3, 1, 2)						) for h in hidden_state						)

       if not return_dict:
         return (last_hidden_state, pooled_output) + hidden_states

       _UpperCAmelCase		:       Tuple       =    hidden_states if output_hidden_states else None

       return TFBaseModelOutputWithPoolingAndNoAttention(
           last_hidden_state=A_       ,				pooler_output=A_       ,				hidden_states=A_       ,				)


@add_start_docstrings(
    '''The bare ResNet model outputting raw features without any specific head on top.'''       ,  A       ,  )
class 		_SCREAMING_SNAKE_CASE    (	A    ):


     def __init__(     self       ,				A_       ,				**A_						):
       super().__init__(A_       ,				**A_						)
       _UpperCAmelCase		:       int       =    TFResNetMainLayer(config=A_       ,				name="""resnet"""						)




     @add_start_docstrings_to_model_forward(A_						)
     @add_code_sample_docstrings(
         checkpoint=_CHECKPOINT_FOR_DOC       ,				output_type=A_       ,				config_class=_CONFIG_FOR_DOC       ,				modality="""vision"""       ,				expected_output=_EXPECTED_OUTPUT_SHAPE       ,				)
     @unpack_inputs
     def 					__snake_case(     self       ,				A_       ,				A_ = None       ,				A_ = None       ,				A_ = False       ,				):
       _UpperCAmelCase		:       int       =    (
           output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
       )
       _UpperCAmelCase		:       Union[str, Any]       =    return_dict if return_dict is not None else self.config.use_return_dict

       _UpperCAmelCase		:       List[str]       =    self.resnet(
           pixel_values=A_       ,				output_hidden_states=A_       ,				return_dict=A_       ,				training=A_       ,				)
       return resnet_outputs


@add_start_docstrings(
    '''
    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    '''       ,  A       ,  )
class 		_SCREAMING_SNAKE_CASE    (	A       ,  A    ):


     def __init__(     self       ,				A_       ,				**A_						):
       super().__init__(A_       ,				**A_						)
       _UpperCAmelCase		:       int       =    config.num_labels
       _UpperCAmelCase		:       str       =    TFResNetMainLayer(A_       ,				name="""resnet"""						)
       # classification head
       _UpperCAmelCase		:       str       =    (
           tf.keras.layers.Dense(config.num_labels       ,				name="""classifier.1"""						)
           if config.num_labels > 0
           else tf.keras.layers.Activation("""linear"""       ,				name="""classifier.1"""						)
       )


     def 					__snake_case(     self       ,				A_						):
       _UpperCAmelCase		:       int       =    tf.keras.layers.Flatten()(A_						)
       _UpperCAmelCase		:       str       =    self.classifier_layer(A_						)
       return logits




     @add_start_docstrings_to_model_forward(A_						)
     @add_code_sample_docstrings(
         checkpoint=_IMAGE_CLASS_CHECKPOINT       ,				output_type=A_       ,				config_class=_CONFIG_FOR_DOC       ,				expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT       ,				)
     @unpack_inputs
     def 					__snake_case(     self       ,				A_ = None       ,				A_ = None       ,				A_ = None       ,				A_ = None       ,				A_ = False       ,				):
       _UpperCAmelCase		:       List[Any]       =    return_dict if return_dict is not None else self.config.use_return_dict

       _UpperCAmelCase		:       List[Any]       =    self.resnet(
           A_       ,				output_hidden_states=A_       ,				return_dict=A_       ,				training=A_						)

       _UpperCAmelCase		:       Optional[int]       =    outputs.pooler_output if return_dict else outputs[1]

       _UpperCAmelCase		:       str       =    self.classifier(A_						)

       _UpperCAmelCase		:       Optional[Any]       =    None if labels is None else self.hf_compute_loss(A_       ,				A_						)

       if not return_dict:
         _UpperCAmelCase		:       Optional[int]       =    (logits,) + outputs[2:]
         return (loss,) + output if loss is not None else output

       return TFImageClassifierOutputWithNoAttention(loss=A_       ,				logits=A_       ,				hidden_states=outputs.hidden_states						)









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__				:    Dict         =    logging.getLogger(__name__)


class 		_SCREAMING_SNAKE_CASE    (	A    ):


     def 					__snake_case(     self       ,				A_       ,				A_       ,				A_=None       ,				A_=None						):
       if current_layer == 0:
         _UpperCAmelCase		:       Any       =    self.embedding_hidden_mapping_in(A_						)
       else:
         _UpperCAmelCase		:       Optional[int]       =    hidden_states[0]

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

       # Index of the hidden group
       _UpperCAmelCase		:       Union[str, Any]       =    int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)						)

       _UpperCAmelCase		:       int       =    self.albert_layer_groups[group_idx](
           A_       ,				A_       ,				head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]       ,				)
       _UpperCAmelCase		:       Dict       =    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.'''       ,  A       ,  )
class 		_SCREAMING_SNAKE_CASE    (	A    ):


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

       _UpperCAmelCase		:       List[Any]       =    AlbertTransformerWithPabee(A_						)

       self.init_weights()
       _UpperCAmelCase		:       Tuple       =    0
       _UpperCAmelCase		:       Optional[Any]       =    0
       _UpperCAmelCase		:       int       =    0

       _UpperCAmelCase		:       Dict       =    0


     def 					__snake_case(     self       ,				A_						):
       _UpperCAmelCase		:       str       =    threshold


     def 					__snake_case(     self       ,				A_						):
       _UpperCAmelCase		:       Optional[Any]       =    patience


     def 					__snake_case(     self						):
       _UpperCAmelCase		:       Tuple       =    0
       _UpperCAmelCase		:       Tuple       =    0


     def 					__snake_case(     self						):
       _UpperCAmelCase		:       Dict       =    self.inference_layers_num / self.inference_instances_num
       _UpperCAmelCase		:       Optional[int]       =    (
           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(A_						)




     @add_start_docstrings_to_model_forward(A_						)
     def 					__snake_case(     self       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=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[str]       =    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		:       List[str]       =    input_ids.device if input_ids is not None else inputs_embeds.device

       if attention_mask is None:
         _UpperCAmelCase		:       Any       =    torch.ones(A_       ,				device=A_						)
       if token_type_ids is None:
         _UpperCAmelCase		:       List[str]       =    torch.zeros(A_       ,				dtype=torch.long       ,				device=A_						)

       _UpperCAmelCase		:       str       =    attention_mask.unsqueeze(1						).unsqueeze(2						)
       _UpperCAmelCase		:       Tuple       =    extended_attention_mask.to(dtype=self.dtype						)  # fp16 compatibility
       _UpperCAmelCase		:       Tuple       =    (1.0 - extended_attention_mask) * -1_0_0_0_0.0
       _UpperCAmelCase		:       Union[str, Any]       =    self.get_head_mask(A_       ,				self.config.num_hidden_layers						)

       _UpperCAmelCase		:       Dict       =    self.embeddings(
           A_       ,				position_ids=A_       ,				token_type_ids=A_       ,				inputs_embeds=A_						)
       _UpperCAmelCase		:       str       =    embedding_output

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

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

           _UpperCAmelCase		:       Tuple       =    self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]						)						)
           _UpperCAmelCase		:       List[str]       =    output_layers[i](A_						)
           if regression:
             _UpperCAmelCase		:       Optional[Any]       =    logits.detach()
             if patient_result is not None:
               _UpperCAmelCase		:       str       =    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		:       Dict       =    logits.detach().argmax(dim=1						)
             if patient_result is not None:
               _UpperCAmelCase		:       Optional[Any]       =    patient_result.detach().argmax(dim=1						)
             if (patient_result is not None) and torch.all(labels.eq(A_						)						):
               patient_counter += 1
             else:
               _UpperCAmelCase		:       Dict       =    0

           _UpperCAmelCase		:       int       =    logits
           if patient_counter == self.patience:
             break
         _UpperCAmelCase		:       Tuple       =    [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. '''       ,  A       ,  )
class 		_SCREAMING_SNAKE_CASE    (	A    ):


     def __init__(     self       ,				A_						):
       super().__init__(A_						)
       _UpperCAmelCase		:       List[str]       =    config.num_labels

       _UpperCAmelCase		:       List[Any]       =    AlbertModelWithPabee(A_						)
       _UpperCAmelCase		:       Optional[int]       =    nn.Dropout(config.classifier_dropout_prob						)
       _UpperCAmelCase		:       Tuple       =    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(A_						)
     def 					__snake_case(     self       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				):
       _UpperCAmelCase		:       Optional[Any]       =    self.albert(
           input_ids=A_       ,				attention_mask=A_       ,				token_type_ids=A_       ,				position_ids=A_       ,				head_mask=A_       ,				inputs_embeds=A_       ,				output_dropout=self.dropout       ,				output_layers=self.classifiers       ,				regression=self.num_labels == 1       ,				)

       _UpperCAmelCase		:       str       =    (logits[-1],)

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

       return outputs





