


'''simple docstring'''


import collections.abc
from typing import Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_poolformer import PoolFormerConfig


_UpperCamelCase 							=  logging.get_logger(__name__)

# General docstring
_UpperCamelCase 							=  """PoolFormerConfig"""

# Base docstring
_UpperCamelCase 							=  """sail/poolformer_s12"""
_UpperCamelCase 							=  [1, 512, 7, 7]

# Image classification docstring
_UpperCamelCase 							=  """sail/poolformer_s12"""
_UpperCamelCase 							=  """tabby, tabby cat"""

_UpperCamelCase 							=  [
    """sail/poolformer_s12""",
    # See all PoolFormer models at https://huggingface.co/models?filter=poolformer
]


def 		_lowercase (SCREAMING_SNAKE_CASE       ,					SCREAMING_SNAKE_CASE = 0.0       ,					SCREAMING_SNAKE_CASE = False							):




     '''simple docstring'''





     if drop_prob == 0.0 or not training:
          return input
     __A	:       int      =							1 - drop_prob
     __A	:       List[str]      =							(input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
     __A	:       List[str]      =							keep_prob + torch.rand(SCREAMING_SNAKE_CASE       ,					dtype=input.dtype       ,					device=input.device							)
     random_tensor.floor_()  # binarize
     __A	:       Optional[int]      =							input.div(SCREAMING_SNAKE_CASE							) * random_tensor
     return output







class 							__magic_name__	(       nn.Module ):

     """simple docstring"""






     def __init__(   self			,     lowerCamelCase = None		):


          '''simple docstring'''





          super().__init__()
          __A	:       int      =							drop_prob


     def lowerCAmelCase__  (   self			,     lowerCamelCase		):


          '''simple docstring'''





          return drop_path(lowerCamelCase			,     self.drop_prob			,     self.training		)


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          return "p={}".format(self.drop_prob		)







class 							__magic_name__	(       nn.Module ):

     """simple docstring"""






     def __init__(   self			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase=None		):


          '''simple docstring'''





          super().__init__()
          __A	:       Optional[Any]      =							patch_size if isinstance(lowerCamelCase			,     collections.abc.Iterable		) else (patch_size, patch_size)
          __A	:       Union[str, Any]      =							stride if isinstance(lowerCamelCase			,     collections.abc.Iterable		) else (stride, stride)
          __A	:       Union[str, Any]      =							padding if isinstance(lowerCamelCase			,     collections.abc.Iterable		) else (padding, padding)

          __A	:       Optional[Any]      =							nn.Convad(lowerCamelCase			,     lowerCamelCase			,     kernel_size=lowerCamelCase			,     stride=lowerCamelCase			,     padding=lowerCamelCase		)
          __A	:       Optional[int]      =							norm_layer(lowerCamelCase		) if norm_layer else nn.Identity()


     def lowerCAmelCase__  (   self			,     lowerCamelCase		):


          '''simple docstring'''





          __A	:       str      =							self.projection(lowerCamelCase		)
          __A	:       Any      =							self.norm(lowerCamelCase		)
          return embeddings







class 							__magic_name__	(       nn.GroupNorm ):

     """simple docstring"""






     def __init__(   self			,     lowerCamelCase			,     **lowerCamelCase		):


          '''simple docstring'''





          super().__init__(1			,     lowerCamelCase			,     **lowerCamelCase		)







class 							__magic_name__	(       nn.Module ):

     """simple docstring"""






     def __init__(   self			,     lowerCamelCase		):


          '''simple docstring'''





          super().__init__()
          __A	:       Dict      =							nn.AvgPoolad(lowerCamelCase			,     stride=1			,     padding=pool_size // 2			,     count_include_pad=lowerCamelCase		)


     def lowerCAmelCase__  (   self			,     lowerCamelCase		):


          '''simple docstring'''





          return self.pool(lowerCamelCase		) - hidden_states







class 							__magic_name__	(       nn.Module ):

     """simple docstring"""






     def __init__(   self			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase		):


          '''simple docstring'''





          super().__init__()
          __A	:       int      =							nn.Convad(lowerCamelCase			,     lowerCamelCase			,     1		)
          __A	:       Union[str, Any]      =							nn.Convad(lowerCamelCase			,     lowerCamelCase			,     1		)
          __A	:       str      =							PoolFormerDropPath(lowerCamelCase		)
          if isinstance(config.hidden_act			,     lowerCamelCase		):
               __A	:       List[str]      =							ACTaFN[config.hidden_act]
          else:
               __A	:       List[Any]      =							config.hidden_act


     def lowerCAmelCase__  (   self			,     lowerCamelCase		):


          '''simple docstring'''





          __A	:       Union[str, Any]      =							self.conva(lowerCamelCase		)
          __A	:       Optional[Any]      =							self.act_fn(lowerCamelCase		)
          __A	:       Optional[Any]      =							self.drop(lowerCamelCase		)
          __A	:       List[Any]      =							self.conva(lowerCamelCase		)
          __A	:       Optional[Any]      =							self.drop(lowerCamelCase		)

          return hidden_states







class 							__magic_name__	(       nn.Module ):

     """simple docstring"""






     def __init__(   self			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase		):


          '''simple docstring'''





          super().__init__()
          __A	:       Optional[Any]      =							PoolFormerPooling(lowerCamelCase		)
          __A	:       List[Any]      =							PoolFormerOutput(lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase		)
          __A	:       int      =							PoolFormerGroupNorm(lowerCamelCase		)
          __A	:       List[str]      =							PoolFormerGroupNorm(lowerCamelCase		)

          # Useful for training neural nets
          __A	:       Union[str, Any]      =							PoolFormerDropPath(lowerCamelCase		) if drop_path > 0.0 else nn.Identity()
          __A	:       Optional[Any]      =							config.use_layer_scale
          if config.use_layer_scale:
               __A	:       List[str]      =							nn.Parameter(
                   config.layer_scale_init_value * torch.ones((lowerCamelCase)		)			,     requires_grad=lowerCamelCase		)
               __A	:       Tuple      =							nn.Parameter(
                   config.layer_scale_init_value * torch.ones((lowerCamelCase)		)			,     requires_grad=lowerCamelCase		)


     def lowerCAmelCase__  (   self			,     lowerCamelCase		):


          '''simple docstring'''





          if self.use_layer_scale:
               __A	:       Union[str, Any]      =							self.pooling(self.before_norm(lowerCamelCase		)		)
               __A	:       Tuple      =							self.layer_scale_a.unsqueeze(-1		).unsqueeze(-1		) * pooling_output
               # First residual connection
               __A	:       Dict      =							hidden_states + self.drop_path(lowerCamelCase		)
               __A	:       Dict      =							()

               __A	:       int      =							self.output(self.after_norm(lowerCamelCase		)		)
               __A	:       List[Any]      =							self.layer_scale_a.unsqueeze(-1		).unsqueeze(-1		) * layer_output
               # Second residual connection
               __A	:       Dict      =							hidden_states + self.drop_path(lowerCamelCase		)

               __A	:       Union[str, Any]      =							(output,) + outputs
               return outputs

          else:
               __A	:       Any      =							self.drop_path(self.pooling(self.before_norm(lowerCamelCase		)		)		)
               # First residual connection
               __A	:       List[str]      =							pooling_output + hidden_states
               __A	:       Tuple      =							()

               # Second residual connection inside the PoolFormerOutput block
               __A	:       int      =							self.drop_path(self.output(self.after_norm(lowerCamelCase		)		)		)
               __A	:       int      =							hidden_states + layer_output

               __A	:       str      =							(output,) + outputs
               return outputs







class 							__magic_name__	(       nn.Module ):

     """simple docstring"""






     def __init__(   self			,     lowerCamelCase		):


          '''simple docstring'''





          super().__init__()
          __A	:       Dict      =							config
          # stochastic depth decay rule
          __A	:       Optional[int]      =							[x.item() for x in torch.linspace(0			,     config.drop_path_rate			,     sum(config.depths		)		)]

          # patch embeddings
          __A	:       Dict      =							[]
          for i in range(config.num_encoder_blocks		):
               embeddings.append(
                   PoolFormerEmbeddings(
                       patch_size=config.patch_sizes[i]			,     stride=config.strides[i]			,     padding=config.padding[i]			,     num_channels=config.num_channels if i == 0 else config.hidden_sizes[i - 1]			,     hidden_size=config.hidden_sizes[i]			,     )		)
          __A	:       str      =							nn.ModuleList(lowerCamelCase		)

          # Transformer blocks
          __A	:       str      =							[]
          __A	:       Any      =							0
          for i in range(config.num_encoder_blocks		):
               # each block consists of layers
               __A	:       Union[str, Any]      =							[]
               if i != 0:
                    cur += config.depths[i - 1]
               for j in range(config.depths[i]		):
                    layers.append(
                        PoolFormerLayer(
                            lowerCamelCase			,     num_channels=config.hidden_sizes[i]			,     pool_size=config.pool_size			,     hidden_size=config.hidden_sizes[i]			,     intermediate_size=int(config.hidden_sizes[i] * config.mlp_ratio		)			,     drop_path=dpr[cur + j]			,     )		)
               blocks.append(nn.ModuleList(lowerCamelCase		)		)

          __A	:       List[str]      =							nn.ModuleList(lowerCamelCase		)


     def lowerCAmelCase__  (   self			,     lowerCamelCase			,     lowerCamelCase=False			,     lowerCamelCase=True		):


          '''simple docstring'''





          __A	:       str      =							() if output_hidden_states else None

          __A	:       Optional[int]      =							pixel_values
          for idx, layers in enumerate(zip(self.patch_embeddings			,     self.block		)		):
               __A  ,__A	:       int      =							layers
               # Get patch embeddings from hidden_states
               __A	:       str      =							embedding_layer(lowerCamelCase		)
               # Send the embeddings through the blocks
               for _, blk in enumerate(lowerCamelCase		):
                    __A	:       Optional[int]      =							blk(lowerCamelCase		)
                    __A	:       Tuple      =							layer_outputs[0]

               if output_hidden_states:
                    __A	:       Tuple      =							all_hidden_states + (hidden_states,)

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

          return BaseModelOutputWithNoAttention(last_hidden_state=lowerCamelCase			,     hidden_states=lowerCamelCase		)







class 							__magic_name__	(       lowerCAmelCase ):

     """simple docstring"""






     lowerCamelCase__       						=				PoolFormerConfig
     lowerCamelCase__       						=				'poolformer'
     lowerCamelCase__       						=				'pixel_values'
     lowerCamelCase__       						=				True


     def lowerCAmelCase__  (   self			,     lowerCamelCase		):


          '''simple docstring'''





          if isinstance(lowerCamelCase			,     (nn.Linear, nn.Convad)		):
               module.weight.data.normal_(mean=0.0			,     std=self.config.initializer_range		)
               if module.bias is not None:
                    module.bias.data.zero_()
          elif isinstance(lowerCamelCase			,     nn.LayerNorm		):
               module.bias.data.zero_()
               module.weight.data.fill_(1.0		)


     def lowerCAmelCase__  (   self			,     lowerCamelCase			,     lowerCamelCase=False		):


          '''simple docstring'''





          if isinstance(lowerCamelCase			,     lowerCamelCase		):
               __A	:       List[Any]      =							value


_UpperCamelCase 							=  r"""
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`PoolFormerConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

_UpperCamelCase 							=  r"""
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`PoolFormerImageProcessor.__call__`] for details.
"""







@add_start_docstrings(
    'The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top.'    ,		lowerCAmelCase    ,		)
class 							__magic_name__	(       lowerCAmelCase ):

     """simple docstring"""






     def __init__(   self			,     lowerCamelCase		):


          '''simple docstring'''





          super().__init__(lowerCamelCase		)
          __A	:       Tuple      =							config

          __A	:       Tuple      =							PoolFormerEncoder(lowerCamelCase		)

          # Initialize weights and apply final processing
          self.post_init()


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          return self.embeddings.patch_embeddings


     @add_start_docstrings_to_model_forward(lowerCamelCase		)
     @add_code_sample_docstrings(
         checkpoint=_CHECKPOINT_FOR_DOC			,     output_type=lowerCamelCase			,     config_class=_CONFIG_FOR_DOC			,     modality="vision"			,     expected_output=_EXPECTED_OUTPUT_SHAPE			,     )
     def lowerCAmelCase__  (   self			,     lowerCamelCase = None			,     lowerCamelCase = None			,     lowerCamelCase = None			,     ):


          '''simple docstring'''





          __A	:       str      =							(
              output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
          )
          __A	:       int      =							return_dict if return_dict is not None else self.config.use_return_dict

          if pixel_values is None:
               raise ValueError("You have to specify pixel_values"		)

          __A	:       str      =							self.encoder(
              lowerCamelCase			,     output_hidden_states=lowerCamelCase			,     return_dict=lowerCamelCase			,     )
          __A	:       str      =							encoder_outputs[0]

          if not return_dict:
               return (sequence_output, None) + encoder_outputs[1:]

          return BaseModelOutputWithNoAttention(
              last_hidden_state=lowerCamelCase			,     hidden_states=encoder_outputs.hidden_states			,     )







class 							__magic_name__	(       nn.Module ):

     """simple docstring"""






     def __init__(   self			,     lowerCamelCase		):


          '''simple docstring'''





          super().__init__()
          __A	:       Any      =							nn.Linear(config.hidden_size			,     config.hidden_size		)


     def lowerCAmelCase__  (   self			,     lowerCamelCase		):


          '''simple docstring'''





          __A	:       List[str]      =							self.dense(lowerCamelCase		)
          return output







@add_start_docstrings(
    '\n    PoolFormer Model transformer with an image classification head on top\n    '    ,		lowerCAmelCase    ,		)
class 							__magic_name__	(       lowerCAmelCase ):

     """simple docstring"""






     def __init__(   self			,     lowerCamelCase		):


          '''simple docstring'''





          super().__init__(lowerCamelCase		)
          __A	:       List[Any]      =							config.num_labels
          __A	:       Dict      =							PoolFormerModel(lowerCamelCase		)

          # Final norm
          __A	:       int      =							PoolFormerGroupNorm(config.hidden_sizes[-1]		)
          # Classifier head
          __A	:       Union[str, Any]      =							(
              nn.Linear(config.hidden_sizes[-1]			,     config.num_labels		) if config.num_labels > 0 else nn.Identity()
          )

          # Initialize weights and apply final processing
          self.post_init()


     @add_start_docstrings_to_model_forward(lowerCamelCase		)
     @add_code_sample_docstrings(
         checkpoint=_IMAGE_CLASS_CHECKPOINT			,     output_type=lowerCamelCase			,     config_class=_CONFIG_FOR_DOC			,     expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT			,     )
     def lowerCAmelCase__  (   self			,     lowerCamelCase = None			,     lowerCamelCase = None			,     lowerCamelCase = None			,     lowerCamelCase = None			,     ):


          '''simple docstring'''





          __A	:       Optional[int]      =							return_dict if return_dict is not None else self.config.use_return_dict

          __A	:       int      =							self.poolformer(
              lowerCamelCase			,     output_hidden_states=lowerCamelCase			,     return_dict=lowerCamelCase			,     )

          __A	:       int      =							outputs[0]

          __A	:       str      =							self.classifier(self.norm(lowerCamelCase		).mean([-2, -1]		)		)

          __A	:       Union[str, Any]      =							None
          if labels is not None:
               if self.config.problem_type is None:
                    if self.num_labels == 1:
                         __A	:       Tuple      =							"regression"
                    elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                         __A	:       Dict      =							"single_label_classification"
                    else:
                         __A	:       Any      =							"multi_label_classification"

               if self.config.problem_type == "regression":
                    __A	:       Tuple      =							MSELoss()
                    if self.num_labels == 1:
                         __A	:       Optional[int]      =							loss_fct(logits.squeeze()			,     labels.squeeze()		)
                    else:
                         __A	:       Optional[Any]      =							loss_fct(lowerCamelCase			,     lowerCamelCase		)
               elif self.config.problem_type == "single_label_classification":
                    __A	:       Dict      =							CrossEntropyLoss()
                    __A	:       Optional[Any]      =							loss_fct(logits.view(-1			,     self.num_labels		)			,     labels.view(-1		)		)
               elif self.config.problem_type == "multi_label_classification":
                    __A	:       Dict      =							BCEWithLogitsLoss()
                    __A	:       Union[str, Any]      =							loss_fct(lowerCamelCase			,     lowerCamelCase		)

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

          return ImageClassifierOutputWithNoAttention(loss=lowerCamelCase			,     logits=lowerCamelCase			,     hidden_states=outputs.hidden_states		)





'''simple docstring'''


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.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEncoder,
    BertModel,
    BertPreTrainedModel,
)


_UpperCamelCase 							=  logging.getLogger(__name__)







class 							__magic_name__	(       lowerCAmelCase ):

     """simple docstring"""






     def lowerCAmelCase__  (   self			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase=None			,     lowerCamelCase=None		):


          '''simple docstring'''





          __A	:       List[Any]      =							self.layer[current_layer](lowerCamelCase			,     lowerCamelCase			,     head_mask[current_layer]		)

          __A	:       Tuple      =							layer_outputs[0]

          return hidden_states







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

     """simple docstring"""






     def __init__(   self			,     lowerCamelCase		):


          '''simple docstring'''





          super().__init__(lowerCamelCase		)

          __A	:       Any      =							BertEncoderWithPabee(lowerCamelCase		)

          self.init_weights()
          __A	:       Optional[Any]      =							0
          __A	:       Union[str, Any]      =							0
          __A	:       Dict      =							0

          __A	:       List[Any]      =							0


     def lowerCAmelCase__  (   self			,     lowerCamelCase		):


          '''simple docstring'''





          __A	:       Any      =							threshold


     def lowerCAmelCase__  (   self			,     lowerCamelCase		):


          '''simple docstring'''





          __A	:       List[Any]      =							patience


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       List[str]      =							0
          __A	:       int      =							0


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Optional[Any]      =							self.inference_layers_num / self.inference_instances_num
          __A	:       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(lowerCamelCase		)


     @add_start_docstrings_to_model_forward(lowerCamelCase		)
     def lowerCAmelCase__  (   self			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=False			,     ):


          '''simple docstring'''





          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:
               __A	:       Optional[int]      =							input_ids.size()
          elif inputs_embeds is not None:
               __A	:       Any      =							inputs_embeds.size()[:-1]
          else:
               raise ValueError("You have to specify either input_ids or inputs_embeds"		)

          __A	:       List[str]      =							input_ids.device if input_ids is not None else inputs_embeds.device

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

          # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
          # ourselves in which case we just need to make it broadcastable to all heads.
          __A	:       torch.Tensor      =							self.get_extended_attention_mask(lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase		)

          # If a 2D ou 3D attention mask is provided for the cross-attention
          # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
          if self.config.is_decoder and encoder_hidden_states is not None:
               __A  ,__A  ,__A	:       Tuple      =							encoder_hidden_states.size()
               __A	:       Any      =							(encoder_batch_size, encoder_sequence_length)
               if encoder_attention_mask is None:
                    __A	:       List[str]      =							torch.ones(lowerCamelCase			,     device=lowerCamelCase		)
               __A	:       Dict      =							self.invert_attention_mask(lowerCamelCase		)
          else:
               __A	:       Union[str, Any]      =							None

          # Prepare head mask if needed
          # 1.0 in head_mask indicate we keep the head
          # attention_probs has shape bsz x n_heads x N x N
          # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
          # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
          __A	:       Any      =							self.get_head_mask(lowerCamelCase			,     self.config.num_hidden_layers		)

          __A	:       Tuple      =							self.embeddings(
              input_ids=lowerCamelCase			,     position_ids=lowerCamelCase			,     token_type_ids=lowerCamelCase			,     inputs_embeds=lowerCamelCase		)
          __A	:       str      =							embedding_output

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

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

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

                    __A	:       int      =							logits
                    if patient_counter == self.patience:
                         break
               __A	:       List[str]      =							[patient_result]
               self.inference_layers_num += calculated_layer_num
               self.inference_instances_num += 1

          return res







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

     """simple docstring"""






     def __init__(   self			,     lowerCamelCase		):


          '''simple docstring'''





          super().__init__(lowerCamelCase		)
          __A	:       Union[str, Any]      =							config.num_labels

          __A	:       Union[str, Any]      =							BertModelWithPabee(lowerCamelCase		)
          __A	:       Union[str, Any]      =							nn.Dropout(config.hidden_dropout_prob		)
          __A	:       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(lowerCamelCase		)
     def lowerCAmelCase__  (   self			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     ):


          '''simple docstring'''





          __A	:       Optional[Any]      =							self.bert(
              input_ids=lowerCamelCase			,     attention_mask=lowerCamelCase			,     token_type_ids=lowerCamelCase			,     position_ids=lowerCamelCase			,     head_mask=lowerCamelCase			,     inputs_embeds=lowerCamelCase			,     output_dropout=self.dropout			,     output_layers=self.classifiers			,     regression=self.num_labels == 1			,     )

          __A	:       Any      =							(logits[-1],)

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

          return outputs

