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


snake_case__				= logging.getLogger(__name__)



class  UpperCamelCase				( __lowercase ):

     '''simple docstring'''

     def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_=None				,				A_=None						) ->				Any:






           """simple docstring"""
           _lowerCamelCase			    =				self.layer[current_layer](A_				,				A_				,				head_mask[current_layer]						)

           _lowerCamelCase			    =				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.'   ,  __lowercase   ,  )
class  UpperCamelCase				( __lowercase ):

     '''simple docstring'''

     def __init__(					self				,				A_						) ->				int:






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

           _lowerCamelCase			    =				BertEncoderWithPabee(A_						)

           self.init_weights()
           _lowerCamelCase			    =				0
           _lowerCamelCase			    =				0
           _lowerCamelCase			    =				0

           _lowerCamelCase			    =				0



     def 		UpperCamelCase_						(					self				,				A_						) ->				Tuple:






           """simple docstring"""
           _lowerCamelCase			    =				threshold



     def 		UpperCamelCase_						(					self				,				A_						) ->				Union[str, Any]:






           """simple docstring"""
           _lowerCamelCase			    =				patience



     def 		UpperCamelCase_						(					self						) ->				List[Any]:






           """simple docstring"""
           _lowerCamelCase			    =				0
           _lowerCamelCase			    =				0



     def 		UpperCamelCase_						(					self						) ->				int:






           """simple docstring"""
           _lowerCamelCase			    =				self.inference_layers_num / self.inference_instances_num
           _lowerCamelCase			    =				(
               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 		UpperCamelCase_						(					self				,				A_=None				,				A_=None				,				A_=None				,				A_=None				,				A_=None				,				A_=None				,				A_=None				,				A_=None				,				A_=None				,				A_=None				,				A_=False				,				) ->				Tuple:






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

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

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

           # 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.
           _lowerCamelCase			    =				self.get_extended_attention_mask(A_				,				A_				,				A_						)

           # 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:
                 _lowerCamelCase    ,					_lowerCamelCase    ,					_lowerCamelCase			    =				encoder_hidden_states.size()
                 _lowerCamelCase			    =				(encoder_batch_size, encoder_sequence_length)
                 if encoder_attention_mask is None:
                       _lowerCamelCase			    =				torch.ones(A_				,				device=A_						)
                 _lowerCamelCase			    =				self.invert_attention_mask(A_						)
           else:
                 _lowerCamelCase			    =				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]
           _lowerCamelCase			    =				self.get_head_mask(A_				,				self.config.num_hidden_layers						)

           _lowerCamelCase			    =				self.embeddings(
               input_ids=A_				,				position_ids=A_				,				token_type_ids=A_				,				inputs_embeds=A_						)
           _lowerCamelCase			    =				embedding_output

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

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

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

                       _lowerCamelCase			    =				logits
                       if patient_counter == self.patience:
                             break
                 _lowerCamelCase			    =				[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. '   ,  __lowercase   ,  )
class  UpperCamelCase				( __lowercase ):

     '''simple docstring'''

     def __init__(					self				,				A_						) ->				Union[str, Any]:






           """simple docstring"""
           super().__init__(A_						)
           _lowerCamelCase			    =				config.num_labels

           _lowerCamelCase			    =				BertModelWithPabee(A_						)
           _lowerCamelCase			    =				nn.Dropout(config.hidden_dropout_prob						)
           _lowerCamelCase			    =				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 		UpperCamelCase_						(					self				,				A_=None				,				A_=None				,				A_=None				,				A_=None				,				A_=None				,				A_=None				,				A_=None				,				) ->				str:






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

           _lowerCamelCase			    =				(logits[-1],)

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

           return outputs
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


snake_case__				= {
    'configuration_whisper': ['WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP', 'WhisperConfig', 'WhisperOnnxConfig'],
    'feature_extraction_whisper': ['WhisperFeatureExtractor'],
    'processing_whisper': ['WhisperProcessor'],
    'tokenization_whisper': ['WhisperTokenizer'],
}

try:
     if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case__				= ['WhisperTokenizerFast']

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case__				= [
         'WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST',
         'WhisperForConditionalGeneration',
         'WhisperModel',
         'WhisperPreTrainedModel',
         'WhisperForAudioClassification',
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case__				= [
         'TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST',
         'TFWhisperForConditionalGeneration',
         'TFWhisperModel',
         'TFWhisperPreTrainedModel',
     ]

try:
     if not is_flax_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case__				= [
         'FlaxWhisperForConditionalGeneration',
         'FlaxWhisperModel',
         'FlaxWhisperPreTrainedModel',
         'FlaxWhisperForAudioClassification',
     ]


if TYPE_CHECKING:
     from .configuration_whisper import WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP, WhisperConfig, WhisperOnnxConfig
     from .feature_extraction_whisper import WhisperFeatureExtractor
     from .processing_whisper import WhisperProcessor
     from .tokenization_whisper import WhisperTokenizer

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_whisper_fast import WhisperTokenizerFast

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_whisper import (
              WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
              WhisperForAudioClassification,
              WhisperForConditionalGeneration,
              WhisperModel,
              WhisperPreTrainedModel,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_whisper import (
              TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
              TFWhisperForConditionalGeneration,
              TFWhisperModel,
              TFWhisperPreTrainedModel,
          )

     try:
          if not is_flax_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_flax_whisper import (
              FlaxWhisperForAudioClassification,
              FlaxWhisperForConditionalGeneration,
              FlaxWhisperModel,
              FlaxWhisperPreTrainedModel,
          )

else:
     import sys

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