





'''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__	:					Tuple	     =       logging.getLogger(__name__)


class 			SCREAMING_SNAKE_CASE__	(					snake_case__						):







					"""simple docstring"""






					def 						lowerCamelCase_					(					self    :	Optional[Any]						,     UpperCAmelCase_    :	Tuple						,     UpperCAmelCase_    :	Any						,     UpperCAmelCase_    :	Optional[int]=None						,     UpperCAmelCase_    :	str=None						):







							"""simple docstring"""




							if current_layer == 0:
									__UpperCAmelCase     :				Optional[Any]					   =			self.embedding_hidden_mapping_in(UpperCAmelCase_						)
							else:
									__UpperCAmelCase     :				Optional[int]					   =			hidden_states[0]

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

							# Index of the hidden group
							__UpperCAmelCase     :				Union[str, Any]					   =			int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)						)

							__UpperCAmelCase     :				Tuple					   =			self.albert_layer_groups[group_idx](
							    UpperCAmelCase_						,     UpperCAmelCase_						,     head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]						,     )
							__UpperCAmelCase     :				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.'''				,snake_case__				,)
class 			SCREAMING_SNAKE_CASE__	(					snake_case__						):







					"""simple docstring"""






					def __init__(					self    :	Optional[int]						,     UpperCAmelCase_    :	Union[str, Any]						):







							"""simple docstring"""
							super().__init__(UpperCAmelCase_						)

							__UpperCAmelCase     :				List[Any]					   =			AlbertTransformerWithPabee(UpperCAmelCase_						)

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

							__UpperCAmelCase     :				Optional[Any]					   =			0
					def 						lowerCamelCase_					(					self    :	Union[str, Any]						,     UpperCAmelCase_    :	int						):







							"""simple docstring"""
							__UpperCAmelCase     :				List[str]					   =			threshold
					def 						lowerCamelCase_					(					self    :	Union[str, Any]						,     UpperCAmelCase_    :	List[str]						):







							"""simple docstring"""
							__UpperCAmelCase     :				Union[str, Any]					   =			patience
					def 						lowerCamelCase_					(					self    :	int						):







							"""simple docstring"""
							__UpperCAmelCase     :				Tuple					   =			0
							__UpperCAmelCase     :				Union[str, Any]					   =			0
					def 						lowerCamelCase_					(					self    :	Any						):







							"""simple docstring"""
							__UpperCAmelCase     :				Any					   =			self.inference_layers_num / self.inference_instances_num
							__UpperCAmelCase     :				Any					   =			(
							    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(UpperCAmelCase_						)






					@add_start_docstrings_to_model_forward(UpperCAmelCase_						)
					def 						lowerCamelCase_					(					self    :	Optional[int]						,     UpperCAmelCase_    :	int=None						,     UpperCAmelCase_    :	List[str]=None						,     UpperCAmelCase_    :	str=None						,     UpperCAmelCase_    :	Optional[int]=None						,     UpperCAmelCase_    :	List[Any]=None						,     UpperCAmelCase_    :	int=None						,     UpperCAmelCase_    :	int=None						,     UpperCAmelCase_    :	Any=None						,     UpperCAmelCase_    :	Any=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:
									__UpperCAmelCase     :				Tuple					   =			input_ids.size()
							elif inputs_embeds is not None:
									__UpperCAmelCase     :				List[str]					   =			inputs_embeds.size()[:-1]
							else:
									raise ValueError("You have to specify either input_ids or inputs_embeds"						)

							__UpperCAmelCase     :				Tuple					   =			input_ids.device if input_ids is not None else inputs_embeds.device

							if attention_mask is None:
									__UpperCAmelCase     :				List[str]					   =			torch.ones(UpperCAmelCase_						,     device=UpperCAmelCase_						)
							if token_type_ids is None:
									__UpperCAmelCase     :				List[str]					   =			torch.zeros(UpperCAmelCase_						,     dtype=torch.long						,     device=UpperCAmelCase_						)

							__UpperCAmelCase     :				int					   =			attention_mask.unsqueeze(1						).unsqueeze(2						)
							__UpperCAmelCase     :				Optional[int]					   =			extended_attention_mask.to(dtype=self.dtype						)  # fp16 compatibility
							__UpperCAmelCase     :				Tuple					   =			(1.0 - extended_attention_mask) * -10_000.0
							__UpperCAmelCase     :				List[str]					   =			self.get_head_mask(UpperCAmelCase_						,     self.config.num_hidden_layers						)

							__UpperCAmelCase     :				List[str]					   =			self.embeddings(
							    UpperCAmelCase_						,     position_ids=UpperCAmelCase_						,     token_type_ids=UpperCAmelCase_						,     inputs_embeds=UpperCAmelCase_						)
							__UpperCAmelCase     :				str					   =			embedding_output

							if self.training:
									__UpperCAmelCase     :				Optional[Any]					   =			[]
									for i in range(self.config.num_hidden_layers						):
											__UpperCAmelCase     :				Tuple					   =			self.encoder.adaptive_forward(
											    UpperCAmelCase_						,     current_layer=UpperCAmelCase_						,     attention_mask=UpperCAmelCase_						,     head_mask=UpperCAmelCase_						,     )

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

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

											__UpperCAmelCase     :				List[Any]					   =			logits
											if patient_counter == self.patience:
													break
									__UpperCAmelCase     :				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. '''				,snake_case__				,)
class 			SCREAMING_SNAKE_CASE__	(					snake_case__						):







					"""simple docstring"""






					def __init__(					self    :	List[Any]						,     UpperCAmelCase_    :	Optional[int]						):







							"""simple docstring"""
							super().__init__(UpperCAmelCase_						)
							__UpperCAmelCase     :				Any					   =			config.num_labels

							__UpperCAmelCase     :				Dict					   =			AlbertModelWithPabee(UpperCAmelCase_						)
							__UpperCAmelCase     :				Optional[int]					   =			nn.Dropout(config.classifier_dropout_prob						)
							__UpperCAmelCase     :				List[str]					   =			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(UpperCAmelCase_						)
					def 						lowerCamelCase_					(					self    :	List[Any]						,     UpperCAmelCase_    :	List[Any]=None						,     UpperCAmelCase_    :	Dict=None						,     UpperCAmelCase_    :	Union[str, Any]=None						,     UpperCAmelCase_    :	Tuple=None						,     UpperCAmelCase_    :	Tuple=None						,     UpperCAmelCase_    :	Any=None						,     UpperCAmelCase_    :	Tuple=None						,     ):







							"""simple docstring"""
							__UpperCAmelCase     :				int					   =			self.albert(
							    input_ids=UpperCAmelCase_						,     attention_mask=UpperCAmelCase_						,     token_type_ids=UpperCAmelCase_						,     position_ids=UpperCAmelCase_						,     head_mask=UpperCAmelCase_						,     inputs_embeds=UpperCAmelCase_						,     output_dropout=self.dropout						,     output_layers=self.classifiers						,     regression=self.num_labels == 1						,     )

							__UpperCAmelCase     :				Dict					   =			(logits[-1],)

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

							return outputs












"""simple docstring"""


from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


_UpperCamelCase				:    List[Any]     =						{
    'configuration_openai': ['OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'OpenAIGPTConfig'],
    'tokenization_openai': ['OpenAIGPTTokenizer'],
}

try:
   if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _UpperCamelCase				:    Any     =						['OpenAIGPTTokenizerFast']

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _UpperCamelCase				:    int     =						[
       'OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST',
       'OpenAIGPTDoubleHeadsModel',
       'OpenAIGPTForSequenceClassification',
       'OpenAIGPTLMHeadModel',
       'OpenAIGPTModel',
       'OpenAIGPTPreTrainedModel',
       'load_tf_weights_in_openai_gpt',
   ]

try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _UpperCamelCase				:    List[str]     =						[
       'TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST',
       'TFOpenAIGPTDoubleHeadsModel',
       'TFOpenAIGPTForSequenceClassification',
       'TFOpenAIGPTLMHeadModel',
       'TFOpenAIGPTMainLayer',
       'TFOpenAIGPTModel',
       'TFOpenAIGPTPreTrainedModel',
   ]


if TYPE_CHECKING:
   from .configuration_openai import OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP, OpenAIGPTConfig
   from .tokenization_openai import OpenAIGPTTokenizer

   try:
      if not is_tokenizers_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_openai_fast import OpenAIGPTTokenizerFast

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_openai import (
          OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
          OpenAIGPTDoubleHeadsModel,
          OpenAIGPTForSequenceClassification,
          OpenAIGPTLMHeadModel,
          OpenAIGPTModel,
          OpenAIGPTPreTrainedModel,
          load_tf_weights_in_openai_gpt,
      )

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_openai import (
          TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFOpenAIGPTDoubleHeadsModel,
          TFOpenAIGPTForSequenceClassification,
          TFOpenAIGPTLMHeadModel,
          TFOpenAIGPTMainLayer,
          TFOpenAIGPTModel,
          TFOpenAIGPTPreTrainedModel,
      )

else:
   import sys

   _UpperCamelCase				:    Union[str, Any]     =						_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




