LLM 模板注册机制与 QwenTemplate 源码解析

与 llama-factory 定义 Qwen 模板的对比

def register_template(template_type: str, template: Template, *, exist_ok: bool = False, **kwargs) -> None:
    if not exist_ok and template_type in TEMPLATE_MAPPING:
        raise ValueError(f'The `{template_type}` has already been registered in the TEMPLATE_MAPPING.')
    template.template_type = template_type
    template_info = {'template': template, **kwargs}
    TEMPLATE_MAPPING[template_type] = template_info

register_template(
    TemplateType.default,
    Template([], ['### Human:\n{{QUERY}}\n\n### Assistant:\n'], ['\n\n'], [['eos_token_id']], DEFAULT_SYSTEM, ['{{SYSTEM}}\n\n'])
)

该函数用于注册一个新的模板。它接受模板类型（字符串）、Template 对象、可选的 exist_ok 标志，以及任何额外的关键字参数。

1. 检查重复：如果模板类型已存在于 TEMPLATE_MAPPING 中且 exist_ok 为 False，则抛出 ValueError。
2. 设置属性：设置模板的 template_type 属性。
3. 注册映射：创建包含模板和额外参数的字典，并将其添加到 TEMPLATE_MAPPING 中。

代码示例展示了默认模板的注册过程，为人类和助手消息使用特定格式，并包含查询和系统消息的占位符。

class DefaultGenerationTemplate(Template):
    def __init__(self):
        super().__init__([], ['{{QUERY}}'], None, [['eos_token_id']], auto_add_bos=True)

register_template(TemplateType.default_generation, DefaultGenerationTemplate(), is_generation=)

register_template(
    TemplateType.default_generation_bos,
    Template([[]], [], , [[]]),
    is_generation=
)