如何将Pandas DataFrame或Hugging Face数据集转为.bin和.idx格式？

Megatron-LM要求的.bin和.idx格式属于IndexedDataset，是为大规模数据高效加载设计的二进制格式。针对你的多列数据集场景，有两种可靠的转换路径：

路径1：通过文本中间文件转换（快速上手） #

步骤1：合并多列为训练用文本 #

不管用Pandas还是Hugging Face Dataset，先把多列数据拼接成模型训练所需的单条文本（比如对话数据可格式化为"用户：{prompt}\n助手：{response}"，格式按任务调整）。

Pandas示例：

import pandas as pd

df = pd.read_csv("your_data.csv")
# 合并多列到text字段
df["text"] = df.apply(lambda row: f"用户：{row['prompt']}\n助手：{row['response']}", axis=1)
# 导出为每行一条样本的纯文本
df["text"].to_csv("processed_text.txt", index=False, header=False)

Hugging Face Dataset示例：

from datasets import load_from_disk

dataset = load_from_disk("your_hf_dataset_dir")
# 定义文本格式化函数
def format_text(example):
    return {"text": f"用户：{example['prompt']}\n助手：{example['response']}"}

dataset = dataset.map(format_text)
# 导出纯文本文件
dataset.to_csv("processed_text.txt", columns=["text"], index=False, header=False)

步骤2：用Megatron-LM脚本生成二进制文件 #

使用Hugging Face版Megatron-LM仓库中的tools/preprocess_data.py脚本，执行以下命令（替换参数为你的实际配置）：

python tools/preprocess_data.py \
    --input processed_text.txt \
    --output_prefix your_dataset \
    --tokenizer_type GPT2Tokenizer \  # 匹配你的模型分词器类型
    --vocab_file path/to/vocab.json \  # 按需指定自定义词表
    --merge_file path/to/merges.txt \
    --do_lower_case \  # 按需开启
    --workers 8

执行后会生成your_dataset.bin和your_dataset.idx，可直接被Megatron-LM的IndexedDataset加载。

路径2：直接用IndexedDataset API生成（大规模数据友好） #

如果不想生成中间文本文件，可直接用Megatron-LM的API处理：

步骤1：分词处理数据集 #

用模型对应的分词器把多列文本转换为token序列：

from datasets import load_from_disk
from transformers import GPT2Tokenizer

tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
dataset = load_from_disk("your_hf_dataset_dir")

# 分词并格式化
def tokenize_example(example):
    text = f"用户：{example['prompt']}\n助手：{example['response']}"
    return tokenizer(text, truncation=True, max_length=512)

tokenized_dataset = dataset.map(tokenize_example, batched=True, remove_columns=dataset.column_names)

步骤2：写入IndexedDataset #

导入Megatron-LM的相关类，直接生成二进制文件：

from megatron.data.indexed_dataset import MMapIndexedDatasetWriter

output_prefix = "your_dataset"
# 初始化写入器
writer = MMapIndexedDatasetWriter(output_prefix, dtype="int64")
# 逐个写入分词后的样本token
for example in tokenized_dataset:
    writer.add_item(example["input_ids"])
# 完成写入
writer.finalize()

执行后同样会生成your_dataset.bin和your_dataset.idx文件。

关键注意事项 #

• 多列数据必须先转换为单条文本：IndexedDataset仅存储一维token序列，无法直接保留多列结构。
• 分词器需与训练模型一致：避免词表不匹配导致的训练错误。
• 规模适配：10万条数据用两种方法都高效；千万级以上数据推荐路径2，减少IO开销。

内容的提问来源于stack exchange，提问作者chriscrutt