# MiniMind数据集代码详解（NLP初学者版）
这段代码定义了训练语言模型的三种不同数据集类，分别对应NLP领域中常见的三种训练范式：预训练、监督微调和偏好优化。我将逐类详细讲解。
## 1. 导入和环境设置
```python
import json
import random
import re
import pandas as pd
import numpy as np
from torch.utils.data import Dataset, DataLoader
import torch
from sklearn.model_selection import train_test_split
import os
import ast
os.environ["TOKENIZERS_PARALLELISM"] = "false"
```
- 导入处理JSON文件的库，用于读取训练数据
- 导入PyTorch的Dataset类，这是自定义数据集的基础
- 设置环境变量关闭tokenizers并行处理，避免多线程加载数据时的问题
## 2. 预训练数据集 (PretrainDataset)
```python
class PretrainDataset(Dataset):
def __init__(self, data_path, tokenizer, max_length=512):
super().__init__()
self.tokenizer = tokenizer # 分词器，将文本转为token ID
self.max_length = max_length # 序列最大长度
self.samples = self.load_data(data_path) # 加载数据
```
初始化函数：接收数据路径、分词器和最大序列长度参数
```python
def load_data(self, path):
samples = []
with open(path, 'r', encoding='utf-8') as f:
for line_num, line in enumerate(f, 1):
data = json.loads(line.strip())
samples.append(data)
return samples
```
加载JSONL格式数据（每行一个JSON对象）
```python
def __len__(self):
return len(self.samples)
```
返回数据集长度
```python
def __getitem__(self, index):
sample = self.samples[index]
# 构建输入文本
text = f"{self.tokenizer.bos_token}{str(sample['text'])}{self.tokenizer.eos_token}"
encoding = self.tokenizer(
text,
max_length=self.max_length,
padding='max_length', # 填充到最大长度
truncation=True, # 截断过长文本
return_tensors='pt' # 返回PyTorch张量
)
input_ids = encoding.input_ids.squeeze()
loss_mask = (input_ids != self.tokenizer.pad_token_id) # 生成掩码：不计算填充位置的损失
X = torch.tensor(input_ids[:-1], dtype=torch.long) # 输入序列（最后一个token不作为输入）
Y = torch.tensor(input_ids[1:], dtype=torch.long) # 目标序列（第一个token不作为目标）
loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long) # 损失掩码（与Y对齐）
return X, Y, loss_mask
```
这个方法实现了语言模型预训练的经典模式：
- 添加开始和结束标记
- 生成滑动窗口式的输入输出对（输入是前n-1个token，输出是后n-1个token）
- 创建损失掩码，忽略填充位置的损失贡献
## 3. 监督微调数据集 (SFTDataset)
```python
class SFTDataset(Dataset):
def __init__(self, jsonl_path, tokenizer, max_length=1024):
super().__init__()
self.tokenizer = tokenizer
self.max_length = max_length
self.samples = self.load_data(jsonl_path)
self.bos_id = tokenizer('assistant\n', add_special_tokens=False).input_ids # 助手回复开始标记\n', add_special_tokens=False).input_ids # 对话结束标记
self.eos_id = tokenizer('
```
初始化函数：注意这里特别定义了助手回复的开始和结束标记ID
```python
def _create_chat_prompt(self, conversations):
"""构建符合ChatML格式的对话"""
messages = []
for i, turn in enumerate(conversations):
role = 'user' if i % 2 == 0 else 'assistant' # 交替的用户和助手角色
messages.append({"role": role, "content": turn['content']})
return self.tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=False
)
```
这个函数将对话数据转换为ChatML格式：
- 根据索引自动判定角色（偶数为用户，奇数为助手）
- 使用tokenizer的chat_template将消息格式化为模型期望的格式
```python
def _generate_loss_mask(self, input_ids):
loss_mask = [0] * len(input_ids) # 初始化全0掩码
i = 0
while i < len(input_ids):
# 寻找助手回复的开始标记
if input_ids[i:i + len(self.bos_id)] == self.bos_id:
start = i + len(self.bos_id)
end = start
# 寻找助手回复的结束标记
while end < len(input_ids):
if input_ids[end:end + len(self.eos_id)] == self.eos_id:
break
end += 1
# 标记助手回复部分为需要计算损失的区域（设为1）
for j in range(start + 1, min(end + len(self.eos_id) + 1, self.max_length)):
loss_mask[j] = 1
i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)
else:
i += 1
return loss_mask
```
这是SFT中的核心函数，它创建了动态损失掩码：
- **关键思想**：只为模型生成的部分（助手回复）计算损失，不为用户输入计算损失
- 遍历序列找到所有 `assistant\n` 到 `\n` 之间的文本
- 将这些区域在掩码中标记为1（参与损失计算）
```python
def __getitem__(self, index):
sample = self.samples[index]
# 构建对话提示
prompt = self._create_chat_prompt(sample['conversations'])
input_ids = self.tokenizer(prompt).input_ids[:self.max_length]
input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids)) # 手动填充
# 生成动态损失掩码
loss_mask = self._generate_loss_mask(input_ids)
# 构建训练数据（与预训练方式相同）
X = torch.tensor(input_ids[:-1], dtype=torch.long)
Y = torch.tensor(input_ids[1:], dtype=torch.long)
loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long) # 对齐预测位置
return X, Y, loss_mask
```
这里实现了SFT的核心数据处理逻辑：
- 将对话转换为模型可理解的格式
- 生成只关注于助手回复部分的损失掩码
- 构建滑动窗口式的输入输出对
## 4. 偏好优化数据集 (DPODataset)
```python
class DPODataset(Dataset):
def __init__(self, file_path, tokenizer, max_length=4096):
super().__init__()
self.tokenizer = tokenizer
self.max_length = max_length
self.padding = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
self.bos_id = tokenizer('assistant\n', add_special_tokens=False).input_ids\n', add_special_tokens=False).input_ids
self.eos_id = tokenizer('
with open(file_path, 'r', encoding='utf-8') as f:
self.data = []
for line in f:
line = line.strip()
obj = json.loads(line)
self.data.append(obj)
```
初始化函数：类似于SFT数据集，但直接加载数据而不是调用单独的方法
```python
def __getitem__(self, index):
item = self.data[index]
chosen = item['chosen'] # 偏好数据：被选中的回复（更好的回复）
rejected = item['rejected'] # 偏好数据：被拒绝的回复（较差的回复）
# 为选中和拒绝的回复分别创建提示
chosen_prompt = self.tokenizer.apply_chat_template(
chosen, tokenize=False, add_generation_prompt=False
)
rejected_prompt = self.tokenizer.apply_chat_template(
rejected, tokenize=False, add_generation_prompt=False
)
# 分词和填充
chosen_encoding = self.tokenizer(
chosen_prompt, truncation=True, max_length=self.max_length, padding='max_length'
)
rejected_encoding = self.tokenizer(
rejected_prompt, truncation=True, max_length=self.max_length, padding='max_length'
)
# 获取token IDs并生成损失掩码
chosen_input_ids = chosen_encoding['input_ids']
chosen_loss_mask = self._generate_loss_mask(chosen_input_ids)
rejected_input_ids = rejected_encoding['input_ids']
rejected_loss_mask = self._generate_loss_mask(rejected_input_ids)
# 创建输入输出对和掩码，为选中和拒绝的回复各创建一份
x_chosen = torch.tensor(chosen_input_ids[:-1], dtype=torch.long)
y_chosen = torch.tensor(chosen_input_ids[1:], dtype=torch.long)
mask_chosen = torch.tensor(chosen_loss_mask[1:], dtype=torch.long)
x_rejected = torch.tensor(rejected_input_ids[:-1], dtype=torch.long)
y_rejected = torch.tensor(rejected_input_ids[1:], dtype=torch.long)
mask_rejected = torch.tensor(rejected_loss_mask[1:], dtype=torch.long)
# 返回字典，包含两组数据
return {
'x_chosen': x_chosen,
'y_chosen': y_chosen,
'mask_chosen': mask_chosen,
'x_rejected': x_rejected,
'y_rejected': y_rejected,
'mask_rejected': mask_rejected
}
```
这个方法实现了DPO算法所需的数据准备：
- 同时处理"好的回复"和"差的回复"
- 为每组回复创建输入、目标和掩码
- 返回字典格式，包含两组完整的训练数据
## NLP初学者的重点解释
1. **三种训练范式的数据处理区别**：
- **预训练**：简单文本，预测下一个词
- **SFT**：对话格式，只对AI回复部分计算损失
- **DPO**：成对的好/差回复，用于学习偏好排序
2. **损失掩码的作用**：
- 控制哪些位置的预测会影响损失计算
- 在SFT/DPO中只对AI助手部分计算损失，不对用户输入计算损失
- 忽略填充位置的损失贡献
3. **滑动窗口训练**：
- `X = input_ids[:-1]` 和 `Y = input_ids[1:]` 创建错位的输入/输出对
- 模型看到前n-1个token，预测后n-1个token
- 这是标准的语言模型训练方式，本质是"下一个词预测"任务
4. **特殊标记处理**：
- 开始标记(`assistant`)和结束标记(``)用于识别助手回答区域
- 在DPO和SFT中，只对这些标记之间的内容计算损失
通过这三个数据集类，MiniMind系统实现了完整的语言模型训练范式，从预训练到微调再到偏好学习，逐步提升模型能力和对齐程度。
Comments NOTHING