preprocess_data.py 的主要功能。这是 Megatron-LM 的数据预处理脚本,主要用于将原始文本数据转换为模型训练所需的格式。
核心功能:
1. 数据预处理流程:
- 输入:原始文本文件(JSON格式)
- 处理:分句 -> 分词 -> 转换为token IDs
- 输出:二进制格式(.bin和.idx文件)
2. 主要组件
class Encoder:
# 负责文本编码的核心类
def split(self, json_line):
# 将文本分割成句子
def encode(self, json_line):
# 将文本转换为token IDs
3. 关键参数:
--input # 输入文件路径
--output-prefix # 输出文件前缀
--vocab-file # 词表文件
--tokenizer-type # 分词器类型
--split-sentences # 是否分句
--workers # 并行处理的工作进程数
--partitions # 文件分区数
4. 特殊功能:
- 支持多进程并行处理
- 支持大文件分区处理
- 支持句子分割(使用NLTK)
- 支持多种tokenizer(GPT2、BERT等)
- 支持压缩文件(.gz)
5. 输出格式:
- .bin文件:包含token IDs的二进制数据
- .idx文件:索引文件,用于快速访问数据
6. 使用示例:
python tools/preprocess_data.py \
--input your_text.json \
--output-prefix output_name \
--vocab-file vocab.json \
--tokenizer-type GPT2BPETokenizer \
--workers 4
7. 这个脚本的主要目的是:
- 将原始文本转换为模型可以直接使用的格式
- 提高数据加载效率(二进制格式)
- 支持大规模数据处理(并行和分区)
8. 输入json文件内容格式要求
输入的 JSON 文件格式有特定要求。让我详细解释:
基本格式要求:
- 每行必须是一个独立的 JSON 对象(jsonl 格式)
- 默认情况下,JSON 对象中需要有 text 字段(可以通过 --json-keys 参数修改)
示例格式:
{"text": "这是第一行文本"}
{"text": "这是第二行文本"}
支持多字段
{"text": "主要文本", "title": "标题", "summary": "摘要"}
如果要处理多个字段,需要在命令中指定 --json-keys 参数,例如:
--json-keys text title summary
小节8 信息主要被来源于对源码以下部分的解读(如有错误欢迎纠正):
JSON格式要求来自 Encoder 类的处理逻辑:
def encode(self, json_line):
data = json.loads(json_line) # 每行必须是合法的JSON
ids = {}
lens = {}
for key in self.args.json_keys: # 默认 json_keys=['text']
text = data[key] # 从JSON中获取指定key的文本
默认的 text 字段来自参数解析部分
def get_args():
parser = argparse.ArgumentParser()
group.add_argument('--json-keys', nargs='+', default=['text'],
help='space separate listed of keys to extract from json')
每行一个JSON对象的要求来自文件处理逻辑:
def process_json_file(self, file_name):
fin = open(input_file_name, 'r', encoding='utf-8')
encoded_docs = pool.imap(encoder.encode, fin, 32) # 逐行处理
多字段支持的功能来自参数定义和处理逻辑:
group.add_argument('--json-keys', nargs='+', default=['text'],
help='space separate listed of keys to extract from json')