
"""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,
)


snake_case			       =		logging.getLogger(__name__)




class 			UpperCamelCase     (	snake_case_				):

				"""simple docstring"""



				def 							A (    self		,     lowercase__		,     lowercase__		,     lowercase__=None		,     lowercase__=None      )   ->   List[str]:
									"""simple docstring"""







									if current_layer == 0:
														SCREAMING_SNAKE_CASE           =     self.embedding_hidden_mapping_in(lowercase__      )
									else:
														SCREAMING_SNAKE_CASE           =     hidden_states[0]

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

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

									SCREAMING_SNAKE_CASE           =     self.albert_layer_groups[group_idx](
									    lowercase__		,     lowercase__		,     head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]		,     )
									SCREAMING_SNAKE_CASE           =     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 			UpperCamelCase     (	snake_case_				):

				"""simple docstring"""



				def __init__(    self		,     lowercase__      )   ->   Union[str, Any]:
									"""simple docstring"""







									super().__init__(lowercase__      )

									SCREAMING_SNAKE_CASE           =     AlbertTransformerWithPabee(lowercase__      )

									self.init_weights()
									SCREAMING_SNAKE_CASE           =     0
									SCREAMING_SNAKE_CASE           =     0
									SCREAMING_SNAKE_CASE           =     0

									SCREAMING_SNAKE_CASE           =     0


				def 							A (    self		,     lowercase__      )   ->   Optional[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     threshold


				def 							A (    self		,     lowercase__      )   ->   Union[str, Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     patience


				def 							A (    self      )   ->   Union[str, Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     0
									SCREAMING_SNAKE_CASE           =     0


				def 							A (    self      )   ->   List[str]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.inference_layers_num / self.inference_instances_num
									SCREAMING_SNAKE_CASE           =     (
									    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 							A (    self		,     lowercase__=None		,     lowercase__=None		,     lowercase__=None		,     lowercase__=None		,     lowercase__=None		,     lowercase__=None		,     lowercase__=None		,     lowercase__=None		,     lowercase__=False		,     )   ->   List[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           =     input_ids.size()
									elif inputs_embeds is not None:
														SCREAMING_SNAKE_CASE           =     inputs_embeds.size()[:-1]
									else:
														raise ValueError('You have to specify either input_ids or inputs_embeds'      )

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

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

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

									SCREAMING_SNAKE_CASE           =     self.embeddings(
									    lowercase__		,     position_ids=lowercase__		,     token_type_ids=lowercase__		,     inputs_embeds=lowercase__      )
									SCREAMING_SNAKE_CASE           =     embedding_output

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

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

																			SCREAMING_SNAKE_CASE           =     self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]      )      )
																			SCREAMING_SNAKE_CASE           =     output_layers[i](lowercase__      )
																			if regression:
																								SCREAMING_SNAKE_CASE           =     logits.detach()
																								if patient_result is not None:
																													SCREAMING_SNAKE_CASE           =     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           =     0
																			else:
																								SCREAMING_SNAKE_CASE           =     logits.detach().argmax(dim=1      )
																								if patient_result is not None:
																													SCREAMING_SNAKE_CASE           =     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           =     0

																			SCREAMING_SNAKE_CASE           =     logits
																			if patient_counter == self.patience:
																								break
														SCREAMING_SNAKE_CASE           =     [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. "  ,				snake_case_  ,				)
class 			UpperCamelCase     (	snake_case_				):

				"""simple docstring"""



				def __init__(    self		,     lowercase__      )   ->   Union[str, Any]:
									"""simple docstring"""







									super().__init__(lowercase__      )
									SCREAMING_SNAKE_CASE           =     config.num_labels

									SCREAMING_SNAKE_CASE           =     AlbertModelWithPabee(lowercase__      )
									SCREAMING_SNAKE_CASE           =     nn.Dropout(config.classifier_dropout_prob      )
									SCREAMING_SNAKE_CASE           =     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 							A (    self		,     lowercase__=None		,     lowercase__=None		,     lowercase__=None		,     lowercase__=None		,     lowercase__=None		,     lowercase__=None		,     lowercase__=None		,     )   ->   Optional[int]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     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           =     (logits[-1],)

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

									return outputs


"""simple docstring"""




import argparse
import logging

import datasets
from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers, processors
from tokenizers.models import Unigram
from tokenizers.trainers import UnigramTrainer

from transformers import AlbertTokenizerFast


snake_case			       =		logging.getLogger(__name__)

def   UpperCamelCase_  (		):
					SCREAMING_SNAKE_CASE           =     argparse.ArgumentParser(description='Train a unigram tokenizer on the wikitext dataset.'      )
					parser.add_argument(
					    '--dataset_name',					type=SCREAMING_SNAKE_CASE_,					default='wikitext',					help='Name of the training. Explore datasets at: hf.co/datasets.',					)
					parser.add_argument(
					    '--dataset_config',					type=SCREAMING_SNAKE_CASE_,					default='wikitext-103-raw-v1',					help='Configuration name of the dataset.'      )
					parser.add_argument(
					    '--batch_size',					type=SCREAMING_SNAKE_CASE_,					default=1_0_0_0,					help='Batch size during training.',					)
					parser.add_argument(
					    '--vocab_size',					type=SCREAMING_SNAKE_CASE_,					default=1_0_0_4_8,					help='Size of the desired vocabulary.',					)
					parser.add_argument(
					    '--limit',					default=SCREAMING_SNAKE_CASE_,					type=SCREAMING_SNAKE_CASE_,					help='Limit the number of shards (used for debugging).',					)
					parser.add_argument(
					    '--export_to_hub',					action='store_true',					)

					SCREAMING_SNAKE_CASE           =     parser.parse_args()
					return args

def   UpperCamelCase_  (		SCREAMING_SNAKE_CASE_      ):
					SCREAMING_SNAKE_CASE           =     datasets.load_dataset(args.dataset_name,					args.dataset_config,					split='train'      )

					if args.limit is not None:
										SCREAMING_SNAKE_CASE           =     min(len(SCREAMING_SNAKE_CASE_      ),					args.limit      )
										SCREAMING_SNAKE_CASE           =     dataset.select(range(SCREAMING_SNAKE_CASE_      )      )
										logger.info(f'''Limiting the dataset to {args.limit} entries.'''      )

					def batch_iterator():
										for i in range(0,					len(SCREAMING_SNAKE_CASE_      ),					args.batch_size      ):
															yield dataset[i : i + args.batch_size]["text"]

    # Prepare the tokenizer.
					SCREAMING_SNAKE_CASE           =     Tokenizer(Unigram()      )
					SCREAMING_SNAKE_CASE           =     normalizers.Sequence([normalizers.Replace('``',					'"'      ), normalizers.Replace('\'\'',					'"'      )]      )
					SCREAMING_SNAKE_CASE           =     pre_tokenizers.Metaspace()

					# Prepare the trainer.
					SCREAMING_SNAKE_CASE           =     UnigramTrainer(
					    unk_token='<unk>',					special_tokens=['[CLS]', '[SEP]', '<unk>', '<pad>', '[MASK]'],					vocab_size=args.vocab_size,					)

					logger.info('Training the tokenizer.'      )
					tokenizer.train_from_iterator(batch_iterator(),					trainer=SCREAMING_SNAKE_CASE_      )
					logger.info('Tokenizer training complete!'      )

					SCREAMING_SNAKE_CASE           =     tokenizer.token_to_id('[CLS]'      )
					SCREAMING_SNAKE_CASE           =     tokenizer.token_to_id('[SEP]'      )
					SCREAMING_SNAKE_CASE           =     processors.TemplateProcessing(
					    single='[CLS]:0 $A:0 [SEP]:0',					pair='[CLS]:0 $A:0 [SEP]:0 $B:1 [SEP]:1',					special_tokens=[
					        ('[CLS]', cls_token_id),
					        ('[SEP]', sep_token_id),
					    ],					)
					SCREAMING_SNAKE_CASE           =     decoders.Metaspace()

					if args.export_to_hub:
										logger.info('Exporting the trained tokenzier to Hub.'      )
										SCREAMING_SNAKE_CASE           =     AlbertTokenizerFast(tokenizer_object=SCREAMING_SNAKE_CASE_      )
										new_tokenizer.push_to_hub('unigram-tokenizer-dataset'      )


if __name__ == "__main__":
	snake_case			       =		parse_args()
	main(args)
