#!/usr/bin/env python3

import argparse
from collections.abc import Iterator

from datasets import load_dataset
from tokenizers import Tokenizer
from tokenizers.models import WordLevel
from tokenizers.normalizers import Sequence, NFC, Strip, Lowercase
from tokenizers.pre_tokenizers import Whitespace
from tokenizers.trainers import WordLevelTrainer
from tqdm.auto import tqdm


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument('--vocabulary', type=int, default=75000, help='Vocabulary size')
    parser.add_argument('--batch', type=int, default=1024, help='Batch size')
    args = parser.parse_args()

    dataset = load_dataset('wikitext', 'wikitext-103-raw-v1', split='train+validation+test')

    tokenizer = Tokenizer(WordLevel(unk_token='<unk>'))
    tokenizer.normalizer = Sequence([NFC(), Strip(), Lowercase()])
    tokenizer.pre_tokenizer = Whitespace()

    def batches(batch_size: int) -> Iterator[str]:
        for batch in tqdm(dataset.iter(batch_size=batch_size), desc='Tokenization'):
            yield batch['text']

    trainer = WordLevelTrainer(vocab_size=args.vocabulary,
                               special_tokens=['<s>', '</s>', '<unk>'])

    tokenizer.train_from_iterator(batches(args.batch), trainer=trainer, length=len(dataset))

    tokenizer.save('tokenizer.json', pretty=True)


if __name__ == '__main__':
    main()