PromptTemplate¶
- class torchtune.data.PromptTemplate(template: Dict[Literal['system', 'user', 'assistant', 'ipython'], Tuple[str, str]])[source]¶
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 andAPPEND_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.