


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 							UpperCAmelCase_   (				)	->     Union[str, Any]:







 """simple docstring"""


 _lowercase				 =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=1000	,		help='''Batch size during training.'''	,		)
 parser.add_argument(
     '''--vocab_size'''	,		type=UpperCamelCase_	,		default=10048	,		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'''	,		)

 _lowercase				 =parser.parse_args()
 return args
def 							UpperCAmelCase_   (				__snake_case		)	->     Optional[Any]:







 """simple docstring"""


 _lowercase				 =datasets.load_dataset(args.dataset_name	,		args.dataset_config	,		split='''train'''		)

 if args.limit is not None:
  _lowercase				 =min(len(UpperCamelCase_		)	,		args.limit		)
  _lowercase				 =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.
 _lowercase				 =Tokenizer(Unigram()		)
 _lowercase				 =normalizers.Sequence([normalizers.Replace('''``'''	,		'''"'''		), normalizers.Replace('''\'\''''	,		'''"'''		)]		)
 _lowercase				 =pre_tokenizers.Metaspace()

 # Prepare the trainer.
 _lowercase				 =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!'''		)

 _lowercase				 =tokenizer.token_to_id('''[CLS]'''		)
 _lowercase				 =tokenizer.token_to_id('''[SEP]'''		)
 _lowercase				 =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),
     ]	,		)
 _lowercase				 =decoders.Metaspace()

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


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









"""simple docstring"""


import importlib.metadata
import subprocess
import sys




def 					a			(  UpperCamelCase_				:					bool = False		)	->		Optional[Any]:
 __A						:							List[Any]      =						False
 if "IPython" in sys.modules:
  __A						:							List[Any]      =						'google.colab' in str(sys.modules['IPython'].get_ipython()		)

 if in_colab:
  if upgrade:
   __A						:							Dict      =						['pip', 'install', '-U', 'torch']
   subprocess.run(UpperCamelCase_  ,						check=UpperCamelCase_		)
  # get the current version of torch
  __A						:							int      =						importlib.metadata.version('torch'		)
  __A						:							int      =						torch_version[: torch_version.rindex('.'		)]
  __A						:							Any      =						F"https://storage.googleapis.com/tpu-pytorch/wheels/colab/torch_xla-{torch_version_trunc}-cp37-cp37m-linux_x86_64.whl"
  __A						:							int      =						['pip', 'install', xla_wheel]
  subprocess.run(UpperCamelCase_  ,						check=UpperCamelCase_		)
 else:
  raise RuntimeError('`install_xla` utility works only on google colab.'		)





