wehub-resource-sync
86 lines
2.9 KiB
Python
86 lines
2.9 KiB
Python
"""The conversation template registry and presets in MLC LLM"""
|
|
|
|
from typing import ClassVar, Dict, Optional # noqa: UP035
|
|
|
|
from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
|
|
|
|
|
|
class ConvTemplateRegistry:
|
|
"""Global conversation template registry for preset templates."""
|
|
|
|
_conv_templates: ClassVar[Dict[str, Conversation]] = {} # noqa: UP006
|
|
|
|
@staticmethod
|
|
def register_conv_template(conv_template: Conversation, override: bool = False) -> None:
|
|
"""Register a new conversation template in the global registry.
|
|
Using `override = True` to override the previously registered
|
|
template with the same name.
|
|
"""
|
|
name = conv_template.name
|
|
if name is None:
|
|
raise ValueError("The template to register should have non-None name.")
|
|
if name in ConvTemplateRegistry._conv_templates and not override:
|
|
raise ValueError(
|
|
"The name of the template has been registered "
|
|
f"for {ConvTemplateRegistry._conv_templates[name].model_dump_json(by_alias=True)}"
|
|
)
|
|
ConvTemplateRegistry._conv_templates[name] = conv_template
|
|
|
|
@staticmethod
|
|
def get_conv_template(name: str) -> Optional[Conversation]:
|
|
"""Return the conversation template specified by the given name,
|
|
or None if the template is not registered.
|
|
"""
|
|
return ConvTemplateRegistry._conv_templates.get(name, None)
|
|
|
|
|
|
# ChatML
|
|
ConvTemplateRegistry.register_conv_template(
|
|
Conversation(
|
|
name="chatml",
|
|
system_template=f"<|im_start|>system\n{MessagePlaceholders.SYSTEM.value}<|im_end|>\n",
|
|
system_message=(
|
|
"A conversation between a user and an LLM-based AI assistant. The "
|
|
"assistant gives helpful and honest answers."
|
|
),
|
|
roles={"user": "<|im_start|>user", "assistant": "<|im_start|>assistant"},
|
|
seps=["<|im_end|>\n"],
|
|
role_content_sep="\n",
|
|
role_empty_sep="\n",
|
|
stop_str=["<|im_end|>"],
|
|
stop_token_ids=[2],
|
|
)
|
|
)
|
|
|
|
# ChatML without a system prompt
|
|
ConvTemplateRegistry.register_conv_template(
|
|
Conversation(
|
|
name="chatml_nosystem",
|
|
system_template=f"{MessagePlaceholders.SYSTEM.value}",
|
|
system_message="",
|
|
roles={"user": "<|im_start|>user", "assistant": "<|im_start|>assistant"},
|
|
seps=["<|im_end|>\n"],
|
|
role_content_sep="\n",
|
|
role_empty_sep="\n",
|
|
stop_str=["<|im_end|>"],
|
|
stop_token_ids=[2],
|
|
)
|
|
)
|
|
|
|
|
|
# Vanilla LM
|
|
ConvTemplateRegistry.register_conv_template(
|
|
Conversation(
|
|
name="LM",
|
|
system_template=f"{MessagePlaceholders.SYSTEM.value}",
|
|
system_message="",
|
|
roles={"user": "", "assistant": ""},
|
|
seps=[""],
|
|
role_content_sep="",
|
|
role_empty_sep="",
|
|
stop_str=[],
|
|
stop_token_ids=[2],
|
|
system_prefix_token_ids=[1],
|
|
)
|
|
)
|