在Megatron-Deepspeed项目中如何下载和预处理Wikipedia数据集
更详细的查看 https://github.com/marsggbo/Megatron-DeepSpeed/blob/main/tutorials/gpt2_wikipedia.md
-
下载Wikipedia压缩数据集(enwiki-latest-pages-articles.xml.bz2)
-
再使用wikiextractor工具将数据集解压缩
pip install wikiextractor
python -m wikiextractor.WikiExtractor --json enwiki-latest-pages-articles.xml.bz2
解压缩后会得到一个文件夹text
,结构如下:
text
├── AA
├── wiki_00
├── wiki__01
├── ...
├── AB
├── AC
├── AD
├── AE
├── ...
├── GD
└── GE
文件夹包含多个子文件夹,每个子文件夹包含多个json格式的数据集,即wiki_00
其实是json格式的文件
- 对解压后的数据集做预处理
我们在训练GPT的时候,解压后的数据集还不能直接拿来用,我们还需要用Megatron-Deepspeed提供的tools/preprocess_data.py
对text
目录下数据集做预处理,最终会得到两个二进制文件,后缀分别是bin
和idx
。
不过tools/preprocess_data.py
只能对单个的json文件做处理,而第二步中我们有几十万个json文件,这个该怎么办呢?一种处理办法就是把第三步中的所有json文件合并到一个json文件中去,最后再对后并后的文件做预处理就可以了。在于处理之前,你需要先运行下面大命令下载GPT相关的文件,这主要是用来预处理的
wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json
wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt
下载好后,执行下面的代码即可
#!/bin/bash
# 设置ROOT路径
ROOT="/data/personal/nus-hx/Wikipedia/text"
# 检查是否存在 wiki_all.json 文件,如果存在则删除
if [ -f "$ROOT/wiki_all.json" ]; then
rm "$ROOT/wiki_all.json"
fi
# 创建一个空的 wiki_all.json 文件
touch "$ROOT/wiki_all.json"
# 遍历ROOT路径下所有的文件
find $ROOT -type f -name "*" -print0 | while IFS= read -r -d $'\0' file; do
# 将所有文件内容追加到wiki_all.json文件中
cat "$file" >> "$ROOT/wiki_all.json"
done
cd /path/to/Megatron-Deepspeed
python tools/preprocess_data.py \
--input "$ROOT/wiki_all.json" \
--output-prefix my-gpt2 \
--dataset-impl mmap \
--tokenizer-type GPT2BPETokenizer \
--append-eod \
--vocab-file gpt2-vocab.json \
--merge-file gpt2-merges.txt \
--workers 16 \
--partitions 16