Adding a new token to a transformer model without breaking tokenization of subwords

Question

I'm running an experiment investigating the internal structure of large pre-trained models (BERT and RoBERTa, to be specific). Part of this experiment involves fine-tuning the models on a made-up new word in a specific sentential context and observing its predictions for that novel word in other contexts post-tuning. Because I am just trying to teach it a new word, we freeze the embeddings for the other words during fine-tuning so that only the weights for the new word are updated. This means that I would like for everything to be treated as if it were "normal," except for adding the new word to the model's vocabulary.

I've added the new word to the model and tokenizer like in this MWE (in the case of BERT):

from transformers import BertTokenizer, BertForMaskedLM
new_words = ['myword1', 'myword2']
model = BertForMaskedLM.from_pretrained('bert-base-uncased')
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', do_basic_tokenize = False)
tokenizer.tokenize('myword1 myword2') 
# verify the words do not already exist in the vocabulary
# result: ['my', '##word', '##1', 'my', '##word', '##2']
    
tokenizer.add_tokens(new_words)
model.resize_token_embeddings(len(tokenizer))
    
tokenizer.tokenize('myword1 myword2')
# result: ['myword1', 'myword2']
   
new_tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', do_basic_tokenize = False)
    
new_tokenizer.tokenize('the period is a subword.')
# result: ['the', 'period', 'is', 'a', 'sub', '##word', '##.']
    
tokenizer.tokenize('but not when it follows myword1.')
# result: ['but', 'not', 'when', 'it', 'follows', 'myword1', '.']

How can I add a new token and have it behave correctly (i.e., by preserving the correct subword tokenization of adjacent strings)? Similar issues happen with RoBERTa, where the following word does not appear to be tokenized correctly (it is tokenized without the 'Ġ' that indicates a preceding space, which is present when the new word is replaced with an existing token). (This Ġ is also not present on the added token, but I assume that as long as the added token never occurs at the beginning of a string, it wouldn't matter.)

Edit: After poking around a bit more, I've found that this is related to my setting do_basic_tokenize=False. If it is not set to false, the results come out as expected. Nevertheless, I'd prefer to keep that set to false if there's a way to fix this.

Answer 1

You can add the tokens as special tokens, similar to [SEP] or [CLS] using the add_special_tokens method. There will be separated during pre-tokenization and not passed further for tokenization.

Answer 2

Reposting the solution I came up with here after first posting it on Stack Overflow, in case anyone else finds it helpful. I originally posted this here.

———————————————————————————————————————————————————————

After continuing to try and figure this out, I seem to have found something that might work. It's not necessarily generalizable, but one can load a tokenizer from a vocabulary file (+ a merges file for RoBERTa). If you manually edit those files to add the new tokens in the right way, everything seems to work as expected. Here's an example for BERT:

from transformers import BertTokenizer

bert = BertTokenizer.from_pretrained('bert-base-uncased', do_basic_tokenize=False)
bert.tokenize('testing.') # ['testing', '##.']
bert.tokenize('mynewword') # ['my', '##ne', '##w', '##word']

bert_vocab = bert.get_vocab() # get the pretrained tokenizer's vocabulary
bert_vocab.update({'mynewword' : len(bert_vocab)}) # add the new word to the end

with open('vocab.tmp', 'w', encoding = 'utf-8') as tmp_vocab_file:
    tmp_vocab_file.write('\n'.join(bert_vocab))
    
new_bert = BertTokenizer(name_or_path = 'bert-base-uncased', vocab_file = 'vocab.tmp', do_basic_tokenize=False)
new_bert.max_model_length = 512 # for identity to this setting on the pretrained one

new_bert.tokenize('mynewword') # ['mynewword']
new_bert.tokenize('mynewword.') # ['mynewword', '##.']

import os
os.remove('vocab.tmp') # cleanup

RoBERTa is much harder since we also have to add the pairs to merges.txt. I have a way of doing this that works for the new tokens, but unfortunately it can affect tokenization of words that are subparts of the new tokens, so it's not perfect—if one is using this to add made up words (as in my use case), you can just choose strings that are unlikely to cause problems (unlike the example here of 'mynewword'), but in other cases it is likely to cause problems. (While it's not a perfect solution, hopefully it might get others to see a better one.)

import re
import json
import requests
from transformers import RobertaTokenizer

roberta = RobertaTokenizer.from_pretrained('roberta-base')
roberta.tokenize('testing a') # ['testing', 'Ġa']
roberta.tokenize('mynewword') # ['my', 'new', 'word']

# update the vocabulary with the new token and the 'Ġ'' version
roberta_vocab = roberta.get_vocab()
roberta_vocab.update({'mynewword' : len(roberta_vocab)}) 
roberta_vocab.update({chr(288) + 'mynewword' : len(roberta_vocab)}) # chr(288) = 'Ġ'
with open('vocab.tmp', 'w', encoding = 'utf-8') as tmp_vocab_file:
    json.dump(roberta_vocab, tmp_vocab_file, ensure_ascii=False)

# get and modify the merges file so that the new token will always be tokenized as a single word
url = 'https://huggingface.co/roberta-base/resolve/main/merges.txt'
roberta_merges = requests.get(url).content.decode().split('\n')

# this is a helper function to loop through a list of new tokens and get the byte-pair encodings
# such that the new token will be treated as a single unit always
def get_roberta_merges_for_new_tokens(new_tokens):
    merges = [gen_roberta_pairs(new_token) for new_token in new_tokens]
    merges = [pair for token in merges for pair in token]
    return merges

def gen_roberta_pairs(new_token, highest = True):
    # highest is used to determine whether we are dealing with the Ġ version or not. 
    # we add those pairs at the end, which is only if highest = True
    
    # this is the hard part...
    chrs = [c for c in new_token] # list of characters in the new token, which we will recursively iterate through to find the BPEs
    
    # the simplest case: add one pair
    if len(chrs) == 2:
        if not highest: 
            return tuple([chrs[0], chrs[1]])
        else:
            return [' '.join([chrs[0], chrs[1]])]
    
    # add the tokenization of the first letter plus the other two letters as an already merged pair
    if len(chrs) == 3:
        if not highest:
            return tuple([chrs[0], ''.join(chrs[1:])])
        else:
            return gen_roberta_pairs(chrs[1:]) + [' '.join([chrs[0], ''.join(chrs[1:])])]
    
    if len(chrs) % 2 == 0:
        pairs = gen_roberta_pairs(''.join(chrs[:-2]), highest = False)
        pairs += gen_roberta_pairs(''.join(chrs[-2:]), highest = False)
        pairs += tuple([''.join(chrs[:-2]), ''.join(chrs[-2:])])
        if not highest:
            return pairs
    else:
        # for new tokens with odd numbers of characters, we need to add the final two tokens before the
        # third-to-last token
        pairs = gen_roberta_pairs(''.join(chrs[:-3]), highest = False)
        pairs += gen_roberta_pairs(''.join(chrs[-2:]), highest = False)
        pairs += gen_roberta_pairs(''.join(chrs[-3:]), highest = False)
        pairs += tuple([''.join(chrs[:-3]), ''.join(chrs[-3:])])
        if not highest:
            return pairs
    
    pairs = tuple(zip(pairs[::2], pairs[1::2]))
    pairs = [' '.join(pair) for pair in pairs]
    
    # pairs with the preceding special token
    g_pairs = []
    for pair in pairs:
        if re.search(r'^' + ''.join(pair.split(' ')), new_token):
            g_pairs.append(chr(288) + pair)
    
    pairs = g_pairs + pairs
    pairs = [chr(288) + ' ' + new_token[0]] + pairs
    
    pairs = list(dict.fromkeys(pairs)) # remove any duplicates
    
    return pairs

# first line of this file is a comment; add the new pairs after it
roberta_merges = roberta_merges[:1] + get_roberta_merges_for_new_tokens(['mynewword']) + roberta_merges[1:]
roberta_merges = list(dict.fromkeys(roberta_merges))
with open('merges.tmp', 'w', encoding = 'utf-8') as tmp_merges_file:
    tmp_merges_file.write('\n'.join(roberta_merges))

new_roberta = RobertaTokenizer(name_or_path='roberta-base', vocab_file='vocab.tmp', merges_file='merges.tmp')

# for some reason, we have to re-add the <mask> token to roberta if we are using it, since
# loading the tokenizer from a file will cause it to be tokenized as separate parts
# the weight matrix is identical, and once re-added, a fill-mask pipeline still identifies
# the mask token correctly (not shown here)
new_roberta.add_tokens(new_roberta.mask_token, special_tokens=True)
new_roberta.model_max_length = 512

new_roberta.tokenize('mynewword') # ['mynewword']
new_roberta.tokenize('mynewword a') # ['mynewword', 'Ġa']
new_roberta.tokenize(' mynewword') # ['Ġmynewword']

# however, this does not guarantee that tokenization of other words will not be affected
roberta.tokenize('mynew') # ['my', 'new']
new_roberta.tokenize('mynew') # ['myne', 'w']

import os
os.remove('vocab.tmp')
os.remove('merges.tmp') # cleanup