





"""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_					:						Any      = logging.getLogger(__name__)

class 			_UpperCAmelCase  (       UpperCAmelCase__		):



      '''simple docstring'''

      def  lowerCamelCase_      (					self ,				snake_case_ ,				snake_case_ ,				snake_case_=None ,				snake_case_=None  ):

        """simple docstring"""

        if current_layer == 0:
          A_  :	Union[str, Any]    						=  self.embedding_hidden_mapping_in(snake_case_  )
        else:
          A_  :	List[str]    						=  hidden_states[0]

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

        # Index of the hidden group
        A_  :	Optional[Any]    						=  int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)  )

        A_  :	List[Any]    						=  self.albert_layer_groups[group_idx](
            snake_case_ ,				snake_case_ ,				head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group] ,				)
        A_  :	str    						=  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 			_UpperCAmelCase  (       UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(					self ,				snake_case_  ):

        """simple docstring"""

        super().__init__(snake_case_  )

        A_  :	Tuple    						=  AlbertTransformerWithPabee(snake_case_  )

        self.init_weights()
        A_  :	str    						=  0
        A_  :	str    						=  0
        A_  :	int    						=  0

        A_  :	Optional[Any]    						=  0

      def  lowerCamelCase_      (					self ,				snake_case_  ):

        """simple docstring"""

        A_  :	List[str]    						=  threshold

      def  lowerCamelCase_      (					self ,				snake_case_  ):

        """simple docstring"""

        A_  :	Any    						=  patience

      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

        A_  :	Optional[int]    						=  0
        A_  :	Union[str, Any]    						=  0

      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

        A_  :	Tuple    						=  self.inference_layers_num / self.inference_instances_num
        A_  :	List[Any]    						=  (
            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  lowerCamelCase_      (					self ,				snake_case_=None ,				snake_case_=None ,				snake_case_=None ,				snake_case_=None ,				snake_case_=None ,				snake_case_=None ,				snake_case_=None ,				snake_case_=None ,				snake_case_=False ,				):

        """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_  :	Tuple    						=  input_ids.size()
        elif inputs_embeds is not None:
          A_  :	Union[str, Any]    						=  inputs_embeds.size()[:-1]
        else:
          raise ValueError('You have to specify either input_ids or inputs_embeds'  )

        A_  :	Tuple    						=  input_ids.device if input_ids is not None else inputs_embeds.device

        if attention_mask is None:
          A_  :	Tuple    						=  torch.ones(snake_case_ ,				device=snake_case_  )
        if token_type_ids is None:
          A_  :	Optional[int]    						=  torch.zeros(snake_case_ ,				dtype=torch.long ,				device=snake_case_  )

        A_  :	List[Any]    						=  attention_mask.unsqueeze(1  ).unsqueeze(2  )
        A_  :	Union[str, Any]    						=  extended_attention_mask.to(dtype=self.dtype  )  # fp16 compatibility
        A_  :	Tuple    						=  (1.0 - extended_attention_mask) * -1_00_00.0
        A_  :	Union[str, Any]    						=  self.get_head_mask(snake_case_ ,				self.config.num_hidden_layers  )

        A_  :	Union[str, Any]    						=  self.embeddings(
            snake_case_ ,				position_ids=snake_case_ ,				token_type_ids=snake_case_ ,				inputs_embeds=snake_case_  )
        A_  :	List[str]    						=  embedding_output

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

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

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

            A_  :	int    						=  logits
            if patient_counter == self.patience:
              break
          A_  :	List[str]    						=  [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 			_UpperCAmelCase  (       UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(					self ,				snake_case_  ):

        """simple docstring"""

        super().__init__(snake_case_  )
        A_  :	int    						=  config.num_labels

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

        """simple docstring"""

        A_  :	List[Any]    						=  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 ,				)

        A_  :	Union[str, Any]    						=  (logits[-1],)

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

        return outputs






"""simple docstring"""




from functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  if not inplace:
    A_  :	Optional[int]    						=  ma + ma
  else:
    ma += ma

  return ma





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  A_  :	Tuple    						=  -1 * len(_UpperCAmelCase )
  A_  :	List[str]    						=  list(range(len(tensor.shape[:zero_index] ) ) )
  return tensor.permute(first_inds + [zero_index + i for i in inds] )





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  return t.reshape(t.shape[:-no_dims] + (-1,) )





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase = 1E-4 ):






  """simple docstring"""





  A_  :	Optional[Any]    						=  mask.expand(*value.shape )
  return torch.sum(mask * value			, dim=_UpperCAmelCase ) / (eps + torch.sum(_UpperCAmelCase			, dim=_UpperCAmelCase ))





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase = 2.3_125			, _UpperCAmelCase = 21.6_875			, _UpperCAmelCase = 64 ):






  """simple docstring"""





  A_  :	List[str]    						=  torch.linspace(_UpperCAmelCase			, _UpperCAmelCase			, no_bins - 1			, device=pts.device )
  A_  :	Optional[Any]    						=  torch.sqrt(torch.sum((pts.unsqueeze(-2 ) - pts.unsqueeze(-3 )) ** 2			, dim=-1 ) )
  return torch.bucketize(_UpperCAmelCase			, _UpperCAmelCase )





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  A_  :	Any    						=  dicts[0]
  A_  :	int    						=  {}
  for k, v in first.items():
    A_  :	Dict    						=  [d[k] for d in dicts]
    if isinstance(_UpperCAmelCase			, _UpperCAmelCase ):
      A_  :	List[Any]    						=  dict_multimap(_UpperCAmelCase			, _UpperCAmelCase )
    else:
      A_  :	Union[str, Any]    						=  fn(_UpperCAmelCase )

  return new_dict





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  A_  :	Tuple    						=  v_bins.view(((1,) * len(x.shape )) + (len(_UpperCAmelCase ),) )
  A_  :	Any    						=  x[..., None] - reshaped_bins
  A_  :	List[str]    						=  torch.argmin(torch.abs(_UpperCAmelCase )			, dim=-1 )
  return nn.functional.one_hot(_UpperCAmelCase			, num_classes=len(_UpperCAmelCase ) ).float()





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase = 0			, _UpperCAmelCase = 0 ):






  """simple docstring"""





  A_  :	List[Union[slice, torch.Tensor]]    						=  []
  for i, s in enumerate(data.shape[:no_batch_dims] ):
    A_  :	List[Any]    						=  torch.arange(_UpperCAmelCase )
    A_  :	Dict    						=  r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape ) - i - 1))) )
    ranges.append(_UpperCAmelCase )

  A_  :	List[Union[slice, torch.Tensor]]    						=  [slice(_UpperCAmelCase ) for _ in range(len(data.shape ) - no_batch_dims )]
  A_  :	Optional[Any]    						=  inds
  ranges.extend(_UpperCAmelCase )
  # Matt note: Editing this to get around the behaviour of using a list as an array index changing
  # in recent Numpy versions
  return data[tuple(_UpperCAmelCase )]


lowerCamelCase_					:						Union[str, Any]      = TypeVar('T')





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  A_  :	Dict[Any, Union[dict, list, tuple, Any]]    						=  {}
  for k, v in dic.items():
    if isinstance(_UpperCAmelCase			, _UpperCAmelCase ):
      A_  :	int    						=  dict_map(_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase )
    else:
      A_  :	Any    						=  tree_map(_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase )

  return new_dict





@overload
def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  ...





@overload
def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  ...





@overload
def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  ...





@overload
def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  ...





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  if isinstance(_UpperCAmelCase			, _UpperCAmelCase ):
    return dict_map(_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase )
  elif isinstance(_UpperCAmelCase			, _UpperCAmelCase ):
    return [tree_map(_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ) for x in tree]
  elif isinstance(_UpperCAmelCase			, _UpperCAmelCase ):
    return tuple(tree_map(_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ) for x in tree )
  elif isinstance(_UpperCAmelCase			, _UpperCAmelCase ):
    return fn(_UpperCAmelCase )
  else:
    print(type(_UpperCAmelCase ) )
    raise ValueError('Not supported' )





lowerCamelCase_					:						Optional[int]      = partial(tree_map, leaf_type=torch.Tensor)