


"""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_											=     logging.getLogger(__name__)
class snake_case   (	_UpperCamelCase):







		def 			a_			(       self   :							Optional[int]     , a__   :							int     , a__   :							str     , a__   :							Union[str, Any]=None     , a__   :							List[str]=None		)							->    Optional[int]:





						'''simple docstring'''







						if current_layer == 0:
										_A       	=      self.embedding_hidden_mapping_in(a__		)
						else:
										_A       	=      hidden_states[0]

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

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

						_A       	=      self.albert_layer_groups[group_idx](
						    a__     , a__     , head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]     , )
						_A       	=      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 snake_case   (	_UpperCamelCase):







		def __init__(       self   :							List[Any]     , a__   :							int		)							->    Optional[int]:





						'''simple docstring'''







						super().__init__(a__		)

						_A       	=      AlbertTransformerWithPabee(a__		)

						self.init_weights()
						_A       	=      0
						_A       	=      0
						_A       	=      0

						_A       	=      0







		def 			a_			(       self   :							int     , a__   :							Optional[Any]		)							->    Dict:





						'''simple docstring'''







						_A       	=      threshold







		def 			a_			(       self   :							Dict     , a__   :							List[str]		)							->    Tuple:





						'''simple docstring'''







						_A       	=      patience







		def 			a_			(       self   :							str		)							->    List[Any]:





						'''simple docstring'''







						_A       	=      0
						_A       	=      0







		def 			a_			(       self   :							int		)							->    Tuple:





						'''simple docstring'''







						_A       	=      self.inference_layers_num / self.inference_instances_num
						_A       	=      (
						    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 			a_			(       self   :							List[str]     , a__   :							Optional[int]=None     , a__   :							str=None     , a__   :							List[Any]=None     , a__   :							List[Any]=None     , a__   :							List[str]=None     , a__   :							List[str]=None     , a__   :							List[Any]=None     , a__   :							str=None     , a__   :							int=False     , )							->    Tuple:





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

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

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

						_A       	=      attention_mask.unsqueeze(1		).unsqueeze(2		)
						_A       	=      extended_attention_mask.to(dtype=self.dtype		)  # fp16 compatibility
						_A       	=      (1.0 - extended_attention_mask) * -1_0_0_0_0.0
						_A       	=      self.get_head_mask(a__     , self.config.num_hidden_layers		)

						_A       	=      self.embeddings(
						    a__     , position_ids=a__     , token_type_ids=a__     , inputs_embeds=a__		)
						_A       	=      embedding_output

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

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

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

														_A       	=      logits
														if patient_counter == self.patience:
																		break
										_A       	=      [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. '  ,					_UpperCamelCase  ,					)
class snake_case   (	_UpperCamelCase):







		def __init__(       self   :							List[str]     , a__   :							Dict		)							->    Tuple:





						'''simple docstring'''







						super().__init__(a__		)
						_A       	=      config.num_labels

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





						'''simple docstring'''







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

						_A       	=      (logits[-1],)

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

						return outputs



"""simple docstring"""


from .hash_table import HashTable
class snake_case   (	_UpperCamelCase):







		def __init__(       self   :							Any     , *a__   :							Any     , **a__   :							Tuple		)							->    Tuple:





						'''simple docstring'''







						super().__init__(*a__     , **a__		)







		def 			a_			(       self   :							Tuple     , a__   :							int     , a__   :							str=None		)							->    Optional[int]:





						'''simple docstring'''







						_A       	=      1
						_A       	=      self.hash_function(key + i * i		)

						while self.values[new_key] is not None and self.values[new_key] != key:
										i += 1
										_A       	=      (
										    self.hash_function(key + i * i		)
										    if not self.balanced_factor() >= self.lim_charge
										    else None
										)

										if new_key is None:
														break

						return new_key