


'''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.albert.modeling_albert import (
    ALBERT_INPUTS_DOCSTRING,
    ALBERT_START_DOCSTRING,
    AlbertModel,
    AlbertPreTrainedModel,
    AlbertTransformer,
)


lowerCAmelCase  =				logging.getLogger(__name__)






class 					lowerCAmelCase_     (    UpperCAmelCase  ):






       def      UpperCamelCase	(   self ,    _UpperCamelCase ,    _UpperCamelCase ,    _UpperCamelCase=None ,    _UpperCamelCase=None      )->			str:
            if current_layer == 0:
                 _A				=       self.embedding_hidden_mapping_in(_UpperCamelCase      )
            else:
                 _A				=       hidden_states[0]

            _A				=       int(self.config.num_hidden_layers / self.config.num_hidden_groups      )

            # Index of the hidden group
            _A				=       int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)      )

            _A				=       self.albert_layer_groups[group_idx](
                _UpperCamelCase ,    _UpperCamelCase ,    head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group] ,    )
            _A				=       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.'			,       UpperCAmelCase			,       )
class 					lowerCAmelCase_     (    UpperCAmelCase  ):






       def __init__(   self ,    _UpperCamelCase      )->			Optional[int]:
            super().__init__(_UpperCamelCase      )

            _A				=       AlbertTransformerWithPabee(_UpperCamelCase      )

            self.init_weights()
            _A				=       0
            _A				=       0
            _A				=       0

            _A				=       0






       def      UpperCamelCase	(   self ,    _UpperCamelCase      )->			int:
            _A				=       threshold






       def      UpperCamelCase	(   self ,    _UpperCamelCase      )->			List[str]:
            _A				=       patience






       def      UpperCamelCase	(   self      )->			Any:
            _A				=       0
            _A				=       0






       def      UpperCamelCase	(   self      )->			Optional[int]:
            _A				=       self.inference_layers_num / self.inference_instances_num
            _A				=       (
                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      UpperCamelCase	(   self ,    _UpperCamelCase=None ,    _UpperCamelCase=None ,    _UpperCamelCase=None ,    _UpperCamelCase=None ,    _UpperCamelCase=None ,    _UpperCamelCase=None ,    _UpperCamelCase=None ,    _UpperCamelCase=None ,    _UpperCamelCase=False ,    )->			int:

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

            _A				=       input_ids.device if input_ids is not None else inputs_embeds.device

            if attention_mask is None:
                 _A				=       torch.ones(_UpperCamelCase ,    device=_UpperCamelCase      )
            if token_type_ids is None:
                 _A				=       torch.zeros(_UpperCamelCase ,    dtype=torch.long ,    device=_UpperCamelCase      )

            _A				=       attention_mask.unsqueeze(1      ).unsqueeze(2      )
            _A				=       extended_attention_mask.to(dtype=self.dtype      )  # fp16 compatibility
            _A				=       (1.0 - extended_attention_mask) * -1_0000.0
            _A				=       self.get_head_mask(_UpperCamelCase ,    self.config.num_hidden_layers      )

            _A				=       self.embeddings(
                _UpperCamelCase ,    position_ids=_UpperCamelCase ,    token_type_ids=_UpperCamelCase ,    inputs_embeds=_UpperCamelCase      )
            _A				=       embedding_output

            if self.training:
                 _A				=       []
                 for i in range(self.config.num_hidden_layers      ):
                      _A				=       self.encoder.adaptive_forward(
                          _UpperCamelCase ,    current_layer=_UpperCamelCase ,    attention_mask=_UpperCamelCase ,    head_mask=_UpperCamelCase ,    )

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

                      _A				=       self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]      )      )
                      _A				=       output_layers[i](_UpperCamelCase      )
                      if regression:
                           _A				=       logits.detach()
                           if patient_result is not None:
                                _A				=       patient_result.detach()
                           if (patient_result is not None) and torch.abs(patient_result - labels      ) < self.regression_threshold:
                                patient_counter += 1
                           else:
                                _A				=       0
                      else:
                           _A				=       logits.detach().argmax(dim=1      )
                           if patient_result is not None:
                                _A				=       patient_result.detach().argmax(dim=1      )
                           if (patient_result is not None) and torch.all(labels.eq(_UpperCamelCase      )      ):
                                patient_counter += 1
                           else:
                                _A				=       0

                      _A				=       logits
                      if patient_counter == self.patience:
                           break
                 _A				=       [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\n    the pooled output) e.g. for GLUE tasks. '			,       UpperCAmelCase			,       )
class 					lowerCAmelCase_     (    UpperCAmelCase  ):






       def __init__(   self ,    _UpperCamelCase      )->			Union[str, Any]:
            super().__init__(_UpperCamelCase      )
            _A				=       config.num_labels

            _A				=       AlbertModelWithPabee(_UpperCamelCase      )
            _A				=       nn.Dropout(config.classifier_dropout_prob      )
            _A				=       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      UpperCamelCase	(   self ,    _UpperCamelCase=None ,    _UpperCamelCase=None ,    _UpperCamelCase=None ,    _UpperCamelCase=None ,    _UpperCamelCase=None ,    _UpperCamelCase=None ,    _UpperCamelCase=None ,    )->			Dict:
            _A				=       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 ,    )

            _A				=       (logits[-1],)

            if labels is not None:
                 _A				=       None
                 _A				=       0
                 for ix, logits_item in enumerate(_UpperCamelCase      ):
                      if self.num_labels == 1:
                           #  We are doing regression
                           _A				=       MSELoss()
                           _A				=       loss_fct(logits_item.view(-1      ) ,    labels.view(-1      )      )
                      else:
                           _A				=       CrossEntropyLoss()
                           _A				=       loss_fct(logits_item.view(-1 ,    self.num_labels      ) ,    labels.view(-1      )      )
                      if total_loss is None:
                           _A				=       loss
                      else:
                           total_loss += loss * (ix + 1)
                      total_weights += ix + 1
                 _A				=       (total_loss / total_weights,) + outputs

            return outputs









'''simple docstring'''
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_torch_available,
    is_vision_available,
)


lowerCAmelCase  =				{
    """configuration_segformer""": ["""SEGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP""", """SegformerConfig""", """SegformerOnnxConfig"""]
}

try:
    if not is_vision_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    lowerCAmelCase  =				["""SegformerFeatureExtractor"""]
    lowerCAmelCase  =				["""SegformerImageProcessor"""]

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    lowerCAmelCase  =				[
        """SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
        """SegformerDecodeHead""",
        """SegformerForImageClassification""",
        """SegformerForSemanticSegmentation""",
        """SegformerLayer""",
        """SegformerModel""",
        """SegformerPreTrainedModel""",
    ]

try:
    if not is_tf_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    lowerCAmelCase  =				[
        """TF_SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
        """TFSegformerDecodeHead""",
        """TFSegformerForImageClassification""",
        """TFSegformerForSemanticSegmentation""",
        """TFSegformerModel""",
        """TFSegformerPreTrainedModel""",
    ]


if TYPE_CHECKING:
    from .configuration_segformer import SEGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, SegformerConfig, SegformerOnnxConfig

    try:
        if not is_vision_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .feature_extraction_segformer import SegformerFeatureExtractor
        from .image_processing_segformer import SegformerImageProcessor

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_segformer import (
            SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
            SegformerDecodeHead,
            SegformerForImageClassification,
            SegformerForSemanticSegmentation,
            SegformerLayer,
            SegformerModel,
            SegformerPreTrainedModel,
        )
    try:
        if not is_tf_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_tf_segformer import (
            TF_SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
            TFSegformerDecodeHead,
            TFSegformerForImageClassification,
            TFSegformerForSemanticSegmentation,
            TFSegformerModel,
            TFSegformerPreTrainedModel,
        )

else:
    import sys

    lowerCAmelCase  =				_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)





