'''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 		SCREAMING_SNAKE_CASE_					(_UpperCAmelCase      ):



			"""simple docstring"""






			def   lowerCamelCase (							self				,			snake_case_				,			snake_case_				,			snake_case_=None				,			snake_case_=None					):
									if current_layer == 0:
															snake_case :       Optional[int]				=    self.embedding_hidden_mapping_in(snake_case_					)
									else:
															snake_case :       Tuple				=    hidden_states[0]

									snake_case :       Optional[Any]				=    int(self.config.num_hidden_layers / self.config.num_hidden_groups					)

									# Index of the hidden group
									snake_case :       List[Any]				=    int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)					)

									snake_case :       Union[str, Any]				=    self.albert_layer_groups[group_idx](
									    snake_case_				,			snake_case_				,			head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]				,			)
									snake_case :       Tuple				=    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 		SCREAMING_SNAKE_CASE_					(_UpperCAmelCase      ):



			"""simple docstring"""






			def __init__(							self				,			snake_case_					):
									super().__init__(snake_case_					)

									snake_case :       int				=    AlbertTransformerWithPabee(snake_case_					)

									self.init_weights()
									snake_case :       str				=    0
									snake_case :       Dict				=    0
									snake_case :       int				=    0

									snake_case :       Optional[Any]				=    0



			def   lowerCamelCase (							self				,			snake_case_					):
									snake_case :       List[str]				=    threshold



			def   lowerCamelCase (							self				,			snake_case_					):
									snake_case :       int				=    patience



			def   lowerCamelCase (							self					):
									snake_case :       Tuple				=    0
									snake_case :       Dict				=    0



			def   lowerCamelCase (							self					):
									snake_case :       Optional[Any]				=    self.inference_layers_num / self.inference_instances_num
									snake_case :       List[str]				=    (
									    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(snake_case_					)






			@add_start_docstrings_to_model_forward(snake_case_					)
			def   lowerCamelCase (							self				,			snake_case_=None				,			snake_case_=None				,			snake_case_=None				,			snake_case_=None				,			snake_case_=None				,			snake_case_=None				,			snake_case_=None				,			snake_case_=None				,			snake_case_=False				,			):

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

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

									if attention_mask is None:
															snake_case :       Dict				=    torch.ones(snake_case_				,			device=snake_case_					)
									if token_type_ids is None:
															snake_case :       Dict				=    torch.zeros(snake_case_				,			dtype=torch.long				,			device=snake_case_					)

									snake_case :       str				=    attention_mask.unsqueeze(1					).unsqueeze(2					)
									snake_case :       int				=    extended_attention_mask.to(dtype=self.dtype					)  # fp16 compatibility
									snake_case :       List[Any]				=    (1.0 - extended_attention_mask) * -10_000.0
									snake_case :       Dict				=    self.get_head_mask(snake_case_				,			self.config.num_hidden_layers					)

									snake_case :       str				=    self.embeddings(
									    snake_case_				,			position_ids=snake_case_				,			token_type_ids=snake_case_				,			inputs_embeds=snake_case_					)
									snake_case :       Tuple				=    embedding_output

									if self.training:
															snake_case :       Any				=    []
															for i in range(self.config.num_hidden_layers					):
																					snake_case :       Dict				=    self.encoder.adaptive_forward(
																					    snake_case_				,			current_layer=snake_case_				,			attention_mask=snake_case_				,			head_mask=snake_case_				,			)

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

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

																					snake_case :       Optional[int]				=    logits
																					if patient_counter == self.patience:
																											break
															snake_case :       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. """							,				_UpperCAmelCase							,				)
class 		SCREAMING_SNAKE_CASE_					(_UpperCAmelCase      ):



			"""simple docstring"""






			def __init__(							self				,			snake_case_					):
									super().__init__(snake_case_					)
									snake_case :       Tuple				=    config.num_labels

									snake_case :       Union[str, Any]				=    AlbertModelWithPabee(snake_case_					)
									snake_case :       int				=    nn.Dropout(config.classifier_dropout_prob					)
									snake_case :       Dict				=    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(snake_case_					)
			def   lowerCamelCase (							self				,			snake_case_=None				,			snake_case_=None				,			snake_case_=None				,			snake_case_=None				,			snake_case_=None				,			snake_case_=None				,			snake_case_=None				,			):
									snake_case :       List[Any]				=    self.albert(
									    input_ids=snake_case_				,			attention_mask=snake_case_				,			token_type_ids=snake_case_				,			position_ids=snake_case_				,			head_mask=snake_case_				,			inputs_embeds=snake_case_				,			output_dropout=self.dropout				,			output_layers=self.classifiers				,			regression=self.num_labels == 1				,			)

									snake_case :       Union[str, Any]				=    (logits[-1],)

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

									return outputs






'''simple docstring'''

import warnings

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding




class 		SCREAMING_SNAKE_CASE_					(_UpperCAmelCase      ):



			"""simple docstring"""






			__magic_name__ :				Dict   							=       ["""image_processor""", """tokenizer"""]
			__magic_name__ :				Union[str, Any]   							=       """VideoMAEImageProcessor"""
			__magic_name__ :				List[str]   							=       ("""CLIPTokenizer""", """CLIPTokenizerFast""")



			def __init__(							self				,			snake_case_=None				,			snake_case_=None				,			**snake_case_					):
									snake_case :       int				=    None
									if "feature_extractor" in kwargs:
															warnings.warn(
															    "The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"
															    " instead."				,			snake_case_				,			)
															snake_case :       Optional[int]				=    kwargs.pop("feature_extractor"					)

									snake_case :       Union[str, Any]				=    image_processor if image_processor is not None else feature_extractor
									if image_processor is None:
															raise ValueError("You need to specify an `image_processor`."					)
									if tokenizer is None:
															raise ValueError("You need to specify a `tokenizer`."					)

									super().__init__(snake_case_				,			snake_case_					)
									snake_case :       List[Any]				=    self.image_processor



			def __call__(							self				,			snake_case_=None				,			snake_case_=None				,			snake_case_=None				,			**snake_case_					):

									if text is None and videos is None:
															raise ValueError("You have to specify either text or videos. Both cannot be none."					)

									if text is not None:
															snake_case :       List[Any]				=    self.tokenizer(snake_case_				,			return_tensors=snake_case_				,			**snake_case_					)

									if videos is not None:
															snake_case :       List[Any]				=    self.image_processor(snake_case_				,			return_tensors=snake_case_				,			**snake_case_					)

									if text is not None and videos is not None:
															snake_case :       Tuple				=    image_features.pixel_values
															return encoding
									elif text is not None:
															return encoding
									else:
															return BatchEncoding(data=dict(**snake_case_					)				,			tensor_type=snake_case_					)



			def   lowerCamelCase (							self				,			*snake_case_				,			**snake_case_					):
									return self.tokenizer.batch_decode(*snake_case_				,			**snake_case_					)



			def   lowerCamelCase (							self				,			*snake_case_				,			**snake_case_					):
									return self.tokenizer.decode(*snake_case_				,			**snake_case_					)



			@property
			def   lowerCamelCase (							self					):
									return ["input_ids", "attention_mask", "position_ids", "pixel_values"]



			@property
			def   lowerCamelCase (							self					):
									warnings.warn(
									    "`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."				,			snake_case_				,			)
									return self.image_processor_class






			@property
			def   lowerCamelCase (							self					):
									warnings.warn(
									    "`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."				,			snake_case_				,			)
									return self.image_processor





