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:

1. 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.

2. 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.

3. 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