wehub-resource-sync
44 lines
2.6 KiB
Python
44 lines
2.6 KiB
Python
import pytest
|
|
|
|
from mlc_llm.conversation_template import ConvTemplateRegistry
|
|
|
|
pytestmark = [pytest.mark.runtime_unittest]
|
|
|
|
|
|
# From the official Llama-3 example:
|
|
# https://llama.meta.com/docs/model-cards-and-prompt-formats/meta-llama-3/
|
|
def test_llama3_prompt():
|
|
conversation = ConvTemplateRegistry.get_conv_template("llama-3")
|
|
system_msg = "You are a helpful AI assistant for travel tips and recommendations"
|
|
user_msg1 = "What is France's capital?"
|
|
assistant_msg1 = "Bonjour! The capital of France is Paris!"
|
|
user_msg2 = "What can I do there?"
|
|
assistant_msg2 = "Paris, the City of Light, offers a romantic getaway with must-see attractions like the Eiffel Tower and Louvre Museum, romantic experiences like river cruises and charming neighborhoods, and delicious food and drink options, with helpful tips for making the most of your trip." # noqa: E501
|
|
prompt = "Give me a detailed list of the attractions I should visit, and time it takes in each one, to plan my trip accordingly." # noqa: E501
|
|
|
|
conversation.system_message = system_msg
|
|
conversation.messages.append(("user", user_msg1))
|
|
conversation.messages.append(("assistant", assistant_msg1))
|
|
conversation.messages.append(("user", user_msg2))
|
|
conversation.messages.append(("assistant", assistant_msg2))
|
|
conversation.messages.append(("user", prompt))
|
|
conversation.messages.append(("assistant", None))
|
|
res = conversation.as_prompt()
|
|
|
|
expected = (
|
|
"<|start_header_id|>system<|end_header_id|>\n\n"
|
|
"You are a helpful AI assistant for travel tips and recommendations<|eot_id|>\n"
|
|
"<|start_header_id|>user<|end_header_id|>\n\n"
|
|
"What is France's capital?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
|
|
"Bonjour! The capital of France is Paris!<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n" # noqa: E501
|
|
"What can I do there?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
|
|
"Paris, the City of Light, offers a romantic getaway with must-see attractions like the Eiffel Tower and Louvre Museum, romantic experiences like river cruises and charming neighborhoods, and delicious food and drink options, with helpful tips for making the most of your trip.<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n" # noqa: E501
|
|
"Give me a detailed list of the attractions I should visit, and time it takes in each one, to plan my trip accordingly.<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" # noqa: E501
|
|
)
|
|
|
|
assert res[0] == expected
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_llama3_prompt()
|