


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_(			)->						List[Any]:
			UpperCamelCase__   =    argparse.ArgumentParser(description="Train a unigram tokenizer on the wikitext dataset."					)
			parser.add_argument(
			    "--dataset_name"			,       type=_A			,       default="wikitext"			,       help="Name of the training. Explore datasets at: hf.co/datasets."			,       )
			parser.add_argument(
			    "--dataset_config"			,       type=_A			,       default="wikitext-103-raw-v1"			,       help="Configuration name of the dataset."					)
			parser.add_argument(
			    "--batch_size"			,       type=_A			,       default=10_00			,       help="Batch size during training."			,       )
			parser.add_argument(
			    "--vocab_size"			,       type=_A			,       default=1_00_48			,       help="Size of the desired vocabulary."			,       )
			parser.add_argument(
			    "--limit"			,       default=_A			,       type=_A			,       help="Limit the number of shards (used for debugging)."			,       )
			parser.add_argument(
			    "--export_to_hub"			,       action="store_true"			,       )

			UpperCamelCase__   =    parser.parse_args()
			return args
def     UpperCamelCase_(			_A				:Dict					)->						List[str]:
			UpperCamelCase__   =    datasets.load_dataset(args.dataset_name			,       args.dataset_config			,       split="train"					)

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

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

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

			# Prepare the trainer.
			UpperCamelCase__   =    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=_A					)
			logger.info("Tokenizer training complete!"					)

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

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


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






from __future__ import annotations
def     UpperCamelCase_(			_A				:list					)->						int | float:
			UpperCamelCase__   =    sorted(_A					)
			UpperCamelCase__   =    len(_A					)
			UpperCamelCase__   =    length >> 1
			return (
			    (sorted_list[mid_index] + sorted_list[mid_index - 1]) / 2
			    if length % 2 == 0
			    else sorted_list[mid_index]
			)
def     UpperCamelCase_(			)->						List[Any]:
			import doctest

			doctest.testmod()


if __name__ == "__main__":
						main()


