Hiding a prompt in a tokenizer

The Hugging Face tokenizer.json spec allows you to hide a prompt in a post processor. In this post, I show you how that works, and how it improves portability for tokenizers.

Prefix prompts

Prompting, or maybe more technically, setting a system prompt, has become ubiquitous ever since OpenAI allowed users to send their own system prompts in API calls to their models. Maybe less well known is that many regular transformers are trained with simple prompts prefixed to their documents, not as separate messages. For example, intfloat/multilingual-e5-base asks you to prepend query : to queries or general documents, and passage : to documents to be retrieved. These prompts are not inherently part of the model artifact; you need to know to use them, and forgetting to use them reduces performance. ¹ I think this last part is especially salient: forgetting to add a prompt will reduce performance, but in a way that is difficult to notice without running an evaluation. Your model will work fine, it will just not work optimally.

sentence-transformers has good support for managing prompts together with their models: prompts are stored as strings in a sentence-transformers-config.json and then prepended to the document, after which the whole string is tokenized. One potential issue with string concatenation is that tokenizers are allowed to tokenize across the prompt boundary, potentially changing results. ²

However, I think it is common to train models using sentence-transformers, and then serve them using a specialized serving framework, such as fastembed. When doing this, you have to remember to prepend the same prompts in the same way.

I show that it is possible and easy to hide a single prompt in a tokenizer, so that it survives to any framework that correctly ingests the tokenizer.json used by Hugging Face tokenizers. This makes it so that you can directly reuse your model in vLLM, sentence-transformers, fastembed, or a framework of your choosing, without having to think about prompt support at all. I have successfully used this technique to transfer models from training to eval to inference without having to think about the prompt.

How does it work?

In the Hugging Face framework, tokenizers may have a post-processing module. This module is responsible for adding special tokens, such as [CLS] and [SEP] to the token sequence after regular tokenization. While there are a few different post-processors, most can be expressed as variants of a TemplatePostProcessor. ³ As the name implies, a TemplatePostProcessor defines a template, which often looks something like SPECIAL_START $A SPECIAL_END. In this format, $A stands for the sequence that was already tokenized, and to which the template is applied. SPECIAL_START and SPECIAL_END are placeholders, and stand for special tokens which are defined in the TemplatePostProcessor (see below). Note that, because the post-processor is applied after tokenization, any special tokens are always inserted in front of the prompt.

Here’s a concrete example of what the special tokens array looks like:

{
    "SPECIAL_START": {
    "id": "SPECIAL_START",
    "ids": [
        101
    ],
    "tokens": [
        "[CLS]"
    ]
    },
    "SPECIAL_END": {
    "id": "SPECIAL_END",
    "ids": [
        102
    ],
    "tokens": [
        "[SEP]"
    ]
    }
}

As you can see, special tokens can actually consist of multiple IDs and multiple tokens, which is exactly the mechanism we’ll be using to insert a prompt. So concretely, to insert a prompt into a tokenizer, we’ll do the following:

1. Add a new special token to the special tokens array with the tokenized prompt tokens. Let’s say we add it with the key PROMPT
2. Add the PROMPT key to the template: SPECIAL_START $A SPECIAL_END -> SPECIAL_START PROMPT $A SPECIAL_END.

And, that’s it. Now your tokenizer has an embedded prompt that is always embedded. Zero chances of users accidentally ingesting a document without a prompt, or the wrong prompt.

A couple of things to note:

• The prompt is pre-tokenized. This means that you will always get the same IDs for your prompt, regardless of what your document is. It is impossible for tokens to “bleed” into the prompt. This probably also saves you some compute, but that’s likely irrelevant.
• There is no way to disable the prompt and insert special tokens. You can set add_special_tokens to False when encoding, but this also disables the regular special tokens.
• You can only load a single prompt in a single tokenizer. Each different prompt requires an additional tokenizer in memory.

Skeletoken

skeletoken added support for prompt addition in v0.6.0. It works as follows:

from skeletoken import TokenizerModel

t = TokenizerModel.from_pretrained("bert-base-uncased")
t.prompt = "query: "

encoded = t.tokenizer.encode("hello")
print(encoded.tokens)
# ['[CLS]', 'query', ':', 'hello', '[SEP]']

This prompt is now stored on the tokenizer.json. Using it does not require skeletoken at all:

from tokenizers import Tokenizer

t.tokenizer.save("query.json")

tokenizer = Tokenizer.from_file("query.json")
print(tokenizer.encode("hello").tokens)
# ['[CLS]', 'query', ':', 'hello', '[SEP]']

Note that, unlike regular special tokens, the prompt survives decoding. This is because special tokens get removed during decoding. ⁴

from skeletoken import TokenizerModel

t = TokenizerModel.from_pretrained("bert-base-uncased")
t.prompt = "query: "

decoded = t.tokenizer.decode(t.tokenizer.encode("hello").ids)
print(decoded)
# 'query : hello'

Setting a new prompt removes the old one, setting it to None erases the prompt.

from skeletoken import TokenizerModel

t = TokenizerModel.from_pretrained("bert-base-uncased")
t.prompt = "query: "
print(t.tokenizer.encode("hello").tokens)
# ['[CLS]', 'query', ':', 'hello', '[SEP]']
t.prompt = "passage: "
print(t.tokenizer.encode("hello").tokens)
# ['[CLS]', 'passage', ':', 'hello', '[SEP]']
t.prompt = None
print(t.tokenizer.encode("hello").tokens)
# ['[CLS]', 'hello', '[SEP]']

Switching a prompt incurs a full reload of the tokenizer, so whether this is a good strategy to perform on the fly depends on the size of your tokenizer. For small tokenizers (e.g., bert-base-uncased), loading only takes about 30ms on my machine, but for large multilingual tokenizers (e.g., intfloat/multilingual-e5-base), this takes about 500ms.

A better strategy is to just save multiple tokenizers separately.

from skeletoken import TokenizerModel

t = TokenizerModel.from_pretrained("bert-base-uncased")
t.prompt = "query: "
t_query = t.to_tokenizer()

t.prompt = "passage: "
t_passage = t.to_tokenizer()

Although I fully concede that this introduces the exact problem we had at the start of this post. It again requires users and frameworks to keep track of multiple tokenizers and prompts. Still, this might be a preferable situation to storing a string somewhere which a user might forget to use.

Conclusion

So is this worth doing? If your model gets consumed by more than one downstream framework, yes: adding the prompt to the tokenizer costs nothing at inference time and protects you from difficult bugs. It also protects you against prompts slicing into your documents. However, if you need multiple prompts or just use a single framework throughout, it might not be that useful to you. For me, it really works well: I can just dump a model somewhere and know it will reproduce the prompt every time.

Footnotes