



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


__lowerCAmelCase							:       Dict          =       logging.getLogger(__name__)



class  _lowerCAmelCase     (					UpperCamelCase__		):





		"""simple docstring"""






		def  UpperCAmelCase__				(       self				,     _lowercase				,     _lowercase				,     _lowercase=None				,     _lowercase=None )			->    str:







								'''simple docstring'''

								snake_case_ :			int            =      self.layer[current_layer](__lowerCamelCase				,     __lowerCamelCase				,     head_mask[current_layer] )

								snake_case_ :			List[Any]            =      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  _lowerCAmelCase     (					UpperCamelCase__		):





		"""simple docstring"""






		def __init__(       self				,     _lowercase )			->    str:







								'''simple docstring'''

								super().__init__(__lowerCamelCase )

								snake_case_ :			str            =      BertEncoderWithPabee(__lowerCamelCase )

								self.init_weights()
								snake_case_ :			List[str]            =      0
								snake_case_ :			List[Any]            =      0
								snake_case_ :			str            =      0

								snake_case_ :			Optional[Any]            =      0





		def  UpperCAmelCase__				(       self				,     _lowercase )			->    List[Any]:







								'''simple docstring'''

								snake_case_ :			List[Any]            =      threshold





		def  UpperCAmelCase__				(       self				,     _lowercase )			->    Union[str, Any]:







								'''simple docstring'''

								snake_case_ :			Union[str, Any]            =      patience





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







								'''simple docstring'''

								snake_case_ :			Optional[Any]            =      0
								snake_case_ :			Union[str, Any]            =      0





		def  UpperCAmelCase__				(       self )			->    Any:







								'''simple docstring'''

								snake_case_ :			Dict            =      self.inference_layers_num / self.inference_instances_num
								snake_case_ :			int            =      (
								    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(__lowerCamelCase )





		@add_start_docstrings_to_model_forward(__lowerCamelCase )
		def  UpperCAmelCase__				(       self				,     _lowercase=None				,     _lowercase=None				,     _lowercase=None				,     _lowercase=None				,     _lowercase=None				,     _lowercase=None				,     _lowercase=None				,     _lowercase=None				,     _lowercase=None				,     _lowercase=None				,     _lowercase=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:
														snake_case_ :			Optional[Any]            =      input_ids.size()
								elif inputs_embeds is not None:
														snake_case_ :			Union[str, Any]            =      inputs_embeds.size()[:-1]
								else:
														raise ValueError("""You have to specify either input_ids or inputs_embeds""" )

								snake_case_ :			str            =      input_ids.device if input_ids is not None else inputs_embeds.device

								if attention_mask is None:
														snake_case_ :			str            =      torch.ones(__lowerCamelCase				,     device=__lowerCamelCase )
								if token_type_ids is None:
														snake_case_ :			List[str]            =      torch.zeros(__lowerCamelCase				,     dtype=torch.long				,     device=__lowerCamelCase )

								# 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.
								snake_case_ :			Any            =      self.get_extended_attention_mask(__lowerCamelCase				,     __lowerCamelCase				,     __lowerCamelCase )

								# 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:
														snake_case_      ,		snake_case_      ,		snake_case_ :			int            =      encoder_hidden_states.size()
														snake_case_ :			int            =      (encoder_batch_size, encoder_sequence_length)
														if encoder_attention_mask is None:
																				snake_case_ :			Optional[int]            =      torch.ones(__lowerCamelCase				,     device=__lowerCamelCase )
														snake_case_ :			int            =      self.invert_attention_mask(__lowerCamelCase )
								else:
														snake_case_ :			Union[str, Any]            =      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]
								snake_case_ :			Dict            =      self.get_head_mask(__lowerCamelCase				,     self.config.num_hidden_layers )

								snake_case_ :			Union[str, Any]            =      self.embeddings(
								    input_ids=__lowerCamelCase				,     position_ids=__lowerCamelCase				,     token_type_ids=__lowerCamelCase				,     inputs_embeds=__lowerCamelCase )
								snake_case_ :			Optional[Any]            =      embedding_output

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

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

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

																				snake_case_ :			Dict            =      logits
																				if patient_counter == self.patience:
																										break
														snake_case_ :			List[str]            =      [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
    the pooled output) e.g. for GLUE tasks. '''	,					UpperCamelCase__	,					)
class  _lowerCAmelCase     (					UpperCamelCase__		):





		"""simple docstring"""






		def __init__(       self				,     _lowercase )			->    Dict:







								'''simple docstring'''

								super().__init__(__lowerCamelCase )
								snake_case_ :			List[Any]            =      config.num_labels

								snake_case_ :			Union[str, Any]            =      BertModelWithPabee(__lowerCamelCase )
								snake_case_ :			Optional[int]            =      nn.Dropout(config.hidden_dropout_prob )
								snake_case_ :			Union[str, Any]            =      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(__lowerCamelCase )
		def  UpperCAmelCase__				(       self				,     _lowercase=None				,     _lowercase=None				,     _lowercase=None				,     _lowercase=None				,     _lowercase=None				,     _lowercase=None				,     _lowercase=None				,     )			->    Optional[Any]:







								'''simple docstring'''

								snake_case_ :			Optional[Any]            =      self.bert(
								    input_ids=__lowerCamelCase				,     attention_mask=__lowerCamelCase				,     token_type_ids=__lowerCamelCase				,     position_ids=__lowerCamelCase				,     head_mask=__lowerCamelCase				,     inputs_embeds=__lowerCamelCase				,     output_dropout=self.dropout				,     output_layers=self.classifiers				,     regression=self.num_labels == 1				,     )

								snake_case_ :			Optional[Any]            =      (logits[-1],)

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

								return outputs








from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


a__: List[Any]							   =						{
    'configuration_x_clip': [
        'XCLIP_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'XCLIPConfig',
        'XCLIPTextConfig',
        'XCLIPVisionConfig',
    ],
    'processing_x_clip': ['XCLIPProcessor'],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       a__: Any							   =						[
           'XCLIP_PRETRAINED_MODEL_ARCHIVE_LIST',
           'XCLIPModel',
           'XCLIPPreTrainedModel',
           'XCLIPTextModel',
           'XCLIPVisionModel',
       ]

if TYPE_CHECKING:
       from .configuration_x_clip import (
           XCLIP_PRETRAINED_CONFIG_ARCHIVE_MAP,
           XCLIPConfig,
           XCLIPTextConfig,
           XCLIPVisionConfig,
       )
       from .processing_x_clip import XCLIPProcessor

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_x_clip import (
                  XCLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
                  XCLIPModel,
                  XCLIPPreTrainedModel,
                  XCLIPTextModel,
                  XCLIPVisionModel,
              )

else:
       import sys

       a__: Optional[Any]							   =						_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)





