import json, os
from transformers import PreTrainedTokenizer

class FelaTokenizer(PreTrainedTokenizer):
    vocab_files_names = {"vocab_file": "vocab.json"}
    def __init__(self, vocab_file=None, **kwargs):
        self.vocab = json.load(open(vocab_file)) if vocab_file else {}
        self.itos = {v: k for k, v in self.vocab.items()}
        super().__init__(**kwargs)
    def _tokenize(self, text, **kwargs):
        return list(text)
    def _convert_token_to_id(self, token):
        return self.vocab.get(token, self.vocab.get("<unk>", 23))
    def _convert_id_to_token(self, index):
        return self.itos.get(index, "<unk>")
    def get_vocab(self):
        return dict(self.vocab)
    def vocab_size(self):
        return len(self.vocab)
    def save_vocabulary(self, save_directory, filename_prefix=None):
        fname = os.path.join(save_directory, (filename_prefix + "-" if filename_prefix else "") + "vocab.json")
        with open(fname, "w") as f:
            json.dump(self.get_vocab(), f)
        return (fname,)