






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


__A     						= logging.getLogger(__name__)




class a_     (     UpperCamelCase_				):







   def      SCREAMING_SNAKE_CASE__       (self       , __a       , __a       , __a=None       , __a=None)      ->   Optional[Any]:
      """simple docstring"""






      if current_layer == 0:
         __snake_case					:   List[Any]  				=			self.embedding_hidden_mapping_in(__a)
      else:
         __snake_case					:   Optional[Any]  				=			hidden_states[0]

      __snake_case					:   Tuple  				=			int(self.config.num_hidden_layers / self.config.num_hidden_groups)

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

      __snake_case					:   Union[str, Any]  				=			self.albert_layer_groups[group_idx](
          __a       , __a       , head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]       , )
      __snake_case					:   Any  				=			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 a_     (     UpperCamelCase_				):







   def __init__(self       , __a)      ->   List[str]:
      """simple docstring"""






      super().__init__(__a)

      __snake_case					:   Optional[Any]  				=			AlbertTransformerWithPabee(__a)

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

      __snake_case					:   str  				=			0







   def      SCREAMING_SNAKE_CASE__       (self       , __a)      ->   Dict:
      """simple docstring"""






      __snake_case					:   List[Any]  				=			threshold







   def      SCREAMING_SNAKE_CASE__       (self       , __a)      ->   Any:
      """simple docstring"""






      __snake_case					:   Union[str, Any]  				=			patience







   def      SCREAMING_SNAKE_CASE__       (self)      ->   Tuple:
      """simple docstring"""






      __snake_case					:   int  				=			0
      __snake_case					:   List[Any]  				=			0







   def      SCREAMING_SNAKE_CASE__       (self)      ->   Union[str, Any]:
      """simple docstring"""






      __snake_case					:   int  				=			self.inference_layers_num / self.inference_instances_num
      __snake_case					:   Tuple  				=			(
          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(__a)







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

      __snake_case					:   Optional[int]  				=			input_ids.device if input_ids is not None else inputs_embeds.device

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

      __snake_case					:   Optional[int]  				=			attention_mask.unsqueeze(1).unsqueeze(2)
      __snake_case					:   List[Any]  				=			extended_attention_mask.to(dtype=self.dtype)  # fp16 compatibility
      __snake_case					:   Optional[Any]  				=			(1.0 - extended_attention_mask) * -10_000.0
      __snake_case					:   int  				=			self.get_head_mask(__a       , self.config.num_hidden_layers)

      __snake_case					:   str  				=			self.embeddings(
          __a       , position_ids=__a       , token_type_ids=__a       , inputs_embeds=__a)
      __snake_case					:   int  				=			embedding_output

      if self.training:
         __snake_case					:   int  				=			[]
         for i in range(self.config.num_hidden_layers):
            __snake_case					:   Any  				=			self.encoder.adaptive_forward(
                __a       , current_layer=__a       , attention_mask=__a       , head_mask=__a       , )

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

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

            __snake_case					:   str  				=			logits
            if patient_counter == self.patience:
               break
         __snake_case					:   Union[str, Any]  				=			[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
    the pooled output) e.g. for GLUE tasks. """		,   UpperCamelCase_		,   )
class a_     (     UpperCamelCase_				):







   def __init__(self       , __a)      ->   List[str]:
      """simple docstring"""






      super().__init__(__a)
      __snake_case					:   int  				=			config.num_labels

      __snake_case					:   str  				=			AlbertModelWithPabee(__a)
      __snake_case					:   List[str]  				=			nn.Dropout(config.classifier_dropout_prob)
      __snake_case					:   List[Any]  				=			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(__a)
   def      SCREAMING_SNAKE_CASE__       (self       , __a=None       , __a=None       , __a=None       , __a=None       , __a=None       , __a=None       , __a=None       , )      ->   List[Any]:
      """simple docstring"""






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

      __snake_case					:   Tuple  				=			(logits[-1],)

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

      return outputs







'''simple docstring'''






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__A     						= {
    '''configuration_prophetnet''': ['''PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ProphetNetConfig'''],
    '''tokenization_prophetnet''': ['''ProphetNetTokenizer'''],
}

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __A     						= [
          '''PROPHETNET_PRETRAINED_MODEL_ARCHIVE_LIST''',
          '''ProphetNetDecoder''',
          '''ProphetNetEncoder''',
          '''ProphetNetForCausalLM''',
          '''ProphetNetForConditionalGeneration''',
          '''ProphetNetModel''',
          '''ProphetNetPreTrainedModel''',
      ]


if TYPE_CHECKING:
      from .configuration_prophetnet import PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP, ProphetNetConfig
      from .tokenization_prophetnet import ProphetNetTokenizer

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_prophetnet import (
                PROPHETNET_PRETRAINED_MODEL_ARCHIVE_LIST,
                ProphetNetDecoder,
                ProphetNetEncoder,
                ProphetNetForCausalLM,
                ProphetNetForConditionalGeneration,
                ProphetNetModel,
                ProphetNetPreTrainedModel,
            )

else:
      import sys

      __A     						= _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)