




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


__SCREAMING_SNAKE_CASE     : Union[str, Any]      =logging.getLogger(__name__)

def       lowerCAmelCase_      (	)       ->      int:


      '''simple docstring'''




      lowerCamelCase__   	=					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'  ,					)

      lowerCamelCase__   	=					parser.parse_args()
      return args

def       lowerCAmelCase_      (	lowerCamelCase:							Union[str, Any]    )       ->      List[str]:


      '''simple docstring'''




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

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

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

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

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


if __name__ == "__main__":
      __SCREAMING_SNAKE_CASE     : Tuple      =parse_args()
      main(args)





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


__SCREAMING_SNAKE_CASE     : Tuple      ={
    """configuration_open_llama""": ["""OPEN_LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP""", """OpenLlamaConfig"""],
}

try:
      if not is_sentencepiece_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __SCREAMING_SNAKE_CASE     : Dict      =["""LlamaTokenizer"""]

try:
      if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __SCREAMING_SNAKE_CASE     : int      =["""LlamaTokenizerFast"""]

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __SCREAMING_SNAKE_CASE     : Optional[Any]      =[
          """OpenLlamaForCausalLM""",
          """OpenLlamaModel""",
          """OpenLlamaPreTrainedModel""",
          """OpenLlamaForSequenceClassification""",
      ]


if TYPE_CHECKING:
      from .configuration_open_llama import OPEN_LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP, OpenLlamaConfig

      try:
            if not is_sentencepiece_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from transformers import LlamaTokenizer

      try:
            if not is_tokenizers_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from transformers import LlamaTokenizerFast

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_open_llama import (
                OpenLlamaForCausalLM,
                OpenLlamaForSequenceClassification,
                OpenLlamaModel,
                OpenLlamaPreTrainedModel,
            )


else:
      import sys

      __SCREAMING_SNAKE_CASE     : Union[str, Any]      =_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)