minimind/notebook/2_dataset.ipynb

563 lines
24 KiB
Plaintext
Raw Blame History

This file contains ambiguous Unicode characters!

This file contains ambiguous Unicode characters that may be confused with others in your current locale. If your use case is intentional and legitimate, you can safely ignore this warning. Use the Escape button to highlight these characters.

{
"cells": [
{
"cell_type": "markdown",
"id": "44db347d-fef7-4e56-aeab-ce55bdab8eab",
"metadata": {},
"source": [
"# 2-Dataset\n",
"\n",
"到这里我们便完成了对于 MiniMind Tokenizer 和 Model 部分的全部了解,我们所熟悉的大语言模型正是由这个组件构成的,接下来,我们需要对大模型训练所使用的数据集结构有个基本的认识。\n",
"\n",
"想要训练一个能够正常对话,并且符合人类对话偏好的大模型一般需要经过以下几个训练阶段:\n",
"\n",
"- 预训练(Pre-training)\n",
"- 有监督微调(Supervised Fine-tuning,SFT)\n",
"- 人类反馈强化学习(Reinforcement Learning from Human Feedback,RLHF)\n",
"\n",
"在不同训练阶段使用的数据集有所不同,下面会从 MiniMind 代码出发进行介绍和解读。"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "e41156fd-ca0e-4083-a7a4-a5f98019ab88",
"metadata": {},
"outputs": [],
"source": [
"import json\n",
"import random\n",
"import re\n",
"\n",
"import pandas as pd\n",
"import numpy as np\n",
"from torch.utils.data import Dataset, DataLoader\n",
"import torch\n",
"from sklearn.model_selection import train_test_split\n",
"import os\n",
"import ast\n",
"from transformers import AutoTokenizer\n",
"\n",
"os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\""
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "27d18ad2-3f71-435c-af03-ad9d347aa048",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"6400\n"
]
}
],
"source": [
"# 从 ../model 目录加载分词器\n",
"\n",
"tokenizer = AutoTokenizer.from_pretrained('../model/minimind_tokenizer')\n",
"print(tokenizer.vocab_size)"
]
},
{
"cell_type": "markdown",
"id": "ad08fe92-1d04-4b9c-b689-6ad9d763a7c2",
"metadata": {},
"source": [
"## 预训练数据集\n",
"\n",
"预训练是模型在大规模语料上进行无监督学习的训练阶段,在该阶段,模型主要学习下一词预测的能力,简单的来说就是学会说话,而不是胡言乱语。因此,该阶段训练的模型不会具有问答能力,而是根据用户输入进行简单的词语接龙。\n",
"\n",
"我们可以看一看预训练的数据集格式:\n",
"\n",
"```\n",
"{\"text\": \"如何才能摆脱拖延症? 治愈拖延症并不容易,但以下建议可能有所帮助...\"}\n",
"```\n",
"\n",
"为了降低该 demo 的运行门槛,在 `./demo` 文件夹下提供了包含两条训练数据的 `pretrain_data.jsonl` 文件作为熟悉训练流程的数据集 demo。"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "cdcac4aa-0ced-4a74-bba0-bc248923c7bc",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Row 1: {'text': 'LLM首先要学习的并非直接与人交流,而是让网络参数中充满知识的墨水,“墨水” 理论上喝的越饱越好,产生大量的对世界的知识积累。 预训练就是让Model先埋头苦学大量基本的知识,例如从Wiki百科、新闻、书籍整理大规模的高质量训练数据。 这个过程是“无监督”的,即人类不需要在过程中做任何“有监督”的校正,而是由模型自己从大量文本中总结规律学习知识点。 模型此阶段目的只有一个:学会词语接龙。例如我们输入“秦始皇”四个字,它可以接龙“是中国的第一位皇帝”。'}\n",
"\n",
"Row 2: {'text': '经过预训练,LLM此时已经掌握了大量知识,然而此时它只会无脑地词语接龙,还不会与人聊天。 SFT阶段就需要把半成品LLM施加一个自定义的聊天模板进行微调。 例如模型遇到这样的模板【问题->回答,问题->回答】后不再无脑接龙,而是意识到这是一段完整的对话结束。 称这个过程为指令微调,就如同让已经学富五车的「牛顿」先生适应21世纪智能手机的聊天习惯,学习屏幕左侧是对方消息,右侧是本人消息这个规律。 在训练时,MiniMind的指令和回答长度被截断在512,是为了节省显存空间。就像我们学习时,会先从短的文章开始,当学会写作200字作文后,800字文章也可以手到擒来。 在需要长度拓展时,只需要准备少量的2k/4k/8k长度对话数据进行进一步微调即可(此时最好配合RoPE-NTK的基准差值)。'}\n",
"\n"
]
}
],
"source": [
"# 我们可以查看一下 demo 中提供的数据\n",
"path_pretrain = './demo/pretrain_data.jsonl'\n",
"\n",
"with open(path_pretrain, 'r', encoding='utf-8') as f:\n",
" for line_num, line in enumerate(f, 1):\n",
" data = json.loads(line.strip())\n",
" print(f'Row {line_num}: {data}\\n')"
]
},
{
"cell_type": "markdown",
"id": "0c0707e3-3ff9-424b-a5be-dab9db55aa44",
"metadata": {},
"source": [
"我们知道,构建一个深度学习数据集需要继承 `torch.utils.data.dataset`,并构建 DataLoader 数据迭代器进行迭代访问。下面,我们来看看 MiniMind 是如何抽象一个预训练数据集的。"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "fb2458e3-3661-40de-96be-aa036291a8fb",
"metadata": {},
"outputs": [],
"source": [
"class PretrainDataset(Dataset):\n",
" def __init__(self, data_path, tokenizer, max_length=512):\n",
" super().__init__()\n",
" self.tokenizer = tokenizer\n",
" self.max_length = max_length\n",
" self.samples = self.load_data(data_path)\n",
"\n",
" def load_data(self, path):\n",
" \"\"\"按行读取 jsonl 文件,并存储在列表中\"\"\"\n",
" samples = []\n",
" with open(path, 'r', encoding='utf-8') as f:\n",
" for line_num, line in enumerate(f, 1):\n",
" data = json.loads(line.strip())\n",
" samples.append(data)\n",
" return samples\n",
"\n",
" def __len__(self):\n",
" return len(self.samples)\n",
"\n",
" def __getitem__(self, index):\n",
" sample = self.samples[index]\n",
"\n",
" text = f\"{self.tokenizer.bos_token}{str(sample['text'])}{self.tokenizer.eos_token}\"\n",
" # print(text) # uncomment to see formating prompt\n",
" encoding = self.tokenizer(\n",
" text,\n",
" max_length=self.max_length,\n",
" padding='max_length',\n",
" truncation=True,\n",
" return_tensors='pt'\n",
" )\n",
" # print(encoding) # uncomment to see encoding result\n",
" input_ids = encoding.input_ids.squeeze()\n",
" loss_mask = (input_ids != self.tokenizer.pad_token_id) # mask to ignore loss on pad token\n",
"\n",
" X = torch.tensor(input_ids[:-1], dtype=torch.long) # <eos> token not included\n",
" Y = torch.tensor(input_ids[1:], dtype=torch.long) # <bos> token not included\n",
" loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long) # align to tensor X\n",
" return X, Y, loss_mask"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "677a1bbb-0df7-45e5-a3ed-d4431cddf43e",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"预训练数据集长度2\n"
]
}
],
"source": [
"pretrain_dataset = PretrainDataset(path_pretrain, tokenizer)\n",
"print(f'预训练数据集长度{len(pretrain_dataset)}')\n",
"# x, y, lm = pretrain_dataset[0]\n",
"# print(x.shape, y.shape, lm.shape)"
]
},
{
"cell_type": "markdown",
"id": "ebb160cb-0e96-45a7-a1c0-be145895e70f",
"metadata": {},
"source": [
"## 有监督微调数据集\n",
"\n",
"有监督微调(Supervised Fine Tuning,SFT)对预训练后得到的基座 LLM 施加一个自定义聊天模板进行微调,由于在这一阶段,模型训练的目标是根据用户指令生成响应(构建问答体系),故又称为指令微调。\n",
"\n",
"我们可以看一看有监督微调的数据集格式:\n",
"\n",
"```\n",
"{\n",
" \"conversations\": [\n",
" {\"role\": \"user\", \"content\": \"你好\"},\n",
" {\"role\": \"assistant\", \"content\": \"你好!\"},\n",
" {\"role\": \"user\", \"content\": \"再见\"},\n",
" {\"role\": \"assistant\", \"content\": \"再见!\"}\n",
" ]\n",
"}\n",
"```\n",
"\n",
"为了降低该 demo 的运行门槛,在 `./demo` 文件夹下提供了包含两条 conversation 问答数据的 `sft_data.jsonl` 文件作为熟悉训练流程的数据集 demo。"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "7e025e25-c937-4468-ae66-b32447440658",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Row 1: {'conversations': [{'role': 'user', 'content': '你好吗?'}, {'role': 'assistant', 'content': '我很好,谢谢!你呢?'}, {'role': 'user', 'content': '我也很好,谢谢!'}, {'role': 'assistant', 'content': '太好了!祝你今天愉快!'}]}\n",
"\n",
"Row 2: {'conversations': [{'role': 'user', 'content': '你喜欢什么运动?'}, {'role': 'assistant', 'content': '我喜欢跑步和游泳。你呢?'}, {'role': 'user', 'content': '我喜欢打篮球!'}, {'role': 'assistant', 'content': '篮球很棒!是一个很好的团队运动。'}]}\n",
"\n"
]
}
],
"source": [
"# 我们可以查看一下 demo 中提供的数据\n",
"path_sft = './demo/sft_data.jsonl'\n",
"\n",
"with open(path_sft, 'r', encoding='utf-8') as f:\n",
" for line_num, line in enumerate(f, 1):\n",
" data = json.loads(line.strip())\n",
" print(f'Row {line_num}: {data}\\n')"
]
},
{
"cell_type": "markdown",
"id": "5f988233-530f-4ddf-8d64-9532442a2535",
"metadata": {},
"source": [
"接下来,我们尝试构造一个数据集对象,实现对 sft 格式数据的读取与处理。"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "8b25f0c6-b5b5-47b1-974f-15a2f0d2760e",
"metadata": {},
"outputs": [],
"source": [
"class SFTDataset(Dataset):\n",
" def __init__(self, jsonl_path, tokenizer, max_length=512):\n",
" super().__init__()\n",
" self.tokenizer = tokenizer\n",
" self.max_length = max_length\n",
" self.samples = self.load_data(jsonl_path)\n",
" self.bos_id = tokenizer('<s>assistant\\n', add_special_tokens=False).input_ids # set bos token\n",
" self.eos_id = tokenizer('</s>\\n', add_special_tokens=False).input_ids # set eos token\n",
"\n",
" def load_data(self, path):\n",
" samples = []\n",
" with open(path, 'r', encoding='utf-8') as f:\n",
" for line_num, line in enumerate(f, 1):\n",
" data = json.loads(line.strip())\n",
" samples.append(data)\n",
" return samples\n",
" \n",
" def __len__(self):\n",
" return len(self.samples)\n",
"\n",
" def _create_chat_prompt(self, conversations):\n",
" \"\"\"构建符合 ChatML 格式的对话\"\"\"\n",
" messages = []\n",
" for i, turn in enumerate(conversations): # for each speaker in one conversation\n",
" role = 'user' if i % 2 == 0 else 'assistant'\n",
" messages.append({\"role\": role, \"content\": turn['content']})\n",
" return self.tokenizer.apply_chat_template(\n",
" messages,\n",
" tokenize=False,\n",
" add_generation_prompt=False\n",
" )\n",
"\n",
" def _generate_loss_mask(self, input_ids):\n",
" loss_mask = [0] * len(input_ids)\n",
" i = 0\n",
" while i < len(input_ids):\n",
" if input_ids[i:i + len(self.bos_id)] == self.bos_id: # check if reach bos token\n",
" ########### find content start & end point ###########\n",
" start = i + len(self.bos_id) # set start point at the end of bos token\n",
" end = start\n",
" while end < len(input_ids):\n",
" if input_ids[end:end + len(self.eos_id)] == self.eos_id: # check if reach eos token\n",
" break\n",
" end += 1\n",
" ######################################################\n",
" for j in range(start + 1, min(end + len(self.eos_id) + 1, self.max_length)): # ignore tokens that reach max input length\n",
" loss_mask[j] = 1\n",
" i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids) # update i to exit current conversation turn\n",
" else:\n",
" i += 1\n",
" return loss_mask\n",
" \n",
" def __getitem__(self, index):\n",
" sample = self.samples[index]\n",
" prompt = self._create_chat_prompt(sample['conversations'])\n",
" # print(prompt) # uncomment to see formating prompt\n",
" input_ids = self.tokenizer(prompt).input_ids[:self.max_length] # encode input\n",
" input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids)) # encode padding\n",
" # print(input_ids) # uncomment to see encoded prompt\n",
" \n",
" loss_mask = self._generate_loss_mask(input_ids)\n",
"\n",
" X = torch.tensor(input_ids[:-1], dtype=torch.long)\n",
" Y = torch.tensor(input_ids[1:], dtype=torch.long)\n",
" loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long) # align with pred pos\n",
"\n",
" return X, Y, loss_mask"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "67a62dbf-6a50-4faf-8ad1-5228eb8aa7c9",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"2\n",
"样本 shape = torch.Size([511]), 标签 shape = torch.Size([511]), loss_mask shape torch.Size([511])\n"
]
}
],
"source": [
"sft_dataset = SFTDataset(path_sft, tokenizer)\n",
"print(len(sft_dataset))\n",
"x, y, lm = sft_dataset[0]\n",
"print(f'样本 shape = {x.shape}, 标签 shape = {y.shape}, loss_mask shape {lm.shape}')\n",
"# print(lm) # 打印 loss mask,你会发现在序列中有两处以 1 填充的序列,这是因为我们在一个 conversation 中开展了两轮对话,其中只有 assistant 回复计算损失"
]
},
{
"cell_type": "markdown",
"id": "a7094ff6-4657-4483-8c34-785ffe9adbd4",
"metadata": {},
"source": [
"## 人类反馈强化学习数据集\n",
"\n",
"在 MiniMind 项目中,采用直接偏好优化(Direct Parameter Optimization,DPO)训练大模型对齐人类偏好。在这一训练阶段,模型将会根据提供的问答正反例进行偏好优化,从而降低让人类不满意的答案出现的几率。\n",
"\n",
"与PPO(Proximal Policy Optimization)这种需要奖励模型、价值模型的RL算法不同; DPO通过推导PPO奖励模型的显式解,把在线奖励模型换成离线数据,Ref模型输出可以提前保存。 DPO性能几乎不变,只用跑 actor_model 和 ref_model 两个模型,大大节省显存开销和增加训练稳定性。\n",
"\n",
"我们可以看一看有监督微调的数据集格式:\n",
"\n",
"```\n",
"{\n",
" \"chosen\": [\n",
" {\"content\": \"Query\", \"role\": \"user\"}, \n",
" {\"content\": \"good answer\", \"role\": \"assistant\"}\n",
" ], \n",
" \"rejected\": [\n",
" {\"content\": \"Query\", \"role\": \"user\"}, \n",
" {\"content\": \"bad answer\", \"role\": \"assistant\"}\n",
" ]\n",
"}\n",
"```\n",
"\n",
"为了降低该 demo 的运行门槛,在 ./demo 文件夹下提供了包含两条 conversation 问答数据的 sft_data.jsonl 文件作为熟悉训练流程的数据集 demo。"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "81682b5f-0d3e-4fa2-aec4-0be2b19d5c05",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Row 1: {'conversations': [{'role': 'user', 'content': '你好吗?'}, {'role': 'assistant', 'content': '我很好,谢谢!你呢?'}, {'role': 'user', 'content': '我也很好,谢谢!'}, {'role': 'assistant', 'content': '太好了!祝你今天愉快!'}]}\n",
"\n",
"Row 2: {'conversations': [{'role': 'user', 'content': '你喜欢什么运动?'}, {'role': 'assistant', 'content': '我喜欢跑步和游泳。你呢?'}, {'role': 'user', 'content': '我喜欢打篮球!'}, {'role': 'assistant', 'content': '篮球很棒!是一个很好的团队运动。'}]}\n",
"\n"
]
}
],
"source": [
"# 我们可以查看一下 demo 中提供的数据\n",
"path_dpo = './demo/dpo_data.jsonl'\n",
"\n",
"with open(path_sft, 'r', encoding='utf-8') as f:\n",
" for line_num, line in enumerate(f, 1):\n",
" data = json.loads(line.strip())\n",
" print(f'Row {line_num}: {data}\\n')"
]
},
{
"cell_type": "markdown",
"id": "ebb54557-54f4-4004-8a7b-443b8671002b",
"metadata": {},
"source": [
"接下来,我们尝试构造 json 对象,实现对 dpo 格式数据的读取和处理"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "efa98777-5d3a-4b3f-a9fb-e0f05e4af81d",
"metadata": {},
"outputs": [],
"source": [
"class DPODataset(Dataset):\n",
" def __init__(self, file_path, tokenizer, max_length=512):\n",
" super().__init__()\n",
" self.tokenizer = tokenizer\n",
" self.max_length = max_length\n",
" self.padding = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0\n",
" self.bos_id = tokenizer('<s>assistant\\n', add_special_tokens=False).input_ids # content prefix\n",
" self.eos_id = tokenizer('</s>\\n', add_special_tokens=False).input_ids # content suffix\n",
" with open(file_path, 'r', encoding='utf-8') as f: # load data\n",
" self.data = []\n",
" for line in f:\n",
" line = line.strip()\n",
" obj = json.loads(line)\n",
" self.data.append(obj)\n",
"\n",
" def __len__(self):\n",
" return len(self.data)\n",
"\n",
" def _generate_loss_mask(self, input_ids):\n",
" \"\"\"此处的损失掩码生成函数与 SFT 阶段逻辑一致\"\"\"\n",
" loss_mask = [0] * len(input_ids)\n",
" i = 0\n",
" while i < len(input_ids):\n",
" if input_ids[i:i + len(self.bos_id)] == self.bos_id:\n",
" start = i + len(self.bos_id)\n",
" end = start\n",
" while end < len(input_ids):\n",
" if input_ids[end:end + len(self.eos_id)] == self.eos_id:\n",
" break\n",
" end += 1\n",
" for j in range(start + 1, min(end + len(self.eos_id) + 1, self.max_length)):\n",
" loss_mask[j] = 1\n",
" i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)\n",
" else:\n",
" i += 1\n",
" return loss_mask\n",
" \n",
" def __getitem__(self, index):\n",
" item = self.data[index]\n",
" chosen = item['chosen'] # 一个 list,里面包含若干 {role, content}\n",
" rejected = item['rejected'] # 同上\n",
" # format prompt\n",
" chosen_prompt = self.tokenizer.apply_chat_template(\n",
" chosen, tokenize=False, add_generation_prompt=False\n",
" )\n",
" print(chosen_prompt) # uncomment to see formating prompt\n",
" rejected_prompt = self.tokenizer.apply_chat_template(\n",
" rejected, tokenize = False, add_gerneration_prompt=False\n",
" )\n",
" print(rejected_prompt) # uncomment to see formating prompt\n",
" # tokenize\n",
" chosen_encoding = self.tokenizer(\n",
" chosen_prompt, truncation=True, max_length=self.max_length, padding='max_length'\n",
" )\n",
" rejected_encoding = self.tokenizer(\n",
" rejected_prompt, truncation=True, max_length=self.max_length, padding='max_length'\n",
" )\n",
" # generate loss mask\n",
" chosen_input_ids = chosen_encoding['input_ids']\n",
" chosen_loss_mask = self._generate_loss_mask(chosen_input_ids)\n",
" rejected_input_ids = rejected_encoding['input_ids']\n",
" rejected_loss_mask = self._generate_loss_mask(rejected_input_ids)\n",
" # same as sft / pretrain\n",
" x_chosen = torch.tensor(chosen_input_ids[:-1], dtype=torch.long)\n",
" y_chosen = torch.tensor(chosen_input_ids[1:], dtype=torch.long)\n",
" mask_chosen = torch.tensor(chosen_loss_mask[1:], dtype=torch.long)\n",
" x_rejected = torch.tensor(rejected_input_ids[:-1], dtype=torch.long)\n",
" y_rejected = torch.tensor(rejected_input_ids[1:], dtype=torch.long)\n",
" mask_rejected = torch.tensor(rejected_loss_mask[1:], dtype=torch.long)\n",
"\n",
" return {\n",
" 'x_chosen': x_chosen,\n",
" 'y_chosen': y_chosen,\n",
" 'mask_chosen': mask_chosen,\n",
" 'x_rejected': x_rejected,\n",
" 'y_rejected': y_rejected,\n",
" 'mask_rejected': mask_rejected\n",
" }"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "5ba205c4-0dd2-48fd-8aa2-26a42afc90d8",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"DPO 数据集长度:2\n",
"<s>system\n",
"你是 MiniMind,是一个有用的人工智能助手。</s>\n",
"<s>user\n",
"你好吗?</s>\n",
"<s>assistant\n",
"我很好,谢谢!你呢?</s>\n",
"<s>user\n",
"今天过得怎么样?</s>\n",
"<s>assistant\n",
"挺好的,去跑步了,心情不错。</s>\n",
"\n",
"<s>system\n",
"你是 MiniMind,是一个有用的人工智能助手。</s>\n",
"<s>user\n",
"你好吗?</s>\n",
"<s>assistant\n",
"不好,我很累。</s>\n",
"<s>user\n",
"你喜欢什么运动?</s>\n",
"<s>assistant\n",
"我不喜欢运动,没兴趣。</s>\n",
"\n"
]
}
],
"source": [
"dpo_dataset = DPODataset(path_dpo, tokenizer)\n",
"print(f'DPO 数据集长度:{len(dpo_dataset)}')\n",
"res = dpo_dataset[0]\n",
"# 如有需要,请自主查看 res 中的元素"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.14"
}
},
"nbformat": 4,
"nbformat_minor": 5
}