
'''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 :    List[str]       =      logging.getLogger(__name__)





class      A							(    __UpperCAmelCase			):



			def A__	(   self    ,       lowerCamelCase__    ,       lowerCamelCase__    ,       lowerCamelCase__=None    ,       lowerCamelCase__=None    )     ->						Optional[Any]:




									'''simple docstring'''



									if current_layer == 0:
															lowercase__								=    self.embedding_hidden_mapping_in(lowerCamelCase__    )
									else:
															lowercase__								=    hidden_states[0]

									lowercase__								=    int(self.config.num_hidden_layers / self.config.num_hidden_groups    )

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

									lowercase__								=    self.albert_layer_groups[group_idx](
									    lowerCamelCase__    ,       lowerCamelCase__    ,       head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]    ,       )
									lowercase__								=    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."""						,  __UpperCAmelCase						,  )
class      A							(    __UpperCAmelCase			):



			def __init__(   self    ,       lowerCamelCase__    )     ->						str:




									'''simple docstring'''



									super().__init__(lowerCamelCase__    )

									lowercase__								=    AlbertTransformerWithPabee(lowerCamelCase__    )

									self.init_weights()
									lowercase__								=    0
									lowercase__								=    0
									lowercase__								=    0

									lowercase__								=    0



			def A__	(   self    ,       lowerCamelCase__    )     ->						Any:




									'''simple docstring'''



									lowercase__								=    threshold



			def A__	(   self    ,       lowerCamelCase__    )     ->						str:




									'''simple docstring'''



									lowercase__								=    patience



			def A__	(   self    )     ->						Tuple:




									'''simple docstring'''



									lowercase__								=    0
									lowercase__								=    0



			def A__	(   self    )     ->						Union[str, Any]:




									'''simple docstring'''



									lowercase__								=    self.inference_layers_num / self.inference_instances_num
									lowercase__								=    (
									    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(lowerCamelCase__    )



			@add_start_docstrings_to_model_forward(lowerCamelCase__    )
			def A__	(   self    ,       lowerCamelCase__=None    ,       lowerCamelCase__=None    ,       lowerCamelCase__=None    ,       lowerCamelCase__=None    ,       lowerCamelCase__=None    ,       lowerCamelCase__=None    ,       lowerCamelCase__=None    ,       lowerCamelCase__=None    ,       lowerCamelCase__=False    ,       )     ->						Optional[int]:




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

									lowercase__								=    input_ids.device if input_ids is not None else inputs_embeds.device

									if attention_mask is None:
															lowercase__								=    torch.ones(lowerCamelCase__    ,       device=lowerCamelCase__    )
									if token_type_ids is None:
															lowercase__								=    torch.zeros(lowerCamelCase__    ,       dtype=torch.long    ,       device=lowerCamelCase__    )

									lowercase__								=    attention_mask.unsqueeze(1    ).unsqueeze(2    )
									lowercase__								=    extended_attention_mask.to(dtype=self.dtype    )  # fp16 compatibility
									lowercase__								=    (1.0 - extended_attention_mask) * -10_000.0
									lowercase__								=    self.get_head_mask(lowerCamelCase__    ,       self.config.num_hidden_layers    )

									lowercase__								=    self.embeddings(
									    lowerCamelCase__    ,       position_ids=lowerCamelCase__    ,       token_type_ids=lowerCamelCase__    ,       inputs_embeds=lowerCamelCase__    )
									lowercase__								=    embedding_output

									if self.training:
															lowercase__								=    []
															for i in range(self.config.num_hidden_layers    ):
																					lowercase__								=    self.encoder.adaptive_forward(
																					    lowerCamelCase__    ,       current_layer=lowerCamelCase__    ,       attention_mask=lowerCamelCase__    ,       head_mask=lowerCamelCase__    ,       )

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

																					lowercase__								=    self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]    )    )
																					lowercase__								=    output_layers[i](lowerCamelCase__    )
																					if regression:
																											lowercase__								=    logits.detach()
																											if patient_result is not None:
																																	lowercase__								=    patient_result.detach()
																											if (patient_result is not None) and torch.abs(patient_result - labels    ) < self.regression_threshold:
																																	patient_counter += 1
																											else:
																																	lowercase__								=    0
																					else:
																											lowercase__								=    logits.detach().argmax(dim=1    )
																											if patient_result is not None:
																																	lowercase__								=    patient_result.detach().argmax(dim=1    )
																											if (patient_result is not None) and torch.all(labels.eq(lowerCamelCase__    )    ):
																																	patient_counter += 1
																											else:
																																	lowercase__								=    0

																					lowercase__								=    logits
																					if patient_counter == self.patience:
																											break
															lowercase__								=    [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. """						,  __UpperCAmelCase						,  )
class      A							(    __UpperCAmelCase			):



			def __init__(   self    ,       lowerCamelCase__    )     ->						Tuple:




									'''simple docstring'''



									super().__init__(lowerCamelCase__    )
									lowercase__								=    config.num_labels

									lowercase__								=    AlbertModelWithPabee(lowerCamelCase__    )
									lowercase__								=    nn.Dropout(config.classifier_dropout_prob    )
									lowercase__								=    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(lowerCamelCase__    )
			def A__	(   self    ,       lowerCamelCase__=None    ,       lowerCamelCase__=None    ,       lowerCamelCase__=None    ,       lowerCamelCase__=None    ,       lowerCamelCase__=None    ,       lowerCamelCase__=None    ,       lowerCamelCase__=None    ,       )     ->						Dict:




									'''simple docstring'''



									lowercase__								=    self.albert(
									    input_ids=lowerCamelCase__    ,       attention_mask=lowerCamelCase__    ,       token_type_ids=lowerCamelCase__    ,       position_ids=lowerCamelCase__    ,       head_mask=lowerCamelCase__    ,       inputs_embeds=lowerCamelCase__    ,       output_dropout=self.dropout    ,       output_layers=self.classifiers    ,       regression=self.num_labels == 1    ,       )

									lowercase__								=    (logits[-1],)

									if labels is not None:
															lowercase__								=    None
															lowercase__								=    0
															for ix, logits_item in enumerate(lowerCamelCase__    ):
																					if self.num_labels == 1:
																											#  We are doing regression
																											lowercase__								=    MSELoss()
																											lowercase__								=    loss_fct(logits_item.view(-1    )    ,       labels.view(-1    )    )
																					else:
																											lowercase__								=    CrossEntropyLoss()
																											lowercase__								=    loss_fct(logits_item.view(-1    ,       self.num_labels    )    ,       labels.view(-1    )    )
																					if total_loss is None:
																											lowercase__								=    loss
																					else:
																											total_loss += loss * (ix + 1)
																					total_weights += ix + 1
															lowercase__								=    (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


__A       =      logging.get_logger(__name__)

__A       =      {
    "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							(    __UpperCAmelCase			):
			lowerCamelCase							:							List[str]    						=					"""vit-hybrid"""



			def __init__(   self    ,       lowerCamelCase__=None    ,       lowerCamelCase__=768    ,       lowerCamelCase__=12    ,       lowerCamelCase__=12    ,       lowerCamelCase__=3_072    ,       lowerCamelCase__="gelu"    ,       lowerCamelCase__=0.0    ,       lowerCamelCase__=0.0    ,       lowerCamelCase__=0.02    ,       lowerCamelCase__=1e-12    ,       lowerCamelCase__=224    ,       lowerCamelCase__=1    ,       lowerCamelCase__=3    ,       lowerCamelCase__=[1, 1_024, 24, 24]    ,       lowerCamelCase__=True    ,       **lowerCamelCase__    ,       )     ->						Optional[Any]:




									'''simple docstring'''



									super().__init__(**lowerCamelCase__    )

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

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

									lowercase__								=    backbone_featmap_shape
									lowercase__								=    backbone_config
									lowercase__								=    hidden_size
									lowercase__								=    num_hidden_layers
									lowercase__								=    num_attention_heads
									lowercase__								=    intermediate_size
									lowercase__								=    hidden_act
									lowercase__								=    hidden_dropout_prob
									lowercase__								=    attention_probs_dropout_prob
									lowercase__								=    initializer_range
									lowercase__								=    layer_norm_eps
									lowercase__								=    image_size
									lowercase__								=    patch_size
									lowercase__								=    num_channels
									lowercase__								=    qkv_bias



			def A__	(   self    )     ->						Dict[str, any]:




									'''simple docstring'''



									lowercase__								=    copy.deepcopy(self.__dict__    )
									lowercase__								=    self.backbone_config.to_dict()
									lowercase__								=    self.__class__.model_type
									return output






