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


__lowerCamelCase :     int          =  logging.getLogger(__name__)






class 		a__						( A__ ):

					def     __UpperCamelCase       (   self				:	Optional[Any],_A				:	Optional[Any],_A				:	Union[str, Any],_A				:	Optional[int]=None,_A				:	int=None   ):







							"""simple docstring"""






							if current_layer == 0:
									SCREAMING_SNAKE_CASE_  :		Tuple    				=   self.embedding_hidden_mapping_in(_A   )
							else:
									SCREAMING_SNAKE_CASE_  :		List[Any]    				=   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_  :		int    				=   self.albert_layer_groups[group_idx](
							    _A,_A,head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group],)
							SCREAMING_SNAKE_CASE_  :		Optional[int]    				=   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.'   ,		A__   ,		)
class 		a__						( A__ ):

					def __init__(   self				:	Dict,_A				:	Dict   ):







							"""simple docstring"""






							super().__init__(_A   )

							SCREAMING_SNAKE_CASE_  :		int    				=   AlbertTransformerWithPabee(_A   )

							self.init_weights()
							SCREAMING_SNAKE_CASE_  :		Tuple    				=   0
							SCREAMING_SNAKE_CASE_  :		Dict    				=   0
							SCREAMING_SNAKE_CASE_  :		List[Any]    				=   0

							SCREAMING_SNAKE_CASE_  :		Optional[Any]    				=   0

					def     __UpperCamelCase       (   self				:	Any,_A				:	int   ):







							"""simple docstring"""






							SCREAMING_SNAKE_CASE_  :		List[Any]    				=   threshold

					def     __UpperCamelCase       (   self				:	Optional[int],_A				:	Optional[Any]   ):







							"""simple docstring"""






							SCREAMING_SNAKE_CASE_  :		List[Any]    				=   patience

					def     __UpperCamelCase       (   self				:	Union[str, Any]   ):







							"""simple docstring"""






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

					def     __UpperCamelCase       (   self				:	Union[str, Any]   ):







							"""simple docstring"""






							SCREAMING_SNAKE_CASE_  :		Optional[Any]    				=   self.inference_layers_num / self.inference_instances_num
							SCREAMING_SNAKE_CASE_  :		List[str]    				=   (
							    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				:	Any,_A				:	Any=None,_A				:	List[str]=None,_A				:	List[str]=None,_A				:	Union[str, Any]=None,_A				:	Tuple=None,_A				:	Tuple=None,_A				:	List[Any]=None,_A				:	int=None,_A				:	List[str]=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:
									SCREAMING_SNAKE_CASE_  :		Union[str, Any]    				=   input_ids.size()
							elif inputs_embeds is not None:
									SCREAMING_SNAKE_CASE_  :		Dict    				=   inputs_embeds.size()[:-1]
							else:
									raise ValueError("You have to specify either input_ids or inputs_embeds"   )

							SCREAMING_SNAKE_CASE_  :		Any    				=   input_ids.device if input_ids is not None else inputs_embeds.device

							if attention_mask is None:
									SCREAMING_SNAKE_CASE_  :		str    				=   torch.ones(_A,device=_A   )
							if token_type_ids is None:
									SCREAMING_SNAKE_CASE_  :		Dict    				=   torch.zeros(_A,dtype=torch.long,device=_A   )

							SCREAMING_SNAKE_CASE_  :		Dict    				=   attention_mask.unsqueeze(1   ).unsqueeze(2   )
							SCREAMING_SNAKE_CASE_  :		Any    				=   extended_attention_mask.to(dtype=self.dtype   )  # fp16 compatibility
							SCREAMING_SNAKE_CASE_  :		Optional[Any]    				=   (1.0 - extended_attention_mask) * -10000.0
							SCREAMING_SNAKE_CASE_  :		int    				=   self.get_head_mask(_A,self.config.num_hidden_layers   )

							SCREAMING_SNAKE_CASE_  :		Optional[int]    				=   self.embeddings(
							    _A,position_ids=_A,token_type_ids=_A,inputs_embeds=_A   )
							SCREAMING_SNAKE_CASE_  :		Union[str, Any]    				=   embedding_output

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

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

											SCREAMING_SNAKE_CASE_  :		int    				=   self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]   )   )
											SCREAMING_SNAKE_CASE_  :		Tuple    				=   output_layers[i](_A   )
											if regression:
													SCREAMING_SNAKE_CASE_  :		Tuple    				=   logits.detach()
													if patient_result is not None:
															SCREAMING_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:
															SCREAMING_SNAKE_CASE_  :		Optional[Any]    				=   0
											else:
													SCREAMING_SNAKE_CASE_  :		Dict    				=   logits.detach().argmax(dim=1   )
													if patient_result is not None:
															SCREAMING_SNAKE_CASE_  :		List[Any]    				=   patient_result.detach().argmax(dim=1   )
													if (patient_result is not None) and torch.all(labels.eq(_A   )   ):
															patient_counter += 1
													else:
															SCREAMING_SNAKE_CASE_  :		str    				=   0

											SCREAMING_SNAKE_CASE_  :		List[str]    				=   logits
											if patient_counter == self.patience:
													break
									SCREAMING_SNAKE_CASE_  :		Dict    				=   [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. '   ,		A__   ,		)
class 		a__						( A__ ):

					def __init__(   self				:	Any,_A				:	str   ):







							"""simple docstring"""






							super().__init__(_A   )
							SCREAMING_SNAKE_CASE_  :		int    				=   config.num_labels

							SCREAMING_SNAKE_CASE_  :		str    				=   AlbertModelWithPabee(_A   )
							SCREAMING_SNAKE_CASE_  :		str    				=   nn.Dropout(config.classifier_dropout_prob   )
							SCREAMING_SNAKE_CASE_  :		int    				=   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				:	Dict,_A				:	Any=None,_A				:	int=None,_A				:	List[str]=None,_A				:	List[Any]=None,_A				:	Tuple=None,_A				:	Optional[Any]=None,_A				:	List[Any]=None,):







							"""simple docstring"""






							SCREAMING_SNAKE_CASE_  :		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,)

							SCREAMING_SNAKE_CASE_  :		str    				=   (logits[-1],)

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

							return outputs

import warnings

from ...utils import logging
from .image_processing_vit import ViTImageProcessor


__lowerCamelCase :     Tuple          =  logging.get_logger(__name__)






class 		a__						( A__ ):

					def __init__(   self				:	List[Any],*_A				:	Union[str, Any],**_A				:	List[Any]   ):







							"""simple docstring"""






							warnings.warn(
							    "The class ViTFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please"
							    " use ViTImageProcessor instead.",_A,)
							super().__init__(*_A,**_A   )
