from tokenizers import Tokenizer
from tokenizers.models import WordPiece
from tokenizers.pre_tokenizers import Whitespace
from tokenizers.trainers import WordPieceTrainer


def train_tokenizer(files, vocab_size=50000):

    tokenizer = Tokenizer(WordPiece(unk_token=":OOV:"))

    tokenizer.pre_tokenizer = Whitespace()

    trainer = WordPieceTrainer(
        vocab_size=vocab_size,
        special_tokens=[":OOV:"],
        continuing_subword_prefix="##",
    )

    tokenizer.train(files, trainer)

    return tokenizer


def save_tokenizer(tokenizer, path):

    tokenizer.save(path)


def load_tokenizer(path):

    return Tokenizer.from_file(path)