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  :					Optional[int]          =  logging.getLogger(__name__)




class        _UpperCamelCase				(						_UpperCAmelCase			):


			"""simple docstring"""






			def 	_SCREAMING_SNAKE_CASE      (       self				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__=None				,       lowerCAmelCase__=None						)					-> List[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        _UpperCamelCase				(						_UpperCAmelCase			):


			"""simple docstring"""






			def __init__(       self				,       lowerCAmelCase__						)					-> Tuple:

							'''simple docstring'''


							super().__init__(lowerCAmelCase__						)

							__lowercase			      =							AlbertTransformerWithPabee(lowerCAmelCase__						)

							self.init_weights()
							__lowercase			      =							0
							__lowercase			      =							0
							__lowercase			      =							0

							__lowercase			      =							0





			def 	_SCREAMING_SNAKE_CASE      (       self				,       lowerCAmelCase__						)					-> Tuple:

							'''simple docstring'''


							__lowercase			      =							threshold





			def 	_SCREAMING_SNAKE_CASE      (       self				,       lowerCAmelCase__						)					-> Any:

							'''simple docstring'''


							__lowercase			      =							patience





			def 	_SCREAMING_SNAKE_CASE      (       self						)					-> Dict:

							'''simple docstring'''


							__lowercase			      =							0
							__lowercase			      =							0





			def 	_SCREAMING_SNAKE_CASE      (       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 	_SCREAMING_SNAKE_CASE      (       self				,       lowerCAmelCase__=None				,       lowerCAmelCase__=None				,       lowerCAmelCase__=None				,       lowerCAmelCase__=None				,       lowerCAmelCase__=None				,       lowerCAmelCase__=None				,       lowerCAmelCase__=None				,       lowerCAmelCase__=None				,       lowerCAmelCase__=False				,       )					-> 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) * -1_00_00.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        _UpperCamelCase				(						_UpperCAmelCase			):


			"""simple docstring"""






			def __init__(       self				,       lowerCAmelCase__						)					-> List[Any]:

							'''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 	_SCREAMING_SNAKE_CASE      (       self				,       lowerCAmelCase__=None				,       lowerCAmelCase__=None				,       lowerCAmelCase__=None				,       lowerCAmelCase__=None				,       lowerCAmelCase__=None				,       lowerCAmelCase__=None				,       lowerCAmelCase__=None				,       )					-> Tuple:

							'''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
from typing import TypeVar

from .arrow_dataset import Dataset, _split_by_node_map_style_dataset
from .iterable_dataset import IterableDataset, _split_by_node_iterable_dataset


__a  :					Union[str, Any]          =  TypeVar("""DatasetType""", Dataset, IterableDataset)





def     UpperCAmelCase						(   lowercase							,      lowercase							,      lowercase							):







				"""simple docstring"""



				if isinstance(lowercase							,      lowercase							):
								return _split_by_node_map_style_dataset(lowercase							,      rank=lowercase							,      world_size=lowercase							)
				else:
								return _split_by_node_iterable_dataset(lowercase							,      rank=lowercase							,      world_size=lowercase							)