



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


A__						:						Tuple          =      logging.getLogger(__name__)


def     UpperCamelCase(						):
		lowerCAmelCase_    :  Any             =				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'''  ,)

		lowerCAmelCase_    :  Dict             =				parser.parse_args()
		return args


def     UpperCamelCase(						__UpperCamelCase							:	int				):
		lowerCAmelCase_    :  List[Any]             =				datasets.load_dataset(args.dataset_name  ,args.dataset_config  ,split='''train'''				)

		if args.limit is not None:
				lowerCAmelCase_    :  Union[str, Any]             =				min(len(__UpperCamelCase				)  ,args.limit				)
				lowerCAmelCase_    :  str             =				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.
		lowerCAmelCase_    :  List[str]             =				Tokenizer(Unigram()				)
		lowerCAmelCase_    :  Dict             =				normalizers.Sequence([normalizers.Replace('''``'''  ,'''"'''				), normalizers.Replace('''\'\''''  ,'''"'''				)]				)
		lowerCAmelCase_    :  Union[str, Any]             =				pre_tokenizers.Metaspace()

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

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

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


if __name__ == "__main__":
						A__						:						Dict          =      parse_args()
						main(args)










from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_torch_available,
    is_vision_available,
)


A__						:						List[str]          =      {
    '''configuration_segformer''': ['''SEGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''SegformerConfig''', '''SegformerOnnxConfig''']
}

try:
						if not is_vision_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						A__						:						int          =      ['''SegformerFeatureExtractor''']
						A__						:						Dict          =      ['''SegformerImageProcessor''']

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						A__						:						Any          =      [
						    '''SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST''',
						    '''SegformerDecodeHead''',
						    '''SegformerForImageClassification''',
						    '''SegformerForSemanticSegmentation''',
						    '''SegformerLayer''',
						    '''SegformerModel''',
						    '''SegformerPreTrainedModel''',
						]

try:
						if not is_tf_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						A__						:						List[Any]          =      [
						    '''TF_SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST''',
						    '''TFSegformerDecodeHead''',
						    '''TFSegformerForImageClassification''',
						    '''TFSegformerForSemanticSegmentation''',
						    '''TFSegformerModel''',
						    '''TFSegformerPreTrainedModel''',
						]


if TYPE_CHECKING:
						from .configuration_segformer import SEGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, SegformerConfig, SegformerOnnxConfig

						try:
												if not is_vision_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .feature_extraction_segformer import SegformerFeatureExtractor
												from .image_processing_segformer import SegformerImageProcessor

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_segformer import (
												    SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
												    SegformerDecodeHead,
												    SegformerForImageClassification,
												    SegformerForSemanticSegmentation,
												    SegformerLayer,
												    SegformerModel,
												    SegformerPreTrainedModel,
												)
						try:
												if not is_tf_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_tf_segformer import (
												    TF_SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
												    TFSegformerDecodeHead,
												    TFSegformerForImageClassification,
												    TFSegformerForSemanticSegmentation,
												    TFSegformerModel,
												    TFSegformerPreTrainedModel,
												)

else:
						import sys

						A__						:						Dict          =      _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)





