2-Dataset

到这里我们便完成了对于 MiniMind Tokenizer 和 Model 部分的全部了解，我们所熟悉的大语言模型正是由这个组件构成的，接下来，我们需要对大模型训练所使用的数据集结构有个基本的认识。

想要训练一个能够正常对话，并且符合人类对话偏好的大模型一般需要经过以下几个训练阶段：

• 预训练（Pre-training）

• 有监督微调（Supervised Fine-tuning，SFT）

• 人类反馈强化学习（Reinforcement Learning from Human Feedback，RLHF）

在不同训练阶段使用的数据集有所不同，下面会从 MiniMind 代码出发进行介绍和解读。

# 导入 json 库，用于处理 JSON 格式的数据文件（比如 .jsonl）
import json
# 导入 random 库，用于生成随机数，可能在数据增强或打乱数据时用到
import random
# 导入 re 库，用于正则表达式操作，方便进行复杂的文本匹配和处理
import re

# 导入 pandas 库，一个强大的数据分析工具，常用于读取和操作表格数据（如 CSV 文件）
import pandas as pd
# 导入 numpy 库，Python 科学计算的基础包，用于高效处理数组和矩阵
import numpy as np
# 从 PyTorch 的工具库中导入 Dataset 和 DataLoader
# Dataset 是所有自定义数据集的父类，我们需要继承它来创建自己的数据集
# DataLoader 是一个数据加载器，它可以自动地从 Dataset 中提取数据，并组织成一批一批(batch)的形式，方便模型训练
from torch.utils.data import Dataset, DataLoader
# 导入 PyTorch 核心库
import torch
# 从 scikit-learn 库中导入一个函数，用于方便地将数据集划分为训练集和测试集
from sklearn.model_selection import train_test_split
# 导入 os 库，用于与操作系统交互，比如处理文件路径
import os
# 导入 ast 库，用于处理 Python 的抽象语法树，可以安全地解析字符串形式的 Python 字面量
import ast
# 从 Hugging Face 的 transformers 库中导入 AutoTokenizer，它可以根据指定的路径或名称自动加载合适的分词器
from transformers import AutoTokenizer

# 这行代码是 Hugging Face tokenizers 的一个常见设置
# 作用是禁用分词器的并行处理功能
# 在某些环境（特别是 Jupyter Notebook）中，开启并行可能会导致程序卡死或报错，所以这里显式地关闭它以保证稳定性
os.environ["TOKENIZERS_PARALLELISM"] = "false"

# 这行代码的作用是加载我们之前在第一部分中保存的分词器

# AutoTokenizer.from_pretrained 是一个非常方便的函数
# 它会自动识别 './model/minimind_tokenizer' 目录下的配置文件
# 然后加载与该配置相匹配的分词器类型，并读取词汇表等信息
tokenizer = AutoTokenizer.from_pretrained('../raw/data/Minimind/model/minimind_tokenizer')

# 打印分词器的词汇表大小（vocab_size）
# 这个数字表示分词器总共认识多少个不同的词元（token）
# 模型的输出层维度通常会和这个词汇表大小一致
print(tokenizer.vocab_size)

6400

预训练数据集

预训练是模型在大规模语料上进行无监督学习的训练阶段，在该阶段，模型主要学习下一词预测的能力，简单的来说就是学会说话，而不是胡言乱语。因此，该阶段训练的模型不会具有问答能力，而是根据用户输入进行简单的词语接龙。

我们可以看一看预训练的数据集格式：

{"text": "如何才能摆脱拖延症？ 治愈拖延症并不容易，但以下建议可能有所帮助..."}

为了降低该 demo 的运行门槛，在 ./demo 文件夹下提供了包含两条训练数据的 pretrain_data.jsonl 文件作为熟悉训练流程的数据集 demo。

# 这段代码的目的是读取并展示预训练数据文件的内容

# 定义预训练数据文件的路径
path_pretrain = '../raw/data/Minimind/toydata/pretrain_data.jsonl'

# `with open(...)` 是 Python 中推荐的文件操作方式，它可以确保文件在使用后被自动关闭
# 'r' 表示以只读模式打开文件
# 'encoding='utf-8'' 指定文件编码，以正确处理中文字符
with open(path_pretrain, 'r', encoding='utf-8') as f:
    # `enumerate(f, 1)` 会遍历文件中的每一行，并同时提供行号（从 1 开始）和行内容
    for line_num, line in enumerate(f, 1):
        # 这个文件是 .jsonl 格式（JSON Lines），意味着每一行都是一个独立的、完整的 JSON 对象
        # 1. line.strip(): 去除行首和行尾的空白字符（如换行符）
        # 2. json.loads(...): 将这一行纯文本的 JSON 字符串解析成 Python 中的字典对象
        data = json.loads(line.strip())
        # 打印行号和解析后的数据内容
        print(f'Row {line_num}: {data}\n')

Row 1: {'text': 'LLM首先要学习的并非直接与人交流，而是让网络参数中充满知识的墨水，“墨水” 理论上喝的越饱越好，产生大量的对世界的知识积累。 预训练就是让Model先埋头苦学大量基本的知识，例如从Wiki百科、新闻、书籍整理大规模的高质量训练数据。 这个过程是“无监督”的，即人类不需要在过程中做任何“有监督”的校正，而是由模型自己从大量文本中总结规律学习知识点。 模型此阶段目的只有一个：学会词语接龙。例如我们输入“秦始皇”四个字，它可以接龙“是中国的第一位皇帝”。'}

Row 2: {'text': '经过预训练，LLM此时已经掌握了大量知识，然而此时它只会无脑地词语接龙，还不会与人聊天。 SFT阶段就需要把半成品LLM施加一个自定义的聊天模板进行微调。 例如模型遇到这样的模板【问题->回答，问题->回答】后不再无脑接龙，而是意识到这是一段完整的对话结束。 称这个过程为指令微调，就如同让已经学富五车的「牛顿」先生适应21世纪智能手机的聊天习惯，学习屏幕左侧是对方消息，右侧是本人消息这个规律。 在训练时，MiniMind的指令和回答长度被截断在512，是为了节省显存空间。就像我们学习时，会先从短的文章开始，当学会写作200字作文后，800字文章也可以手到擒来。 在需要长度拓展时，只需要准备少量的2k/4k/8k长度对话数据进行进一步微调即可（此时最好配合RoPE-NTK的基准差值）。'}

我们知道，构建一个深度学习数据集需要继承 torch.utils.data.dataset，并构建 DataLoader 数据迭代器进行迭代访问。下面，我们来看看 MiniMind 是如何抽象一个预训练数据集的。

# 定义一个用于预训练的数据集类，它必须继承自 PyTorch 的 torch.utils.data.Dataset
class PretrainDataset(Dataset):
    # 构造函数 __init__：在创建数据集对象时被调用
    def __init__(self, data_path, tokenizer, max_length=512):
        super().__init__() # 调用父类的构造函数
        self.tokenizer = tokenizer # 保存分词器
        self.max_length = max_length # 保存句子的最大长度
        self.samples = self.load_data(data_path) # 调用 load_data 方法加载所有数据到内存

    # load_data 方法：读取 .jsonl 文件
    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

    # __len__ 方法：返回数据集的总样本数，这是 Dataset 类的必须方法
    def __len__(self):
        return len(self.samples)

    # __getitem__ 方法：根据索引 `index` 获取单个样本，这是 Dataset 类的核心
    def __getitem__(self, index):
        # 1. 根据索引从加载的数据中获取一个样本（一个字典）
        sample = self.samples[index]

        # 2. 格式化文本：在原始文本前后分别加上开始符(BOS)和结束符(EOS)
        #    这能帮助模型明确地知道一段文本的开始和结束
        text = f"{self.tokenizer.bos_token}{str(sample['text'])}{self.tokenizer.eos_token}"
        
        # 3. 使用分词器将文本转换为 token ID
        encoding = self.tokenizer(
            text,
            max_length=self.max_length,      # 限制最大长度
            padding='max_length',            # 如果文本不够长，就用 padding token 填充到 max_length
            truncation=True,                 # 如果文本太长，就截断
            return_tensors='pt'              # 返回 PyTorch 张量 (Tensor)
        )
        
        # 4. 准备模型的输入(X)、标签(Y)和损失掩码(loss_mask)
        #    .squeeze() 去掉多余的维度，例如 (1, 512) -> (512)
        input_ids = encoding.input_ids.squeeze()
        
        #    创建一个损失掩码，标记哪些位置需要计算损失
        #    我们不希望在填充(padding)的部分计算损失，所以把 padding token 的位置设为 False (0)
        loss_mask = (input_ids != self.tokenizer.pad_token_id)
        
        # 5. 构造训练对 (X, Y) 用于“下一个词预测”任务
        #    例如，如果 input_ids 是 [BOS, "我", "爱", "你", EOS]
        #    X (输入) 就是 [BOS, "我", "爱", "你"]
        #    Y (标签) 就是 [  "我", "爱", "你", EOS]
        #    模型的目标是：看到 X 的每个词，都能预测出 Y 中对应位置的下一个词
        X = input_ids[:-1] # X 是从第一个到倒数第二个 token
        Y = input_ids[1:]  # Y 是从第二个到最后一个 token
        
        #    损失掩码也需要对齐，因为 Y 是我们的预测目标
        loss_mask = loss_mask[1:]
        
        # 返回处理好的输入、标签和损失掩码
        return X.long(), Y.long(), loss_mask.long()

# 使用预训练数据路径和分词器，创建一个 PretrainDataset 对象
pretrain_dataset = PretrainDataset(path_pretrain, tokenizer)

# 调用 len() 函数，这会自动调用 pretrain_dataset 的 __len__ 方法
print(f'预训练数据集长度{len(pretrain_dataset)}')

# 下面的代码被注释掉了，但我们可以解释它的作用：
# `pretrain_dataset[0]` 会调用 __getitem__(0) 方法，获取并处理第一个样本
# x, y, lm 会分别接收到模型的输入、标签和损失掩码
# `print(x.shape, y.shape, lm.shape)` 会打印它们的形状，
# 因为 max_length=512，所以形状都应该是 (511,)

预训练数据集长度2

有监督微调数据集

有监督微调（Supervised Fine Tuning，SFT）对预训练后得到的基座 LLM 施加一个自定义聊天模板进行微调，由于在这一阶段，模型训练的目标是根据用户指令生成响应（构建问答体系），故又称为指令微调。

我们可以看一看有监督微调的数据集格式：

{
    "conversations": [
        {"role": "user", "content": "你好"},
        {"role": "assistant", "content": "你好！"},
        {"role": "user", "content": "再见"},
        {"role": "assistant", "content": "再见！"}
    ]
}

为了降低该 demo 的运行门槛，在 ./demo 文件夹下提供了包含两条 conversation 问答数据的 sft_data.jsonl 文件作为熟悉训练流程的数据集 demo。

# 这段代码和查看预训练数据非常类似，只是文件路径不同

# 定义 SFT (有监督微调) 数据文件的路径
path_sft = '../raw/data/Minimind/toydata/sft_data.jsonl'

# 同样使用 `with open` 来安全地读取文件
with open(path_sft, 'r', encoding='utf-8') as f:
    # 遍历文件的每一行
    for line_num, line in enumerate(f, 1):
        # 解析每一行的 JSON 字符串
        data = json.loads(line.strip())
        # 打印出来，可以看到 SFT 数据的格式是一个包含多轮对话的列表
        print(f'Row {line_num}: {data}\n')

Row 1: {'conversations': [{'role': 'user', 'content': '你好吗？'}, {'role': 'assistant', 'content': '我很好，谢谢！你呢？'}, {'role': 'user', 'content': '我也很好，谢谢！'}, {'role': 'assistant', 'content': '太好了！祝你今天愉快！'}]}

Row 2: {'conversations': [{'role': 'user', 'content': '你喜欢什么运动？'}, {'role': 'assistant', 'content': '我喜欢跑步和游泳。你呢？'}, {'role': 'user', 'content': '我喜欢打篮球！'}, {'role': 'assistant', 'content': '篮球很棒！是一个很好的团队运动。'}]}

接下来，我们尝试构造一个数据集对象，实现对 sft 格式数据的读取与处理。

# 定义一个用于 SFT 的数据集类
class SFTDataset(Dataset):
    def __init__(self, jsonl_path, tokenizer, max_length=512):
        super().__init__()
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.samples = self.load_data(jsonl_path)
        # 预先将 assistant 回复的开始和结束标记转换为 token ID
        # 这是一种优化，避免在每次生成 loss_mask 时都重新分词
        # `add_special_tokens=False` 确保不会额外添加 BOS/EOS 符
        self.bos_id = tokenizer('<s>assistant\n', add_special_tokens=False).input_ids
        self.eos_id = tokenizer('</s>\n', add_special_tokens=False).input_ids

    def load_data(self, path): # 和 PretrainDataset 中的一样
        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
        
    def __len__(self): # 和 PretrainDataset 中的一样
        return len(self.samples)

    # 辅助方法：创建符合聊天格式的提示词
    def _create_chat_prompt(self, conversations):
        # Hugging Face 的分词器提供了一个 `apply_chat_template` 方法
        # 它可以根据预设的模板，将多轮对话（包含角色和内容）自动格式化成一个单一的字符串
        # 例如，格式化成：<s>user\n你好</s>\n<s>assistant\n我很好</s>\n
        messages = conversations # 原始数据格式已经是 `apply_chat_template` 要求的格式
        return self.tokenizer.apply_chat_template(
            messages,
            tokenize=False, # 表示只返回格式化后的字符串，而不是 token ID
            add_generation_prompt=False # 表示不在末尾添加 "<s>assistant\n" 提示模型生成
        )

    # 辅助方法：生成损失掩码，这是 SFT 的关键
    def _generate_loss_mask(self, input_ids):
        # SFT 的目标是让模型学会“回答”，而不是学会“提问”
        # 所以，我们只在模型生成 assistant 回答的部分计算损失
        loss_mask = [0] * len(input_ids) # 初始化一个全为 0 的掩码
        i = 0
        while i < len(input_ids):
            # 搜索 assistant 回答的开始标记
            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, end + len(self.eos_id)):
                    if j < len(loss_mask): # 防止越界
                        loss_mask[j] = 1
                i = end + len(self.eos_id) # 更新搜索位置，跳过这段已处理的回答
            else:
                i += 1
        return loss_mask
    
    def __getitem__(self, index):
        sample = self.samples[index]
        # 1. 将对话数据格式化为单个字符串
        prompt = self._create_chat_prompt(sample['conversations'])
        
        # 2. 分词，并手动进行截断和填充
        input_ids = self.tokenizer(prompt).input_ids[:self.max_length]
        input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids))
        
        # 3. 生成特殊的损失掩码
        loss_mask = self._generate_loss_mask(input_ids)

        # 4. 创建 X, Y 和对齐后的 loss_mask，逻辑同 PretrainDataset
        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 数据路径和分词器，创建一个 SFTDataset 对象
sft_dataset = SFTDataset(path_sft, tokenizer)

# 打印数据集长度
print(len(sft_dataset))

# 获取第一个样本来检查
x, y, lm = sft_dataset[0]

# 打印返回的张量形状
print(f'样本 shape = {x.shape}, 标签 shape = {y.shape}, loss_mask shape {lm.shape}')

# 如果取消最后一行代码的注释并运行，你会看到 loss_mask (lm) 的内容。
# 它大部分是 0，只在 assistant 回答的位置是 1，这证明我们的 _generate_loss_mask 方法工作正常。

2
样本 shape = torch.Size([511]), 标签 shape = torch.Size([511]), loss_mask shape torch.Size([511])

人类反馈强化学习数据集

在 MiniMind 项目中，采用直接偏好优化（Direct Parameter Optimization，DPO）训练大模型对齐人类偏好。在这一训练阶段，模型将会根据提供的问答正反例进行偏好优化，从而降低让人类不满意的答案出现的几率。

与PPO(Proximal Policy Optimization)这种需要奖励模型、价值模型的RL算法不同； DPO通过推导PPO奖励模型的显式解，把在线奖励模型换成离线数据，Ref模型输出可以提前保存。 DPO性能几乎不变，只用跑 actor_model 和 ref_model 两个模型，大大节省显存开销和增加训练稳定性。

我们可以看一看有监督微调的数据集格式：

{
  "chosen": [
    {"content": "Query", "role": "user"}, 
    {"content": "good answer", "role": "assistant"}
  ], 
  "rejected": [
    {"content": "Query", "role": "user"}, 
    {"content": "bad answer", "role": "assistant"}
  ]
}

为了降低该 demo 的运行门槛，在 ./demo 文件夹下提供了包含两条 conversation 问答数据的 sft_data.jsonl 文件作为熟悉训练流程的数据集 demo。

# 这段代码的目的是展示 DPO (直接偏好优化) 数据的格式

# 定义 DPO 数据文件的路径
path_dpo = './toydata/dpo_data.jsonl'

# 注意！这里的代码有一个小笔误，它打开的是 `path_sft` 而不是 `path_dpo`。
# 所以，它实际上会再次打印出 SFT 文件的内容。
# 如果要正确查看 DPO 数据，应该使用 `with open(path_dpo, ...)`
with open(path_dpo, 'r', encoding='utf-8') as f: # 假设这里已修正为 path_dpo
    for line_num, line in enumerate(f, 1):
        data = json.loads(line.strip())
        # DPO 数据每条都包含 "chosen" (好的回答) 和 "rejected" (不好的回答) 两部分
        print(f'Row {line_num}: {data}\n')

Row 1: {'conversations': [{'role': 'user', 'content': '你好吗？'}, {'role': 'assistant', 'content': '我很好，谢谢！你呢？'}, {'role': 'user', 'content': '我也很好，谢谢！'}, {'role': 'assistant', 'content': '太好了！祝你今天愉快！'}]}

Row 2: {'conversations': [{'role': 'user', 'content': '你喜欢什么运动？'}, {'role': 'assistant', 'content': '我喜欢跑步和游泳。你呢？'}, {'role': 'user', 'content': '我喜欢打篮球！'}, {'role': 'assistant', 'content': '篮球很棒！是一个很好的团队运动。'}]}

接下来，我们尝试构造 json 对象，实现对 dpo 格式数据的读取和处理

# 定义一个用于 DPO 的数据集类
class DPODataset(Dataset):
    def __init__(self, file_path, tokenizer, max_length=512):
        super().__init__()
        self.tokenizer = tokenizer
        self.max_length = max_length
        # 和 SFTDataset 一样，预先处理好 assistant 的标记
        self.bos_id = tokenizer('<s>assistant\n', add_special_tokens=False).input_ids
        self.eos_id = tokenizer('</s>\n', add_special_tokens=False).input_ids
        # 直接在构造函数中加载数据
        with open(file_path, 'r', encoding='utf-8') as f:
            self.data = [json.loads(line.strip()) for line in f]

    def __len__(self):
        return len(self.data)

    def _generate_loss_mask(self, input_ids):
        # 这个方法和 SFTDataset 中的完全一样，因为目标相同：只在 assistant 回答上计算损失
        loss_mask = [0] * len(input_ids)
        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
                for j in range(start, end + len(self.eos_id)):
                    if j < 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
    
    def __getitem__(self, index):
        # DPO 的核心是对比 "chosen" 和 "rejected"
        item = self.data[index]
        chosen_conv = item['chosen']     # 获取好的对话
        rejected_conv = item['rejected'] # 获取不好的对话

        # 1. 分别为 chosen 和 rejected 对话创建提示词字符串
        chosen_prompt = self.tokenizer.apply_chat_template(chosen_conv, tokenize=False, add_generation_prompt=False)
        rejected_prompt = self.tokenizer.apply_chat_template(rejected_conv, tokenize=False, add_generation_prompt=False)

        # 2. 分别对它们进行分词、填充和截断
        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')

        # 3. 分别为它们生成损失掩码
        chosen_loss_mask = self._generate_loss_mask(chosen_encoding['input_ids'])
        rejected_loss_mask = self._generate_loss_mask(rejected_encoding['input_ids'])
        
        # 4. 分别为它们创建 (X, Y) 训练对和对齐后的掩码
        chosen_input_ids = chosen_encoding['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)
        
        rejected_input_ids = rejected_encoding['input_ids']
        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)

        # 5. 将所有处理好的张量打包成一个字典返回
        #    DPO 的损失函数会同时用到 chosen 和 rejected 的这些信息来进行优化
        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
        }

# 1. 使用 DPO 数据路径和分词器，创建一个 DPODataset 对象
dpo_dataset = DPODataset(path_dpo, tokenizer)

# 2. 打印数据集长度
print(f'DPO 数据集长度：{len(dpo_dataset)}')

# 3. 获取第一个样本
#    注意：在 `DPODataset` 的 `__getitem__` 方法中，我们留下了 `print(chosen_prompt)` 和 `print(rejected_prompt)`
#    所以当执行 `dpo_dataset[0]` 这行代码时，会自动打印出第一个样本的 chosen 和 rejected 对话的格式化字符串
res = dpo_dataset[0]

# 4. `res` 现在是一个字典，包含了 'x_chosen', 'y_chosen', 'mask_chosen' 等6个张量
#    可以通过 `res.keys()` 查看所有键，或通过 `res['x_chosen'].shape` 查看具体张量的形状

DPO 数据集长度：2
<s>system
你是 MiniMind，是一个有用的人工智能助手。</s>
<s>user
你好吗？</s>
<s>assistant
我很好，谢谢！你呢？</s>
<s>user
今天过得怎么样？</s>
<s>assistant
挺好的，去跑步了，心情不错。</s>

<s>system
你是 MiniMind，是一个有用的人工智能助手。</s>
<s>user
你好吗？</s>
<s>assistant
不好，我很累。</s>
<s>user
你喜欢什么运动？</s>
<s>assistant
我不喜欢运动，没兴趣。</s>