File size: 1,035 Bytes
bd4f849 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 |
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,)
|