core.tokenizers.text.libraries.chat_template#

Module Contents#

Classes#

MegatronTokenizerChatTemplate

Chat template class for Megatron text tokenizers.

API#

class core.tokenizers.text.libraries.chat_template.MegatronTokenizerChatTemplate#

Chat template class for Megatron text tokenizers.

apply_chat_template(
conversation: List[Dict[str, str]],
chat_template: str,
tokenize: Optional[bool] = True,
truncation: Optional[bool] = False,
max_length: Optional[int] = None,
add_generation_prompt: Optional[bool] = False,
) → Union[str, List[int]]#

Applies tokenizer’s chat template to the conversation.

Parameters:

  • conversation (List[Dict[str, str]]) – a list of dicts with “role” and “content” keys, representing the chat history so far. Conversation example: [ {“role”: “system”, “content”: “You are a witty and helpful assistant.”}, {“role”: “user”, “content”: “Hey, what’s a fun fact about octopuses?”}, {“role”: “assistant”, “content”: “Octopuses blood is blue!”}, {“role”: “user”, “content”: “Whoa, why is their blood blue?”}, ]

  • tokenize (bool) – whether to tokenize the output. If False, the output will be a string.

  • truncation (bool) – whether to truncate sequences at the maximum length. Has no effect if tokenize is False.

  • max_length (int) – maximum length to use fro truncation. Has no effect if tokenize is False.

  • add_generation_prompt (bool) – If this is set, a prompt with the token(s) that indicate the start of an assistant message will be appended to the formatted output. This is useful when you want to generate a response from the model. Note that this argument will be passed to the chat template, and so it must be supported in the template for this argument to have any effect.