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





class 				lowercase    (    snake_case__			):


   def __UpperCAmelCase							( self					:						List[str] ,		__lowerCAmelCase					:						List[Any] ,		__lowerCAmelCase					:						Any ,		__lowerCAmelCase					:						Optional[Any]=None ,		__lowerCAmelCase					:						str=None)   ->       Union[str, Any]:
         if current_layer == 0:
               lowercase_     			=   self.embedding_hidden_mapping_in(lowercase_)
         else:
               lowercase_     			=   hidden_states[0]

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

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

         lowercase_     			=   self.albert_layer_groups[group_idx](
             lowercase_ ,		lowercase_ ,		head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group] ,		)
         lowercase_     			=   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.'			,					snake_case__			,					)
class 				lowercase    (    snake_case__			):


   def __init__( self					:						Optional[Any] ,		__lowerCAmelCase					:						List[Any])   ->       int:
         super().__init__(lowercase_)

         lowercase_     			=   AlbertTransformerWithPabee(lowercase_)

         self.init_weights()
         lowercase_     			=   0
         lowercase_     			=   0
         lowercase_     			=   0

         lowercase_     			=   0


   def __UpperCAmelCase							( self					:						List[Any] ,		__lowerCAmelCase					:						Union[str, Any])   ->       str:
         lowercase_     			=   threshold


   def __UpperCAmelCase							( self					:						List[str] ,		__lowerCAmelCase					:						Tuple)   ->       int:
         lowercase_     			=   patience


   def __UpperCAmelCase							( self					:						List[Any])   ->       Dict:
         lowercase_     			=   0
         lowercase_     			=   0


   def __UpperCAmelCase							( self					:						Union[str, Any])   ->       List[str]:
         lowercase_     			=   self.inference_layers_num / self.inference_instances_num
         lowercase_     			=   (
             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(lowercase_)


   @add_start_docstrings_to_model_forward(lowercase_)
   def __UpperCAmelCase							( self					:						List[str] ,		__lowerCAmelCase					:						List[str]=None ,		__lowerCAmelCase					:						List[str]=None ,		__lowerCAmelCase					:						str=None ,		__lowerCAmelCase					:						Dict=None ,		__lowerCAmelCase					:						Tuple=None ,		__lowerCAmelCase					:						int=None ,		__lowerCAmelCase					:						Optional[Any]=None ,		__lowerCAmelCase					:						Dict=None ,		__lowerCAmelCase					:						Dict=False ,		)   ->       List[Any]:

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

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

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

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

         lowercase_     			=   self.embeddings(
             lowercase_ ,		position_ids=lowercase_ ,		token_type_ids=lowercase_ ,		inputs_embeds=lowercase_)
         lowercase_     			=   embedding_output

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

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

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

                     lowercase_     			=   logits
                     if patient_counter == self.patience:
                           break
               lowercase_     			=   [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. '			,					snake_case__			,					)
class 				lowercase    (    snake_case__			):


   def __init__( self					:						List[Any] ,		__lowerCAmelCase					:						Any)   ->       Optional[int]:
         super().__init__(lowercase_)
         lowercase_     			=   config.num_labels

         lowercase_     			=   AlbertModelWithPabee(lowercase_)
         lowercase_     			=   nn.Dropout(config.classifier_dropout_prob)
         lowercase_     			=   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(lowercase_)
   def __UpperCAmelCase							( self					:						Tuple ,		__lowerCAmelCase					:						Optional[Any]=None ,		__lowerCAmelCase					:						List[Any]=None ,		__lowerCAmelCase					:						List[str]=None ,		__lowerCAmelCase					:						Optional[int]=None ,		__lowerCAmelCase					:						int=None ,		__lowerCAmelCase					:						Tuple=None ,		__lowerCAmelCase					:						Dict=None ,		)   ->       Any:
         lowercase_     			=   self.albert(
             input_ids=lowercase_ ,		attention_mask=lowercase_ ,		token_type_ids=lowercase_ ,		position_ids=lowercase_ ,		head_mask=lowercase_ ,		inputs_embeds=lowercase_ ,		output_dropout=self.dropout ,		output_layers=self.classifiers ,		regression=self.num_labels == 1 ,		)

         lowercase_     			=   (logits[-1],)

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

         return outputs








'''simple docstring'''
from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...file_utils import TensorType, is_torch_available
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import logging


a_           =    logging.get_logger(__name__)

a_           =    {
    "facebook/blenderbot-3B": "https://huggingface.co/facebook/blenderbot-3B/resolve/main/config.json",
    # See all Blenderbot models at https://huggingface.co/models?filter=blenderbot
}






class  UpperCAmelCase_					(						snake_case__       ):
 UpperCAmelCase_		     =    """blenderbot"""
 UpperCAmelCase_		     =    ["""past_key_values"""]
 UpperCAmelCase_		     =    {"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}





 def __init__(     self			,							lowercase_=80_08			,							lowercase_=1_28			,							lowercase_=2			,							lowercase_=1_02_40			,							lowercase_=32			,							lowercase_=24			,							lowercase_=1_02_40			,							lowercase_=32			,							lowercase_=0.0			,							lowercase_=0.0			,							lowercase_=True			,							lowercase_=True			,							lowercase_="gelu"			,							lowercase_=25_60			,							lowercase_=0.1			,							lowercase_=0.0			,							lowercase_=0.0			,							lowercase_=0.02			,							lowercase_=1			,							lowercase_=False			,							lowercase_=0			,							lowercase_=1			,							lowercase_=2			,							lowercase_=3			,							lowercase_=2			,							**lowercase_			,							):
     snake_case_			:     str        =    vocab_size
     snake_case_			:     str        =    max_position_embeddings
     snake_case_			:     Optional[Any]        =    d_model
     snake_case_			:     Any        =    encoder_ffn_dim
     snake_case_			:     Optional[Any]        =    encoder_layers
     snake_case_			:     List[Any]        =    encoder_attention_heads
     snake_case_			:     List[Any]        =    decoder_ffn_dim
     snake_case_			:     Optional[int]        =    decoder_layers
     snake_case_			:     Optional[int]        =    decoder_attention_heads
     snake_case_			:     Any        =    dropout
     snake_case_			:     List[str]        =    attention_dropout
     snake_case_			:     int        =    activation_dropout
     snake_case_			:     Dict        =    activation_function
     snake_case_			:     Any        =    init_std
     snake_case_			:     Union[str, Any]        =    encoder_layerdrop
     snake_case_			:     Optional[int]        =    decoder_layerdrop
     snake_case_			:     Any        =    use_cache
     snake_case_			:     List[str]        =    encoder_layers
     snake_case_			:     str        =    scale_embedding  # scale factor will be sqrt(d_model) if True

     super().__init__(
         pad_token_id=lowercase_			,							bos_token_id=lowercase_			,							eos_token_id=lowercase_			,							is_encoder_decoder=lowercase_			,							decoder_start_token_id=lowercase_			,							encoder_no_repeat_ngram_size=lowercase_			,							forced_eos_token_id=lowercase_			,							**lowercase_			,							)






class  UpperCAmelCase_					(						snake_case__       ):





 @property
 def 							snake_case__		(     self):
     if self.task in ["default", "seq2seq-lm"]:
         snake_case_			:     int        =    OrderedDict(
             [
                 ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                 ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
             ])
         if self.use_past:
             snake_case_			:     Any        =    {0: "batch"}
             snake_case_			:     Optional[int]        =    {0: "batch", 1: "past_decoder_sequence + sequence"}
         else:
             snake_case_			:     int        =    {0: "batch", 1: "decoder_sequence"}
             snake_case_			:     int        =    {0: "batch", 1: "decoder_sequence"}
         if self.use_past:
             self.fill_with_past_key_values_(lowercase_			,							direction="inputs")
     elif self.task == "causal-lm":
         snake_case_			:     Tuple        =    OrderedDict(
             [
                 ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                 ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
             ])
         if self.use_past:
             snake_case_      ,			snake_case_			:     Any        =    self.num_layers
             for i in range(lowercase_):
                 snake_case_			:     Dict        =    {0: "batch", 2: "past_sequence + sequence"}
                 snake_case_			:     Dict        =    {0: "batch", 2: "past_sequence + sequence"}
     else:
         snake_case_			:     str        =    OrderedDict(
             [
                 ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                 ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
                 ("decoder_input_ids", {0: "batch", 1: "decoder_sequence"}),
                 ("decoder_attention_mask", {0: "batch", 1: "decoder_sequence"}),
             ])

     return common_inputs





 @property
 # Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.outputs
 def 							snake_case__		(     self):
     if self.task in ["default", "seq2seq-lm"]:
         snake_case_			:     Optional[int]        =    super().outputs
     else:
         snake_case_			:     Dict        =    super(lowercase_			,							self).outputs
         if self.use_past:
             snake_case_      ,			snake_case_			:     Any        =    self.num_layers
             for i in range(lowercase_):
                 snake_case_			:     str        =    {0: "batch", 2: "past_sequence + sequence"}
                 snake_case_			:     int        =    {0: "batch", 2: "past_sequence + sequence"}
     return common_outputs





 def 							snake_case__		(     self			,							lowercase_			,							lowercase_ = -1			,							lowercase_ = -1			,							lowercase_ = False			,							lowercase_ = None			,							):
     snake_case_			:     List[Any]        =    self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
         lowercase_			,							lowercase_			,							lowercase_			,							lowercase_			,							lowercase_)
     # Generate decoder inputs
     snake_case_			:     int        =    seq_length if not self.use_past else 1
     snake_case_			:     List[str]        =    self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
         lowercase_			,							lowercase_			,							lowercase_			,							lowercase_			,							lowercase_)
     snake_case_			:     str        =    {F'decoder_{name}': tensor for name, tensor in decoder_inputs.items()}
     snake_case_			:     Any        =    dict(**lowercase_			,							**lowercase_)

     if self.use_past:
         if not is_torch_available():
             raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed.")
         else:
             import torch
         snake_case_      ,			snake_case_			:     int        =    common_inputs["input_ids"].shape
         snake_case_			:     str        =    common_inputs["decoder_input_ids"].shape[1]
         snake_case_      ,			snake_case_			:     Tuple        =    self.num_attention_heads
         snake_case_			:     str        =    (
             batch,
             num_encoder_attention_heads,
             encoder_seq_length,
             self._config.hidden_size // num_encoder_attention_heads,
         )
         snake_case_			:     Dict        =    decoder_seq_length
         snake_case_			:     List[Any]        =    (
             batch,
             num_decoder_attention_heads,
             decoder_past_length,
             self._config.hidden_size // num_decoder_attention_heads,
         )
         snake_case_			:     Optional[Any]        =    torch.cat(
             [common_inputs["decoder_attention_mask"], torch.ones(lowercase_			,							lowercase_)]			,							dim=1)
         snake_case_			:     Tuple        =    []
         snake_case_      ,			snake_case_			:     Tuple        =    self.num_layers

         for _ in range(lowercase_):
             common_inputs["past_key_values"].append(
                 (
                     torch.zeros(lowercase_),
                     torch.zeros(lowercase_),
                     torch.zeros(lowercase_),
                     torch.zeros(lowercase_),
                 ))
     return common_inputs





 def 							snake_case__		(     self			,							lowercase_			,							lowercase_ = -1			,							lowercase_ = -1			,							lowercase_ = False			,							lowercase_ = None			,							):
     snake_case_			:     str        =    self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
         lowercase_			,							lowercase_			,							lowercase_			,							lowercase_			,							lowercase_)

     if self.use_past:
         if not is_torch_available():
             raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed.")
         else:
             import torch
         snake_case_      ,			snake_case_			:     Optional[Any]        =    common_inputs["input_ids"].shape
         snake_case_			:     List[str]        =    seqlen
         snake_case_      ,			snake_case_			:     Optional[int]        =    self.num_layers
         snake_case_      ,			snake_case_			:     Union[str, Any]        =    self.num_attention_heads
         snake_case_			:     Optional[int]        =    (
             batch,
             num_encoder_attention_heads,
             past_key_values_length,
             self._config.hidden_size // num_encoder_attention_heads,
         )
         snake_case_			:     Dict        =    common_inputs["attention_mask"].dtype
         snake_case_			:     Optional[Any]        =    torch.cat(
             [common_inputs["attention_mask"], torch.ones(lowercase_			,							lowercase_			,							dtype=lowercase_)]			,							dim=1)
         snake_case_			:     Tuple        =    [
             (torch.zeros(lowercase_), torch.zeros(lowercase_)) for _ in range(lowercase_)
         ]
     return common_inputs





 def 							snake_case__		(     self			,							lowercase_			,							lowercase_ = -1			,							lowercase_ = -1			,							lowercase_ = False			,							lowercase_ = None			,							):
     # Copied from OnnxConfig.generate_dummy_inputs
     # Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
     # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
     snake_case_			:     Any        =    compute_effective_axis_dimension(
         lowercase_			,							fixed_dimension=OnnxConfig.default_fixed_batch			,							num_token_to_add=0)

     # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
     snake_case_			:     Tuple        =    tokenizer.num_special_tokens_to_add(lowercase_)
     snake_case_			:     Tuple        =    compute_effective_axis_dimension(
         lowercase_			,							fixed_dimension=OnnxConfig.default_fixed_sequence			,							num_token_to_add=lowercase_)

     # Generate dummy inputs according to compute batch and sequence
     snake_case_			:     Optional[int]        =    [" ".join([tokenizer.unk_token]) * seq_length] * batch_size
     snake_case_			:     Tuple        =    dict(tokenizer(lowercase_			,							return_tensors=lowercase_))
     return common_inputs





 def 							snake_case__		(     self			,							lowercase_			,							lowercase_ = -1			,							lowercase_ = -1			,							lowercase_ = False			,							lowercase_ = None			,							):
     if self.task in ["default", "seq2seq-lm"]:
         snake_case_			:     Optional[Any]        =    self._generate_dummy_inputs_for_default_and_seqaseq_lm(
             lowercase_			,							batch_size=lowercase_			,							seq_length=lowercase_			,							is_pair=lowercase_			,							framework=lowercase_)

     elif self.task == "causal-lm":
         snake_case_			:     Dict        =    self._generate_dummy_inputs_for_causal_lm(
             lowercase_			,							batch_size=lowercase_			,							seq_length=lowercase_			,							is_pair=lowercase_			,							framework=lowercase_)
     else:
         snake_case_			:     Dict        =    self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
             lowercase_			,							batch_size=lowercase_			,							seq_length=lowercase_			,							is_pair=lowercase_			,							framework=lowercase_)

     return common_inputs





 def 							snake_case__		(     self			,							lowercase_			,							lowercase_			,							lowercase_			,							lowercase_):
     if self.task in ["default", "seq2seq-lm"]:
         snake_case_			:     int        =    super()._flatten_past_key_values_(lowercase_			,							lowercase_			,							lowercase_			,							lowercase_)
     else:
         snake_case_			:     List[Any]        =    super(lowercase_			,							self)._flatten_past_key_values_(
             lowercase_			,							lowercase_			,							lowercase_			,							lowercase_)





 def 							snake_case__		(     self			,							lowercase_			,							lowercase_):
     if direction not in ["inputs", "outputs"]:
         raise ValueError(F'direction must either be "inputs" or "outputs", but {direction} was given')

     snake_case_			:     Union[str, Any]        =    "past_key_values" if direction == "inputs" else "present"
     snake_case_      ,			snake_case_			:     Union[str, Any]        =    self.num_layers

     snake_case_			:     List[Any]        =    "past_encoder_sequence"
     snake_case_			:     str        =    "past_decoder_sequence" if direction == "inputs" else "past_decoder_sequence + sequence"

     for i in range(lowercase_):
         snake_case_			:     Union[str, Any]        =    {0: "batch", 2: decoder_sequence}
         snake_case_			:     str        =    {0: "batch", 2: decoder_sequence}
         snake_case_			:     Any        =    {0: "batch", 2: encoder_sequence}
         snake_case_			:     Dict        =    {0: "batch", 2: encoder_sequence}




