



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


SCREAMING_SNAKE_CASE__          =  logging.getLogger(__name__)


class 					snake_case			(UpperCamelCase ):



       def 				_a   (					self  ,UpperCAmelCase_  ,UpperCAmelCase_  ,UpperCAmelCase_=None  ,UpperCAmelCase_=None   )	->				List[Any]:
           lowercase__									=    self.layer[current_layer](UpperCAmelCase_  ,UpperCAmelCase_  ,head_mask[current_layer]   )

           lowercase__									=    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."		,   UpperCamelCase		,   )
class 					snake_case			(UpperCamelCase ):



       def __init__(					self  ,UpperCAmelCase_   )	->				int:
           super().__init__(UpperCAmelCase_   )

           lowercase__									=    BertEncoderWithPabee(UpperCAmelCase_   )

           self.init_weights()
           lowercase__									=    0
           lowercase__									=    0
           lowercase__									=    0

           lowercase__									=    0



       def 				_a   (					self  ,UpperCAmelCase_   )	->				Any:
           lowercase__									=    threshold



       def 				_a   (					self  ,UpperCAmelCase_   )	->				Union[str, Any]:
           lowercase__									=    patience



       def 				_a   (					self   )	->				str:
           lowercase__									=    0
           lowercase__									=    0



       def 				_a   (					self   )	->				List[Any]:
           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(UpperCAmelCase_   )







       @add_start_docstrings_to_model_forward(UpperCAmelCase_   )
       def 				_a   (					self  ,UpperCAmelCase_=None  ,UpperCAmelCase_=None  ,UpperCAmelCase_=None  ,UpperCAmelCase_=None  ,UpperCAmelCase_=None  ,UpperCAmelCase_=None  ,UpperCAmelCase_=None  ,UpperCAmelCase_=None  ,UpperCAmelCase_=None  ,UpperCAmelCase_=None  ,UpperCAmelCase_=False  ,)	->				int:

           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(UpperCAmelCase_  ,device=UpperCAmelCase_   )
           if token_type_ids is None:
               lowercase__									=    torch.zeros(UpperCAmelCase_  ,dtype=torch.long  ,device=UpperCAmelCase_   )

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

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

           lowercase__									=    self.embeddings(
               input_ids=UpperCAmelCase_  ,position_ids=UpperCAmelCase_  ,token_type_ids=UpperCAmelCase_  ,inputs_embeds=UpperCAmelCase_   )
           lowercase__									=    embedding_output

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

                   lowercase__									=    self.pooler(UpperCAmelCase_   )
                   lowercase__									=    output_layers[i](output_dropout(UpperCAmelCase_   )   )
                   res.append(UpperCAmelCase_   )
           elif self.patience == 0:  # Use all layers for inference
               lowercase__									=    self.encoder(
                   UpperCAmelCase_  ,attention_mask=UpperCAmelCase_  ,head_mask=UpperCAmelCase_  ,encoder_hidden_states=UpperCAmelCase_  ,encoder_attention_mask=UpperCAmelCase_  ,)
               lowercase__									=    self.pooler(encoder_outputs[0]   )
               lowercase__									=    [output_layers[self.config.num_hidden_layers - 1](UpperCAmelCase_   )]
           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(
                       UpperCAmelCase_  ,current_layer=UpperCAmelCase_  ,attention_mask=UpperCAmelCase_  ,head_mask=UpperCAmelCase_   )

                   lowercase__									=    self.pooler(UpperCAmelCase_   )
                   lowercase__									=    output_layers[i](UpperCAmelCase_   )
                   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(UpperCAmelCase_   )   ):
                           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(
    "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. "		,   UpperCamelCase		,   )
class 					snake_case			(UpperCamelCase ):



       def __init__(					self  ,UpperCAmelCase_   )	->				Optional[Any]:
           super().__init__(UpperCAmelCase_   )
           lowercase__									=    config.num_labels

           lowercase__									=    BertModelWithPabee(UpperCAmelCase_   )
           lowercase__									=    nn.Dropout(config.hidden_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(UpperCAmelCase_   )
       def 				_a   (					self  ,UpperCAmelCase_=None  ,UpperCAmelCase_=None  ,UpperCAmelCase_=None  ,UpperCAmelCase_=None  ,UpperCAmelCase_=None  ,UpperCAmelCase_=None  ,UpperCAmelCase_=None  ,)	->				Dict:
           lowercase__									=    self.bert(
               input_ids=UpperCAmelCase_  ,attention_mask=UpperCAmelCase_  ,token_type_ids=UpperCAmelCase_  ,position_ids=UpperCAmelCase_  ,head_mask=UpperCAmelCase_  ,inputs_embeds=UpperCAmelCase_  ,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(UpperCAmelCase_   ):
                   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 typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


SCREAMING_SNAKE_CASE__          =  {
    "configuration_clap": [
        "CLAP_PRETRAINED_MODEL_ARCHIVE_LIST",
        "ClapAudioConfig",
        "ClapConfig",
        "ClapTextConfig",
    ],
    "processing_clap": ["ClapProcessor"],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 SCREAMING_SNAKE_CASE__          =  [
     "CLAP_PRETRAINED_MODEL_ARCHIVE_LIST",
     "ClapModel",
     "ClapPreTrainedModel",
     "ClapTextModel",
     "ClapTextModelWithProjection",
     "ClapAudioModel",
     "ClapAudioModelWithProjection",
 ]
 SCREAMING_SNAKE_CASE__          =  ["ClapFeatureExtractor"]

if TYPE_CHECKING:
 from .configuration_clap import (
     CLAP_PRETRAINED_MODEL_ARCHIVE_LIST,
     ClapAudioConfig,
     ClapConfig,
     ClapTextConfig,
 )
 from .processing_clap import ClapProcessor

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .feature_extraction_clap import ClapFeatureExtractor
  from .modeling_clap import (
      CLAP_PRETRAINED_MODEL_ARCHIVE_LIST,
      ClapAudioModel,
      ClapAudioModelWithProjection,
      ClapModel,
      ClapPreTrainedModel,
      ClapTextModel,
      ClapTextModelWithProjection,
  )


else:
 import sys

 SCREAMING_SNAKE_CASE__          =  _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)




