| import json
|
| import os
|
| from transformers import PreTrainedTokenizer
|
|
|
|
|
| class NomTokenizer(PreTrainedTokenizer):
|
| vocab_files_names = {'vocab_file': 'vocab.json'}
|
|
|
| def __init__(
|
| self,
|
| vocab_file,
|
| unk_token='<UNK>',
|
| unk_token_id=0,
|
| id_start=1,
|
| **kwargs
|
| ):
|
| self.vocab_file = vocab_file
|
| self.id_start = id_start
|
| self.unk_token = unk_token
|
| self.unk_token_id = unk_token_id
|
| self.pad_token = unk_token
|
| self.pad_token_id = unk_token_id
|
|
|
| with open(vocab_file, 'r', encoding='utf-8') as f:
|
| self.vocab_dict = json.load(f)
|
|
|
| self.char2id = {}
|
| self.id2char = {}
|
| for i, char in enumerate(self.vocab_dict.keys(), start=id_start):
|
| self.char2id[char] = i
|
| self.id2char[i] = char
|
| self.id_to_options = {idx: v for idx, v in enumerate(self.vocab_dict.values(), start=id_start)}
|
|
|
| super().__init__(**kwargs)
|
|
|
| def _tokenize(self, text):
|
| return list(text)
|
|
|
| def _convert_token_to_id(self, token):
|
| return self.char2id.get(token, self.unk_token_id)
|
|
|
| def _convert_id_to_token(self, index):
|
| if index == self.unk_token_id:
|
| return self.unk_token
|
| return self.id2char.get(index, self.unk_token)
|
|
|
| @property
|
| def vocab_size(self):
|
| return len(self.char2id) + 1
|
|
|
| def get_vocab(self):
|
| vocab = {**self.char2id, **self.added_tokens_encoder}
|
| return vocab
|
|
|
| def save_vocabulary(self, save_directory, filename_prefix=None):
|
| if filename_prefix:
|
| vocab_file = os.path.join(save_directory, f'{filename_prefix}-vocab.json')
|
| else:
|
| vocab_file = os.path.join(save_directory, 'vocab.json')
|
|
|
| with open(vocab_file, 'w', encoding='utf-8') as f:
|
| json.dump(self.vocab_dict, f, ensure_ascii=False)
|
|
|
| return (vocab_file,)
|
|
|
