从零手搓中文大模型｜Day03｜数据预处理

走过路过不要错过，先关注一下，第一时间获取最新进度（或催更）

从零手搓中文大模型｜🚀Day03

数据预处理

虽然省略了数据清洗的逻辑，但是我们还是需要对数据进行预处理，以便于后续的模型训练。

包括以下两个细节：

1. 在每个文本后添加eos标记，以便于模型识别句子的结束。

2. 将文本转换为数字序列，以便于模型处理。

   这一步其实也可以放到模型训练的时候进行，但提前处理可以减少训练时的计算量。

数据集划分

解压数据集，得到48个jsonl文件，共计3952863行json数据。

我之前已经解压过了，并且将原始数据和处理过后的数据分别存在了不同路径下。

这里把命令贴出来以供参考。

# !mkdir -p ../../Data/TinyStoriesChinese/raw_data/train
# !mkdir -p ../../Data/TinyStoriesChinese/raw_data/val
# !mkdir -p ../../Data/TinyStoriesChinese/processed_data/train
# !mkdir -p ../../Data/TinyStoriesChinese/processed_data/val

# !tar zxvf ../../Data/TinyStoriesChinese/TinyStories_all_data_zh.tar.gz -C ../../Data/TinyStoriesChinese/raw_data/train

我把最后一个文件data47_zh.jsonl（共计78538行）里切分出来4w行作为eval数据。

# !mv ../../Data/TinyStoriesChinese/raw_data/train/data47_zh.jsonl ../../Data/TinyStoriesChinese/raw_data/val/
# !head -n 40000 ../../Data/TinyStoriesChinese/raw_data/val/data47_zh.jsonl > ../../Data/TinyStoriesChinese/raw_data/val/val.jsonl
# !tail -n +40000 ../../Data/TinyStoriesChinese/raw_data/val/data47_zh.jsonl > ../../Data/TinyStoriesChinese/raw_data/train/data47_zh.jsonl
# !rm ../../Data/TinyStoriesChinese/raw_data/val/data47_zh.jsonl

先看一条数据

（都打印出来太长了，所以只输出前100个字符）

import json

with open("../../Data/TinyStoriesChinese/raw_data/train/data00_zh.jsonl", "r") as f:
    for line in f.readlines():
        js = json.loads(line)
        print(js["story_zh"][:100])
        break

莉莉和本是朋友。他们喜欢在公园里玩。有一天，他们在一棵大树下看到了一个秋千。莉莉想试试那个秋千。她跑到树下，爬上了秋千。
"推我，本！"她说。本轻轻地推了她一下。莉莉感到很开心。她越荡越高，笑着喊叫。

适配框架API

由于选择了使用⚡️litgpt框架进行训练，所以需要引入框架相关的Class和API来封装我们的数据准备逻辑。

这里我们可以参考源码里集成的Tinyllama的数据预处理代码里的代码，稍作修改。

主要是需要将Day02里的line处理逻辑封装到ligtgpt的API中。

但在此之前我们先熟悉一下litgpt的Tokenizer的使用方法：

先安装一下litgpt以及它所以赖的litdata:

# !pip install litgpt
# !pip install litdata

import torch
from litgpt import Tokenizer

litgpt_tokenizer = Tokenizer("../../References/chatglm3-6b")

这里也实验了一下结果，对比发现和咱们之前Day02里用原生Tokenizer处理的结果一致。

结果这里就不贴出来了，有兴趣的可以自己试一下。

⚠️不过需要注意litgpt的Tokenizer.encode返回的是一个torch的Tensor

import numpy as np

litgpt_encoded = litgpt_tokenizer.encode(
    json.loads(line)["story_zh"][:100], eos=True
)  # 记得设置eos=True
print(litgpt_encoded)
# print(np.array(litgpt_encoded, dtype=np.uint16))
print(litgpt_tokenizer.decode(litgpt_encoded))

tensor([30910, 56623, 56623, 54542, 50154, 31761, 31155, 31633, 31815, 54534,
        32693, 54662, 55409, 31155, 35632, 31123, 31633, 34383, 57427, 47658,
        54578, 34518, 31623, 55567, 55226, 31155, 56623, 56623, 54695, 39887,
        32437, 55567, 55226, 31155, 54790, 41309, 52624, 31123, 56856, 32660,
        55567, 55226, 31155,    13, 30955, 54834, 54546, 31123, 54613, 31404,
        30955, 36213, 31155, 54613, 36660, 54563, 54834, 43881, 32024, 31155,
        56623, 56623, 32707, 54657, 33436, 31155, 54790, 54937, 56567, 40714,
        31123, 38502, 56653, 55483, 31155,     2], dtype=torch.int32)
莉莉和本是朋友。他们喜欢在公园里玩。有一天，他们在一棵大树下看到了一个秋千。莉莉想试试那个秋千。她跑到树下，爬上了秋千。
"推我，本！"她说。本轻轻地推了她一下。莉莉感到很开心。她越荡越高，笑着喊叫。

数据处理代码

数据处理直接参考了上面给出的litgpt samples，我们需要仿照prepare_slimpajama.py实现里面相关函数（之前Day 02里实现的函数需要稍加改造一下）。

# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.

import json
import os
import time
import numpy as np
from pathlib import Path

from litgpt.tokenizer import Tokenizer
from litgpt.data.prepare_starcoder import DataChunkRecipe
from litdata import TokensLoader
from litgpt.utils import extend_checkpoint_dir


class TinyStoriesZhDataRecipe(DataChunkRecipe):
    is_generator = True

    def __init__(self, tokenizer: Tokenizer, chunk_size: int):
        super().__init__(chunk_size)
        self.tokenizer = tokenizer

    def prepare_structure(self, input_dir):
        files = Path(input_dir).rglob("*.jsonl")
        return [str(file) for file in files]

    def prepare_item(self, filepath):

        with open(filepath, "rb") as f:
            for line in f.readlines():
                js = json.loads(line)
                story = js["story_zh"]
                # 注意这里要添加eos
                # 还记得吗：我们的vocab size在int16范围内，所以可以转换为uint16来节省内存
                # story_ids = np.array(
                #     self.tokenizer.encode(story, eos=True), dtype=np.uint16
                # )
                # 很遗憾，实际使用的时候发现如果按照上面这样写，
                # litdata反序列化数据的时候会错误地得到torch.int64且超界的Tensor，
                # 但直接存torch.Tensor没问题（加上litdata不支持torch.uint16），
                # 所以最后实际使用的时候还是用下面这种写法
                story_ids = self.tokenizer.encode(story, eos=True)
                yield story_ids


def prepare(
    input_dir: Path = Path("../../Data/TinyStoriesChinese/raw_data/train"),
    output_dir: Path = Path("../../Data/TinyStoriesChinese/processed_data/train"),
    tokenizer_path: Path = Path("../../References/chatglm3-6b"),
    chunk_size: int = (2049 * 8012),
    fast_dev_run: bool = False,
) -> None:
    from litdata.processing.data_processor import DataProcessor

    tokenizer_path = extend_checkpoint_dir(tokenizer_path)
    tokenizer = Tokenizer(tokenizer_path)
    data_recipe = TinyStoriesZhDataRecipe(tokenizer=tokenizer, chunk_size=chunk_size)
    data_processor = DataProcessor(
        input_dir=str(input_dir),
        output_dir=str(output_dir),
        fast_dev_run=fast_dev_run,
        num_workers=os.cpu_count(),
        num_downloaders=1,
        # 这里有个「巨坑」，如果不加这一行，处理好的数据配对的index.json里
        # 有一个名为"dim"的key值会为null，导致后续有一个无法规避的报错
        # 但是官方的例子里是没有这一行的，很奇怪为何会有这个问题
        item_loader=TokensLoader(),
    )

    start_time = time.time()
    data_processor.run(data_recipe)
    elapsed_time = time.time() - start_time
    print(f"Time taken: {elapsed_time:.2f} seconds")

首先，我这里主要就是把之前实现的line处理逻辑封装到litgpt的DataChunkRecipe中：

• prepare_structure函数给定路径返回符合我们期望的数据文件的路径列表
• prepare_item函数给定一个上面的数据文件的路径，根据我们自定义的tokenization处理逻辑返回一个np.array对象

然后，定义了一个prepare函数，指定我们数据的输入路径和输出路径以及一些其它参数配置（其实用默认的即可），其余的都交给了litdata的DataProcessor，它基于我前面定义的DataChunkRecipe来处理数据。

感兴趣的可以看看DataProcessor的源码，里面做了很多并行之类的数据处理优化。

先用eval数据集测试

prepare(
    input_dir=Path("../../Data/TinyStoriesChinese/raw_data/val"),
    output_dir=Path("../../Data/TinyStoriesChinese/processed_data/val"),
    tokenizer_path=Path("../../References/chatglm3-6b"),
)

（也可以设置fast_dev_run=True来处理更少的数据，尤其是debug时十分有用）

执行完可以在processed_data/eval目录下看到生成的.bin文件以及记录了每个chunk文件信息的index.json。

比较一下可以发现从原先的83m的.jsonl文件压缩到了13m的.bin，压缩比（83/13≈6.385）还是很可观的。

处理train数据集

在32核的CPU上处理train数据集耗时不到1min。

prepare(
    input_dir=Path("../../Data/TinyStoriesChinese/raw_data/train"),
    output_dir=Path("../../Data/TinyStoriesChinese/processed_data/train"),
    tokenizer_path=Path("../../References/chatglm3-6b"),
)

小结

1. 数据预处理的逻辑主要是将文本转换为数字序列，以便于模型处理。
2. 通过litgpt的Tokenizer可以方便的实现文本到数字序列的转换。
3. litdata提供了数据处理的API，可以方便的封装我们的数据处理逻辑。
4. 基于上面的开发，将TinyStoriesChinese数据集做了数据划分并完成了预处理。

走之前点个赞和关注吧，你的支持是我坚持更新的动力！
我们下期见！