



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


UpperCAmelCase     		=   logging.getLogger(__name__)

class __snake_case(				lowerCamelCase__ ):


						'''simple docstring'''



						def     __snake_case   ( self     ,	A_     ,	A_     ,	A_=None     ,	A_=None				)     ->     Optional[int]:
							lowerCAmelCase   =     self.layer[current_layer](lowercase__     ,	lowercase__     ,	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."		,							lowerCamelCase__		,							)
class __snake_case(				lowerCamelCase__ ):


						'''simple docstring'''



						def __init__( self     ,	A_				)     ->     Any:
							super().__init__(lowercase__				)

							lowerCAmelCase   =     BertEncoderWithPabee(lowercase__				)

							self.init_weights()
							lowerCAmelCase   =     0
							lowerCAmelCase   =     0
							lowerCAmelCase   =     0

							lowerCAmelCase   =     0


						def     __snake_case   ( self     ,	A_				)     ->     int:
							lowerCAmelCase   =     threshold


						def     __snake_case   ( self     ,	A_				)     ->     List[str]:
							lowerCAmelCase   =     patience


						def     __snake_case   ( self				)     ->     Optional[Any]:
							lowerCAmelCase   =     0
							lowerCAmelCase   =     0


						def     __snake_case   ( self				)     ->     Union[str, Any]:
							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(lowercase__				)




						@add_start_docstrings_to_model_forward(lowercase__				)
						def     __snake_case   ( self     ,	A_=None     ,	A_=None     ,	A_=None     ,	A_=None     ,	A_=None     ,	A_=None     ,	A_=None     ,	A_=None     ,	A_=None     ,	A_=None     ,	A_=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:
								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(lowercase__     ,	device=lowercase__				)
							if token_type_ids is None:
								lowerCAmelCase   =     torch.zeros(lowercase__     ,	dtype=torch.long     ,	device=lowercase__				)

							# 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(lowercase__     ,	lowercase__     ,	lowercase__				)

							# 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(lowercase__     ,	device=lowercase__				)
								lowerCAmelCase   =     self.invert_attention_mask(lowercase__				)
							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(lowercase__     ,	self.config.num_hidden_layers				)

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

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

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

									lowerCAmelCase   =     self.pooler(lowercase__				)
									lowerCAmelCase   =     output_layers[i](lowercase__				)
									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(lowercase__				)				):
											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. "		,							lowerCamelCase__		,							)
class __snake_case(				lowerCamelCase__ ):


						'''simple docstring'''



						def __init__( self     ,	A_				)     ->     Optional[int]:
							super().__init__(lowercase__				)
							lowerCAmelCase   =     config.num_labels

							lowerCAmelCase   =     BertModelWithPabee(lowercase__				)
							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(lowercase__				)
						def     __snake_case   ( self     ,	A_=None     ,	A_=None     ,	A_=None     ,	A_=None     ,	A_=None     ,	A_=None     ,	A_=None     ,	)     ->     Optional[Any]:
							lowerCAmelCase   =     self.bert(
							    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     ,	)

							lowerCAmelCase   =     (logits[-1],)

							if labels is not None:
								lowerCAmelCase   =     None
								lowerCAmelCase   =     0
								for ix, logits_item in enumerate(lowercase__				):
									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




'''simple docstring'''

from torch import nn






def 		_snake_case		( _SCREAMING_SNAKE_CASE					:						Optional[int]							)      ->						Union[str, Any]:







	"""simple docstring"""
	if act_fn in ["swish", "silu"]:
		return nn.SiLU()
	elif act_fn == "mish":
		return nn.Mish()
	elif act_fn == "gelu":
		return nn.GELU()
	else:
		raise ValueError(f'Unsupported activation function: {act_fn}'							)