from tokenizers import ByteLevelBPETokenizer
from datasets import load_from_disk

import os
import argparse


def parse_arguments(parser):
    
    parser.add_argument(
        "--dataset-dir",
        required=True,
        help="Define dataset folder",
    )

    parser.add_argument(
        "--out",
        default="./",
        type=str,
        help="Path to the output directory, where the files will be saved",
    )

    parser.add_argument(
        "--name", 
        default="bpe-bytelevel", 
        type=str, 
        help="The name of the output vocab files"
    )

    return parser.parse_args()


def main(args):

    if not os.path.isdir(args.dataset_dir):
        raise Exception('--dataset-dir not exists')

    os.makedirs(args.out, exist_ok=True)
   
    print('LOADING %s...' % (args.dataset_dir))
    raw_dataset = load_from_disk(args.dataset_dir)
    print(raw_dataset)

    def batch_iterator(batch_size=10000):
        for i in range(0, len(raw_dataset), batch_size):
            yield raw_dataset[i : i + batch_size]['text']

    tokenizer = ByteLevelBPETokenizer(
        lowercase=False,
    )

    tokenizer.enable_truncation(max_length=512)
    tokenizer.train_from_iterator(batch_iterator(), vocab_size=50265, min_frequency=2, special_tokens=[
        "<s>",
        "<pad>",
        "</s>",
        "<unk>",
        "<mask>",
    ])

    print("SAVING TOKENIZER CONFIG INTO...")
    tokenizer.save("{}/tokenizer.json".format(args.out))
    tokenizer.save_model(args.out, args.name)
    
    print("SO FAR SO GOOD...")


if __name__ == '__main__':
    args = parse_arguments(argparse.ArgumentParser())
    print(args)
    
    try:
        main(args)
    except Exception as excp:
        print(excp)