






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


_UpperCamelCase :    Tuple      =			logging.getLogger(__name__)






class __a ( snake_case_					):



  def     A_					(self	:		Optional[int]					,      lowercase_	:		Tuple					,      lowercase_	:		int					,      lowercase_	:		List[str]=None					,      lowercase_	:		Optional[int]=None )		->				List[str]:
   if current_layer == 0:
    __A						:							int      =						self.embedding_hidden_mapping_in(lowercase_ )
   else:
    __A						:							int      =						hidden_states[0]

   __A						:							int      =						int(self.config.num_hidden_layers / self.config.num_hidden_groups )

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

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



  def __init__(self	:		Tuple					,      lowercase_	:		List[str] )		->				Any:
   super().__init__(lowercase_ )

   __A						:							str      =						AlbertTransformerWithPabee(lowercase_ )

   self.init_weights()
   __A						:							List[str]      =						0
   __A						:							Union[str, Any]      =						0
   __A						:							List[str]      =						0

   __A						:							Union[str, Any]      =						0



  def     A_					(self	:		Tuple					,      lowercase_	:		Union[str, Any] )		->				int:
   __A						:							List[Any]      =						threshold



  def     A_					(self	:		str					,      lowercase_	:		Any )		->				Tuple:
   __A						:							Optional[int]      =						patience



  def     A_					(self	:		Optional[Any] )		->				Union[str, Any]:
   __A						:							List[Any]      =						0
   __A						:							Any      =						0



  def     A_					(self	:		List[Any] )		->				int:
   __A						:							Dict      =						self.inference_layers_num / self.inference_instances_num
   __A						:							Tuple      =						(
       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	:		str					,      lowercase_	:		Optional[int]=None					,      lowercase_	:		Optional[int]=None					,      lowercase_	:		Any=None					,      lowercase_	:		Any=None					,      lowercase_	:		Union[str, Any]=None					,      lowercase_	:		Dict=None					,      lowercase_	:		str=None					,      lowercase_	:		Optional[Any]=None					,      lowercase_	:		List[str]=False					,      )		->				Dict:

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

   __A						:							Optional[Any]      =						input_ids.device if input_ids is not None else inputs_embeds.device

   if attention_mask is None:
    __A						:							List[Any]      =						torch.ones(lowercase_					,      device=lowercase_ )
   if token_type_ids is None:
    __A						:							Dict      =						torch.zeros(lowercase_					,      dtype=torch.long					,      device=lowercase_ )

   __A						:							int      =						attention_mask.unsqueeze(1 ).unsqueeze(2 )
   __A						:							List[str]      =						extended_attention_mask.to(dtype=self.dtype )  # fp16 compatibility
   __A						:							Optional[int]      =						(1.0 - extended_attention_mask) * -1_0000.0
   __A						:							Optional[int]      =						self.get_head_mask(lowercase_					,      self.config.num_hidden_layers )

   __A						:							List[str]      =						self.embeddings(
       lowercase_					,      position_ids=lowercase_					,      token_type_ids=lowercase_					,      inputs_embeds=lowercase_ )
   __A						:							str      =						embedding_output

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

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

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

     __A						:							int      =						logits
     if patient_counter == self.patience:
      break
    __A						:							str      =						[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 __a ( snake_case_					):



  def __init__(self	:		Union[str, Any]					,      lowercase_	:		str )		->				List[Any]:
   super().__init__(lowercase_ )
   __A						:							Tuple      =						config.num_labels

   __A						:							Optional[Any]      =						AlbertModelWithPabee(lowercase_ )
   __A						:							Any      =						nn.Dropout(config.classifier_dropout_prob )
   __A						:							int      =						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	:		List[Any]					,      lowercase_	:		Union[str, Any]=None					,      lowercase_	:		Dict=None					,      lowercase_	:		Optional[Any]=None					,      lowercase_	:		Optional[Any]=None					,      lowercase_	:		Union[str, Any]=None					,      lowercase_	:		Union[str, Any]=None					,      lowercase_	:		Union[str, Any]=None					,      )		->				List[str]:
   __A						:							str      =						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					,      )

   __A						:							int      =						(logits[-1],)

   if labels is not None:
    __A						:							Union[str, Any]      =						None
    __A						:							Any      =						0
    for ix, logits_item in enumerate(lowercase_ ):
     if self.num_labels == 1:
      #  We are doing regression
      __A						:							List[Any]      =						MSELoss()
      __A						:							str      =						loss_fct(logits_item.view(-1 )					,      labels.view(-1 ) )
     else:
      __A						:							Dict      =						CrossEntropyLoss()
      __A						:							Any      =						loss_fct(logits_item.view(-1					,      self.num_labels )					,      labels.view(-1 ) )
     if total_loss is None:
      __A						:							List[str]      =						loss
     else:
      total_loss += loss * (ix + 1)
     total_weights += ix + 1
    __A						:							List[Any]      =						(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


_UpperCamelCase :    List[str]      =			logging.getLogger(__name__)




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

 __A						:							List[Any]      =						parser.parse_args()
 return args




def 					a			(  UpperCamelCase_				:					Optional[Any]		)	->		Optional[Any]:
 __A						:							Union[str, Any]      =						datasets.load_dataset(args.dataset_name  ,						args.dataset_config  ,						split='train'		)

 if args.limit is not None:
  __A						:							Tuple      =						min(len(UpperCamelCase_		)  ,						args.limit		)
  __A						:							Optional[int]      =						dataset.select(range(UpperCamelCase_		)		)
  logger.info(F"Limiting the dataset to {args.limit} entries."		)

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

    # Prepare the tokenizer.
 __A						:							Any      =						Tokenizer(Unigram()		)
 __A						:							List[Any]      =						normalizers.Sequence([normalizers.Replace('``'  ,						'"'		), normalizers.Replace('\'\''  ,						'"'		)]		)
 __A						:							Optional[int]      =						pre_tokenizers.Metaspace()

 # Prepare the trainer.
 __A						:							int      =						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=UpperCamelCase_		)
 logger.info('Tokenizer training complete!'		)

 __A						:							Union[str, Any]      =						tokenizer.token_to_id('[CLS]'		)
 __A						:							str      =						tokenizer.token_to_id('[SEP]'		)
 __A						:							Tuple      =						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),
     ]  ,						)
 __A						:							List[str]      =						decoders.Metaspace()

 if args.export_to_hub:
  logger.info('Exporting the trained tokenzier to Hub.'		)
  __A						:							Union[str, Any]      =						AlbertTokenizerFast(tokenizer_object=UpperCamelCase_		)
  new_tokenizer.push_to_hub('unigram-tokenizer-dataset'		)


if __name__ == "__main__":
   _UpperCamelCase :    Optional[Any]      =			parse_args()
   main(args)





