chore: import upstream snapshot with attribution

2026-07-13 13:23:58 +08:00
commit 770d92cb1f
694 changed files with 114634 additions and 0 deletions
@@ -0,0 +1,83 @@
+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")