



"""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   UpperCAmelCase__		(   UpperCamelCase  ):



		def        A_		(			self	:     List[str]   ,				snake_case	:     Union[str, Any]   ,				snake_case	:     str   ,				snake_case	:     List[str]=None   ,				snake_case	:     Any=None					)      ->    int:






									'''simple docstring'''

									if current_layer == 0:
																A    =       self.embedding_hidden_mapping_in(snake_case					)
									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](
									    snake_case   ,				snake_case   ,				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   UpperCAmelCase__		(   UpperCamelCase  ):



		def __init__(			self	:     Union[str, Any]   ,				snake_case	:     Any					)      ->    Any:






									'''simple docstring'''

									super().__init__(snake_case					)

									A    =       AlbertTransformerWithPabee(snake_case					)

									self.init_weights()
									A    =       0
									A    =       0
									A    =       0

									A    =       0



		def        A_		(			self	:     List[str]   ,				snake_case	:     Optional[Any]					)      ->    List[Any]:






									'''simple docstring'''

									A    =       threshold



		def        A_		(			self	:     str   ,				snake_case	:     List[Any]					)      ->    str:






									'''simple docstring'''

									A    =       patience



		def        A_		(			self	:     List[str]					)      ->    int:






									'''simple docstring'''

									A    =       0
									A    =       0



		def        A_		(			self	:     int					)      ->    int:






									'''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(snake_case					)



		@add_start_docstrings_to_model_forward(snake_case					)
		def        A_		(			self	:     Optional[Any]   ,				snake_case	:     Tuple=None   ,				snake_case	:     List[str]=None   ,				snake_case	:     List[Any]=None   ,				snake_case	:     List[str]=None   ,				snake_case	:     str=None   ,				snake_case	:     Optional[int]=None   ,				snake_case	:     List[Any]=None   ,				snake_case	:     int=None   ,				snake_case	:     List[Any]=False   ,				)      ->    Optional[Any]:






									'''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(snake_case   ,				device=snake_case					)
									if token_type_ids is None:
																A    =       torch.zeros(snake_case   ,				dtype=torch.long   ,				device=snake_case					)

									A    =       attention_mask.unsqueeze(1					).unsqueeze(2					)
									A    =       extended_attention_mask.to(dtype=self.dtype					)  # fp16 compatibility
									A    =       (1.0 - extended_attention_mask) * -10000.0
									A    =       self.get_head_mask(snake_case   ,				self.config.num_hidden_layers					)

									A    =       self.embeddings(
									    snake_case   ,				position_ids=snake_case   ,				token_type_ids=snake_case   ,				inputs_embeds=snake_case					)
									A    =       embedding_output

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

																							A    =       self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]					)					)
																							A    =       output_layers[i](output_dropout(snake_case					)					)
																							res.append(snake_case					)
									elif self.patience == 0:  # Use all layers for inference
																A    =       self.encoder(snake_case   ,				snake_case   ,				head_mask=snake_case					)
																A    =       self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]					)					)
																A    =       [output_layers[self.config.num_hidden_layers - 1](snake_case					)]
									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(
																							    snake_case   ,				current_layer=snake_case   ,				attention_mask=snake_case   ,				head_mask=snake_case   ,				)

																							A    =       self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]					)					)
																							A    =       output_layers[i](snake_case					)
																							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(snake_case					)					):
																																					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
    the pooled output) e.g. for GLUE tasks. """ ,UpperCamelCase ,)
class   UpperCAmelCase__		(   UpperCamelCase  ):



		def __init__(			self	:     Any   ,				snake_case	:     str					)      ->    Optional[int]:






									'''simple docstring'''

									super().__init__(snake_case					)
									A    =       config.num_labels

									A    =       AlbertModelWithPabee(snake_case					)
									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(snake_case					)
		def        A_		(			self	:     List[str]   ,				snake_case	:     List[str]=None   ,				snake_case	:     List[str]=None   ,				snake_case	:     List[str]=None   ,				snake_case	:     str=None   ,				snake_case	:     Tuple=None   ,				snake_case	:     List[Any]=None   ,				snake_case	:     Tuple=None   ,				)      ->    List[str]:






									'''simple docstring'''

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

									A    =       (logits[-1],)

									if labels is not None:
																A    =       None
																A    =       0
																for ix, logits_item in enumerate(snake_case					):
																							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"""







import os

from ...utils import logging
from .utils import DataProcessor, InputExample


A     =      logging.get_logger(__name__)






class   UpperCAmelCase__		(   UpperCamelCase  ):



		def __init__(			self	:     List[str]   ,				snake_case	:     List[Any]   ,				snake_case	:     str=None					)      ->    Tuple:






									'''simple docstring'''

									A    =       language
									A    =       train_language



		def        A_		(			self	:     Optional[int]   ,				snake_case	:     Any					)      ->    Union[str, Any]:






									'''simple docstring'''

									A    =       self.language if self.train_language is None else self.train_language
									A    =       self._read_tsv(os.path.join(snake_case   ,				f"""XNLI-MT-1.0/multinli/multinli.train.{lg}.tsv"""					)					)
									A    =       []
									for i, line in enumerate(snake_case					):
																if i == 0:
																							continue
																A    =       f"""train-{i}"""
																A    =       line[0]
																A    =       line[1]
																A    =       'contradiction' if line[2] == 'contradictory' else line[2]
																if not isinstance(snake_case   ,				snake_case					):
																							raise ValueError(f"""Training input {text_a} is not a string"""					)
																if not isinstance(snake_case   ,				snake_case					):
																							raise ValueError(f"""Training input {text_b} is not a string"""					)
																if not isinstance(snake_case   ,				snake_case					):
																							raise ValueError(f"""Training label {label} is not a string"""					)
																examples.append(InputExample(guid=snake_case   ,				text_a=snake_case   ,				text_b=snake_case   ,				label=snake_case					)					)
									return examples



		def        A_		(			self	:     int   ,				snake_case	:     int					)      ->    Optional[int]:






									'''simple docstring'''

									A    =       self._read_tsv(os.path.join(snake_case   ,				'XNLI-1.0/xnli.test.tsv'					)					)
									A    =       []
									for i, line in enumerate(snake_case					):
																if i == 0:
																							continue
																A    =       line[0]
																if language != self.language:
																							continue
																A    =       f"""test-{i}"""
																A    =       line[6]
																A    =       line[7]
																A    =       line[1]
																if not isinstance(snake_case   ,				snake_case					):
																							raise ValueError(f"""Training input {text_a} is not a string"""					)
																if not isinstance(snake_case   ,				snake_case					):
																							raise ValueError(f"""Training input {text_b} is not a string"""					)
																if not isinstance(snake_case   ,				snake_case					):
																							raise ValueError(f"""Training label {label} is not a string"""					)
																examples.append(InputExample(guid=snake_case   ,				text_a=snake_case   ,				text_b=snake_case   ,				label=snake_case					)					)
									return examples



		def        A_		(			self	:     Any					)      ->    List[Any]:






									'''simple docstring'''

									return ["contradiction", "entailment", "neutral"]



A     =      {
    'xnli': XnliProcessor,
}

A     =      {
    'xnli': 'classification',
}

A     =      {
    'xnli': 3,
}



