



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_					(		):

				"""simple docstring"""
				a__ 					=						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=1_000       ,						help="""Batch size during training."""       ,						)
				parser.add_argument(
				    """--vocab_size"""       ,						type=A__       ,						default=10_048       ,						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"""       ,						)

				a__ 					=						parser.parse_args()
				return args

def 	A_					(		__a   : int						):

				"""simple docstring"""
				a__ 					=						datasets.load_dataset(args.dataset_name       ,						args.dataset_config       ,						split="""train"""						)

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

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

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

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


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



import os
import shutil
import unittest
from unittest.mock import patch

from transformers.testing_utils import CaptureStd, is_pt_tf_cross_test, require_torch







class      lowerCamelCase		(  unittest.TestCase  ):





     @patch("sys.argv"				,    ["fakeprogrampath", "env"]   )
     def 		UpperCAmelCase(self      :    Optional[Any]   )						->						Dict:
         # test transformers-cli env
         import transformers.commands.transformers_cli

         with CaptureStd() as cs:
             transformers.commands.transformers_cli.main()
         self.assertIn("Python version"				,    cs.out   )
         self.assertIn("Platform"				,    cs.out   )
         self.assertIn("Using distributed or parallel set-up in script?"				,    cs.out   )





     @is_pt_tf_cross_test
     @patch(
         "sys.argv"				,    ["fakeprogrampath", "pt-to-tf", "--model-name", "hf-internal-testing/tiny-random-gptj", "--no-pr"]   )
     def 		UpperCAmelCase(self      :    Any   )						->						int:
         import transformers.commands.transformers_cli

         shutil.rmtree("/tmp/hf-internal-testing/tiny-random-gptj"				,    ignore_errors=_A   )  # cleans potential past runs
         transformers.commands.transformers_cli.main()

         # The original repo has no TF weights -- if they exist, they were created by the CLI
         self.assertTrue(os.path.exists("/tmp/hf-internal-testing/tiny-random-gptj/tf_model.h5"   )   )





     @require_torch
     @patch("sys.argv"				,    ["fakeprogrampath", "download", "hf-internal-testing/tiny-random-gptj", "--cache-dir", "/tmp"]   )
     def 		UpperCAmelCase(self      :    Dict   )						->						List[Any]:
         import transformers.commands.transformers_cli

         # # remove any previously downloaded model to start clean
         shutil.rmtree("/tmp/models--hf-internal-testing--tiny-random-gptj"				,    ignore_errors=_A   )

         # run the command
         transformers.commands.transformers_cli.main()

         # check if the model files are downloaded correctly on /tmp/models--hf-internal-testing--tiny-random-gptj
         self.assertTrue(os.path.exists("/tmp/models--hf-internal-testing--tiny-random-gptj/blobs"   )   )
         self.assertTrue(os.path.exists("/tmp/models--hf-internal-testing--tiny-random-gptj/refs"   )   )
         self.assertTrue(os.path.exists("/tmp/models--hf-internal-testing--tiny-random-gptj/snapshots"   )   )





     @require_torch
     @patch(
         "sys.argv"				,    [
             "fakeprogrampath",
             "download",
             "hf-internal-testing/test_dynamic_model_with_tokenizer",
             "--trust-remote-code",
             "--cache-dir",
             "/tmp",
         ]				,    )
     def 		UpperCAmelCase(self      :    Any   )						->						List[str]:
         import transformers.commands.transformers_cli

         # # remove any previously downloaded model to start clean
         shutil.rmtree("/tmp/models--hf-internal-testing--test_dynamic_model_with_tokenizer"				,    ignore_errors=_A   )

         # run the command
         transformers.commands.transformers_cli.main()

         # check if the model files are downloaded correctly on /tmp/models--hf-internal-testing--test_dynamic_model_with_tokenizer
         self.assertTrue(os.path.exists("/tmp/models--hf-internal-testing--test_dynamic_model_with_tokenizer/blobs"   )   )
         self.assertTrue(os.path.exists("/tmp/models--hf-internal-testing--test_dynamic_model_with_tokenizer/refs"   )   )
         self.assertTrue(
             os.path.exists("/tmp/models--hf-internal-testing--test_dynamic_model_with_tokenizer/snapshots"   )   )

