|
|
| 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,) |
|
|
