



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


A_	:						Tuple				=			logging.getLogger(__name__)
class 						lowerCAmelCase__	(  _lowerCamelCase   ):


	'''simple docstring'''







	def  _lowerCAmelCase    (			self							:					List[str]    ,					_SCREAMING_SNAKE_CASE							:					int    ,					_SCREAMING_SNAKE_CASE							:					Tuple    ,					_SCREAMING_SNAKE_CASE							:					Optional[int]=None    ,					_SCREAMING_SNAKE_CASE							:					List[Any]=None						)	->  List[Any]:





					"""simple docstring"""


					if current_layer == 0:
									SCREAMING_SNAKE_CASE	:       str 							=    self.embedding_hidden_mapping_in(_SCREAMING_SNAKE_CASE						)
					else:
									SCREAMING_SNAKE_CASE	:       int 							=    hidden_states[0]

					SCREAMING_SNAKE_CASE	:       Optional[int] 							=    int(self.config.num_hidden_layers / self.config.num_hidden_groups						)

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

					SCREAMING_SNAKE_CASE	:       Optional[int] 							=    self.albert_layer_groups[group_idx](
					    _SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]    ,					)
					SCREAMING_SNAKE_CASE	:       Optional[Any] 							=    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.'''       ,       _lowerCamelCase       ,       )
class 						lowerCAmelCase__	(  _lowerCamelCase   ):


	'''simple docstring'''







	def __init__(			self							:					Optional[Any]    ,					_SCREAMING_SNAKE_CASE							:					Dict						)	->  List[str]:





					"""simple docstring"""


					super().__init__(_SCREAMING_SNAKE_CASE						)

					SCREAMING_SNAKE_CASE	:       Tuple 							=    AlbertTransformerWithPabee(_SCREAMING_SNAKE_CASE						)

					self.init_weights()
					SCREAMING_SNAKE_CASE	:       Dict 							=    0
					SCREAMING_SNAKE_CASE	:       Any 							=    0
					SCREAMING_SNAKE_CASE	:       Any 							=    0

					SCREAMING_SNAKE_CASE	:       List[str] 							=    0





	def  _lowerCAmelCase    (			self							:					Any    ,					_SCREAMING_SNAKE_CASE							:					Dict						)	->  Dict:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE	:       str 							=    threshold





	def  _lowerCAmelCase    (			self							:					int    ,					_SCREAMING_SNAKE_CASE							:					str						)	->  Tuple:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE	:       Dict 							=    patience





	def  _lowerCAmelCase    (			self							:					Tuple						)	->  Tuple:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE	:       List[Any] 							=    0
					SCREAMING_SNAKE_CASE	:       Optional[Any] 							=    0





	def  _lowerCAmelCase    (			self							:					Optional[Any]						)	->  str:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE	:       int 							=    self.inference_layers_num / self.inference_instances_num
					SCREAMING_SNAKE_CASE	:       Union[str, Any] 							=    (
					    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(_SCREAMING_SNAKE_CASE						)







	@add_start_docstrings_to_model_forward(_SCREAMING_SNAKE_CASE						)
	def  _lowerCAmelCase    (			self							:					int    ,					_SCREAMING_SNAKE_CASE							:					int=None    ,					_SCREAMING_SNAKE_CASE							:					Tuple=None    ,					_SCREAMING_SNAKE_CASE							:					Any=None    ,					_SCREAMING_SNAKE_CASE							:					int=None    ,					_SCREAMING_SNAKE_CASE							:					List[Any]=None    ,					_SCREAMING_SNAKE_CASE							:					List[Any]=None    ,					_SCREAMING_SNAKE_CASE							:					List[str]=None    ,					_SCREAMING_SNAKE_CASE							:					List[Any]=None    ,					_SCREAMING_SNAKE_CASE							:					str=False    ,					)	->  List[str]:





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

					SCREAMING_SNAKE_CASE	:       Optional[int] 							=    input_ids.device if input_ids is not None else inputs_embeds.device

					if attention_mask is None:
									SCREAMING_SNAKE_CASE	:       Optional[Any] 							=    torch.ones(_SCREAMING_SNAKE_CASE    ,					device=_SCREAMING_SNAKE_CASE						)
					if token_type_ids is None:
									SCREAMING_SNAKE_CASE	:       List[str] 							=    torch.zeros(_SCREAMING_SNAKE_CASE    ,					dtype=torch.long    ,					device=_SCREAMING_SNAKE_CASE						)

					SCREAMING_SNAKE_CASE	:       List[Any] 							=    attention_mask.unsqueeze(1						).unsqueeze(2						)
					SCREAMING_SNAKE_CASE	:       Optional[int] 							=    extended_attention_mask.to(dtype=self.dtype						)  # fp16 compatibility
					SCREAMING_SNAKE_CASE	:       List[str] 							=    (1.0 - extended_attention_mask) * -1_0_0_0_0.0
					SCREAMING_SNAKE_CASE	:       Dict 							=    self.get_head_mask(_SCREAMING_SNAKE_CASE    ,					self.config.num_hidden_layers						)

					SCREAMING_SNAKE_CASE	:       List[Any] 							=    self.embeddings(
					    _SCREAMING_SNAKE_CASE    ,					position_ids=_SCREAMING_SNAKE_CASE    ,					token_type_ids=_SCREAMING_SNAKE_CASE    ,					inputs_embeds=_SCREAMING_SNAKE_CASE						)
					SCREAMING_SNAKE_CASE	:       Optional[int] 							=    embedding_output

					if self.training:
									SCREAMING_SNAKE_CASE	:       List[str] 							=    []
									for i in range(self.config.num_hidden_layers						):
													SCREAMING_SNAKE_CASE	:       Any 							=    self.encoder.adaptive_forward(
													    _SCREAMING_SNAKE_CASE    ,					current_layer=_SCREAMING_SNAKE_CASE    ,					attention_mask=_SCREAMING_SNAKE_CASE    ,					head_mask=_SCREAMING_SNAKE_CASE    ,					)

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

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

													SCREAMING_SNAKE_CASE	:       Any 							=    logits
													if patient_counter == self.patience:
																	break
									SCREAMING_SNAKE_CASE	:       Any 							=    [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
    the pooled output) e.g. for GLUE tasks. '''       ,       _lowerCamelCase       ,       )
class 						lowerCAmelCase__	(  _lowerCamelCase   ):


	'''simple docstring'''







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





					"""simple docstring"""


					super().__init__(_SCREAMING_SNAKE_CASE						)
					SCREAMING_SNAKE_CASE	:       Optional[Any] 							=    config.num_labels

					SCREAMING_SNAKE_CASE	:       Optional[int] 							=    AlbertModelWithPabee(_SCREAMING_SNAKE_CASE						)
					SCREAMING_SNAKE_CASE	:       List[str] 							=    nn.Dropout(config.classifier_dropout_prob						)
					SCREAMING_SNAKE_CASE	:       List[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(_SCREAMING_SNAKE_CASE						)
	def  _lowerCAmelCase    (			self							:					Optional[Any]    ,					_SCREAMING_SNAKE_CASE							:					Union[str, Any]=None    ,					_SCREAMING_SNAKE_CASE							:					Tuple=None    ,					_SCREAMING_SNAKE_CASE							:					Optional[int]=None    ,					_SCREAMING_SNAKE_CASE							:					Optional[Any]=None    ,					_SCREAMING_SNAKE_CASE							:					Any=None    ,					_SCREAMING_SNAKE_CASE							:					Optional[int]=None    ,					_SCREAMING_SNAKE_CASE							:					List[Any]=None    ,					)	->  List[Any]:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE	:       List[str] 							=    self.albert(
					    input_ids=_SCREAMING_SNAKE_CASE    ,					attention_mask=_SCREAMING_SNAKE_CASE    ,					token_type_ids=_SCREAMING_SNAKE_CASE    ,					position_ids=_SCREAMING_SNAKE_CASE    ,					head_mask=_SCREAMING_SNAKE_CASE    ,					inputs_embeds=_SCREAMING_SNAKE_CASE    ,					output_dropout=self.dropout    ,					output_layers=self.classifiers    ,					regression=self.num_labels == 1    ,					)

					SCREAMING_SNAKE_CASE	:       Union[str, Any] 							=    (logits[-1],)

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

					return outputs










"""simple docstring"""

def    __snake_case       (  __A			:	int    )    ->					int:






				'''simple docstring'''
				if not isinstance(__A       ,		__A    ):
								raise ValueError('double_factorial() only accepts integral values'    )
				if num < 0:
								raise ValueError('double_factorial() not defined for negative values'    )
				SCREAMING_SNAKE_CASE	:       List[str] 							=    1
				for i in range(__A       ,		0       ,		-2    ):
								value *= i
				return value


if __name__ == "__main__":
							import doctest

							doctest.testmod()





