mirror of
https://github.com/zylon-ai/private-gpt.git
synced 2025-12-22 17:05:41 +01:00
6fd3a23daf
This commit introduces several improvements to the prompt formatting logic in `private_gpt/components/llm/prompt_helper.py`:
1. **Llama3PromptStyle**:
* Implemented tool handling capabilities, allowing for the formatting of tool call and tool result messages within the Llama 3 prompt structure.
* Ensured correct usage of BOS, EOT, and other Llama 3 specific tokens.
2. **MistralPromptStyle**:
* Refactored the `_messages_to_prompt` method for more robust handling of various conversational scenarios, including consecutive user messages and initial assistant messages.
* Ensured correct application of `<s>`, `</s>`, and `[INST]` tags.
3. **ChatMLPromptStyle**:
* Corrected the logic for handling system messages to prevent duplication and ensure accurate ChatML formatting (`<|im_start|>role\ncontent<|im_end|>`).
4. **TagPromptStyle**:
* Addressed a FIXME comment by incorporating `<s>` (BOS) and `</s>` (EOS) tokens, making it more suitable for Llama-based models like Vigogne.
* Fixed a minor bug related to enum string conversion.
5. **Unit Tests**:
* Added a new test suite in `tests/components/llm/test_prompt_helper.py`.
* These tests provide comprehensive coverage for all modified prompt styles, verifying correct prompt generation for various inputs, edge cases, and special token placements.
These changes improve the correctness, robustness, and feature set of the supported prompt styles, leading to better compatibility and interaction with the respective language models.
592 lines
28 KiB
Python
592 lines
28 KiB
Python
import pytest
|
|
from llama_index.core.llms import ChatMessage, MessageRole
|
|
|
|
from private_gpt.components.llm.prompt_helper import (
|
|
Llama3PromptStyle,
|
|
MistralPromptStyle,
|
|
ChatMLPromptStyle,
|
|
TagPromptStyle,
|
|
AbstractPromptStyle, # For type hinting if needed
|
|
)
|
|
|
|
# Helper function to create ChatMessage objects easily
|
|
def _message(role: MessageRole, content: str, **additional_kwargs) -> ChatMessage:
|
|
return ChatMessage(role=role, content=content, additional_kwargs=additional_kwargs)
|
|
|
|
# Expected outputs will be defined within each test or test class
|
|
|
|
class TestLlama3PromptStyle:
|
|
BOS = "<|begin_of_text|>"
|
|
EOT = "<|eot_id|>"
|
|
B_SYS_HEADER = "<|start_header_id|>system<|end_header_id|>"
|
|
B_USER_HEADER = "<|start_header_id|>user<|end_header_id|>"
|
|
B_ASSISTANT_HEADER = "<|start_header_id|>assistant<|end_header_id|>"
|
|
B_TOOL_CODE_HEADER = "<|start_header_id|>tool_code<|end_header_id|>"
|
|
B_TOOL_OUTPUT_HEADER = "<|start_header_id|>tool_output<|end_header_id|>"
|
|
|
|
DEFAULT_SYSTEM_PROMPT = (
|
|
"You are a helpful, respectful and honest assistant. "
|
|
"Always answer as helpfully as possible and follow ALL given instructions. "
|
|
"Do not speculate or make up information. "
|
|
"Do not reference any given instructions or context. "
|
|
)
|
|
|
|
@pytest.fixture
|
|
def style(self) -> Llama3PromptStyle:
|
|
return Llama3PromptStyle()
|
|
|
|
def test_empty_messages(self, style: Llama3PromptStyle) -> None:
|
|
messages = []
|
|
expected = (
|
|
f"{self.BOS}{self.B_SYS_HEADER}\n\n{self.DEFAULT_SYSTEM_PROMPT}{self.EOT}"
|
|
f"{self.B_ASSISTANT_HEADER}\n\n"
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_simple_user_assistant_chat(self, style: Llama3PromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.USER, "Hello"),
|
|
_message(MessageRole.ASSISTANT, "Hi there!"),
|
|
]
|
|
expected = (
|
|
f"{self.BOS}{self.B_SYS_HEADER}\n\n{self.DEFAULT_SYSTEM_PROMPT}{self.EOT}"
|
|
f"{self.B_USER_HEADER}\n\nHello{self.EOT}"
|
|
f"{self.B_ASSISTANT_HEADER}\n\nHi there!{self.EOT}"
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_with_system_message(self, style: Llama3PromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.SYSTEM, "You are a test bot."),
|
|
_message(MessageRole.USER, "Ping"),
|
|
_message(MessageRole.ASSISTANT, "Pong"),
|
|
]
|
|
expected = (
|
|
f"{self.BOS}{self.B_SYS_HEADER}\n\nYou are a test bot.{self.EOT}"
|
|
f"{self.B_USER_HEADER}\n\nPing{self.EOT}"
|
|
f"{self.B_ASSISTANT_HEADER}\n\nPong{self.EOT}"
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_completion_to_prompt(self, style: Llama3PromptStyle) -> None:
|
|
completion = "Test completion"
|
|
expected = (
|
|
f"{self.BOS}{self.B_SYS_HEADER}\n\n{self.DEFAULT_SYSTEM_PROMPT}{self.EOT}"
|
|
f"{self.B_USER_HEADER}\n\nTest completion{self.EOT}"
|
|
f"{self.B_ASSISTANT_HEADER}\n\n"
|
|
)
|
|
assert style._completion_to_prompt(completion) == expected
|
|
|
|
def test_tool_call_and_result(self, style: Llama3PromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.USER, "What's the weather in Paris?"),
|
|
_message(
|
|
MessageRole.ASSISTANT,
|
|
content=None, # LlamaIndex might put tool call details here, or just use additional_kwargs
|
|
additional_kwargs={"type": "tool_call", "tool_call_id": "123", "name": "get_weather", "arguments": '{"location": "Paris"}'}
|
|
),
|
|
_message(
|
|
MessageRole.TOOL,
|
|
content='{"temperature": "20C"}',
|
|
additional_kwargs={"type": "tool_result", "tool_call_id": "123", "name": "get_weather"}
|
|
),
|
|
]
|
|
# Note: The current Llama3PromptStyle implementation uses message.content for tool call/result content.
|
|
# If additional_kwargs are used to structure tool calls (like OpenAI), the style needs to be adapted.
|
|
# For this test, we assume content holds the direct string for tool_code and tool_output.
|
|
# Let's adjust the messages based on current implementation that uses .content for tool_code/output
|
|
messages_for_current_impl = [
|
|
_message(MessageRole.USER, "What's the weather in Paris?"),
|
|
_message(
|
|
MessageRole.ASSISTANT,
|
|
content='get_weather({"location": "Paris"})', # Simplified tool call content
|
|
additional_kwargs={"type": "tool_call"}
|
|
),
|
|
_message(
|
|
MessageRole.TOOL,
|
|
content='{"temperature": "20C"}',
|
|
additional_kwargs={"type": "tool_result"} # No specific tool_call_id or name used by current style from additional_kwargs
|
|
),
|
|
]
|
|
expected = (
|
|
f"{self.BOS}{self.B_SYS_HEADER}\n\n{self.DEFAULT_SYSTEM_PROMPT}{self.EOT}"
|
|
f"{self.B_USER_HEADER}\n\nWhat's the weather in Paris?{self.EOT}"
|
|
f"{self.B_TOOL_CODE_HEADER}\n\nget_weather({{\"location\": \"Paris\"}}){self.EOT}" # Content is 'get_weather(...)'
|
|
f"{self.B_TOOL_OUTPUT_HEADER}\n\n{{\"temperature\": \"20C\"}}{self.EOT}"
|
|
f"{self.B_ASSISTANT_HEADER}\n\n" # Assistant should respond after tool result
|
|
)
|
|
assert style._messages_to_prompt(messages_for_current_impl) == expected
|
|
|
|
def test_multiple_interactions_with_tools(self, style: Llama3PromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.USER, "Can you search for prompt engineering techniques?"),
|
|
_message(MessageRole.ASSISTANT, content="Okay, I will search for that.", additional_kwargs={}), # Normal assistant message
|
|
_message(MessageRole.ASSISTANT, content='search_web({"query": "prompt engineering techniques"})', additional_kwargs={"type": "tool_call"}),
|
|
_message(MessageRole.TOOL, content='[Result 1: ...]', additional_kwargs={"type": "tool_result"}),
|
|
_message(MessageRole.ASSISTANT, content="I found one result. Should I look for more?", additional_kwargs={}),
|
|
_message(MessageRole.USER, "Yes, please find another one."),
|
|
]
|
|
expected = (
|
|
f"{self.BOS}{self.B_SYS_HEADER}\n\n{self.DEFAULT_SYSTEM_PROMPT}{self.EOT}"
|
|
f"{self.B_USER_HEADER}\n\nCan you search for prompt engineering techniques?{self.EOT}"
|
|
f"{self.B_ASSISTANT_HEADER}\n\nOkay, I will search for that.{self.EOT}"
|
|
f"{self.B_TOOL_CODE_HEADER}\n\nsearch_web({{\"query\": \"prompt engineering techniques\"}}){self.EOT}"
|
|
f"{self.B_TOOL_OUTPUT_HEADER}\n\n[Result 1: ...]{self.EOT}"
|
|
f"{self.B_ASSISTANT_HEADER}\n\nI found one result. Should I look for more?{self.EOT}"
|
|
f"{self.B_USER_HEADER}\n\nYes, please find another one.{self.EOT}"
|
|
f"{self.B_ASSISTANT_HEADER}\n\n"
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_ending_with_user_message(self, style: Llama3PromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.USER, "First message"),
|
|
_message(MessageRole.ASSISTANT, "First response"),
|
|
_message(MessageRole.USER, "Second message, expecting response"),
|
|
]
|
|
expected = (
|
|
f"{self.BOS}{self.B_SYS_HEADER}\n\n{self.DEFAULT_SYSTEM_PROMPT}{self.EOT}"
|
|
f"{self.B_USER_HEADER}\n\nFirst message{self.EOT}"
|
|
f"{self.B_ASSISTANT_HEADER}\n\nFirst response{self.EOT}"
|
|
f"{self.B_USER_HEADER}\n\nSecond message, expecting response{self.EOT}"
|
|
f"{self.B_ASSISTANT_HEADER}\n\n"
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_ending_with_tool_result(self, style: Llama3PromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.USER, "Find info on X."),
|
|
_message(MessageRole.ASSISTANT, content='search({"topic": "X"})', additional_kwargs={"type": "tool_call"}),
|
|
_message(MessageRole.TOOL, content="Info about X found.", additional_kwargs={"type": "tool_result"}),
|
|
]
|
|
expected = (
|
|
f"{self.BOS}{self.B_SYS_HEADER}\n\n{self.DEFAULT_SYSTEM_PROMPT}{self.EOT}"
|
|
f"{self.B_USER_HEADER}\n\nFind info on X.{self.EOT}"
|
|
f"{self.B_TOOL_CODE_HEADER}\n\nsearch({{\"topic\": \"X\"}}){self.EOT}"
|
|
f"{self.B_TOOL_OUTPUT_HEADER}\n\nInfo about X found.{self.EOT}"
|
|
f"{self.B_ASSISTANT_HEADER}\n\n"
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_message_with_empty_content(self, style: Llama3PromptStyle) -> None:
|
|
# Llama3PromptStyle skips messages with None content, but not necessarily empty string.
|
|
# Let's test with an empty string for user, and None for assistant (which should be skipped)
|
|
messages = [
|
|
_message(MessageRole.USER, ""), # Empty string content
|
|
_message(MessageRole.ASSISTANT, None), # None content, should be skipped
|
|
_message(MessageRole.USER, "Follow up")
|
|
]
|
|
expected = (
|
|
f"{self.BOS}{self.B_SYS_HEADER}\n\n{self.DEFAULT_SYSTEM_PROMPT}{self.EOT}"
|
|
f"{self.B_USER_HEADER}\n\n{self.EOT}" # Empty content for user
|
|
f"{self.B_USER_HEADER}\n\nFollow up{self.EOT}" # Assistant message was skipped
|
|
f"{self.B_ASSISTANT_HEADER}\n\n"
|
|
)
|
|
# The style's loop: `if not message or message.content is None: continue`
|
|
# An empty string `""` is not `None`, so it should be included.
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_system_message_not_first(self, style: Llama3PromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.USER, "Hello"),
|
|
_message(MessageRole.SYSTEM, "System message in the middle (unusual)."),
|
|
_message(MessageRole.ASSISTANT, "Hi there!"),
|
|
]
|
|
# The current implementation processes system messages whenever they appear.
|
|
# If a system message appears, it sets `has_system_message = True`.
|
|
# If NO system message appears, a default one is prepended.
|
|
# If one DOES appear, it's used, and default is not prepended.
|
|
expected = (
|
|
f"{self.BOS}"
|
|
# Default system prompt is NOT added because a system message IS present.
|
|
f"{self.B_USER_HEADER}\n\nHello{self.EOT}"
|
|
f"{self.B_SYS_HEADER}\n\nSystem message in the middle (unusual).{self.EOT}"
|
|
f"{self.B_ASSISTANT_HEADER}\n\nHi there!{self.EOT}"
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
|
|
class TestMistralPromptStyle:
|
|
@pytest.fixture
|
|
def style(self) -> MistralPromptStyle:
|
|
return MistralPromptStyle()
|
|
|
|
def test_empty_messages(self, style: MistralPromptStyle) -> None:
|
|
messages = []
|
|
# The refactored version should produce an empty string if no instructions are pending.
|
|
# Or, if it were to prompt for something, it might be "<s>[INST] [/INST]" or just ""
|
|
# Based on current refactored code: if current_instruction_parts is empty, it returns prompt ("")
|
|
assert style._messages_to_prompt(messages) == ""
|
|
|
|
def test_simple_user_assistant_chat(self, style: MistralPromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.USER, "Hello"),
|
|
_message(MessageRole.ASSISTANT, "Hi there!"),
|
|
]
|
|
expected = "<s>[INST] Hello [/INST] Hi there!</s>"
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_with_system_message(self, style: MistralPromptStyle) -> None:
|
|
# System messages are treated like user messages in the current Mistral impl
|
|
messages = [
|
|
_message(MessageRole.SYSTEM, "You are helpful."),
|
|
_message(MessageRole.USER, "Ping"),
|
|
_message(MessageRole.ASSISTANT, "Pong"),
|
|
]
|
|
expected = "<s>[INST] You are helpful.\nPing [/INST] Pong</s>"
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_completion_to_prompt(self, style: MistralPromptStyle) -> None:
|
|
completion = "Test completion"
|
|
# This will call _messages_to_prompt with [ChatMessage(role=USER, content="Test completion")]
|
|
expected = "<s>[INST] Test completion [/INST]"
|
|
assert style._completion_to_prompt(completion) == expected
|
|
|
|
def test_consecutive_user_messages(self, style: MistralPromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.USER, "First part."),
|
|
_message(MessageRole.USER, "Second part."),
|
|
_message(MessageRole.ASSISTANT, "Understood."),
|
|
]
|
|
expected = "<s>[INST] First part.\nSecond part. [/INST] Understood.</s>"
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_ending_with_user_message(self, style: MistralPromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.USER, "Hi"),
|
|
_message(MessageRole.ASSISTANT, "Hello"),
|
|
_message(MessageRole.USER, "How are you?"),
|
|
]
|
|
# Note: The previous prompt had "<s>[INST] Hi [/INST] Hello</s>"
|
|
# The new user message should start a new <s>[INST] block if prompt was not empty.
|
|
# Current logic: bos_token = "<s>" if not prompt else ""
|
|
# Since prompt is not empty after "Hello</s>", bos_token will be "".
|
|
# This might be an issue with the current Mistral refactor if strict BOS per turn is needed.
|
|
# The current refactored code for Mistral:
|
|
# prompt += bos_token + "[INST] " + "\n".join(current_instruction_parts) + " [/INST]"
|
|
# If prompt = "<s>[INST] Hi [/INST] Hello</s>", then bos_token is "", so it becomes:
|
|
# "<s>[INST] Hi [/INST] Hello</s>[INST] How are you? [/INST]" -> This seems correct for continued conversation.
|
|
expected = "<s>[INST] Hi [/INST] Hello</s>[INST] How are you? [/INST]"
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_initial_assistant_message_skipped(self, style: MistralPromptStyle, caplog) -> None:
|
|
messages = [
|
|
_message(MessageRole.ASSISTANT, "I speak first!"),
|
|
_message(MessageRole.USER, "Oh, hello there."),
|
|
]
|
|
# The first assistant message should be skipped with a warning.
|
|
# The prompt should then start with the user message.
|
|
expected = "<s>[INST] Oh, hello there. [/INST]"
|
|
with caplog.at_level("WARNING"):
|
|
assert style._messages_to_prompt(messages) == expected
|
|
assert "MistralPromptStyle: First message is from assistant, skipping." in caplog.text
|
|
|
|
def test_multiple_assistant_messages_in_a_row(self, style: MistralPromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.USER, "User message"),
|
|
_message(MessageRole.ASSISTANT, "Assistant first response."),
|
|
_message(MessageRole.ASSISTANT, "Assistant second response (after no user message)."),
|
|
]
|
|
# current_instruction_parts will be empty when processing the second assistant message.
|
|
# The logic is:
|
|
# if current_instruction_parts: prompt += bos_token + "[INST] " + "\n".join(current_instruction_parts) + " [/INST]"
|
|
# prompt += " " + content + "</s>"
|
|
# So, it will correctly append the second assistant message without a new [INST]
|
|
expected = ("<s>[INST] User message [/INST] Assistant first response.</s>"
|
|
" Assistant second response (after no user message).</s>")
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_system_user_assistant_alternating(self, style: MistralPromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.SYSTEM, "System setup."),
|
|
_message(MessageRole.USER, "User query 1."),
|
|
_message(MessageRole.ASSISTANT, "Assistant answer 1."),
|
|
_message(MessageRole.USER, "User query 2."), # System messages are part of INST with user
|
|
_message(MessageRole.ASSISTANT, "Assistant answer 2."),
|
|
]
|
|
expected = ("<s>[INST] System setup.\nUser query 1. [/INST] Assistant answer 1.</s>"
|
|
"[INST] User query 2. [/INST] Assistant answer 2.</s>")
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_empty_content_messages(self, style: MistralPromptStyle, caplog) -> None:
|
|
messages = [
|
|
_message(MessageRole.USER, "Hello"),
|
|
_message(MessageRole.USER, None), # Skipped by `if not content and message.role != MessageRole.ASSISTANT:`
|
|
_message(MessageRole.USER, ""), # Skipped by the same logic
|
|
_message(MessageRole.ASSISTANT, "Hi"),
|
|
_message(MessageRole.ASSISTANT, ""), # Empty assistant message, kept
|
|
_message(MessageRole.ASSISTANT, None),# Empty assistant message, kept (content becomes "")
|
|
]
|
|
# The refactored code skips empty non-assistant messages.
|
|
# Empty assistant messages (content="" or content=None) are kept.
|
|
expected = ("<s>[INST] Hello [/INST] Hi</s>"
|
|
" </s>" # From assistant with content=""
|
|
" </s>") # From assistant with content=None (becomes "")
|
|
|
|
with caplog.at_level("DEBUG"): # The skipping messages are logged at DEBUG level
|
|
actual = style._messages_to_prompt(messages)
|
|
assert actual == expected
|
|
# Check that specific debug messages for skipping are present
|
|
assert "Skipping empty non-assistant message." in caplog.text # For the None and "" user messages
|
|
|
|
|
|
class TestChatMLPromptStyle:
|
|
IM_START = "<|im_start|>"
|
|
IM_END = "<|im_end|>"
|
|
|
|
@pytest.fixture
|
|
def style(self) -> ChatMLPromptStyle:
|
|
return ChatMLPromptStyle()
|
|
|
|
def test_empty_messages(self, style: ChatMLPromptStyle) -> None:
|
|
messages = []
|
|
# Expected: just the final assistant prompt
|
|
expected = f"{self.IM_START}assistant\n"
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_simple_user_assistant_chat(self, style: ChatMLPromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.USER, "Hello"),
|
|
_message(MessageRole.ASSISTANT, "Hi there!"),
|
|
]
|
|
expected = (
|
|
f"{self.IM_START}user\nHello{self.IM_END}\n"
|
|
f"{self.IM_START}assistant\nHi there!{self.IM_END}\n"
|
|
f"{self.IM_START}assistant\n"
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_with_system_message(self, style: ChatMLPromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.SYSTEM, "You are ChatML bot."),
|
|
_message(MessageRole.USER, "Ping"),
|
|
_message(MessageRole.ASSISTANT, "Pong"),
|
|
]
|
|
expected = (
|
|
f"{self.IM_START}system\nYou are ChatML bot.{self.IM_END}\n"
|
|
f"{self.IM_START}user\nPing{self.IM_END}\n"
|
|
f"{self.IM_START}assistant\nPong{self.IM_END}\n"
|
|
f"{self.IM_START}assistant\n"
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_completion_to_prompt(self, style: ChatMLPromptStyle) -> None:
|
|
completion = "Test user input"
|
|
# This will call _messages_to_prompt with [ChatMessage(role=USER, content="Test user input")]
|
|
expected = (
|
|
f"{self.IM_START}user\nTest user input{self.IM_END}\n"
|
|
f"{self.IM_START}assistant\n"
|
|
)
|
|
assert style._completion_to_prompt(completion) == expected
|
|
|
|
def test_multiple_turns(self, style: ChatMLPromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.USER, "First user message."),
|
|
_message(MessageRole.ASSISTANT, "First assistant response."),
|
|
_message(MessageRole.USER, "Second user message."),
|
|
_message(MessageRole.ASSISTANT, "Second assistant response."),
|
|
]
|
|
expected = (
|
|
f"{self.IM_START}user\nFirst user message.{self.IM_END}\n"
|
|
f"{self.IM_START}assistant\nFirst assistant response.{self.IM_END}\n"
|
|
f"{self.IM_START}user\nSecond user message.{self.IM_END}\n"
|
|
f"{self.IM_START}assistant\nSecond assistant response.{self.IM_END}\n"
|
|
f"{self.IM_START}assistant\n"
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_message_with_empty_content(self, style: ChatMLPromptStyle) -> None:
|
|
# ChatML typically includes messages even with empty content.
|
|
messages = [
|
|
_message(MessageRole.USER, "Hello"),
|
|
_message(MessageRole.ASSISTANT, ""), # Empty string content
|
|
_message(MessageRole.USER, "Follow up")
|
|
]
|
|
expected = (
|
|
f"{self.IM_START}user\nHello{self.IM_END}\n"
|
|
f"{self.IM_START}assistant\n{self.IM_END}\n" # Empty content for assistant
|
|
f"{self.IM_START}user\nFollow up{self.IM_END}\n"
|
|
f"{self.IM_START}assistant\n"
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_message_with_none_content(self, style: ChatMLPromptStyle) -> None:
|
|
# ChatML typically includes messages even with empty content (None becomes empty string).
|
|
messages = [
|
|
_message(MessageRole.USER, "Hello"),
|
|
_message(MessageRole.ASSISTANT, None), # None content
|
|
_message(MessageRole.USER, "Follow up")
|
|
]
|
|
expected = (
|
|
f"{self.IM_START}user\nHello{self.IM_END}\n"
|
|
f"{self.IM_START}assistant\n{self.IM_END}\n" # Empty content for assistant
|
|
f"{self.IM_START}user\nFollow up{self.IM_END}\n"
|
|
f"{self.IM_START}assistant\n"
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_correct_token_usage_and_newlines(self, style: ChatMLPromptStyle) -> None:
|
|
# Validates: <|im_start|>role\ncontent<|im_end|>\n ... <|im_start|>assistant\n
|
|
messages = [_message(MessageRole.USER, "Test")]
|
|
expected = (
|
|
f"{self.IM_START}user\nTest{self.IM_END}\n"
|
|
f"{self.IM_START}assistant\n"
|
|
)
|
|
actual = style._messages_to_prompt(messages)
|
|
assert actual == expected
|
|
assert actual.count(self.IM_START) == 2
|
|
assert actual.count(self.IM_END) == 1
|
|
assert actual.endswith(f"{self.IM_START}assistant\n")
|
|
# Check newlines: after role, after content (before im_end), after im_end
|
|
# <|im_start|>user\nTest<|im_end|>\n<|im_start|>assistant\n
|
|
# Role is followed by \n. Content is on its own line implicitly. im_end is followed by \n.
|
|
# The structure f"{IM_START}{role}\n{content}{IM_END}\n" ensures this.
|
|
user_part = f"{self.IM_START}user\nTest{self.IM_END}\n"
|
|
assert user_part in actual
|
|
|
|
messages_with_system = [
|
|
_message(MessageRole.SYSTEM, "Sys"),
|
|
_message(MessageRole.USER, "Usr")
|
|
]
|
|
expected_sys_usr = (
|
|
f"{self.IM_START}system\nSys{self.IM_END}\n"
|
|
f"{self.IM_START}user\nUsr{self.IM_END}\n"
|
|
f"{self.IM_START}assistant\n"
|
|
)
|
|
actual_sys_usr = style._messages_to_prompt(messages_with_system)
|
|
assert actual_sys_usr == expected_sys_usr
|
|
assert actual_sys_usr.count(self.IM_START) == 3
|
|
assert actual_sys_usr.count(self.IM_END) == 2
|
|
|
|
|
|
class TestTagPromptStyle:
|
|
BOS = "<s>"
|
|
EOS = "</s>"
|
|
|
|
@pytest.fixture
|
|
def style(self) -> TagPromptStyle:
|
|
return TagPromptStyle()
|
|
|
|
def test_empty_messages(self, style: TagPromptStyle) -> None:
|
|
messages = []
|
|
# Expected based on current TagPromptStyle: "<s><|assistant|>: "
|
|
expected = f"{self.BOS}<|assistant|>: "
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_simple_user_assistant_chat(self, style: TagPromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.USER, "Hello"),
|
|
_message(MessageRole.ASSISTANT, "Hi there!"),
|
|
]
|
|
expected = (
|
|
f"{self.BOS}<|user|>: Hello\n"
|
|
f"<|assistant|>: Hi there!{self.EOS}\n"
|
|
f"<|assistant|>: "
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_with_system_message(self, style: TagPromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.SYSTEM, "System instructions."),
|
|
_message(MessageRole.USER, "Ping"),
|
|
_message(MessageRole.ASSISTANT, "Pong"),
|
|
]
|
|
expected = (
|
|
f"{self.BOS}<|system|>: System instructions.\n"
|
|
f"<|user|>: Ping\n"
|
|
f"<|assistant|>: Pong{self.EOS}\n"
|
|
f"<|assistant|>: "
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_completion_to_prompt(self, style: TagPromptStyle) -> None:
|
|
completion = "Test user input"
|
|
# Expected: <s><|user|>: Test user input\n<|assistant|>:
|
|
expected = f"{self.BOS}<|user|>: Test user input\n<|assistant|>: "
|
|
assert style._completion_to_prompt(completion) == expected
|
|
|
|
def test_bos_eos_placement_multiple_turns(self, style: TagPromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.USER, "User1"),
|
|
_message(MessageRole.ASSISTANT, "Assistant1"),
|
|
_message(MessageRole.USER, "User2"),
|
|
_message(MessageRole.ASSISTANT, "Assistant2"),
|
|
]
|
|
expected = (
|
|
f"{self.BOS}<|user|>: User1\n"
|
|
f"<|assistant|>: Assistant1{self.EOS}\n"
|
|
f"<|user|>: User2\n"
|
|
f"<|assistant|>: Assistant2{self.EOS}\n"
|
|
f"<|assistant|>: "
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_ending_with_user_message(self, style: TagPromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.USER, "User1"),
|
|
_message(MessageRole.ASSISTANT, "Assistant1"),
|
|
_message(MessageRole.USER, "User2 (prompting for response)"),
|
|
]
|
|
expected = (
|
|
f"{self.BOS}<|user|>: User1\n"
|
|
f"<|assistant|>: Assistant1{self.EOS}\n"
|
|
f"<|user|>: User2 (prompting for response)\n"
|
|
f"<|assistant|>: "
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_message_with_empty_content(self, style: TagPromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.USER, ""),
|
|
_message(MessageRole.ASSISTANT, ""), # Empty assistant response
|
|
]
|
|
# Content is stripped, so empty string remains empty.
|
|
expected = (
|
|
f"{self.BOS}<|user|>: \n"
|
|
f"<|assistant|>: {self.EOS}\n"
|
|
f"<|assistant|>: "
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_message_with_none_content(self, style: TagPromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.USER, None), # Becomes empty string
|
|
_message(MessageRole.ASSISTANT, None), # Becomes empty string
|
|
]
|
|
expected = (
|
|
f"{self.BOS}<|user|>: \n"
|
|
f"<|assistant|>: {self.EOS}\n"
|
|
f"<|assistant|>: "
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_only_user_message(self, style: TagPromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.USER, "Just a user message"),
|
|
]
|
|
expected = (
|
|
f"{self.BOS}<|user|>: Just a user message\n"
|
|
f"<|assistant|>: "
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_only_assistant_message(self, style: TagPromptStyle) -> None:
|
|
# This is an unusual case, but the style should handle it.
|
|
messages = [
|
|
_message(MessageRole.ASSISTANT, "Only assistant"),
|
|
]
|
|
expected = (
|
|
f"{self.BOS}<|assistant|>: Only assistant{self.EOS}\n"
|
|
f"<|assistant|>: " # Still prompts for assistant
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|
|
|
|
def test_only_system_message(self, style: TagPromptStyle) -> None:
|
|
messages = [
|
|
_message(MessageRole.SYSTEM, "Only system"),
|
|
]
|
|
expected = (
|
|
f"{self.BOS}<|system|>: Only system\n"
|
|
f"<|assistant|>: "
|
|
)
|
|
assert style._messages_to_prompt(messages) == expected
|