与 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 标志,以及任何额外的关键字参数。
- 检查重复:如果模板类型已存在于
TEMPLATE_MAPPING中且exist_ok为False,则抛出ValueError。 - 设置属性:设置模板的
template_type属性。 - 注册映射:创建包含模板和额外参数的字典,并将其添加到
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=
)

