usrbinenv python codingutf8 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license script for training a unigram tokenizer import argparse import logging import datasets from tokenizers import tokenizer decoders normalizers pretokenizers processors from tokenizers models import unigram from tokenizers trainers import unigramtrainer from transformers import alberttokenizerfast logger logging getloggername def parseargs parser argparse argumentparserdescriptiontrain a unigram tokenizer on the wikitext dataset parser addargument datasetname typestr defaultwikitext helpname of the training explore datasets at hf codatasets parser addargument datasetconfig typestr defaultwikitext103rawv1 helpconfiguration name of the dataset parser addargument batchsize typeint default1000 helpbatch size during training parser addargument vocabsize typeint default10048 helpsize of the desired vocabulary parser addargument limit defaultnone typeint helplimit the number of shards used for debugging parser addargument exporttohub actionstoretrue args parser parseargs return args def mainargs dataset datasets loaddatasetargs datasetname args datasetconfig splittrain if args limit is not none maxtrainsamples minlendataset args limit dataset dataset selectrangemaxtrainsamples logger infoflimiting the dataset to args limit entries def batchiterator for i in range0 lendataset args batchsize yield dataseti i args batchsizetext prepare the tokenizer tokenizer tokenizerunigram tokenizer normalizer normalizers sequencenormalizers replace normalizers replace tokenizer pretokenizer pretokenizers metaspace prepare the trainer trainer unigramtrainer unktokenunk specialtokenscls sep unk pad mask vocabsizeargs vocabsize logger infotraining the tokenizer tokenizer trainfromiteratorbatchiterator trainertrainer logger infotokenizer training complete clstokenid tokenizer tokentoidcls septokenid tokenizer tokentoidsep tokenizer postprocessor processors templateprocessing singlecls 0 a 0 sep 0 paircls 0 a 0 sep 0 b 1 sep 1 specialtokens cls clstokenid sep septokenid tokenizer decoder decoders metaspace if args exporttohub logger infoexporting the trained tokenzier to hub newtokenizer alberttokenizerfasttokenizerobjecttokenizer newtokenizer pushtohubunigramtokenizerdataset if name main args parseargs mainargs usr bin env python coding utf 8 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license script for training a unigram tokenizer prepare the tokenizer prepare the trainer
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


logger = logging.getLogger(__name__)


def parse_args():
    parser = argparse.ArgumentParser(description="Train a unigram tokenizer on the wikitext dataset.")
    parser.add_argument(
        "--dataset_name",
        type=str,
        default="wikitext",
        help="Name of the training. Explore datasets at: hf.co/datasets.",
    )
    parser.add_argument(
        "--dataset_config", type=str, default="wikitext-103-raw-v1", help="Configuration name of the dataset."
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=1000,
        help="Batch size during training.",
    )
    parser.add_argument(
        "--vocab_size",
        type=int,
        default=10048,
        help="Size of the desired vocabulary.",
    )
    parser.add_argument(
        "--limit",
        default=None,
        type=int,
        help="Limit the number of shards (used for debugging).",
    )
    parser.add_argument(
        "--export_to_hub",
        action="store_true",
    )

    args = parser.parse_args()
    return args


def main(args):
    dataset = datasets.load_dataset(args.dataset_name, args.dataset_config, split="train")

    if args.limit is not None:
        max_train_samples = min(len(dataset), args.limit)
        dataset = dataset.select(range(max_train_samples))
        logger.info(f"Limiting the dataset to {args.limit} entries.")

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

    
    tokenizer = Tokenizer(Unigram())
    tokenizer.normalizer = normalizers.Sequence([normalizers.Replace("``", '"'), normalizers.Replace("''", '"')])
    tokenizer.pre_tokenizer = pre_tokenizers.Metaspace()

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

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

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


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