


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


lowerCAmelCase   :List[str]								=    logging.getLogger(__name__)
class _lowerCamelCase (		lowercase__							):



		'''simple docstring'''







		def      __lowerCAmelCase	(      self	:		Optional[int]						,  _A	:		List[Any]						,  _A	:		Union[str, Any]						,  _A	:		Dict=None						,  _A	:		List[str]=None    )						->				List[Any]:
			if current_layer == 0:
				__magic_name__ :					List[str]  			=  self.embedding_hidden_mapping_in(_A    )
			else:
				__magic_name__ :					Dict  			=  hidden_states[0]

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

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

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



		'''simple docstring'''







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

			__magic_name__ :					Optional[int]  			=  AlbertTransformerWithPabee(_A    )

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

			__magic_name__ :					Optional[int]  			=  0







		def      __lowerCAmelCase	(      self	:		Tuple						,  _A	:		Dict    )						->				Dict:
			__magic_name__ :					Tuple  			=  threshold







		def      __lowerCAmelCase	(      self	:		Any						,  _A	:		Tuple    )						->				List[Any]:
			__magic_name__ :					Dict  			=  patience







		def      __lowerCAmelCase	(      self	:		int    )						->				Dict:
			__magic_name__ :					Optional[Any]  			=  0
			__magic_name__ :					List[str]  			=  0







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

			__magic_name__ :					Dict  			=  input_ids.device if input_ids is not None else inputs_embeds.device

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

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

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

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

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

					__magic_name__ :					List[str]  			=  self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]    )    )
					__magic_name__ :					Tuple  			=  output_layers[i](_A    )
					if regression:
						__magic_name__ :					Any  			=  logits.detach()
						if patient_result is not None:
							__magic_name__ :					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:
							__magic_name__ :					Optional[Any]  			=  0
					else:
						__magic_name__ :					Tuple  			=  logits.detach().argmax(dim=1    )
						if patient_result is not None:
							__magic_name__ :					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:
							__magic_name__ :					Union[str, Any]  			=  0

					__magic_name__ :					List[Any]  			=  logits
					if patient_counter == self.patience:
						break
				__magic_name__ :					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 _lowerCamelCase (		lowercase__							):



		'''simple docstring'''







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

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

			__magic_name__ :					List[Any]  			=  (logits[-1],)

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

			return outputs



'''simple docstring'''

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCAmelCase   :List[Any]								=    logging.get_logger(__name__)

lowerCAmelCase   :Dict								=    {
    '''uw-madison/yoso-4096''': '''https://huggingface.co/uw-madison/yoso-4096/resolve/main/config.json''',
    # See all YOSO models at https://huggingface.co/models?filter=yoso
}
class _lowerCamelCase (		lowercase__							):



		'''simple docstring'''







		A_			:    Optional[Any]   					=       """yoso"""







		def __init__(      self	:		Any						,  _A	:		List[Any]=50265						,  _A	:		Any=768						,  _A	:		List[str]=12						,  _A	:		Dict=12						,  _A	:		Union[str, Any]=3072						,  _A	:		List[Any]="gelu"						,  _A	:		Dict=0.1						,  _A	:		Union[str, Any]=0.1						,  _A	:		Union[str, Any]=4096						,  _A	:		Dict=1						,  _A	:		Tuple=0.02						,  _A	:		Any=1E-12						,  _A	:		Optional[int]="absolute"						,  _A	:		Tuple=True						,  _A	:		Any=9						,  _A	:		Tuple=64						,  _A	:		str=None						,  _A	:		Dict=True						,  _A	:		str=True						,  _A	:		int=1						,  _A	:		Optional[int]=0						,  _A	:		Tuple=2						,  **_A	:		Optional[Any]						,  )						->				Any:
			super().__init__(pad_token_id=_A						,  bos_token_id=_A						,  eos_token_id=_A						,  **_A    )

			__magic_name__ :					List[Any]  			=  vocab_size
			__magic_name__ :					int  			=  max_position_embeddings
			__magic_name__ :					int  			=  hidden_size
			__magic_name__ :					Any  			=  num_hidden_layers
			__magic_name__ :					List[str]  			=  num_attention_heads
			__magic_name__ :					str  			=  intermediate_size
			__magic_name__ :					Any  			=  hidden_act
			__magic_name__ :					Union[str, Any]  			=  hidden_dropout_prob
			__magic_name__ :					str  			=  attention_probs_dropout_prob
			__magic_name__ :					int  			=  initializer_range
			__magic_name__ :					List[Any]  			=  type_vocab_size
			__magic_name__ :					Tuple  			=  layer_norm_eps
			__magic_name__ :					Tuple  			=  position_embedding_type
			__magic_name__ :					Optional[Any]  			=  use_expectation
			__magic_name__ :					Union[str, Any]  			=  hash_code_len
			__magic_name__ :					int  			=  num_hash
			__magic_name__ :					List[Any]  			=  conv_window
			__magic_name__ :					Optional[Any]  			=  use_fast_hash
			__magic_name__ :					int  			=  lsh_backward