1. 初始tokennize设计

instruction = tokenizer(  <|im_start|>system\n{system_prompt}\n<|im_end|>\n<im_start>user\n{example['text']}\n<|im_end|>\n<|im_start|>assistant\n<think>\n{example['reason']}\n</think>\n",add_special_tokens=False)
response = tokenizer(f"{example['output']}", add_special_tokens=False)
input_ids = instruction["input_ids"] + response["input_ids"] + [tokenizer.pad_token_id]
attention_mask = (instruction["attention_mask"] + response["attention_mask"] + [1])
labels = [-100] * len(instruction["input_ids"]) + response["input_ids"] + [tokenizer.pad_token_id]

2. 重新设计

问题：
在instruction中包含了example['reason']（即思考过程），这应该属于assistant的response部分

当前的拼接方式会导致：
input_ids：将example['reason']和example['output']都放在assistant部分

labels：由于instruction部分的标签被设置为-100（忽略loss），而只有response部分的token有loss
因此，example['reason']被放在了instruction的末尾，导致它没有标签（即被忽略），而example['output']有标签

但这样处理是不对的，因为example['reason']（思考过程）也应该是assistant生成的内容的一部分，因此它应该放在response中，并且计算loss。


我们期望的流程：
用户提问（example['text']）后，assistant首先生成思考过程（<thinking>标签包裹），然后生成回复（<answer>标签包裹）

因此，整个assistant的回复应该是：
<think>\n{example['reason']}\n</think>\n{example['output']}


因此，修改思路如下：
将assistant的整个回复内容合并为一个字符串：assistant_response = f"<think>\n{example['reason']}\n</think>\n{example['output']}"

构建instruction时，只包含到assistant开始（即<|im_start|>assistant\n），不要包含assistant的任何输出。

然后，将assistant_response作为assistant的输出，并tokenize。

最后将instruction的token和assistant_response的token拼接，并设置labels：


instruction部分（包括system、user和assistant的开始标记）的token对应的label为-100
assistant_response的token对应的label为实际token


修改后代码：
# 构建instruction部分（到assistant开始标记为止）
instruction_text = (
    f"<|im_start|>system\n{system_prompt}\n<|im_end|>\n"
    f"<|im_start|>user\n{example['text']}\n<|im_end|>\n"
    f"<|im_start|>assistant\n"
)
instruction = tokenizer(
    instruction_text,
    add_special_tokens=False
)

# 构建assistant的整个回复（包括<thinking>和<answer>部分）
assistant_response = f"<think>\n{example['reason']}\n</think>\n{example['output']}"

# Tokenize assistant的回复
response = tokenizer(assistant_response, add_special_tokens=False)

# 合并input_ids和attention_mask
input_ids = instruction["input_ids"] + response["input_ids"] + [tokenizer.eos_token_id]
attention_mask = instruction["attention_mask"] + response["attention_mask"] + [1]

# 设置labels：instruction部分为-100，assistant_response部分为实际token
labels = [-100] * len(instruction["input_ids"]) + response["input_ids"] + [tokenizer.eos_token_id]

# 注意：这里添加了EOS token，并且它的label也要计算loss（所以不是-100）

解释：
这样设计后，模型在训练时会学习生成assistant的整个回复（包括<thinking>和<answer>两部分）
在生成时，可以引导模型首先生成<thinking>思考过程，然后生成<answer>最终答案


注意：这里使用了tokenizer.eos_token_id作为结束符，并计算其loss。


另外，关于标签：
在训练时，我们只计算assistant生成部分的loss（即从<|im_start|>assistant后面的所有内容）

因此，instruction部分（包括system、user和assistant开始标记）都不计算loss（用-100屏蔽）


这样修改后，数据放置就正确了。
</think>
以下是修正后的代码设计，重点解决了 
<think>/</think>
 数据位置问题 并优化了整体逻辑：


def format_example(example):
    # 重构instruction：保证system和user部分不包含assistant输出
    # 注意：这里移除了对 example['reason'] 的错误包含
    instruction = tokenizer(
        f"<|im_start|>system\n{system_prompt}\n<|im_end|>\n"
        f"<|im_start|>user\n{example['text']}\n<|im_end|>\n"
        f"<|im_start|>assistant\n",  # 仅包含assistant起始标记
        add_special_tokens=False,
        return_tensors="pt"
    )
    
    # 正确拼接assistant响应：完整包含思考+回答的完整输出
    # 使用特殊标签明确区分思考/回答区间
    full_response = (
        f"<think>\n{example['reason']}\n</think>\n"  # 思考区间
        f"{example['output']}"              # 最终回答
    )
    
    # Tokenize完整响应
    response = tokenizer(
        full_response,
        add_special_tokens=False,
        return_tensors="pt"
    )
    
    # 合并输入序列（添加EOS标记）
    input_ids = torch.cat([
        instruction.input_ids[0],
        response.input_ids[0],
        torch.tensor([tokenizer.eos_token_id])
    ])
    
    # 创建attention mask
    attention_mask = torch.ones_like(input_ids)
    
    # 创建labels: 仅对assistant输出计算损失
    labels = torch.cat([
        torch.full(instruction.input_ids.shape, -100),  # 忽略系统/用户部分
        response.input_ids[0],                          # 学习整个响应
        torch.tensor([tokenizer.eos_token_id])           # 学习结束符
    ])

    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels
    }

关键修正说明：

问题诊断：
原代码错误将 example['reason'] (思考过程) 放在 instruction 中

导致思考过程被包含在输入序列但未计入训练目标

修正方案：
# 错误：思考过程放错位置 (instruction内)
f"...assistant\n<think>\n{example['reason']}..."

# 正确：思考过程作为响应的一部分
full_response = f"<think>\n{example['reason']}...{example['output']}"

损失计算优化：
labels = [
    [-100] * len(instruction),  # 忽略系统/用户输入
    response_ids,                # 学习完整响应（含思考）
    [eos_token]                  # 学习结束符
]


保证模型学习整个 
<think>...思考...</think>...输出
 序列

与推理时生成流程保持一致

序列结构可视化：
[系统提示][用户问题][助理开头标记] | <think> 思考内容 </think> 最终回答 <EOS>
└─── 忽略损失 ───┘ └───────────── 计算损失的目标区域 ─────────────┘


使用验证示例：

# 样本结构
example = {
    "text": "如何计算圆的面积？",
    "reason": "圆面积公式为πr²，需先确定半径",
    "output": "面积 = π × 半径的平方"
}

# 处理后模型学习目标：
# 输入: [系统][用户问题][assistant开始]
# 目标: <think>\n如何计算圆的面积？\n</think>\n面积=π×半径²<EOS>
#        ^^^^^^^^^^^^^^^^^^^^^^^^
#        模型需要完整生成的响应部分


核心原则：在微调任务中，思考过程必须包含在模型需要生成的响应序列内，不能作为输入内容的一部分，这样才能让模型学习到"思考-回答"的完整生成逻辑。