Prompt Templates¶
Prompt templates are structured text templates which are used to format user prompts to optimize model performance on specific tasks. They can serve many purposes:
Model-specific templates that are required whenever the model is prompted, such as the [INST] tags in the instruct-tuned Llama2 and Mistral models. These models were pre-trained with these tags and using them in inference can help ensure optimal performance.
Task-specific templates to gear models for a particular task that it will expect after training. Example include grammar correction (
GrammarErrorCorrectionTemplate), summarization (SummarizeTemplate), question answering (QuestionAnswerTemplate), and more.Community standardized templates, such as
ChatMLTemplate
For example, if I wanted to fine-tune a model to perform a grammar correction task, I could use the GrammarErrorCorrectionTemplate
to add the text “Correct this to standard English: {prompt} — Corrected: {response}” to all my data samples.
from torchtune.data import GrammarErrorCorrectionTemplate, Message
sample = {
"incorrect": "This are a cat",
"correct": "This is a cat.",
}
msgs = [
Message(role="user", content=sample["incorrect"]),
Message(role="assistant", content=sample["correct"]),
]
gec_template = GrammarErrorCorrectionTemplate()
templated_msgs = gec_template(msgs)
for msg in templated_msgs:
print(msg.text_content)
# Correct this to standard English: This are a cat
# ---
# Corrected:
# This is a cat.
The added text is different from special tokens that are added by the model tokenizer. For an extended discussion on the different between prompt templates and special tokens, see Tokenizing prompt templates & special tokens.
Using prompt templates¶
Prompt templates are passed into the tokenizer and will be automatically applied for the dataset you are fine-tuning on. You can pass it in two ways:
A string dotpath to a prompt template class, i.e., “torchtune.models.mistral.MistralChatTemplate” or “path.to.my.CustomPromptTemplate”
A dictionary that maps role to a tuple of strings indicating the text to add before and after the message content
Defining via dotpath string¶
# In code
from torchtune.models.mistral import mistral_tokenizer
m_tokenizer = mistral_tokenizer(
path="/tmp/Mistral-7B-v0.1/tokenizer.model"
prompt_template="torchtune.models.mistral.MistralChatTemplate"
)
# In config
tokenizer:
_component_: torchtune.models.mistral.mistral_tokenizer
path: /tmp/Mistral-7B-v0.1/tokenizer.model
prompt_template: torchtune.models.mistral.MistralChatTemplate
Defining via dictionary¶
For example to achieve the following prompt template:
System: {content}\\n
User: {content}\\n
Assistant: {content}\\n
Tool: {content}\\n
You need to pass in a tuple for each role, where PREPEND_TAG is the string
added before the text content and APPEND_TAG is the string added after.
template = {role: (PREPEND_TAG, APPEND_TAG)}
Thus, the template would be defined as follows:
template = {
"system": ("System: ", "\\n"),
"user": ("User: ", "\\n"),
"assistant": ("Assistant: ", "\\n"),
"ipython": ("Tool: ", "\\n"),
}
Now we can pass it into the tokenizer as a dictionary:
# In code
from torchtune.models.mistral import mistral_tokenizer
template = {
"system": ("System: ", "\\n"),
"user": ("User: ", "\\n"),
"assistant": ("Assistant: ", "\\n"),
"ipython": ("Tool: ", "\\n"),
}
m_tokenizer = mistral_tokenizer(
path="/tmp/Mistral-7B-v0.1/tokenizer.model"
prompt_template=template,
)
# In config
tokenizer:
_component_: torchtune.models.mistral.mistral_tokenizer
path: /tmp/Mistral-7B-v0.1/tokenizer.model
prompt_template:
system:
- "System: "
- "\\n"
user:
- "User: "
- "\\n"
assistant:
- "Assistant: "
- "\\n"
ipython:
- "Tool: "
- "\\n"
If you don’t want to add a prepend/append tag to a role, you can just pass in an empty string “” where needed.
Using the PromptTemplate class¶
A template dictionary can also be passed into PromptTemplate so you can use it as a standalone custom
prompt template class.
from torchtune.data import PromptTemplate
def my_custom_template() -> PromptTemplate:
return PromptTemplate(
template={
"user": ("User: ", "\\n"),
"assistant": ("Assistant: ", "\\n"),
},
)
template = my_custom_template()
msgs = [
Message(role="user", content="Hello world!"),
Message(role="assistant", content="Is AI overhyped?"),
]
templated_msgs = template(msgs)
for msg in templated_msgs:
print(msg.role, msg.text_content)
# user, User: Hello world!
#
# assistant, Assistant: Is AI overhyped?
#
Custom prompt templates¶
For more advanced configuration that doesn’t neatly fall into the PREPEND_TAG content APPEND_TAG
pattern, you can create a new class that inherits from PromptTemplateInterface
and implements the __call__ method.
from torchtune.data import Message
class PromptTemplateInterface(Protocol):
def __call__(
self,
messages: List[Message],
inference: bool = False,
) -> List[Message]:
"""
Format each role's message(s) according to the prompt template
Args:
messages (List[Message]): a single conversation, structured as a list
of :class:`~torchtune.data.Message` objects
inference (bool): Whether the template is being used for inference or not.
Returns:
The formatted list of messages
"""
pass
# Contrived example - make all assistant prompts say "Eureka!"
class EurekaTemplate(PromptTemplateInterface):
def __call__(
self,
messages: List[Message],
inference: bool = False,
) -> List[Message]:
formatted_dialogue = []
for message in messages:
if message.role == "assistant":
content = "Eureka!"
else:
content = message.content
formatted_dialogue.append(
Message(
role=message.role,
content=content,
masked=message.masked,
ipython=message.ipython,
eot=message.eot,
),
)
return formatted_dialogue
template = EurekaTemplate()
msgs = [
Message(role="user", content="Hello world!"),
Message(role="assistant", content="Is AI overhyped?"),
]
templated_msgs = template(msgs)
for msg in templated_msgs:
print(msg.role, msg.text_content)
# user, Hello world!
# assistant, Eureka!
For more examples, you can look at MistralChatTemplate or
Llama2ChatTemplate.
To use this custom template in the tokenizer, you can pass it in via dotpath string:
# In code
from torchtune.models.mistral import mistral_tokenizer
m_tokenizer = mistral_tokenizer(
path="/tmp/Mistral-7B-v0.1/tokenizer.model",
prompt_template="path.to.template.EurekaTemplate",
)
# In config
tokenizer:
_component_: torchtune.models.mistral.mistral_tokenizer
path: /tmp/Mistral-7B-v0.1/tokenizer.model
prompt_template: path.to.template.EurekaTemplate
