

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


_UpperCamelCase     =    logging.getLogger(__name__)



class 					lowercase		( a_      ):



 '''simple docstring'''


 def 						UpperCamelCase__							(self , __a , __a , __a=None , __a=None				)				->      List[str]:






  """simple docstring"""





  if current_layer == 0:
   UpperCAmelCase__           =       self.embedding_hidden_mapping_in(_snake_case				)
  else:
   UpperCAmelCase__           =       hidden_states[0]

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

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

  UpperCAmelCase__           =       self.albert_layer_groups[group_idx](
      _snake_case , _snake_case , head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group] , )
  UpperCAmelCase__           =       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 					lowercase		( a_      ):



 '''simple docstring'''


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






  """simple docstring"""





  super().__init__(_snake_case				)

  UpperCAmelCase__           =       AlbertTransformerWithPabee(_snake_case				)

  self.init_weights()
  UpperCAmelCase__           =       0
  UpperCAmelCase__           =       0
  UpperCAmelCase__           =       0

  UpperCAmelCase__           =       0




 def 						UpperCamelCase__							(self , __a				)				->      List[str]:






  """simple docstring"""





  UpperCAmelCase__           =       threshold




 def 						UpperCamelCase__							(self , __a				)				->      Any:






  """simple docstring"""





  UpperCAmelCase__           =       patience




 def 						UpperCamelCase__							(self				)				->      int:






  """simple docstring"""





  UpperCAmelCase__           =       0
  UpperCAmelCase__           =       0




 def 						UpperCamelCase__							(self				)				->      int:






  """simple docstring"""





  UpperCAmelCase__           =       self.inference_layers_num / self.inference_instances_num
  UpperCAmelCase__           =       (
      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(_snake_case				)




 @add_start_docstrings_to_model_forward(_snake_case				)
 def 						UpperCamelCase__							(self , __a=None , __a=None , __a=None , __a=None , __a=None , __a=None , __a=None , __a=None , __a=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:
   UpperCAmelCase__           =       input_ids.size()
  elif inputs_embeds is not None:
   UpperCAmelCase__           =       inputs_embeds.size()[:-1]
  else:
   raise ValueError('You have to specify either input_ids or inputs_embeds'				)

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

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

  UpperCAmelCase__           =       attention_mask.unsqueeze(1				).unsqueeze(2				)
  UpperCAmelCase__           =       extended_attention_mask.to(dtype=self.dtype				)  # fp16 compatibility
  UpperCAmelCase__           =       (1.0 - extended_attention_mask) * -10000.0
  UpperCAmelCase__           =       self.get_head_mask(_snake_case , self.config.num_hidden_layers				)

  UpperCAmelCase__           =       self.embeddings(
      _snake_case , position_ids=_snake_case , token_type_ids=_snake_case , inputs_embeds=_snake_case				)
  UpperCAmelCase__           =       embedding_output

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

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

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

    UpperCAmelCase__           =       logits
    if patient_counter == self.patience:
     break
   UpperCAmelCase__           =       [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 					lowercase		( a_      ):



 '''simple docstring'''


 def __init__(self , __a				)				->      Tuple:






  """simple docstring"""





  super().__init__(_snake_case				)
  UpperCAmelCase__           =       config.num_labels

  UpperCAmelCase__           =       AlbertModelWithPabee(_snake_case				)
  UpperCAmelCase__           =       nn.Dropout(config.classifier_dropout_prob				)
  UpperCAmelCase__           =       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(_snake_case				)
 def 						UpperCamelCase__							(self , __a=None , __a=None , __a=None , __a=None , __a=None , __a=None , __a=None , )				->      Tuple:






  """simple docstring"""





  UpperCAmelCase__           =       self.albert(
      input_ids=_snake_case , attention_mask=_snake_case , token_type_ids=_snake_case , position_ids=_snake_case , head_mask=_snake_case , inputs_embeds=_snake_case , output_dropout=self.dropout , output_layers=self.classifiers , regression=self.num_labels == 1 , )

  UpperCAmelCase__           =       (logits[-1],)

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

  return outputs






def 		UpperCamelCase_(    snake_case__:						int     ) ->				int:
 if not isinstance(snake_case__						,					snake_case__     ):
  UpperCAmelCase__           =       f"Input value of [number={number}] must be an integer"
  raise TypeError(snake_case__     )
 if number < 1:
  raise ValueError('Input must be a positive integer'     )
 return number * (2 * number - 1)


if __name__ == "__main__":
 import doctest

 doctest.testmod()



