





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.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEncoder,
    BertModel,
    BertPreTrainedModel,
)


UpperCAmelCase_     							= logging.getLogger(__name__)






class     lowercase__       (	__lowerCamelCase				):

   '''simple docstring'''


   def       UpperCamelCase__  (    self,		__magic_name__,		__magic_name__,		__magic_name__=None,		__magic_name__=None							) ->      Dict:






       """simple docstring"""
       UpperCamelCase__							:				Tuple          =      self.layer[current_layer](__magic_name__,		__magic_name__,		head_mask[current_layer]							)

       UpperCamelCase__							:				str          =      layer_outputs[0]

       return hidden_states






@add_start_docstrings(
    "The bare Bert Model transformer with PABEE outputting raw hidden-states without any specific head on top."			,				__lowerCamelCase			,				)
class     lowercase__       (	__lowerCamelCase				):

   '''simple docstring'''


   def __init__(    self,		__magic_name__							) ->      Tuple:






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

       UpperCamelCase__							:				Dict          =      BertEncoderWithPabee(__magic_name__							)

       self.init_weights()
       UpperCamelCase__							:				Union[str, Any]          =      0
       UpperCamelCase__							:				Optional[Any]          =      0
       UpperCamelCase__							:				List[Any]          =      0

       UpperCamelCase__							:				Any          =      0




   def       UpperCamelCase__  (    self,		__magic_name__							) ->      str:






       """simple docstring"""
       UpperCamelCase__							:				Optional[int]          =      threshold




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






       """simple docstring"""
       UpperCamelCase__							:				Optional[Any]          =      patience




   def       UpperCamelCase__  (    self							) ->      Any:






       """simple docstring"""
       UpperCamelCase__							:				Dict          =      0
       UpperCamelCase__							:				List[Any]          =      0




   def       UpperCamelCase__  (    self							) ->      Any:






       """simple docstring"""
       UpperCamelCase__							:				Optional[Any]          =      self.inference_layers_num / self.inference_instances_num
       UpperCamelCase__							:				List[str]          =      (
           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(__magic_name__							)




   @add_start_docstrings_to_model_forward(__magic_name__							)
   def       UpperCamelCase__  (    self,		__magic_name__=None,		__magic_name__=None,		__magic_name__=None,		__magic_name__=None,		__magic_name__=None,		__magic_name__=None,		__magic_name__=None,		__magic_name__=None,		__magic_name__=None,		__magic_name__=None,		__magic_name__=False,		) ->      Optional[Any]:






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

       UpperCamelCase__							:				List[str]          =      input_ids.device if input_ids is not None else inputs_embeds.device

       if attention_mask is None:
           UpperCamelCase__							:				Union[str, Any]          =      torch.ones(__magic_name__,		device=__magic_name__							)
       if token_type_ids is None:
           UpperCamelCase__							:				int          =      torch.zeros(__magic_name__,		dtype=torch.long,		device=__magic_name__							)

       # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
       # ourselves in which case we just need to make it broadcastable to all heads.
       UpperCamelCase__							:				torch.Tensor          =      self.get_extended_attention_mask(__magic_name__,		__magic_name__,		__magic_name__							)

       # If a 2D ou 3D attention mask is provided for the cross-attention
       # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
       if self.config.is_decoder and encoder_hidden_states is not None:
           UpperCamelCase__   ,UpperCamelCase__   ,UpperCamelCase__							:				Optional[Any]          =      encoder_hidden_states.size()
           UpperCamelCase__							:				Tuple          =      (encoder_batch_size, encoder_sequence_length)
           if encoder_attention_mask is None:
               UpperCamelCase__							:				Any          =      torch.ones(__magic_name__,		device=__magic_name__							)
           UpperCamelCase__							:				Dict          =      self.invert_attention_mask(__magic_name__							)
       else:
           UpperCamelCase__							:				Optional[int]          =      None

       # Prepare head mask if needed
       # 1.0 in head_mask indicate we keep the head
       # attention_probs has shape bsz x n_heads x N x N
       # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
       # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
       UpperCamelCase__							:				Optional[Any]          =      self.get_head_mask(__magic_name__,		self.config.num_hidden_layers							)

       UpperCamelCase__							:				Optional[Any]          =      self.embeddings(
           input_ids=__magic_name__,		position_ids=__magic_name__,		token_type_ids=__magic_name__,		inputs_embeds=__magic_name__							)
       UpperCamelCase__							:				str          =      embedding_output

       if self.training:
           UpperCamelCase__							:				List[str]          =      []
           for i in range(self.config.num_hidden_layers							):
               UpperCamelCase__							:				List[Any]          =      self.encoder.adaptive_forward(
                   __magic_name__,		current_layer=__magic_name__,		attention_mask=__magic_name__,		head_mask=__magic_name__							)

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

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

               UpperCamelCase__							:				List[Any]          =      logits
               if patient_counter == self.patience:
                   break
           UpperCamelCase__							:				Tuple          =      [patient_result]
           self.inference_layers_num += calculated_layer_num
           self.inference_instances_num += 1

       return res






@add_start_docstrings(
    "Bert 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. "			,				__lowerCamelCase			,				)
class     lowercase__       (	__lowerCamelCase				):

   '''simple docstring'''


   def __init__(    self,		__magic_name__							) ->      Tuple:






       """simple docstring"""
       super().__init__(__magic_name__							)
       UpperCamelCase__							:				int          =      config.num_labels

       UpperCamelCase__							:				Optional[Any]          =      BertModelWithPabee(__magic_name__							)
       UpperCamelCase__							:				Optional[int]          =      nn.Dropout(config.hidden_dropout_prob							)
       UpperCamelCase__							:				Tuple          =      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(__magic_name__							)
   def       UpperCamelCase__  (    self,		__magic_name__=None,		__magic_name__=None,		__magic_name__=None,		__magic_name__=None,		__magic_name__=None,		__magic_name__=None,		__magic_name__=None,		) ->      Dict:






       """simple docstring"""
       UpperCamelCase__							:				List[str]          =      self.bert(
           input_ids=__magic_name__,		attention_mask=__magic_name__,		token_type_ids=__magic_name__,		position_ids=__magic_name__,		head_mask=__magic_name__,		inputs_embeds=__magic_name__,		output_dropout=self.dropout,		output_layers=self.classifiers,		regression=self.num_labels == 1,		)

       UpperCamelCase__							:				Optional[int]          =      (logits[-1],)

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

       return outputs








from typing import List, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase_     							= logging.get_logger(__name__)

UpperCAmelCase_     							= {
    'huggingface/informer-tourism-monthly': (
        'https://huggingface.co/huggingface/informer-tourism-monthly/resolve/main/config.json'
    ),
    # See all Informer models at https://huggingface.co/models?filter=informer
}






class     lowercase__       (	__lowerCamelCase				):

   '''simple docstring'''


   a						:							Optional[int]   		=		"informer"
   a						:							Tuple   		=		{
       "hidden_size": "d_model",
       "num_attention_heads": "encoder_attention_heads",
       "num_hidden_layers": "encoder_layers",
   }




   def __init__(    self,		__magic_name__ = None,		__magic_name__ = None,		__magic_name__ = "student_t",		__magic_name__ = "nll",		__magic_name__ = 1,		__magic_name__ = None,		__magic_name__ = "mean",		__magic_name__ = 0,		__magic_name__ = 0,		__magic_name__ = 0,		__magic_name__ = 0,		__magic_name__ = None,		__magic_name__ = None,		__magic_name__ = 64,		__magic_name__ = 32,		__magic_name__ = 32,		__magic_name__ = 2,		__magic_name__ = 2,		__magic_name__ = 2,		__magic_name__ = 2,		__magic_name__ = True,		__magic_name__ = "gelu",		__magic_name__ = 0.05,		__magic_name__ = 0.1,		__magic_name__ = 0.1,		__magic_name__ = 0.1,		__magic_name__ = 0.1,		__magic_name__ = 100,		__magic_name__ = 0.02,		__magic_name__=True,		__magic_name__ = "prob",		__magic_name__ = 5,		__magic_name__ = True,		**__magic_name__,		) ->      Optional[int]:






       """simple docstring"""
       # time series specific configuration
       UpperCamelCase__							:				List[Any]          =      prediction_length
       UpperCamelCase__							:				Any          =      context_length or prediction_length
       UpperCamelCase__							:				Optional[int]          =      distribution_output
       UpperCamelCase__							:				Union[str, Any]          =      loss
       UpperCamelCase__							:				Optional[Any]          =      input_size
       UpperCamelCase__							:				Dict          =      num_time_features
       UpperCamelCase__							:				Dict          =      lags_sequence if lags_sequence is not None else [1, 2, 3, 4, 5, 6, 7]
       UpperCamelCase__							:				Optional[int]          =      scaling
       UpperCamelCase__							:				Any          =      num_dynamic_real_features
       UpperCamelCase__							:				Optional[int]          =      num_static_real_features
       UpperCamelCase__							:				Optional[Any]          =      num_static_categorical_features

       # set cardinality
       if cardinality and num_static_categorical_features > 0:
           if len(__magic_name__							) != num_static_categorical_features:
               raise ValueError(
                   '''The cardinality should be a list of the same length as `num_static_categorical_features`'''							)
           UpperCamelCase__							:				str          =      cardinality
       else:
           UpperCamelCase__							:				int          =      [0]

       # set embedding_dimension
       if embedding_dimension and num_static_categorical_features > 0:
           if len(__magic_name__							) != num_static_categorical_features:
               raise ValueError(
                   '''The embedding dimension should be a list of the same length as `num_static_categorical_features`'''							)
           UpperCamelCase__							:				Any          =      embedding_dimension
       else:
           UpperCamelCase__							:				Union[str, Any]          =      [min(50,		(cat + 1) // 2							) for cat in self.cardinality]

       UpperCamelCase__							:				Optional[Any]          =      num_parallel_samples

       # Transformer architecture configuration
       UpperCamelCase__							:				Union[str, Any]          =      input_size * len(self.lags_sequence							) + self._number_of_features
       UpperCamelCase__							:				Optional[Any]          =      d_model
       UpperCamelCase__							:				Tuple          =      encoder_attention_heads
       UpperCamelCase__							:				Any          =      decoder_attention_heads
       UpperCamelCase__							:				Dict          =      encoder_ffn_dim
       UpperCamelCase__							:				Optional[Any]          =      decoder_ffn_dim
       UpperCamelCase__							:				str          =      encoder_layers
       UpperCamelCase__							:				Optional[int]          =      decoder_layers

       UpperCamelCase__							:				Optional[Any]          =      dropout
       UpperCamelCase__							:				List[Any]          =      attention_dropout
       UpperCamelCase__							:				Any          =      activation_dropout
       UpperCamelCase__							:				Optional[int]          =      encoder_layerdrop
       UpperCamelCase__							:				Union[str, Any]          =      decoder_layerdrop

       UpperCamelCase__							:				Tuple          =      activation_function
       UpperCamelCase__							:				List[str]          =      init_std

       UpperCamelCase__							:				int          =      use_cache

       # Informer
       UpperCamelCase__							:				Optional[int]          =      attention_type
       UpperCamelCase__							:				Optional[int]          =      sampling_factor
       UpperCamelCase__							:				Dict          =      distil

       super().__init__(is_encoder_decoder=__magic_name__,		**__magic_name__							)




   @property
   def       UpperCamelCase__  (    self							) ->      int:






       """simple docstring"""
       return (
           sum(self.embedding_dimension							)
           + self.num_dynamic_real_features
           + self.num_time_features
           + self.num_static_real_features
           + self.input_size * 2  # the log1p(abs(loc)) and log(scale) features
       )

