




"""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,
)


SCREAMING_SNAKE_CASE       = logging.getLogger(__name__)




class 			UpperCAmelCase_						(		a__							):




    def 							__magic_name__							(					self		: Tuple			,				snake_case_		: int			,				snake_case_		: List[Any]			,				snake_case_		: List[Any]=None			,				snake_case_		: List[Any]=None   )   ->							int:




     '''simple docstring'''

     if current_layer == 0:
      A__ =       self.embedding_hidden_mapping_in(SCREAMING_SNAKE_CASE_   )
     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](
         SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				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.''',     a__,     )
class 			UpperCAmelCase_						(		a__							):




    def __init__(					self		: int			,				snake_case_		: Tuple   )   ->							Tuple:




     '''simple docstring'''

     super().__init__(SCREAMING_SNAKE_CASE_   )

     A__ =       AlbertTransformerWithPabee(SCREAMING_SNAKE_CASE_   )

     self.init_weights()
     A__ =       0
     A__ =       0
     A__ =       0

     A__ =       0




    def 							__magic_name__							(					self		: Optional[int]			,				snake_case_		: List[Any]   )   ->							Dict:




     '''simple docstring'''

     A__ =       threshold




    def 							__magic_name__							(					self		: Any			,				snake_case_		: Any   )   ->							Optional[int]:




     '''simple docstring'''

     A__ =       patience




    def 							__magic_name__							(					self		: Optional[Any]   )   ->							List[str]:




     '''simple docstring'''

     A__ =       0
     A__ =       0




    def 							__magic_name__							(					self		: Union[str, Any]   )   ->							List[Any]:




     '''simple docstring'''

     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(SCREAMING_SNAKE_CASE_   )




    @add_start_docstrings_to_model_forward(SCREAMING_SNAKE_CASE_   )
    def 							__magic_name__							(					self		: int			,				snake_case_		: Optional[int]=None			,				snake_case_		: int=None			,				snake_case_		: Dict=None			,				snake_case_		: Tuple=None			,				snake_case_		: Optional[Any]=None			,				snake_case_		: Tuple=None			,				snake_case_		: Optional[int]=None			,				snake_case_		: Any=None			,				snake_case_		: Tuple=False			,				)   ->							str:




     '''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__ =       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(SCREAMING_SNAKE_CASE_			,				device=SCREAMING_SNAKE_CASE_   )
     if token_type_ids is None:
      A__ =       torch.zeros(SCREAMING_SNAKE_CASE_			,				dtype=torch.long			,				device=SCREAMING_SNAKE_CASE_   )

     A__ =       attention_mask.unsqueeze(1   ).unsqueeze(2   )
     A__ =       extended_attention_mask.to(dtype=self.dtype   )  # fp16 compatibility
     A__ =       (1.0 - extended_attention_mask) * -10_000.0
     A__ =       self.get_head_mask(SCREAMING_SNAKE_CASE_			,				self.config.num_hidden_layers   )

     A__ =       self.embeddings(
         SCREAMING_SNAKE_CASE_			,				position_ids=SCREAMING_SNAKE_CASE_			,				token_type_ids=SCREAMING_SNAKE_CASE_			,				inputs_embeds=SCREAMING_SNAKE_CASE_   )
     A__ =       embedding_output

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

       A__ =       self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]   )   )
       A__ =       output_layers[i](output_dropout(SCREAMING_SNAKE_CASE_   )   )
       res.append(SCREAMING_SNAKE_CASE_   )
     elif self.patience == 0:  # Use all layers for inference
      A__ =       self.encoder(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				head_mask=SCREAMING_SNAKE_CASE_   )
      A__ =       self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]   )   )
      A__ =       [output_layers[self.config.num_hidden_layers - 1](SCREAMING_SNAKE_CASE_   )]
     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(
           SCREAMING_SNAKE_CASE_			,				current_layer=SCREAMING_SNAKE_CASE_			,				attention_mask=SCREAMING_SNAKE_CASE_			,				head_mask=SCREAMING_SNAKE_CASE_			,				)

       A__ =       self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]   )   )
       A__ =       output_layers[i](SCREAMING_SNAKE_CASE_   )
       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(SCREAMING_SNAKE_CASE_   )   ):
         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. ''',     a__,     )
class 			UpperCAmelCase_						(		a__							):




    def __init__(					self		: Optional[Any]			,				snake_case_		: Optional[int]   )   ->							Optional[Any]:




     '''simple docstring'''

     super().__init__(SCREAMING_SNAKE_CASE_   )
     A__ =       config.num_labels

     A__ =       AlbertModelWithPabee(SCREAMING_SNAKE_CASE_   )
     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(SCREAMING_SNAKE_CASE_   )
    def 							__magic_name__							(					self		: Dict			,				snake_case_		: str=None			,				snake_case_		: Tuple=None			,				snake_case_		: Tuple=None			,				snake_case_		: Dict=None			,				snake_case_		: Optional[int]=None			,				snake_case_		: Optional[Any]=None			,				snake_case_		: Tuple=None			,				)   ->							Any:




     '''simple docstring'''

     A__ =       self.albert(
         input_ids=SCREAMING_SNAKE_CASE_			,				attention_mask=SCREAMING_SNAKE_CASE_			,				token_type_ids=SCREAMING_SNAKE_CASE_			,				position_ids=SCREAMING_SNAKE_CASE_			,				head_mask=SCREAMING_SNAKE_CASE_			,				inputs_embeds=SCREAMING_SNAKE_CASE_			,				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(SCREAMING_SNAKE_CASE_   ):
       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 transformers import PretrainedConfig




class 			UpperCAmelCase_						(		A_							):
    lowercase__		      =   '''custom'''




    def __init__(					self		: Any			,				snake_case_		: List[Any]=1			,				**snake_case_		: Any   )   ->							str:




     '''simple docstring'''

     A__ =       attribute
     super().__init__(**snake_case_   )






class 			UpperCAmelCase_						(		A_							):
    lowercase__		      =   '''custom'''




    def __init__(					self		: Any			,				snake_case_		: Tuple=1			,				**snake_case_		: Dict   )   ->							Any:




     '''simple docstring'''

     A__ =       attribute





