






"""simple docstring"""


import unittest

import numpy as np

from transformers import ElectraConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
   from transformers.models.electra.modeling_flax_electra import (
       FlaxElectraForCausalLM,
       FlaxElectraForMaskedLM,
       FlaxElectraForMultipleChoice,
       FlaxElectraForPreTraining,
       FlaxElectraForQuestionAnswering,
       FlaxElectraForSequenceClassification,
       FlaxElectraForTokenClassification,
       FlaxElectraModel,
   )






class __a ( unittest.TestCase					):



  def __init__(self	:		List[str]					,      lowercase_	:		List[Any]					,      lowercase_	:		Union[str, Any]=13					,      lowercase_	:		Optional[Any]=7					,      lowercase_	:		List[str]=True					,      lowercase_	:		Optional[int]=True					,      lowercase_	:		Optional[Any]=True					,      lowercase_	:		List[str]=True					,      lowercase_	:		Tuple=99					,      lowercase_	:		List[Any]=24					,      lowercase_	:		str=32					,      lowercase_	:		Any=5					,      lowercase_	:		Any=4					,      lowercase_	:		Dict=37					,      lowercase_	:		int="gelu"					,      lowercase_	:		Dict=0.1					,      lowercase_	:		Any=0.1					,      lowercase_	:		str=512					,      lowercase_	:		Any=16					,      lowercase_	:		str=2					,      lowercase_	:		Optional[int]=0.02					,      lowercase_	:		str=4					,      )		->				Optional[Any]:
   __A						:							List[Any]      =						parent
   __A						:							Any      =						batch_size
   __A						:							Optional[int]      =						seq_length
   __A						:							List[Any]      =						is_training
   __A						:							List[str]      =						use_attention_mask
   __A						:							Union[str, Any]      =						use_token_type_ids
   __A						:							str      =						use_labels
   __A						:							Union[str, Any]      =						vocab_size
   __A						:							Any      =						hidden_size
   __A						:							Dict      =						embedding_size
   __A						:							Tuple      =						num_hidden_layers
   __A						:							Optional[int]      =						num_attention_heads
   __A						:							int      =						intermediate_size
   __A						:							str      =						hidden_act
   __A						:							List[Any]      =						hidden_dropout_prob
   __A						:							Union[str, Any]      =						attention_probs_dropout_prob
   __A						:							Union[str, Any]      =						max_position_embeddings
   __A						:							Union[str, Any]      =						type_vocab_size
   __A						:							Optional[Any]      =						type_sequence_label_size
   __A						:							Any      =						initializer_range
   __A						:							Optional[Any]      =						num_choices



  def     A_					(self	:		List[Any] )		->				Tuple:
   __A						:							int      =						ids_tensor([self.batch_size, self.seq_length]					,      self.vocab_size )

   __A						:							Dict      =						None
   if self.use_attention_mask:
    __A						:							List[Any]      =						random_attention_mask([self.batch_size, self.seq_length] )

   __A						:							Any      =						None
   if self.use_token_type_ids:
    __A						:							Tuple      =						ids_tensor([self.batch_size, self.seq_length]					,      self.type_vocab_size )

   __A						:							int      =						ElectraConfig(
       vocab_size=self.vocab_size					,      hidden_size=self.hidden_size					,      embedding_size=self.embedding_size					,      num_hidden_layers=self.num_hidden_layers					,      num_attention_heads=self.num_attention_heads					,      intermediate_size=self.intermediate_size					,      hidden_act=self.hidden_act					,      hidden_dropout_prob=self.hidden_dropout_prob					,      attention_probs_dropout_prob=self.attention_probs_dropout_prob					,      max_position_embeddings=self.max_position_embeddings					,      type_vocab_size=self.type_vocab_size					,      initializer_range=self.initializer_range					,      )

   return config, input_ids, token_type_ids, attention_mask



  def     A_					(self	:		Any )		->				Tuple:
   __A						:							Any      =						self.prepare_config_and_inputs()
   __A     ,       __A     ,       __A     ,       __A						:							int      =						config_and_inputs
   __A						:							Dict      =						{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': attention_mask}
   return config, inputs_dict






@require_flax
class __a ( snake_case_			,unittest.TestCase					):
  __lowercase     :     Union[str, Any]       						=       True

  __lowercase     :     Dict       						=       (
      (
          FlaxElectraModel,
          FlaxElectraForCausalLM,
          FlaxElectraForMaskedLM,
          FlaxElectraForPreTraining,
          FlaxElectraForTokenClassification,
          FlaxElectraForQuestionAnswering,
          FlaxElectraForMultipleChoice,
          FlaxElectraForSequenceClassification,
      )
      if is_flax_available()
      else ()
  )



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



  @slow
  def     A_					(self	:		Tuple )		->				Union[str, Any]:
   for model_class_name in self.all_model_classes:
    if model_class_name == FlaxElectraForMaskedLM:
     __A						:							Optional[Any]      =						model_class_name.from_pretrained('google/electra-small-generator' )
    else:
     __A						:							str      =						model_class_name.from_pretrained('google/electra-small-discriminator' )
    __A						:							Tuple      =						model(np.ones((1, 1) ) )
    self.assertIsNotNone(lowercase_ )













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





