Transformers chat template 怎么生成 assistant token 掩码

在做对话微调时，真正需要的不是把整段聊天都当成监督信号，而是让 loss 主要落在 assistant 回复上。Transformers 已经提供了现成的路径：让 chat template 用 {% generation %} 标记 assistant 内容，再调用 apply_chat_template() 时打开 return_assistant_tokens_mask=True。返回字典里的 assistant_masks 会与 input_ids 等长，assistant token 为 1，system/user 以及其他上下文 token 为 0。

官方地址：https://huggingface.co/docs/transformers/main/en/chat_templating

先把三个返回条件配齐

最小调用要同时满足三个条件：输出需要被切成 token，返回结构需要是字典，并显式要求 assistant 掩码。掩码字段名是 assistant_masks，不是 labels，也不是模型生成后的新 token。

from transformers import AutoTokenizer

# 这里使用支持 chat_template 的 tokenizer；模型名称可替换为自己的检查点
tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta")
messages = [
    {"role": "user", "content": "用一句话解释向量数据库"},
    {"role": "assistant", "content": "它用向量相似度检索语义相关内容。"},
]

# tokenize 和 return_dict 不能省略，否则无法请求 assistant_masks
encoded = tokenizer.apply_chat_template(
    messages,
    tokenize=True,
    return_dict=True,
    return_assistant_tokens_mask=True,
    add_generation_prompt=False,
)

# 这些数组应当一一对应，便于构造训练 labels 或检查标注范围
input_ids = encoded["input_ids"]
assistant_masks = encoded["assistant_masks"]
print(len(input_ids), len(assistant_masks))
print(sum(assistant_masks))  # 只统计被标成 assistant 的 token 数

这里的 mask 只描述已经出现在消息里的 assistant 内容。如果最后一条消息是 user，并且你把 add_generation_prompt=True 用于推理，它是在末尾添加 assistant 起始提示，并不等于已经生成了 assistant 回复，所以不要把这两个概念混在一起。

模板必须显式包住 assistant 内容

return_assistant_tokens_mask 的实现依赖模板渲染时记录 assistant 内容的字符区间。自定义 Jinja 模板时，把 assistant 分支放进 generation 块；system 和 user 分支不要误包进去。

{%- for message in messages %}
  {{- '\n' }}
  {%- if message['role'] == 'assistant' %}
    {%- generation %}{{- message['content'] }}{%- endgeneration %}
  {%- else %}
    {{- message['content'] }}
  {%- endif %}
  {{- eos_token }}
{%- endfor %}

模板本身仍然负责控制 token、换行和结束标记；generation 块只是告诉 Transformers 哪一段渲染文本属于 assistant。它不会改变对话格式，也不会自动把所有 role 为 assistant 的消息重新改写。

如果使用的是模型自带模板，先查看 tokenizer.chat_template，确认它是否包含 generation 标记。没有标记时，最直接的结果不是“mask 计算错了”，而是当前模板没有提供可追踪的 assistant 区间。

用对齐检查定位全零掩码

排查时不要只打印一串 0/1。把 token、掩码和解码文本放在同一行，能快速看出是模板没有标记，还是某些字符没有映射到 token。

# 逐 token 对照，避免把掩码位置误认为字符下标
tokens = tokenizer.convert_ids_to_tokens(encoded["input_ids"])
for index, (token, flag) in enumerate(zip(tokens, encoded["assistant_masks"])):
    # flag=1 表示该 token 落在模板的 assistant generation 区间
    print(f"{index:03d}  {flag}  {token}")

# 训练时只把非 assistant 区域设为 -100，忽略它们的交叉熵
labels = [token_id if flag else -100
          for token_id, flag in zip(encoded["input_ids"], encoded["assistant_masks"])]

训练和推理不要混用两种意图

训练完整对话时，通常使用 add_generation_prompt=False，让样本保留已有 assistant 回复，再用 mask 选择监督区域。推理准备新回复时，通常使用 add_generation_prompt=True；此时没有真实 assistant 内容可供掩码标 1。若最后一条消息是 assistant 预填内容，则考虑 continue_final_message=True，它与 add_generation_prompt=True 不能同时使用。

还要留意多模态消息。图像占位符和处理器的文本映射可能让 assistant 区间出现全零或偏移，不能因为纯文本样例正常就推断所有 processor 都正常。先用纯文本对话确认模板，再单独验证图文输入的 token 与 mask 对齐。

常见问题

assistant_masks 是模型自动生成的标签吗？

不是。它是 chat template 渲染出的 assistant 字符区间映射到 token 后得到的输入标记，模型不会替你判断业务上的监督范围。

能不能只传 return_assistant_tokens_mask=True？

不能。该选项要求同时返回 token 化结果和字典结构，也就是 tokenize=True、return_dict=True。

模板已有 assistant role，为什么仍然全零？

role 判断和 generation 跟踪是两件事。只有模板显式使用 generation 块，Transformers 才能记录可映射的 assistant 区间。

mask 能直接替代 attention_mask 吗？

不能。attention_mask 控制哪些输入位置参与注意力，assistant_masks 用来选择哪些 token 作为 assistant 监督区域，训练代码里通常还要把它转换成 labels 的 -100。