Hugging Face 之 Datasets使用

原创

一个风轻云淡

发布于 2026-03-01 16:23:52

1360

文章被收录于专栏：深度学习深度学习

搞 NLP 模型训练，数据处理总让你头疼？🤯 今天给你介绍一个神器——Hugging Face Datasets！这个轻量级的库简直就是 NLP 开发者的福音，让数据处理变得像喝水一样简单！💦

Datasets概述

Datasets 是 Hugging Face 生态中的一个轻量级数据处理工具库🧰，专门为自然语言处理任务优化设计。它在模型训练流程中，能够高效完成数据的加载与预处理工作⚡。

它的核心优势体现在以下几个方面✨：

数据加载灵活便捷📂：无论是读取本地常见格式文件（CSV、JSON 等），还是直接加载 Hugging Face Hub 上的海量公开数据集🌐，都能轻松实现。
数据结构直观清晰📊：数据集内部以类似表格的结构进行组织，每条数据样本都由若干个明确定义的字段构成，查看和操作都非常直观👀。
与 HF 生态无缝集成🤝：能够与 Tokenizer 等其他 Hugging Face 组件高度协同，方便快捷地构建模型所需的输入格式🔧。
数据处理功能丰富🛠️：提供了 .map()方法进行高效的批量数据处理，支持按条件筛选样本🔍，也可以通过 .train_test_split()等方法轻松划分训练集、验证集等📈，覆盖了常见的预处理需求。

安装命令

pip install datasets

加载数据集

Datasets 库最棒的功能之一就是提供了统一的接口 load_dataset()🎯。无论你的数据是躺在本地硬盘里，还是托管在 Hugging Face Hub 上，都能用同一行代码轻松加载！🌈

加载本地数据 📁

load_dataset()支持多种常见文件格式，包括 CSV、JSON、Parquet 等，还能一次加载单个或多个文件，超级灵活！🔄

from datasets import load_dataset

dataset = load_dataset(format, data_files=路径或字典)

参数

参数	类型	说明
format	str	文件格式，如 "csv"📄、"json"🔤、"parquet"📊 等
data_files	str 或 dict	文件路径。可传字符串（单个文件）或字典（多个文件）

具体怎么用？

1）加载多个文件 📚

from datasets import load_dataset

# 同时加载训练集和测试集
dataset_dict = load_dataset('csv', data_files={
    'train': './data/train.csv',   # 🏋️ 训练数据
    'test': './data/test.csv'      # 🧪 测试数据
})

此时返回的是一个 DatasetDict，里面包含两个 Dataset（每个称为一个 split）：

print(dataset_dict)
# DatasetDict({
#     train: Dataset(...),  # ✅ 训练集
#     test: Dataset(...)     # ✅ 测试集
# })

2）加载单个文件 📄

from datasets import load_dataset

# 只加载一个文件也很简单
dataset_dict = load_dataset('csv', data_files='./data/dataset.csv')

即使只加载单个文件，返回的也是 DatasetDict，不过默认命名为 "train"：

print(dataset_dict)
# DatasetDict({
#     train: Dataset(...)  # 🎯 默认就叫 train
# })

查看数据集

获取 Dataset 🎯

假设我们有一个电影评论数据集，load_dataset()返回的是一个 DatasetDict 对象，我们可以像操作字典一样通过键名来访问具体的 split：

from datasets import load_dataset

# 加载数据集 - 这次用电影评论
dataset_dict = load_dataset('csv', data_files='data/movie_reviews.csv')

# 获取数据集
dataset = dataset_dict["train"]  # 📌 注意：单个文件时默认 split 名为 "train"
print("🎬 数据集加载完成！")

此时 dataset是一个 Dataset 对象，包含了所有电影评论数据。🍿

访问样本 📄

Dataset 支持类似列表的索引和切片操作来访问具体样本：

# 查看第一条样本
print("📖 第一条评论：")
print(dataset[0])       # 👉 返回第一条完整记录
# 输出：{'text': '这部电影的视觉效果太震撼了！', 'sentiment': 'positive', 'rating': 5}

# 查看前三条样本
print("\n📚 前三部电影评论：")
print(dataset[:3])      # 👉 注意返回的结构是字典，值变成了列表！
# 输出：{'text': ['这部电影的视觉效果太震撼了！', '剧情拖沓，有点失望', '演员表演出色'],
#       'sentiment': ['positive', 'negative', 'positive'],
#       'rating': [5, 2, 4]}

📊 返回结构说明：

访问方式	返回示例	说明
dataset[0]	{'text': '这部电影的视觉效果太震撼了！', 'sentiment': 'positive', 'rating': 5}	返回单条评论，是字典格式 🎥
dataset[:3]	{'text': ['...', '...', '...'], 'sentiment': ['positive','negative','positive'], 'rating': [5,2,4]}	返回多条评论，但按字段分组成了列表 📋

访问某个字段值 🎯

想要只看评论内容或评分？可以进一步通过字段名访问：

# 获取第一条样本的评论内容
print("🎯 第一条评论的文本：")
print(dataset[0]['text'])        
# 输出：'这部电影的视觉效果太震撼了！' ✨

# 获取前三条样本的所有评分
print("\n⭐ 前三部电影的评分：")
print(dataset[:3]['rating'])       
# 输出：[5, 2, 4] 🌟

预处理数据集 🛠️

除了加载数据，datasets库还支持各种数据预处理操作，让你的数据准备好被模型"消化"！🍽️ 从清理无用列到编码文本，再到数据集划分，我们来一步步看看这些实用的功能。

删除列 🗑️

用 .remove_columns()删除不需要的字段，让数据集更整洁：

# 假设我们的数据集有这些字段
print("📊 原始字段：", dataset.column_names)
# 输出：['review', 'label', 'cat', 'date', 'user_id']

# 删除 'cat' 字段
dataset = dataset.remove_columns(["cat"])
print("✨ 清理后字段：", dataset.column_names)
# 输出：['review', 'label', 'date', 'user_id']

💡 小技巧：一次删除多个字段

dataset = dataset.remove_columns(["cat", "date", "user_id"])
# 现在只剩下 'review' 和 'label' 了 🧹

过滤行 🔍

用 .filter()筛选符合条件的样本，确保数据质量：

# 过滤掉空评论和无效标签
original_size = len(dataset)
dataset = dataset.filter(lambda x: 
    x["review"] is not None and           # 评论不为空
    x["review"].strip() != "" and         # 不是纯空格
    x["label"] in [0, 1]                  # 标签只能是0或1
)
new_size = len(dataset)
print(f"✅ 过滤完成！从 {original_size} 条过滤到 {new_size} 条")
print(f"🗑️  移除了 {original_size - new_size} 条无效数据")

🔧 高级过滤：

# 只保留正面评论
positive_only = dataset.filter(lambda x: x["label"] == 1)
print(f"😊 正面评论数量：{len(positive_only)}")

# 保留长评论（假设我们关心详细评论）
detailed_reviews = dataset.filter(lambda x: len(x["review"]) > 20)
print(f"📝 长评论数量：{len(detailed_reviews)}")

划分数据集 📊

用 .train_test_split()轻松将数据分为训练集和测试集：

# 将数据按8:2划分
dataset_dict = dataset.train_test_split(
    test_size=0.2,      # 测试集占20%
    seed=42,           # 随机种子，保证可重复
    shuffle=True       # 打乱数据
)

train_dataset = dataset_dict["train"]  # 🏋️ 训练集
test_dataset = dataset_dict["test"]    # 🧪 测试集

print(f"📈 训练集大小：{len(train_dataset)} 条")
print(f"📉 测试集大小：{len(test_dataset)} 条")
print(f"🔢 总数据量：{len(train_dataset) + len(test_dataset)} 条")

🎯 更多划分选项：

# 如果需要验证集，可以二次划分
temp = train_dataset.train_test_split(test_size=0.1, seed=42)
train_dataset = temp["train"]  # 最终训练集 (72%)
val_dataset = temp["test"]     # 验证集 (8%)
test_dataset = test_dataset    # 测试集 (20%)

print(f"🏋️  训练集: {len(train_dataset)}")
print(f"📊 验证集: {len(val_dataset)}")
print(f"🧪 测试集: {len(test_dataset)}")

编码数据 🔤

这是最重要的部分！用 .map()配合 tokenizer 将文本转换为模型可理解的数字。

🧠 `.map()`方法简介

.map()是 datasets 的核心方法之一，可以批量处理数据：

dataset = dataset.map(function, batched=False, remove_columns=None)

参数	说明	示例
function	处理函数，作用于每条/批样本	lambda x: x*2
batched	是否批量处理（更快！）	batched=True🚀
remove_columns	处理完后删除的字段	remove_columns=['text']

def preprocess_dataset(dataset):
    """一站式预处理函数"""
    
    # 1. 删除不需要的列 🗑️
    dataset = dataset.remove_columns(["cat", "date", "user_id"])
    
    # 2. 过滤无效数据 🔍
    dataset = dataset.filter(lambda x: 
        x["review"] and 
        x["review"].strip() != "" and 
        x["label"] in [0, 1]
    )
    
    # 3. 划分数据集 📊
    split = dataset.train_test_split(test_size=0.2, seed=42)
    
    return split["train"], split["test"]

# 使用示例
train, test = preprocess_dataset(dataset)
print("🎉 预处理完成！")
print(f"🏋️  训练集: {len(train)} 条")
print(f"🧪 测试集: {len(test)} 条")

保存数据集 💾

处理好的数据当然要保存下来！这样下次就不用重复预处理了，既省时又省力。Datasets提供了多种保存格式，适应不同需求：

数据格式	保存方法	适用对象	特点
Arrow 🏹	save_to_disk()	Dataset 或 DatasetDict	Hugging Face 官方推荐，支持所有特性
CSV 📄	to_csv()	仅限 Dataset	通用表格格式，兼容性好
JSON 🔤	to_json()	仅限 Dataset	结构灵活，适合嵌套数据

Arrow格式 🏹（推荐！）

Arrow 是 Hugging Face 官方推荐的格式，速度快、内存效率高，而且能完整保存数据集的所有信息！

💾 保存数据集

# 保存整个 DatasetDict（包含 train/test/val 等）
dataset_dict.save_to_disk("./data/processed")

# 或者只保存单个 Dataset
train_dataset.save_to_disk("./data/processed/train")
test_dataset.save_to_disk("./data/processed/test")

📁 保存后的目录结构：

processed/
├─ dataset_dict.json          # 数据集字典配置
├─ test/                     # 测试集文件夹
│   ├─ data-00000-of-00001.arrow  # 实际数据（Arrow格式）
│   ├─ dataset_info.json     # 数据集信息
│   └─ state.json           # 状态信息
└─ train/                    # 训练集文件夹
    ├─ data-00000-of-00001.arrow
    ├─ dataset_info.json
    └─ state.json

✨ Arrow格式优点：

⚡ 超快加载：比CSV/JSON快10-100倍
🐏 内存友好：支持内存映射，不占用大量RAM
🔄 完整保存：保留所有字段类型、分片信息
🔧 功能完整：支持后续所有datasets操作

📂 加载Arrow格式

from datasets import load_from_disk

# 加载整个 DatasetDict
dataset_dict = load_from_disk("./data/processed")
print(f"✅ 加载完成！包含 {len(dataset_dict)} 个split")

# 或者加载单个 Dataset
train_dataset = load_from_disk("./data/processed/train")
print(f"📊 训练集大小：{len(train_dataset)} 条")

# 验证数据完整性
print("🔍 字段列表：", train_dataset.column_names)
print("📋 第一条数据：", train_dataset[0])

CSV和JSON格式 📁

如果需要和其他工具（如Excel、Pandas）共享数据，可以导出为通用格式。

💾 保存为CSV/JSON

# 保存为 CSV（表格格式）
train_dataset.to_csv("./data/processed/train.csv")
test_dataset.to_csv("./data/processed/test.csv")

# 保存为 JSON（更灵活的结构）
train_dataset.to_json("./data/processed/train.json", indent=2)  # indent 让文件更易读

⚠️ 注意：to_csv()和 to_json()只能用于单个 Dataset，不能直接保存 DatasetDict。如果需要保存多个split，要分别保存。

📂 加载CSV/JSON格式

from datasets import load_dataset

# 加载 CSV 文件
csv_dataset = load_dataset(
    "csv", 
    data_files={
        "train": "./data/processed/train.csv",
        "test": "./data/processed/test.csv"
    }
)

# 加载 JSON 文件
json_dataset = load_dataset(
    "json",
    data_files={
        "train": "./data/processed/train.json",
        "test": "./data/processed/test.json"
    }
)

print(f"📁 CSV数据集：{csv_dataset}")
print(f"📁 JSON数据集：{json_dataset}")

🎯 格式选择建议

场景	推荐格式	原因
继续用 Hugging Face 训练 🤗	Arrow	加载快、功能全、内存友好
用 Pandas 分析 🐼	CSV	直接用 pd.read_csv()读取
给前端/API 用 🌐	JSON	通用性强，各种语言都支持
长期存储 📦	Arrow	性能最好，兼容性有保障
临时查看 👀	CSV	用 Excel 就能打开查看

集成Dataloader 🔄

数据准备好了，现在要喂给模型啦！虽然 datasets.Dataset不是 PyTorch 原生的 Dataset类，但它实现了 __len__()和 __getitem__()接口，所以能和 DataLoader完美配合！

🔧 设置张量格式

# 首先，告诉数据集我们想要PyTorch张量格式
train_dataset.set_format(
    type="torch",  # 🧪 指定输出为 PyTorch 张量
    columns=["input_ids", "attention_mask", "label"]  # 只转换需要的字段
)

# 验证一下
sample = train_dataset[0]
print(f"🔬 转换后格式：")
print(f"  input_ids 类型: {type(sample['input_ids'])}")  # torch.Tensor 🎉
print(f"  input_ids 形状: {sample['input_ids'].shape}")  # torch.Size([128])
print(f"  label 类型: {type(sample['label'])}")  # torch.Tensor 🎉

📌 重要说明：

set_format()只改变访问时的返回格式，不修改底层存储
只有 columns指定的字段会被转换为张量
未指定的字段在访问时会被自动过滤掉（访问不到！）

🎯 创建DataLoader

from torch.utils.data import DataLoader

# 训练集 DataLoader
train_dataloader = DataLoader(
    train_dataset,
    batch_size=32,      # 🧺 每批32个样本
    shuffle=True,       # 🎲 打乱顺序（训练时需要）
    num_workers=4,      # ⚡ 使用4个进程加载数据
    pin_memory=True     # 🚀 加速GPU传输
)

# 验证/测试集 DataLoader
test_dataloader = DataLoader(
    test_dataset,
    batch_size=32,
    shuffle=False,      # 测试时不需要打乱
    num_workers=2,
    pin_memory=True
)

print(f"✅ 创建 DataLoader 成功！")
print(f"📊 训练集批次数：{len(train_dataloader)}")
print(f"📊 测试集批次数：{len(test_dataloader)}")

🔄 批量迭代示例

# 看看第一批数据长什么样
print("🔍 查看第一批数据：")
first_batch = next(iter(train_dataloader))

for key, value in first_batch.items():
    print(f"{key}: {value.shape} | 类型: {value.dtype}")
# 输出：
# input_ids: torch.Size([32, 128]) | 类型: torch.int64
# attention_mask: torch.Size([32, 128]) | 类型: torch.int64  
# label: torch.Size([32]) | 类型: torch.int64

原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。

如有侵权，请联系 cloudcommunity@tencent.com 删除。

NLP技术

原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。

如有侵权，请联系 cloudcommunity@tencent.com 删除。

NLP技术

登录后参与评论

0 条评论

热度