【Hugging Face】的datasets库中的常用方法总结：load_dataset、map、filter、set_format、train_test_split、concatenate_dat

Hugging Face datasets 库的常用方法总结

datasets 库是 Hugging Face 提供的 高效数据处理库，用于 加载、转换、处理、筛选、格式化和保存数据集，特别适用于 自然语言处理（NLP）、计算机视觉（CV）和语音任务。

---

1. 数据加载

1.1. 加载 Hugging Face Hub 上的公开数据集

from datasets import load_dataset

dataset = load_dataset("imdb")
print(dataset)

返回：

DatasetDict({
    train: Dataset({
        features: ['text', 'label'],
        num_rows: 25000
    })
    test: Dataset({
        features: ['text', 'label'],
        num_rows: 25000
    })
})

• load_dataset("imdb")：加载 imdb 电影评论数据集
• dataset["train"]：访问训练集

---

1.2. 加载本地 CSV/JSON/TXT 文件

dataset = load_dataset("csv", data_files="data.csv")
dataset = load_dataset("json", data_files="data.json")
dataset = load_dataset("text", data_files="data.txt")

---

2. 数据集基本操作

2.1. 查看数据

print(dataset["train"][0])  # 访问第一条数据

返回：

{'text': 'This is a great movie!', 'label': 1}

---

2.2. 访问某一列数据

texts = dataset["train"]["text"]
labels = dataset["train"]["label"]

---

2.3. 训练集和测试集拆分

train_test_split = dataset["train"].train_test_split(test_size=0.2)
print(train_test_split)

---

3. 数据转换与预处理

3.1. map() - 对整个数据集进行转换

• 将文本转换为小写

def lowercase(example):
    example["text"] = example["text"].lower()
    return example

dataset = dataset.map(lowercase)

• 添加新字段

def add_length(example):
    example["text_length"] = len(example["text"])
    return example

dataset = dataset.map(add_length)

• 批量处理（提高速度）

def batch_process(batch):
    batch["text_length"] = [len(text) for text in batch["text"]]
    return batch

dataset = dataset.map(batch_process, batched=True)

• 移除不需要的列

dataset = dataset.map(lowercase, remove_columns=["text_length"])

---

3.2. filter() - 按条件筛选数据

• 筛选文本长度大于 100 的数据

dataset = dataset.filter(lambda example: len(example["text"]) > 100)

• 筛选正面情感的评论

dataset = dataset.filter(lambda example: example["label"] == 1)

• 使用 num_proc 进行多进程加速

dataset = dataset.filter(lambda example: len(example["text"]) > 100, num_proc=4)

---

4. 数据格式转换

4.1. 转换为 PyTorch/TensorFlow 格式

dataset.set_format(type="torch", columns=["text", "label"])
dataset.set_format(type="tensorflow", columns=["text", "label"])

---

5. 数据存储

5.1. 保存数据集

dataset.save_to_disk("./my_dataset")

5.2. 重新加载保存的数据

from datasets import load_from_disk

dataset = load_from_disk("./my_dataset")

---

6. 处理超大规模数据

6.1. 流式加载超大数据集

dataset = load_dataset("c4", split="train", streaming=True)
for sample in dataset:
    print(sample)
    break  # 仅打印第一条数据

适用于 大规模 NLP 数据，如 C4、Common Crawl。

---

7. 数据集拆分与合并

7.1. 按比例拆分数据

train_test_split = dataset["train"].train_test_split(test_size=0.2)
train_dataset = train_test_split["train"]
test_dataset = train_test_split["test"]

7.2. 组合多个数据集

combined_dataset = dataset1["train"].concatenate(dataset2["train"])

---

8. datasets 主要方法对比

方法	作用
load_dataset	加载 Hugging Face Hub 或本地数据集
map	转换数据（分词、文本清理、特征添加等）
filter	筛选数据（去除短文本、保留特定类别等）
set_format	转换数据格式（PyTorch/TensorFlow）
train_test_split	拆分数据集（训练集/测试集）
concatenate_datasets	合并多个数据集
save_to_disk	保存数据集到本地
load_from_disk	加载本地数据集

---

9. 总结

1. datasets 提供 高效的数据加载、预处理、筛选、存储 方法。
2. map 适用于批量数据转换，如 文本清理、特征添加。
3. filter 适用于数据筛选，如 删除短文本、筛选特定类别。
4. set_format 支持 PyTorch/TensorFlow 格式转换，适用于深度学习训练。
5. 支持超大数据集流式处理，适用于 Common Crawl、C4 这类大规模 NLP 数据。