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


UpperCamelCase     					=     logging.getLogger(__name__)
class   UpperCAmelCase__				(    lowercase__			):


	"""simple docstring"""







	def       A    (			self       ,     _SCREAMING_SNAKE_CASE       ,     _SCREAMING_SNAKE_CASE       ,     _SCREAMING_SNAKE_CASE=None       ,     _SCREAMING_SNAKE_CASE=None )    ->	List[Any]:
							if current_layer == 0:
													a_							:			List[str]     		=					self.embedding_hidden_mapping_in(_A )
							else:
													a_							:			Dict     		=					hidden_states[0]

							a_							:			str     		=					int(self.config.num_hidden_layers / self.config.num_hidden_groups )

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

							a_							:			Optional[int]     		=					self.albert_layer_groups[group_idx](
							    _A       ,     _A       ,     head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]       ,     )
							a_							:			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.""",      lowercase__,      )
class   UpperCAmelCase__				(    lowercase__			):


	"""simple docstring"""







	def __init__(			self       ,     _SCREAMING_SNAKE_CASE )    ->	List[str]:
							super().__init__(_A )

							a_							:			Optional[int]     		=					AlbertTransformerWithPabee(_A )

							self.init_weights()
							a_							:			Tuple     		=					0
							a_							:			Optional[Any]     		=					0
							a_							:			int     		=					0

							a_							:			Optional[int]     		=					0







	def       A    (			self       ,     _SCREAMING_SNAKE_CASE )    ->	Dict:
							a_							:			Tuple     		=					threshold







	def       A    (			self       ,     _SCREAMING_SNAKE_CASE )    ->	List[Any]:
							a_							:			Dict     		=					patience







	def       A    (			self )    ->	Dict:
							a_							:			Optional[Any]     		=					0
							a_							:			List[str]     		=					0







	def       A    (			self )    ->	Dict:
							a_							:			Tuple     		=					self.inference_layers_num / self.inference_instances_num
							a_							:			Dict     		=					(
							    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       A    (			self       ,     _SCREAMING_SNAKE_CASE=None       ,     _SCREAMING_SNAKE_CASE=None       ,     _SCREAMING_SNAKE_CASE=None       ,     _SCREAMING_SNAKE_CASE=None       ,     _SCREAMING_SNAKE_CASE=None       ,     _SCREAMING_SNAKE_CASE=None       ,     _SCREAMING_SNAKE_CASE=None       ,     _SCREAMING_SNAKE_CASE=None       ,     _SCREAMING_SNAKE_CASE=False       ,     )    ->	List[Any]:

							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_							:			List[Any]     		=					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_							:			Dict     		=					input_ids.device if input_ids is not None else inputs_embeds.device

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

							a_							:			int     		=					attention_mask.unsqueeze(1 ).unsqueeze(2 )
							a_							:			Optional[Any]     		=					extended_attention_mask.to(dtype=self.dtype )  # fp16 compatibility
							a_							:			Tuple     		=					(1.0 - extended_attention_mask) * -1_0_0_0_0.0
							a_							:			Union[str, Any]     		=					self.get_head_mask(_A       ,     self.config.num_hidden_layers )

							a_							:			List[str]     		=					self.embeddings(
							    _A       ,     position_ids=_A       ,     token_type_ids=_A       ,     inputs_embeds=_A )
							a_							:			List[Any]     		=					embedding_output

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

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

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

																			a_							:			List[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
    the pooled output) e.g. for GLUE tasks. """,      lowercase__,      )
class   UpperCAmelCase__				(    lowercase__			):


	"""simple docstring"""







	def __init__(			self       ,     _SCREAMING_SNAKE_CASE )    ->	List[Any]:
							super().__init__(_A )
							a_							:			Union[str, Any]     		=					config.num_labels

							a_							:			Optional[Any]     		=					AlbertModelWithPabee(_A )
							a_							:			List[str]     		=					nn.Dropout(config.classifier_dropout_prob )
							a_							:			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(_A )
	def       A    (			self       ,     _SCREAMING_SNAKE_CASE=None       ,     _SCREAMING_SNAKE_CASE=None       ,     _SCREAMING_SNAKE_CASE=None       ,     _SCREAMING_SNAKE_CASE=None       ,     _SCREAMING_SNAKE_CASE=None       ,     _SCREAMING_SNAKE_CASE=None       ,     _SCREAMING_SNAKE_CASE=None       ,     )    ->	Tuple:
							a_							:			str     		=					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_							:			List[Any]     		=					(logits[-1],)

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

							return outputs





"""simple docstring"""






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


UpperCamelCase     					=     {'configuration_levit': ['LEVIT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'LevitConfig', 'LevitOnnxConfig']}

try:
				if not is_vision_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				UpperCamelCase     					=     ['LevitFeatureExtractor']
				UpperCamelCase     					=     ['LevitImageProcessor']

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				UpperCamelCase     					=     [
				    'LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST',
				    'LevitForImageClassification',
				    'LevitForImageClassificationWithTeacher',
				    'LevitModel',
				    'LevitPreTrainedModel',
				]


if TYPE_CHECKING:
				from .configuration_levit import LEVIT_PRETRAINED_CONFIG_ARCHIVE_MAP, LevitConfig, LevitOnnxConfig

				try:
								if not is_vision_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .feature_extraction_levit import LevitFeatureExtractor
								from .image_processing_levit import LevitImageProcessor

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_levit import (
								    LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
								    LevitForImageClassification,
								    LevitForImageClassificationWithTeacher,
								    LevitModel,
								    LevitPreTrainedModel,
								)
else:
				import sys

				UpperCamelCase     					=     _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




