


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


__lowercase			:		Union[str, Any] =			logging.getLogger(__name__)





class 							__UpperCamelCase       (	lowerCAmelCase_      ):




	def    __UpperCAmelCase							(     self		,       __a		,       __a		,       __a=None		,       __a=None							):

						'''simple docstring'''





						if current_layer == 0:
											__a    : Optional[Any]      					=						self.embedding_hidden_mapping_in(__a							)
						else:
											__a    : Union[str, Any]      					=						hidden_states[0]

						__a    : List[Any]      					=						int(self.config.num_hidden_layers / self.config.num_hidden_groups							)

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

						__a    : Tuple      					=						self.albert_layer_groups[group_idx](
						    __a		,       __a		,       head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]		,       )
						__a    : List[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 							__UpperCamelCase       (	lowerCAmelCase_      ):




	def __init__(     self		,       __a							):

						'''simple docstring'''





						super().__init__(__a							)

						__a    : Any      					=						AlbertTransformerWithPabee(__a							)

						self.init_weights()
						__a    : Optional[Any]      					=						0
						__a    : Union[str, Any]      					=						0
						__a    : Union[str, Any]      					=						0

						__a    : str      					=						0




	def    __UpperCAmelCase							(     self		,       __a							):

						'''simple docstring'''





						__a    : Optional[int]      					=						threshold




	def    __UpperCAmelCase							(     self		,       __a							):

						'''simple docstring'''





						__a    : Any      					=						patience




	def    __UpperCAmelCase							(     self							):

						'''simple docstring'''





						__a    : List[str]      					=						0
						__a    : str      					=						0




	def    __UpperCAmelCase							(     self							):

						'''simple docstring'''





						__a    : Tuple      					=						self.inference_layers_num / self.inference_instances_num
						__a    : Optional[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(__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=False		,       ):

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

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

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

						__a    : Optional[Any]      					=						attention_mask.unsqueeze(1							).unsqueeze(2							)
						__a    : Tuple      					=						extended_attention_mask.to(dtype=self.dtype							)  # fp16 compatibility
						__a    : List[str]      					=						(1.0 - extended_attention_mask) * -10000.0
						__a    : Optional[Any]      					=						self.get_head_mask(__a		,       self.config.num_hidden_layers							)

						__a    : Tuple      					=						self.embeddings(
						    __a		,       position_ids=__a		,       token_type_ids=__a		,       inputs_embeds=__a							)
						__a    : str      					=						embedding_output

						if self.training:
											__a    : Optional[Any]      					=						[]
											for i in range(self.config.num_hidden_layers							):
																__a    : Dict      					=						self.encoder.adaptive_forward(
																    __a		,       current_layer=__a		,       attention_mask=__a		,       head_mask=__a		,       )

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

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

																__a    : Any      					=						logits
																if patient_counter == self.patience:
																					break
											__a    : List[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\n    the pooled output) e.g. for GLUE tasks. "	,   lowerCAmelCase_	,   )
class 							__UpperCamelCase       (	lowerCAmelCase_      ):




	def __init__(     self		,       __a							):

						'''simple docstring'''





						super().__init__(__a							)
						__a    : List[str]      					=						config.num_labels

						__a    : Union[str, Any]      					=						AlbertModelWithPabee(__a							)
						__a    : Optional[int]      					=						nn.Dropout(config.classifier_dropout_prob							)
						__a    : Tuple      					=						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		,       ):

						'''simple docstring'''





						__a    : Tuple      					=						self.albert(
						    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		,       )

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

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

						return outputs







'''simple docstring'''



def 						lowerCamelCase							(_SCREAMING_SNAKE_CASE   :     set		,   _SCREAMING_SNAKE_CASE   :     set     ):
					assert isinstance(_SCREAMING_SNAKE_CASE		,   _SCREAMING_SNAKE_CASE     ), F"""The input value of [set_a={set_a}] is not a set"""
					assert isinstance(_SCREAMING_SNAKE_CASE		,   _SCREAMING_SNAKE_CASE     ), F"""The input value of [set_b={set_b}] is not a set"""

					return {a + b for a in set_a for b in set_b}


if __name__ == "__main__":
	from doctest import testmod

	testmod()



