92 lines
3.1 KiB
Python
92 lines
3.1 KiB
Python
"""ChatML format strategy."""
|
||
|
||
from typing import Dict, List, Tuple
|
||
|
||
from pipeline.tokenize import AutoTokenizer
|
||
from pipeline.strategies.base import PromptStrategy
|
||
from pipeline.strategies.factory import StrategyFactory
|
||
|
||
DEFAULT_CHATML_TEMPLATE = (
|
||
"{% for message in messages %}"
|
||
"{% if message['role'] == 'system' %}"
|
||
"{{ '<|im_start|>system\n' + message['content'] + '<|im_end|>\n' }}"
|
||
"{% elif message['role'] == 'user' %}"
|
||
"{{ '<|im_start|>user\n' + message['content'] + '<|im_end|>\n' }}"
|
||
"{% elif message['role'] == 'assistant' %}"
|
||
"{{ '<|im_start|>assistant\n' + message['content'] + '<|im_end|>\n' }}"
|
||
"{% endif %}"
|
||
"{% endfor %}"
|
||
"{% if add_generation_prompt %}"
|
||
"{{ '<|im_start|>assistant\n' }}"
|
||
"{% endif %}"
|
||
)
|
||
|
||
|
||
@StrategyFactory.register("chatml")
|
||
class ChatMLStrategy(PromptStrategy):
|
||
"""ChatML format strategy.
|
||
|
||
Renders messages using the tokenizer's jinja chat_template from
|
||
``tokenizer_config.json``. Falls back to DEFAULT_CHATML_TEMPLATE
|
||
when no template is configured.
|
||
|
||
The strategy does **not** hard-code any special tokens – all
|
||
formatting is driven by the jinja template.
|
||
"""
|
||
|
||
def __init__(self, tokenizer: AutoTokenizer):
|
||
super().__init__(tokenizer)
|
||
if tokenizer._chat_template is None:
|
||
tokenizer.set_chat_template(DEFAULT_CHATML_TEMPLATE)
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return "chatml"
|
||
|
||
def format_messages(
|
||
self,
|
||
messages: List[Dict[str, str]],
|
||
) -> Tuple[List[int], List[int]]:
|
||
"""Render a single-turn messages conversation.
|
||
|
||
Returns ``(prompt_tokens, response_tokens)`` where
|
||
*prompt_tokens* contains everything up to (and including) the
|
||
last assistant start marker, and *response_tokens* is the
|
||
assistant content plus the closing markers.
|
||
"""
|
||
last_asst = max(
|
||
i for i, m in enumerate(messages) if m["role"] == "assistant"
|
||
)
|
||
|
||
prompt = self.tokenizer.apply_chat_template(
|
||
messages[:last_asst],
|
||
add_generation_prompt=True,
|
||
tokenize=True,
|
||
)
|
||
full = self.tokenizer.apply_chat_template(
|
||
messages[: last_asst + 1],
|
||
add_generation_prompt=False,
|
||
tokenize=True,
|
||
)
|
||
return prompt, full[len(prompt) :]
|
||
|
||
def assemble_prompt(self, query_tokens: List[int]) -> List[int]:
|
||
text = self.tokenizer.decode(query_tokens)
|
||
return self.tokenizer.apply_chat_template(
|
||
[{"role": "user", "content": text}],
|
||
add_generation_prompt=True,
|
||
tokenize=True,
|
||
)
|
||
|
||
def assemble_response(self, response_tokens: List[int]) -> List[int]:
|
||
text = self.tokenizer.decode(response_tokens)
|
||
full = self.tokenizer.apply_chat_template(
|
||
[{"role": "assistant", "content": text}],
|
||
add_generation_prompt=False,
|
||
tokenize=True,
|
||
)
|
||
opening = self.tokenizer.apply_chat_template(
|
||
[], add_generation_prompt=True, tokenize=True
|
||
)
|
||
return full[len(opening) :]
|