
'''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     a (					__magic_name__     ):




				def 						__snake_case				(				self						:					Tuple,     SCREAMING_SNAKE_CASE_						:					Union[str, Any],     SCREAMING_SNAKE_CASE_						:					Optional[int],     SCREAMING_SNAKE_CASE_						:					Optional[Any]=None,     SCREAMING_SNAKE_CASE_						:					Tuple=None   ):
										if current_layer == 0:
																snake_case							:  str			=     self.embedding_hidden_mapping_in(SCREAMING_SNAKE_CASE_   )
										else:
																snake_case							:  Any			=     hidden_states[0]

										snake_case							:  Tuple			=     int(self.config.num_hidden_layers / self.config.num_hidden_groups   )

										# Index of the hidden group
										snake_case							:  List[str]			=     int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)   )

										snake_case							:  Any			=     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],     )
										snake_case							:  Dict			=     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.'''					,__magic_name__					,)
class     a (					__magic_name__     ):




				def __init__(				self						:					List[Any],     SCREAMING_SNAKE_CASE_						:					Optional[int]   ):
										super().__init__(SCREAMING_SNAKE_CASE_   )

										snake_case							:  Union[str, Any]			=     AlbertTransformerWithPabee(SCREAMING_SNAKE_CASE_   )

										self.init_weights()
										snake_case							:  str			=     0
										snake_case							:  Optional[int]			=     0
										snake_case							:  List[Any]			=     0

										snake_case							:  str			=     0




				def 						__snake_case				(				self						:					List[Any],     SCREAMING_SNAKE_CASE_						:					Optional[int]   ):
										snake_case							:  Any			=     threshold




				def 						__snake_case				(				self						:					List[str],     SCREAMING_SNAKE_CASE_						:					Dict   ):
										snake_case							:  Optional[int]			=     patience




				def 						__snake_case				(				self						:					str   ):
										snake_case							:  Union[str, Any]			=     0
										snake_case							:  Tuple			=     0




				def 						__snake_case				(				self						:					Optional[int]   ):
										snake_case							:  Tuple			=     self.inference_layers_num / self.inference_instances_num
										snake_case							:  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(SCREAMING_SNAKE_CASE_   )




				@add_start_docstrings_to_model_forward(SCREAMING_SNAKE_CASE_   )
				def 						__snake_case				(				self						:					List[str],     SCREAMING_SNAKE_CASE_						:					List[Any]=None,     SCREAMING_SNAKE_CASE_						:					Union[str, Any]=None,     SCREAMING_SNAKE_CASE_						:					Tuple=None,     SCREAMING_SNAKE_CASE_						:					Dict=None,     SCREAMING_SNAKE_CASE_						:					Optional[int]=None,     SCREAMING_SNAKE_CASE_						:					Optional[int]=None,     SCREAMING_SNAKE_CASE_						:					Tuple=None,     SCREAMING_SNAKE_CASE_						:					List[Any]=None,     SCREAMING_SNAKE_CASE_						:					List[Any]=False,     ):

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

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

										if attention_mask is None:
																snake_case							:  Optional[int]			=     torch.ones(SCREAMING_SNAKE_CASE_,     device=SCREAMING_SNAKE_CASE_   )
										if token_type_ids is None:
																snake_case							:  Dict			=     torch.zeros(SCREAMING_SNAKE_CASE_,     dtype=torch.long,     device=SCREAMING_SNAKE_CASE_   )

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

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

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

																						snake_case							:  Union[str, Any]			=     self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]   )   )
																						snake_case							:  Union[str, Any]			=     output_layers[i](output_dropout(SCREAMING_SNAKE_CASE_   )   )
																						res.append(SCREAMING_SNAKE_CASE_   )
										elif self.patience == 0:  # Use all layers for inference
																snake_case							:  int			=     self.encoder(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     head_mask=SCREAMING_SNAKE_CASE_   )
																snake_case							:  Optional[int]			=     self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]   )   )
																snake_case							:  Union[str, Any]			=     [output_layers[self.config.num_hidden_layers - 1](SCREAMING_SNAKE_CASE_   )]
										else:
																snake_case							:  Optional[int]			=     0
																snake_case							:  List[Any]			=     None
																snake_case							:  Optional[int]			=     0
																for i in range(self.config.num_hidden_layers   ):
																						calculated_layer_num += 1
																						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_,     )

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

																						snake_case							:  List[str]			=     logits
																						if patient_counter == self.patience:
																												break
																snake_case							:  Optional[int]			=     [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. '''					,__magic_name__					,)
class     a (					__magic_name__     ):




				def __init__(				self						:					Optional[Any],     SCREAMING_SNAKE_CASE_						:					Union[str, Any]   ):
										super().__init__(SCREAMING_SNAKE_CASE_   )
										snake_case							:  List[str]			=     config.num_labels

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

										snake_case							:  List[str]			=     (logits[-1],)

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

										return outputs



'''simple docstring'''





import copy
from typing import Dict

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto.configuration_auto import CONFIG_MAPPING
from ..bit import BitConfig


UpperCAmelCase            = logging.get_logger(__name__)

UpperCAmelCase            = {
    "google/vit-hybrid-base-bit-384": "https://huggingface.co/vit-hybrid-base-bit-384/resolve/main/config.json",
    # See all ViT hybrid models at https://huggingface.co/models?filter=vit
}

class     a (					__magic_name__     ):
				_snake_case		     =			'''vit-hybrid'''




				def __init__(				self						:					Any,     SCREAMING_SNAKE_CASE_						:					Union[str, Any]=None,     SCREAMING_SNAKE_CASE_						:					Tuple=7_68,     SCREAMING_SNAKE_CASE_						:					List[str]=12,     SCREAMING_SNAKE_CASE_						:					Optional[Any]=12,     SCREAMING_SNAKE_CASE_						:					Union[str, Any]=30_72,     SCREAMING_SNAKE_CASE_						:					List[Any]="gelu",     SCREAMING_SNAKE_CASE_						:					Any=0.0,     SCREAMING_SNAKE_CASE_						:					Dict=0.0,     SCREAMING_SNAKE_CASE_						:					List[Any]=0.02,     SCREAMING_SNAKE_CASE_						:					List[Any]=1e-12,     SCREAMING_SNAKE_CASE_						:					int=2_24,     SCREAMING_SNAKE_CASE_						:					int=1,     SCREAMING_SNAKE_CASE_						:					Optional[Any]=3,     SCREAMING_SNAKE_CASE_						:					Any=[1, 10_24, 24, 24],     SCREAMING_SNAKE_CASE_						:					List[str]=True,     **SCREAMING_SNAKE_CASE_						:					List[str],     ):
										super().__init__(**SCREAMING_SNAKE_CASE_   )

										if backbone_config is None:
																logger.info('''`backbone_config` is `None`. Initializing the config with a `BiT` backbone.'''   )
																snake_case							:  int			=     {
																    '''global_padding''': '''same''',
																    '''layer_type''': '''bottleneck''',
																    '''depths''': [3, 4, 9],
																    '''out_features''': ['''stage3'''],
																    '''embedding_dynamic_padding''': True,
																}

										if isinstance(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_   ):
																if "model_type" in backbone_config:
																						snake_case							:  Any			=     CONFIG_MAPPING[backbone_config['''model_type''']]
																else:
																						logger.info(
																						    '''`model_type` is not found in `backbone_config`. Use `Bit` as the backbone configuration class.'''   )
																						snake_case							:  Any			=     BitConfig
																snake_case							:  Optional[Any]			=     backbone_config_class(**SCREAMING_SNAKE_CASE_   )

										snake_case							:  str			=     backbone_featmap_shape
										snake_case							:  Union[str, Any]			=     backbone_config
										snake_case							:  List[str]			=     hidden_size
										snake_case							:  List[Any]			=     num_hidden_layers
										snake_case							:  Dict			=     num_attention_heads
										snake_case							:  List[str]			=     intermediate_size
										snake_case							:  Optional[int]			=     hidden_act
										snake_case							:  List[str]			=     hidden_dropout_prob
										snake_case							:  Tuple			=     attention_probs_dropout_prob
										snake_case							:  Optional[int]			=     initializer_range
										snake_case							:  Optional[int]			=     layer_norm_eps
										snake_case							:  List[Any]			=     image_size
										snake_case							:  Any			=     patch_size
										snake_case							:  str			=     num_channels
										snake_case							:  List[str]			=     qkv_bias




				def 						__snake_case				(				self						:					Any   ):
										snake_case							:  Tuple			=     copy.deepcopy(self.__dict__   )
										snake_case							:  List[Any]			=     self.backbone_config.to_dict()
										snake_case							:  Optional[Any]			=     self.__class__.model_type
										return output

