PromptTemplate

class torchtune.data.PromptTemplate(template: Dict[Literal['system', 'user', 'assistant', 'ipython'], Tuple[str, str]])

Quickly define a custom prompt template by passing in a dictionary mapping role to the prepend and append tags. 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"),
}

Once instantiated, you must call the prompt template on a list of messages. It will return the same list of messages updated with the template.

Note

Any tags prepended/appended to the assistant message will be included in the loss calculation. All other prepend/append tags for other roles (system, user, ipython) are, in most cases, not included in loss. Consider using the append tags for user messages for tags that need to come before the assistant message but should not be included in loss. For more custom masking and prompt templating, you can create your own class based off the PromptTemplate interface.

Parameters:

template (Dict[Role, Tuple[str, str]]) – a dictionary mapping role to the prepend and append tags