在Megatron-Deepspeed项目中如何下载和预处理Wikipedia数据集

更详细的查看 https://github.com/marsggbo/Megatron-DeepSpeed/blob/main/tutorials/gpt2_wikipedia.md

下载Wikipedia压缩数据集（enwiki-latest-pages-articles.xml.bz2）
再使用wikiextractor工具将数据集解压缩

代码语言：javascript复制

pip install wikiextractor
python -m wikiextractor.WikiExtractor --json enwiki-latest-pages-articles.xml.bz2

解压缩后会得到一个文件夹text，结构如下：

代码语言：javascript复制

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相关的文件，这主要是用来预处理的

代码语言：javascript复制

wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json
wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt

下载好后，执行下面的代码即可

代码语言：javascript复制

#!/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 $'' 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

参考：https://github.com/NVIDIA/Megatron-LM/issues/117

json root wiki wikipedia 数据

0 人点赞