AI tokenizer chat template 统一多模型消息格式

要让同一套对话服务切换多个聊天模型，应该统一的是应用层的 messages 数据结构，而不是强行统一模型看到的特殊 token。业务侧只传 role 与 content；每个模型再由自己的 tokenizer.chat_template 把消息转换成训练时使用的 token 序列。这样，模型 A 可以使用 [INST]，模型 B 可以使用 ，上层接口仍保持不变。

官方文档：https://huggingface.co/docs/transformers/chat_templating

我第一次做多模型路由时，最直觉的方案是在网关里拼接“用户：”“助手：”和统一结束符。单模型演示看起来没问题，一换 checkpoint 就出现回复质量下降、模型续写用户内容、BOS/EOS 重复等现象。根因不是消息 JSON 不统一，而是不同模型在微调时学到的序列化契约不同。

一、业务负载：统一消息对象，保留模型契约

本文讨论的是文本聊天服务：请求可能被路由到多个 Hugging Face Transformers 模型，但上层产品不希望为每个模型改一次请求格式。建议把内部协议固定为一个消息列表，每条消息只保留两个核心字段：

• role：文本聊天常见值为 system、user、assistant。
• content：当前范围内限定为字符串，避免把多模态内容误交给文本 tokenizer。

def normalize_messages(raw_messages):
    # 文本聊天只接受这三类角色；工具和多模态消息另设适配器
    allowed_roles = {"system", "user", "assistant"}
    normalized = []

    for index, item in enumerate(raw_messages):
        role = item.get("role")
        content = item.get("content")

        # 在进入模型层前拒绝不完整数据，避免模板渲染后才暴露错误
        if role not in allowed_roles:
            raise ValueError(f"messages[{index}].role 不支持: {role}")
        if not isinstance(content, str):
            raise TypeError(f"messages[{index}].content 必须是字符串")

        normalized.append({"role": role, "content": content})

    if not normalized:
        raise ValueError("messages 不能为空")
    return normalized

这段规范化代码没有添加任何模型控制 token。它负责的是应用边界：字段类型稳定、角色明确、错误尽早暴露。至于每个角色前后需要哪些特殊标记，应交给与 checkpoint 配套的 tokenizer。

二、约束条件：聊天模型最终只会继续 token 序列

聊天接口看起来在处理多轮消息，底层因果语言模型仍然是在继续一段 token 序列。chat template 的作用，就是把结构化消息转换成模型熟悉的序列。两个模型即使来自相同基座，也可能因为聊天微调格式不同而使用完全不同的控制 token。

这带来一个重要架构决定：不要在公共网关里维护一份所谓万能 prompt 模板。那会把模型训练契约泄漏到业务代码中，后续更换 checkpoint 时还容易漏改。更可靠的绑定关系应是：

• 模型 ID 决定加载哪个 tokenizer；
• tokenizer 自带的 chat_template 决定消息如何序列化；
• 应用只决定本次是“开始新回复”“续写最后一条消息”还是“训练预处理”。

如果模型仓库没有 chat template，生产系统最好显式拒绝接入，或者为该模型登记经过确认的专用模板。不要根据模型名称猜测 [INST]、ChatML 或其他格式。

三、方案对比：手写模板、公共模板还是 tokenizer 模板

第三种方案不是把所有差异消灭，而是把差异放回正确的边界。模型注册表保存 checkpoint；通用适配器负责加载 tokenizer 并调用统一 API；模板细节仍由 tokenizer 管理。

四、推荐架构：一个入口，两种生成模式

普通聊天请求通常以 user 消息结尾，需要模型开始一条新的 assistant 回复，此时使用 add_generation_prompt=True。如果最后一条已经是 assistant，而且业务想让模型从指定前缀继续写，则使用 continue_final_message=True。这两个参数语义相反，不能同时启用。

from functools import lru_cache
from transformers import AutoTokenizer

MODEL_REGISTRY = {
    # 注册表只保存模型定位信息，不保存手写控制 token
    "support": "HuggingFaceH4/zephyr-7b-beta",
    "writer": "mistralai/Mistral-7B-Instruct-v0.1",
}


@lru_cache(maxsize=None)
def load_tokenizer(model_key):
    # 缓存 tokenizer，避免每次请求重复读取模型配置与词表
    model_id = MODEL_REGISTRY[model_key]
    tokenizer = AutoTokenizer.from_pretrained(model_id)

    # 缺失模板时显式失败，不用猜测的默认格式替代模型契约
    if not tokenizer.chat_template:
        raise ValueError(f"模型 {model_id} 未提供 chat_template")
    return tokenizer


def build_model_inputs(model_key, raw_messages, mode="reply"):
    messages = normalize_messages(raw_messages)
    tokenizer = load_tokenizer(model_key)

    common = {
        # 直接让模板返回 token，避免二次分词重复添加特殊 token
        "tokenize": True,
        "return_dict": True,
        "return_tensors": "pt",
    }

    if mode == "reply":
        # 新回复模式要求最后一条通常来自用户
        if messages[-1]["role"] != "user":
            raise ValueError("reply 模式要求最后一条消息是 user")
        return tokenizer.apply_chat_template(
            messages,
            add_generation_prompt=True,
            **common,
        )

    if mode == "prefill":
        # 预填模式续写最后一条 assistant 内容，不创建新消息头
        if messages[-1]["role"] != "assistant":
            raise ValueError("prefill 模式要求最后一条消息是 assistant")
        return tokenizer.apply_chat_template(
            messages,
            continue_final_message=True,
            **common,
        )

    raise ValueError(f"不支持的生成模式: {mode}")

返回值可直接移动到模型设备，再传给 generate()。适配层不需要知道 Zephyr、Mistral 或其他模型具体用了什么标记；只要 tokenizer 与模型 checkpoint 匹配，模板就会负责序列化。

def generate_reply(model, model_key, messages):
    # 构建与目标模型模板一致的 input_ids 和 attention_mask
    model_inputs = build_model_inputs(model_key, messages, mode="reply")
    model_inputs = model_inputs.to(model.device)

    # 这里只展示生成调用；采样参数应由业务策略单独管理
    output_ids = model.generate(**model_inputs, max_new_tokens=256)
    return output_ids

五、普通回复、预填续写与训练不能混成一种模式

add_generation_prompt=True 会在末尾添加“现在开始 assistant 消息”的控制标记，适合最后一条是用户问题的常规聊天。continue_final_message=True 会让模型继续最后一条消息，适合预填 JSON 前缀、固定回答开头等场景。官方文档明确说明二者一起使用会报错。

训练预处理又是第三种情况。训练样本已经包含 assistant 答案，末尾不应再追加一个新的 assistant 开始标记，因此通常使用 add_generation_prompt=False：

def format_training_example(tokenizer, messages):
    # 训练样本已有完整回答，不追加下一轮 assistant 起始标记
    return tokenizer.apply_chat_template(
        normalize_messages(messages),
        tokenize=False,
        add_generation_prompt=False,
    )

六、特殊 token 的风险点：优先直接 tokenize

chat template 通常已经包含模型需要的 BOS、EOS 或角色控制标记。如果先用 tokenize=False 得到字符串，随后又让 tokenizer 自动添加特殊 token，就可能重复 BOS/EOS，影响模型表现。对在线推理，优先直接使用 apply_chat_template(..., tokenize=True)。

确实需要先检查格式化文本时，后续分词应关闭自动添加：

def inspect_then_tokenize(tokenizer, messages):
    # 先得到可读字符串，适合调试模板边界
    rendered = tokenizer.apply_chat_template(
        normalize_messages(messages),
        tokenize=False,
        add_generation_prompt=True,
    )

    # 模板已经负责特殊 token，二次分词时禁止重复添加
    encoded = tokenizer(
        rendered,
        add_special_tokens=False,
        return_tensors="pt",
    )
    return rendered, encoded

这里的“inspect”是开发阶段查看序列字符串，不是生产日志要求。实际服务不要把含有用户敏感内容的完整 prompt 无条件写入日志。

七、风险点：模型能力差异不能靠模板抹平

chat template 统一了调用入口，但不会让所有模型自动拥有相同能力。接入模型时还应记录以下约束：

• system 角色：有些模板支持独立 system 消息，有些会把系统指令合并到首条 user 内容，也有模型会拒绝特定角色顺序。
• 连续角色：部分模型不接受两条连续 assistant 消息；预填模式要单独处理。
• 工具调用：工具模型可能需要 tools 参数、JSON Schema 和命名模板，不能只传基础文本消息。
• 多模态：多模态模板通常位于 processor，content 可能是图像、视频和文本字典列表，不适合本文的字符串规范化函数。
• 推理字段：某些推理模型会使用独立的 reasoning 或 thinking 字段，预填时要确认模板实际引用哪个字段。

因此，推荐在注册表旁维护一份能力元数据，例如是否支持 system、tool use、多模态与 prefill。统一接口负责早期拒绝不兼容请求，而不是把所有内容勉强塞进 content。

八、落地清单

1. 业务层固定 messages=[{"role": ..., "content": ...}]，不要混入模型特殊 token。
2. 模型 ID、tokenizer 与 checkpoint 一一绑定，并检查 tokenizer 是否提供 chat template。
3. 在线推理优先 tokenize=True，减少重复特殊 token 的机会。
4. 普通回复使用 add_generation_prompt=True；预填续写使用 continue_final_message=True，二者不同时传。
5. 训练格式使用完整对话，不追加新回复起始标记。
6. 为每个模型保存最小兼容样例：角色顺序、空 system、长对话、预填、工具和多模态边界。
7. 模板缺失或角色不受支持时明确报错，不在公共层猜测格式。

多模型消息格式的“统一”不是让每个模型吃下同一串字符，而是让上层系统遵守同一数据协议，再让模型自己的 tokenizer 把协议翻译成正确 token 序列。把差异放在 tokenizer/template 这一层，模型切换会更稳，业务代码也更容易维护。