


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		:List[Any]											=       logging.getLogger(__name__)


class       __lowerCAmelCase       (	a   ):






  """simple docstring"""





  def    lowerCAmelCase__					(       self    :    Optional[int]      ,			_lowerCAmelCase    :    Union[str, Any]      ,			_lowerCAmelCase    :    Optional[int]      ,			_lowerCAmelCase    :    Union[str, Any]=None      ,			_lowerCAmelCase    :    List[str]=None					)				->	List[str]:




      """simple docstring"""
      if current_layer == 0:
          snake_case_         =   self.embedding_hidden_mapping_in(_lowerCAmelCase					)
      else:
          snake_case_         =   hidden_states[0]

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

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

      snake_case_         =   self.albert_layer_groups[group_idx](
          _lowerCAmelCase      ,			_lowerCAmelCase      ,			head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]      ,			)
      snake_case_         =   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       __lowerCAmelCase       (	a   ):






  """simple docstring"""





  def __init__(       self    :    List[Any]      ,			_lowerCAmelCase    :    Union[str, Any]					)				->	Dict:




      """simple docstring"""
      super().__init__(_lowerCAmelCase					)

      snake_case_         =   AlbertTransformerWithPabee(_lowerCAmelCase					)

      self.init_weights()
      snake_case_         =   0
      snake_case_         =   0
      snake_case_         =   0

      snake_case_         =   0



  def    lowerCAmelCase__					(       self    :    Optional[Any]      ,			_lowerCAmelCase    :    List[Any]					)				->	int:




      """simple docstring"""
      snake_case_         =   threshold



  def    lowerCAmelCase__					(       self    :    Any      ,			_lowerCAmelCase    :    Optional[int]					)				->	Optional[Any]:




      """simple docstring"""
      snake_case_         =   patience



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




      """simple docstring"""
      snake_case_         =   0
      snake_case_         =   0



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




      """simple docstring"""
      snake_case_         =   self.inference_layers_num / self.inference_instances_num
      snake_case_         =   (
          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    :    List[Any]      ,			_lowerCAmelCase    :    Optional[Any]=None      ,			_lowerCAmelCase    :    List[Any]=None      ,			_lowerCAmelCase    :    Tuple=None      ,			_lowerCAmelCase    :    List[Any]=None      ,			_lowerCAmelCase    :    Dict=None      ,			_lowerCAmelCase    :    Optional[int]=None      ,			_lowerCAmelCase    :    str=None      ,			_lowerCAmelCase    :    Tuple=None      ,			_lowerCAmelCase    :    List[str]=False      ,			)				->	Optional[int]:




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

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

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

      snake_case_         =   attention_mask.unsqueeze(1					).unsqueeze(2					)
      snake_case_         =   extended_attention_mask.to(dtype=self.dtype					)  # fp16 compatibility
      snake_case_         =   (1.0 - extended_attention_mask) * -1_0_0_0_0.0
      snake_case_         =   self.get_head_mask(_lowerCAmelCase      ,			self.config.num_hidden_layers					)

      snake_case_         =   self.embeddings(
          _lowerCAmelCase      ,			position_ids=_lowerCAmelCase      ,			token_type_ids=_lowerCAmelCase      ,			inputs_embeds=_lowerCAmelCase					)
      snake_case_         =   embedding_output

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

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

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

              snake_case_         =   logits
              if patient_counter == self.patience:
                  break
          snake_case_         =   [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       __lowerCAmelCase       (	a   ):






  """simple docstring"""





  def __init__(       self    :    int      ,			_lowerCAmelCase    :    Dict					)				->	Tuple:




      """simple docstring"""
      super().__init__(_lowerCAmelCase					)
      snake_case_         =   config.num_labels

      snake_case_         =   AlbertModelWithPabee(_lowerCAmelCase					)
      snake_case_         =   nn.Dropout(config.classifier_dropout_prob					)
      snake_case_         =   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    :    List[str]      ,			_lowerCAmelCase    :    int=None      ,			_lowerCAmelCase    :    Optional[int]=None      ,			_lowerCAmelCase    :    Any=None      ,			_lowerCAmelCase    :    str=None      ,			_lowerCAmelCase    :    Optional[int]=None      ,			_lowerCAmelCase    :    Optional[Any]=None      ,			_lowerCAmelCase    :    Any=None      ,			)				->	Union[str, Any]:




      """simple docstring"""
      snake_case_         =   self.albert(
          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      ,			)

      snake_case_         =   (logits[-1],)

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

      return outputs








import math
from collections import OrderedDict

import torch
from packaging import version
from torch import Tensor, nn

from .utils import logging


SCREAMING_SNAKE_CASE		:Tuple											=       logging.get_logger(__name__)


class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





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




     """simple docstring"""
     super().__init__()
     if version.parse(torch.__version__					) < version.parse("1.12.0"					):
        raise ImportError(
            F'''You are using torch=={torch.__version__}, but torch>=1.12.0 is required to use '''
            "PytorchGELUTanh. Please upgrade torch."					)



  def    lowerCAmelCase__					(       self    :    Optional[int]      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return nn.functional.gelu(_lowerCAmelCase      ,			approximate="tanh"					)






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def    lowerCAmelCase__					(       self    :    int      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi					) * (input + 0.044_715 * torch.pow(_lowerCAmelCase      ,			3.0					))					))






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def __init__(       self    :    int      ,			_lowerCAmelCase    :    bool = False					)				->	int:




     """simple docstring"""
     super().__init__()
     if use_gelu_python:
        snake_case_         =   self._gelu_python
     else:
        snake_case_         =   nn.functional.gelu



  def    lowerCAmelCase__					(       self    :    Optional[Any]      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return input * 0.5 * (1.0 + torch.erf(input / math.sqrt(2.0					)					))



  def    lowerCAmelCase__					(       self    :    Dict      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return self.act(_lowerCAmelCase					)






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def    lowerCAmelCase__					(       self    :    Optional[int]      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return 0.5 * input * (1.0 + torch.tanh(input * 0.7_978_845_608 * (1.0 + 0.044_715 * input * input)					))






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def    lowerCAmelCase__					(       self    :    int      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return input * torch.sigmoid(1.702 * input					)






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def __init__(       self    :    List[str]      ,			_lowerCAmelCase    :    float      ,			_lowerCAmelCase    :    float					)				->	Any:




     """simple docstring"""
     if min > max:
        raise ValueError(F'''min should be < max (got min: {min}, max: {max})'''					)

     super().__init__()
     snake_case_         =   min
     snake_case_         =   max



  def    lowerCAmelCase__					(       self    :    Union[str, Any]      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return torch.clip(gelu(_lowerCAmelCase					)      ,			self.min      ,			self.max					)






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def __init__(       self    :    Any					)				->	str:




     """simple docstring"""
     super().__init__()
     snake_case_         =   math.sqrt(2 / math.pi					)



  def    lowerCAmelCase__					(       self    :    Optional[Any]      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return 0.5 * input * (1 + torch.tanh(self.precomputed_constant * (input + 0.044_715 * torch.pow(_lowerCAmelCase      ,			3					))					))






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def    lowerCAmelCase__					(       self    :    str      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return nn.functional.silu(_lowerCAmelCase					)






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





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




     """simple docstring"""
     super().__init__()
     if version.parse(torch.__version__					) < version.parse("1.9.0"					):
        snake_case_         =   self._mish_python
     else:
        snake_case_         =   nn.functional.mish



  def    lowerCAmelCase__					(       self    :    Tuple      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return input * torch.tanh(nn.functional.softplus(_lowerCAmelCase					)					)



  def    lowerCAmelCase__					(       self    :    Optional[Any]      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return self.act(_lowerCAmelCase					)






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def    lowerCAmelCase__					(       self    :    List[str]      ,			_lowerCAmelCase    :    Tensor					)				->	Tensor:




     """simple docstring"""
     return input






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def    lowerCAmelCase__					(       self    :    List[Any]      ,			_lowerCAmelCase    :    List[Any]      ,			_lowerCAmelCase    :    Union[str, Any]=0.707_107      ,			_lowerCAmelCase    :    int=0.282_095					)				->	List[str]:




     """simple docstring"""
     snake_case_         =   (input - mu).div(sigma * math.sqrt(2.0					)					)
     return 0.5 * (1.0 + torch.erf(_lowerCAmelCase					))






class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def    lowerCAmelCase__					(       self    :    Any      ,			_lowerCAmelCase    :    Union[str, Any]					)				->	Any:




     """simple docstring"""
     snake_case_         =   nn.functional.relu(_lowerCAmelCase					)
     snake_case_         =   torch.square(_lowerCAmelCase					)
     return squared






class       __lowerCAmelCase       (	a   ):






  """simple docstring"""





  def __getitem__(       self    :    Optional[Any]      ,			_lowerCAmelCase    :    str					)				->	Optional[int]:




     """simple docstring"""
     snake_case_         =   super().__getitem__(_lowerCAmelCase					)
     snake_case_      ,      snake_case_         =   content if isinstance(_lowerCAmelCase      ,			_lowerCAmelCase					) else (content, {})
     return cls(**_lowerCAmelCase					)






SCREAMING_SNAKE_CASE		:Optional[int]											=       {
    '''gelu''': GELUActivation,
    '''gelu_10''': (ClippedGELUActivation, {'''min''': -10, '''max''': 10}),
    '''gelu_fast''': FastGELUActivation,
    '''gelu_new''': NewGELUActivation,
    '''gelu_python''': (GELUActivation, {'''use_gelu_python''': True}),
    '''gelu_pytorch_tanh''': PytorchGELUTanh,
    '''gelu_accurate''': AccurateGELUActivation,
    '''laplace''': LaplaceActivation,
    '''linear''': LinearActivation,
    '''mish''': MishActivation,
    '''quick_gelu''': QuickGELUActivation,
    '''relu''': nn.ReLU,
    '''relu2''': ReLUSquaredActivation,
    '''relu6''': nn.ReLUa,
    '''sigmoid''': nn.Sigmoid,
    '''silu''': SiLUActivation,
    '''swish''': SiLUActivation,
    '''tanh''': nn.Tanh,
}
SCREAMING_SNAKE_CASE		:Optional[int]											=       ClassInstantier(ACTaCLS)



def      _lowerCAmelCase  (							lowerCAmelCase_							:Optional[Any]							)->Union[str, Any]:
   '''simple docstring'''






   if activation_string in ACTaFN:
      return ACTaFN[activation_string]
   else:
      raise KeyError(F'''function {activation_string} not found in ACT2FN mapping {list(ACTaFN.keys()							)}'''							)


# For backwards compatibility with: from activations import gelu_python
SCREAMING_SNAKE_CASE		:Tuple											=       get_activation('''gelu_python''')
SCREAMING_SNAKE_CASE		:int											=       get_activation('''gelu_new''')
SCREAMING_SNAKE_CASE		:List[Any]											=       get_activation('''gelu''')
SCREAMING_SNAKE_CASE		:int											=       get_activation('''gelu_fast''')
SCREAMING_SNAKE_CASE		:Tuple											=       get_activation('''quick_gelu''')
SCREAMING_SNAKE_CASE		:Tuple											=       get_activation('''silu''')
SCREAMING_SNAKE_CASE		:Any											=       get_activation('''mish''')



SCREAMING_SNAKE_CASE		:Optional[Any]											=       get_activation('''linear''')




