
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   						=     logging.getLogger(__name__)


def  A__ (       ):
   SCREAMING_SNAKE_CASE_	=      argparse.ArgumentParser(description='''Train a unigram tokenizer on the wikitext dataset.'''						)
   parser.add_argument(
       '''--dataset_name''',  type=__lowerCamelCase,  default='''wikitext''',  help='''Name of the training. Explore datasets at: hf.co/datasets.''',  )
   parser.add_argument(
       '''--dataset_config''',  type=__lowerCamelCase,  default='''wikitext-103-raw-v1''',  help='''Configuration name of the dataset.'''						)
   parser.add_argument(
       '''--batch_size''',  type=__lowerCamelCase,  default=10_00,  help='''Batch size during training.''',  )
   parser.add_argument(
       '''--vocab_size''',  type=__lowerCamelCase,  default=1_00_48,  help='''Size of the desired vocabulary.''',  )
   parser.add_argument(
       '''--limit''',  default=__lowerCamelCase,  type=__lowerCamelCase,  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  A__ (       __lowerCamelCase						):
   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(__lowerCamelCase						),  args.limit						)
      SCREAMING_SNAKE_CASE_	=      dataset.select(range(__lowerCamelCase						)						)
      logger.info(F'''Limiting the dataset to {args.limit} entries.'''						)

   def batch_iterator():
      for i in range(0,  len(__lowerCamelCase						),  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=__lowerCamelCase						)
   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=__lowerCamelCase						)
      new_tokenizer.push_to_hub('''unigram-tokenizer-dataset'''						)


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




from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


__UpperCAmelCase   						=     {
    "configuration_donut_swin": ["DONUT_SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP", "DonutSwinConfig"],
    "processing_donut": ["DonutProcessor"],
}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __UpperCAmelCase   						=     [
       "DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST",
       "DonutSwinModel",
       "DonutSwinPreTrainedModel",
   ]

try:
   if not is_vision_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __UpperCAmelCase   						=     ["DonutFeatureExtractor"]
   __UpperCAmelCase   						=     ["DonutImageProcessor"]


if TYPE_CHECKING:
   from .configuration_donut_swin import DONUT_SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP, DonutSwinConfig
   from .processing_donut import DonutProcessor

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_donut_swin import (
          DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
          DonutSwinModel,
          DonutSwinPreTrainedModel,
      )

   try:
      if not is_vision_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .feature_extraction_donut import DonutFeatureExtractor
      from .image_processing_donut import DonutImageProcessor

else:
   import sys

   __UpperCAmelCase   						=     _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)


