


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_	( snake_case		):




			def 			_snake_case   (      self					,      lowercase					,      lowercase					,      lowercase=None					,      lowercase=None       )			->						Optional[int]:


								'''simple docstring'''






								if current_layer == 0:
													__SCREAMING_SNAKE_CASE					:							Any			   = self.embedding_hidden_mapping_in(lowercase       )
								else:
													__SCREAMING_SNAKE_CASE					:							Dict			   = hidden_states[0]

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

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

								__SCREAMING_SNAKE_CASE					:							int			   = self.albert_layer_groups[group_idx](
								    lowercase					,      lowercase					,      head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]					,      )
								__SCREAMING_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.'''					,					snake_case					,					)
class      SCREAMING_SNAKE_CASE_	( snake_case		):




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


								'''simple docstring'''






								super().__init__(lowercase       )

								__SCREAMING_SNAKE_CASE					:							Any			   = AlbertTransformerWithPabee(lowercase       )

								self.init_weights()
								__SCREAMING_SNAKE_CASE					:							int			   = 0
								__SCREAMING_SNAKE_CASE					:							Dict			   = 0
								__SCREAMING_SNAKE_CASE					:							Any			   = 0

								__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = 0




			def 			_snake_case   (      self					,      lowercase       )			->						Dict:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							str			   = threshold




			def 			_snake_case   (      self					,      lowercase       )			->						str:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = patience




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


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = 0
								__SCREAMING_SNAKE_CASE					:							str			   = 0




			def 			_snake_case   (      self       )			->						Optional[int]:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							Tuple			   = self.inference_layers_num / self.inference_instances_num
								__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = (
								    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(lowercase       )




			@add_start_docstrings_to_model_forward(lowercase       )
			def 			_snake_case   (      self					,      lowercase=None					,      lowercase=None					,      lowercase=None					,      lowercase=None					,      lowercase=None					,      lowercase=None					,      lowercase=None					,      lowercase=None					,      lowercase=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:
													__SCREAMING_SNAKE_CASE					:							Any			   = input_ids.size()
								elif inputs_embeds is not None:
													__SCREAMING_SNAKE_CASE					:							List[Any]			   = inputs_embeds.size()[:-1]
								else:
													raise ValueError('''You have to specify either input_ids or inputs_embeds'''       )

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

								if attention_mask is None:
													__SCREAMING_SNAKE_CASE					:							Tuple			   = torch.ones(lowercase					,      device=lowercase       )
								if token_type_ids is None:
													__SCREAMING_SNAKE_CASE					:							Tuple			   = torch.zeros(lowercase					,      dtype=torch.long					,      device=lowercase       )

								__SCREAMING_SNAKE_CASE					:							Any			   = attention_mask.unsqueeze(1       ).unsqueeze(2       )
								__SCREAMING_SNAKE_CASE					:							str			   = extended_attention_mask.to(dtype=self.dtype       )  # fp16 compatibility
								__SCREAMING_SNAKE_CASE					:							str			   = (1.0 - extended_attention_mask) * -1_0_0_0_0.0
								__SCREAMING_SNAKE_CASE					:							List[str]			   = self.get_head_mask(lowercase					,      self.config.num_hidden_layers       )

								__SCREAMING_SNAKE_CASE					:							Tuple			   = self.embeddings(
								    lowercase					,      position_ids=lowercase					,      token_type_ids=lowercase					,      inputs_embeds=lowercase       )
								__SCREAMING_SNAKE_CASE					:							List[str]			   = embedding_output

								if self.training:
													__SCREAMING_SNAKE_CASE					:							List[Any]			   = []
													for i in range(self.config.num_hidden_layers       ):
																		__SCREAMING_SNAKE_CASE					:							Any			   = self.encoder.adaptive_forward(
																		    lowercase					,      current_layer=lowercase					,      attention_mask=lowercase					,      head_mask=lowercase					,      )

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

																		__SCREAMING_SNAKE_CASE					:							Dict			   = self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]       )       )
																		__SCREAMING_SNAKE_CASE					:							str			   = output_layers[i](lowercase       )
																		if regression:
																							__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = logits.detach()
																							if patient_result is not None:
																												__SCREAMING_SNAKE_CASE					:							List[str]			   = 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					:							Optional[Any]			   = 0
																		else:
																							__SCREAMING_SNAKE_CASE					:							int			   = logits.detach().argmax(dim=1       )
																							if patient_result is not None:
																												__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = patient_result.detach().argmax(dim=1       )
																							if (patient_result is not None) and torch.all(labels.eq(lowercase       )       ):
																												patient_counter += 1
																							else:
																												__SCREAMING_SNAKE_CASE					:							List[Any]			   = 0

																		__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = logits
																		if patient_counter == self.patience:
																							break
													__SCREAMING_SNAKE_CASE					:							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. '''					,					snake_case					,					)
class      SCREAMING_SNAKE_CASE_	( snake_case		):




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


								'''simple docstring'''






								super().__init__(lowercase       )
								__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = config.num_labels

								__SCREAMING_SNAKE_CASE					:							List[Any]			   = AlbertModelWithPabee(lowercase       )
								__SCREAMING_SNAKE_CASE					:							int			   = nn.Dropout(config.classifier_dropout_prob       )
								__SCREAMING_SNAKE_CASE					:							List[Any]			   = 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(lowercase       )
			def 			_snake_case   (      self					,      lowercase=None					,      lowercase=None					,      lowercase=None					,      lowercase=None					,      lowercase=None					,      lowercase=None					,      lowercase=None					,      )			->						List[str]:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = self.albert(
								    input_ids=lowercase					,      attention_mask=lowercase					,      token_type_ids=lowercase					,      position_ids=lowercase					,      head_mask=lowercase					,      inputs_embeds=lowercase					,      output_dropout=self.dropout					,      output_layers=self.classifiers					,      regression=self.num_labels == 1					,      )

								__SCREAMING_SNAKE_CASE					:							Any			   = (logits[-1],)

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

								return outputs






'''simple docstring'''






# this script builds a small sample spm file tests/fixtures/test_sentencepiece_no_bos.model, with features needed by pegasus

# 1. pip install sentencepiece
#
# 2. wget https://raw.githubusercontent.com/google/sentencepiece/master/data/botchan.txt

# 3. build
import sentencepiece as spm

# pegasus:
# 1. no bos
# 2. eos_id is 1
# 3. unk_id is 2
# build a sample spm file accordingly
spm.SentencePieceTrainer.train(
    """--input=botchan.txt --model_prefix=test_sentencepiece_no_bos --bos_id=-1 --unk_id=2  --eos_id=1  --vocab_size=1000"""
)

# 4. now update the fixture
# mv test_sentencepiece_no_bos.model ../../tests/fixtures/


