Using Hugginface Transformers and Tokenizers with a fixed vocabulary? - huggingface-transformers

Question

I have a special, non-language use case using a fixed vocabulary—i.e., a relatively small set of generated tokens that represent the entire vocabulary of our "language." I’d like to be able to use this with any of the different models and I’m wondering what would be the best approach? It’s just a vocab.txt file of short strings, which I don’t think will work with any of the BPE tokenizers. Am I correct in that assumption? Also, is there a way to “force” a vocabulary onto any of the tokenizers?
To clarify, our "language" uses prefixes to identify certain types of tokens, which have certain functions in the overall syntax. We want to be able to mask by type during inference, both on input and as part of the selection process, for example, by limiting top-k or top-p sampling to a give type. With a fixed/hand-tuned vocabulary we can be very specific about which ids, or how many ids we need; i.e., we know which tokens are used by each type, so we can mask/filter accordingly. However, with BPE tokenization a given type may be tokenized with any number of tokens, making this process much less straightforward.
The motivation is just to make life easier by fitting into the Huggingface universe a little better, so we can experiment with off-the-shelf models more fluently. We already have this working using the standard BertTokenizer with both GPT2 and RoBERTa, but it would be nice to be able to experiment with different Huggingface models "out of the box," so to speak (using Trainers, Pipelines, etc.). With the BertTokenizer we just load our vocab.txt and we're done, so I wondered whether there would be some way of doing this with the other tokenizers (really, the BPE ones are the only issue, at this point).
It seems to me that being able specify a vocab for any tokenizer would be more straightforward than getting our tokenizer working with other models. Though perhaps a better approach would be to look at streamlining that process? I suppose I could fork and modify AutoTokenizer... ??
Any help much appreciated.

Answer 1

As far as I understand the solution below might help you, as you can use this tokenizer, like you would the other pre-trained ones.
As I do not really understand all the inner workings of the tokenizer, I may very well be off with this solution, but hopefully it can help someone.
The main idea is to subclass the PreTrainedTokenizer. This way, you should only override some of the key methods like _tokenize, _convert_token_to_id, etc..., which are more straightforward than implementing a whole new tokenizer.
import json
from pathlib import Path
from typing import Optional, Tuple, Dict
from transformers import PreTrainedTokenizer
class FixedVocabTokenizer(PreTrainedTokenizer):
def __init__(self, vocab: Dict[str, int], max_len: int = None):
super().__init__(max_len=max_len)
self.__token_ids = vocab
self.__id_tokens: Dict[int, str] = {value: key for key, value in vocab.items()}
def _tokenize(self, text: str, **kwargs):
return text.split(' ')
def _convert_token_to_id(self, token: str) -> int:
return self.__token_ids[token] if token in self.__token_ids else self.unk_token_id
def _convert_id_to_token(self, index: int) -> str:
return self.__id_tokens[index] if index in self.__id_tokens else self.unk_token
def get_vocab(self) -> Dict[str, int]:
return self.__token_ids.copy()
def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
if filename_prefix is None:
filename_prefix = ''
vocab_path = Path(save_directory, filename_prefix + 'vocab.json')
json.dump(self.__token_ids, open(vocab_path, 'w'))
return str(vocab_path),
#property
def vocab_size(self) -> int:
return len(self.__token_ids)
if __name__ == '__main__':
# your custom, fixed vocabulary
custom_vocab = {
'[UNK]': 0,
'word0': 1,
'word1': 2,
'word2': 3,
'word3': 4,
'word4': 5,
'word5': 6,
'[CLS]': 7,
'[SEP]': 8,
'[PAD]': 9
}
model_max_len = 8
tokenizer = FixedVocabTokenizer(custom_vocab, max_len=model_max_len)
# tell your tokenizer about your special tokens
tokenizer.add_special_tokens({
'unk_token': '[UNK]',
'pad_token': '[PAD]',
'cls_token': '[CLS]',
'sep_token': '[SEP]'
})
res = tokenizer(
[
'word1 word2 word word1 word3',
'word2 word0 word0 word3 word5 word4 word2 word1 word0'
],
padding=True,
truncation=True
)
# the result should look like something like this
# res -> BatchEncoding(
# data: {
# 'input_ids': [[2, 3, 0, 2, 4, 9, 9, 9], [3, 1, 1, 4, 6, 5, 3, 2]],
# 'attention_mask': [[1, 1, 1, 1, 1, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1]],
# ...
# },
# ...
# )
This is the solution I could come up with, however I could not figure out if you could do something similar with PreTrainedTokenizerFast. So one more note being, that you can only use slow tokenizers using this method.