






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  __UpperCAmelCase						(__A       ):







			def 				UpperCamelCase (      self:    int ,      UpperCAmelCase_:    List[Any] ,      UpperCAmelCase_:    Any ,      UpperCAmelCase_:    int=None ,      UpperCAmelCase_:    Tuple=None	):







									'''simple docstring'''




									if current_layer == 0:
															_SCREAMING_SNAKE_CASE       =    self.embedding_hidden_mapping_in(_lowerCamelCase	)
									else:
															_SCREAMING_SNAKE_CASE       =    hidden_states[0]

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

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

									_SCREAMING_SNAKE_CASE       =    self.albert_layer_groups[group_idx](
									    _lowerCamelCase ,      _lowerCamelCase ,      head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group] ,      )
									_SCREAMING_SNAKE_CASE       =    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  __UpperCAmelCase						(__A       ):







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







									'''simple docstring'''




									super().__init__(_lowerCamelCase	)

									_SCREAMING_SNAKE_CASE       =    AlbertTransformerWithPabee(_lowerCamelCase	)

									self.init_weights()
									_SCREAMING_SNAKE_CASE       =    0
									_SCREAMING_SNAKE_CASE       =    0
									_SCREAMING_SNAKE_CASE       =    0

									_SCREAMING_SNAKE_CASE       =    0







			def 				UpperCamelCase (      self:    Tuple ,      UpperCAmelCase_:    Union[str, Any]	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    threshold







			def 				UpperCamelCase (      self:    List[Any] ,      UpperCAmelCase_:    Dict	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    patience







			def 				UpperCamelCase (      self:    str	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    0
									_SCREAMING_SNAKE_CASE       =    0







			def 				UpperCamelCase (      self:    Any	):







									'''simple docstring'''




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

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

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

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

									_SCREAMING_SNAKE_CASE       =    self.embeddings(
									    _lowerCamelCase ,      position_ids=_lowerCamelCase ,      token_type_ids=_lowerCamelCase ,      inputs_embeds=_lowerCamelCase	)
									_SCREAMING_SNAKE_CASE       =    embedding_output

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

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

																					_SCREAMING_SNAKE_CASE       =    self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]	)	)
																					_SCREAMING_SNAKE_CASE       =    output_layers[i](_lowerCamelCase	)
																					if regression:
																											_SCREAMING_SNAKE_CASE       =    logits.detach()
																											if patient_result is not None:
																																	_SCREAMING_SNAKE_CASE       =    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       =    0
																					else:
																											_SCREAMING_SNAKE_CASE       =    logits.detach().argmax(dim=1	)
																											if patient_result is not None:
																																	_SCREAMING_SNAKE_CASE       =    patient_result.detach().argmax(dim=1	)
																											if (patient_result is not None) and torch.all(labels.eq(_lowerCamelCase	)	):
																																	patient_counter += 1
																											else:
																																	_SCREAMING_SNAKE_CASE       =    0

																					_SCREAMING_SNAKE_CASE       =    logits
																					if patient_counter == self.patience:
																											break
															_SCREAMING_SNAKE_CASE       =    [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  __UpperCAmelCase						(__A       ):







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







									'''simple docstring'''




									super().__init__(_lowerCamelCase	)
									_SCREAMING_SNAKE_CASE       =    config.num_labels

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







									'''simple docstring'''




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

									_SCREAMING_SNAKE_CASE       =    (logits[-1],)

									if labels is not None:
															_SCREAMING_SNAKE_CASE       =    None
															_SCREAMING_SNAKE_CASE       =    0
															for ix, logits_item in enumerate(_lowerCamelCase	):
																					if self.num_labels == 1:
																											#  We are doing regression
																											_SCREAMING_SNAKE_CASE       =    MSELoss()
																											_SCREAMING_SNAKE_CASE       =    loss_fct(logits_item.view(-1	) ,      labels.view(-1	)	)
																					else:
																											_SCREAMING_SNAKE_CASE       =    CrossEntropyLoss()
																											_SCREAMING_SNAKE_CASE       =    loss_fct(logits_item.view(-1 ,      self.num_labels	) ,      labels.view(-1	)	)
																					if total_loss is None:
																											_SCREAMING_SNAKE_CASE       =    loss
																					else:
																											total_loss += loss * (ix + 1)
																					total_weights += ix + 1
															_SCREAMING_SNAKE_CASE       =    (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_layoutlm': ['LAYOUTLM_PRETRAINED_CONFIG_ARCHIVE_MAP', 'LayoutLMConfig', 'LayoutLMOnnxConfig'],
    'tokenization_layoutlm': ['LayoutLMTokenizer'],
}

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

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    :     Optional[Any]			  =   [
      'LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST',
      'LayoutLMForMaskedLM',
      'LayoutLMForSequenceClassification',
      'LayoutLMForTokenClassification',
      'LayoutLMForQuestionAnswering',
      'LayoutLMModel',
      'LayoutLMPreTrainedModel',
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    :     str			  =   [
      'TF_LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST',
      'TFLayoutLMForMaskedLM',
      'TFLayoutLMForSequenceClassification',
      'TFLayoutLMForTokenClassification',
      'TFLayoutLMForQuestionAnswering',
      'TFLayoutLMMainLayer',
      'TFLayoutLMModel',
      'TFLayoutLMPreTrainedModel',
  ]


if TYPE_CHECKING:
  from .configuration_layoutlm import LAYOUTLM_PRETRAINED_CONFIG_ARCHIVE_MAP, LayoutLMConfig, LayoutLMOnnxConfig
  from .tokenization_layoutlm import LayoutLMTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_layoutlm_fast import LayoutLMTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_layoutlm import (
        LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
        LayoutLMForMaskedLM,
        LayoutLMForQuestionAnswering,
        LayoutLMForSequenceClassification,
        LayoutLMForTokenClassification,
        LayoutLMModel,
        LayoutLMPreTrainedModel,
    )
  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_layoutlm import (
        TF_LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFLayoutLMForMaskedLM,
        TFLayoutLMForQuestionAnswering,
        TFLayoutLMForSequenceClassification,
        TFLayoutLMForTokenClassification,
        TFLayoutLMMainLayer,
        TFLayoutLMModel,
        TFLayoutLMPreTrainedModel,
    )

else:
  import sys

  UpperCamelCase__    :     List[str]			  =   _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)
