wehub-resource-sync
84 lines
2.6 KiB
Python
84 lines
2.6 KiB
Python
import pytest
|
|
|
|
from mlc_llm.conversation_template import ConvTemplateRegistry
|
|
from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
|
|
|
|
|
|
def get_conv_templates():
|
|
return [
|
|
"llama-3",
|
|
"llama-2",
|
|
"mistral_default",
|
|
"gorilla",
|
|
"gorilla-openfunctions-v2",
|
|
"chatml",
|
|
"phi-2",
|
|
"codellama_completion",
|
|
"codellama_instruct",
|
|
"rwkv_world",
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("conv_template_name", get_conv_templates())
|
|
def test_json(conv_template_name):
|
|
template = ConvTemplateRegistry.get_conv_template(conv_template_name)
|
|
j = template.to_json_dict()
|
|
template_parsed = Conversation.from_json_dict(j)
|
|
assert template == template_parsed
|
|
|
|
|
|
@pytest.mark.parametrize("conv_template_name", get_conv_templates())
|
|
def test_prompt(conv_template_name):
|
|
conversation = ConvTemplateRegistry.get_conv_template(conv_template_name)
|
|
user_msg = "test1"
|
|
assistant_msg = "test2"
|
|
prompt = "test3"
|
|
|
|
expected_user_msg = (
|
|
conversation.role_templates["user"]
|
|
.replace(MessagePlaceholders.USER.value, user_msg)
|
|
.replace(MessagePlaceholders.FUNCTION.value, "")
|
|
)
|
|
|
|
expected_prompt = (
|
|
conversation.role_templates["user"]
|
|
.replace(MessagePlaceholders.USER.value, prompt)
|
|
.replace(MessagePlaceholders.FUNCTION.value, "")
|
|
)
|
|
|
|
conversation.messages.append(("user", user_msg))
|
|
conversation.messages.append(("assistant", assistant_msg))
|
|
conversation.messages.append(("user", prompt))
|
|
conversation.messages.append(("assistant", None))
|
|
res = conversation.as_prompt()
|
|
|
|
system_msg = conversation.system_template.replace(
|
|
MessagePlaceholders.SYSTEM.value, conversation.system_message
|
|
)
|
|
expected_final_prompt = (
|
|
system_msg
|
|
+ (conversation.seps[0] if system_msg != "" else "")
|
|
+ (
|
|
conversation.roles["user"] + conversation.role_content_sep
|
|
if conversation.add_role_after_system_message
|
|
else ""
|
|
)
|
|
+ expected_user_msg
|
|
+ conversation.seps[0 % len(conversation.seps)]
|
|
+ conversation.roles["assistant"]
|
|
+ conversation.role_content_sep
|
|
+ assistant_msg
|
|
+ conversation.seps[1 % len(conversation.seps)]
|
|
+ conversation.roles["user"]
|
|
+ conversation.role_content_sep
|
|
+ expected_prompt
|
|
+ conversation.seps[0 % len(conversation.seps)]
|
|
+ conversation.roles["assistant"]
|
|
+ conversation.role_empty_sep
|
|
)
|
|
assert res == expected_final_prompt
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_json("llama-3")
|