Python实用工具：Datasets库快速上手指南_数据处理与加载必备神器

作者：

在

一、Datasets库核心概述

1.1 用途与工作原理

Hugging Face的Datasets库是一款专为自然语言处理（NLP）以及机器学习领域打造的数据集处理工具，它能够帮助开发者快速加载、预处理、转换和共享各类数据集。其核心工作原理是基于一种统一的数据集抽象结构Dataset和DatasetDict，将不同来源、不同格式的数据集（如CSV、JSON、文本文件、Hugging Face Hub上的开源数据集）进行标准化封装，同时内置高效的数据处理管道，支持并行处理、流式加载和内存映射，大幅降低了数据预处理的门槛。

1.2 优缺点分析

优点

数据集资源丰富：无缝对接Hugging Face Hub，拥有上万种开源数据集，涵盖文本分类、问答、翻译等数十种任务类型。
高效便捷的预处理：内置map()、filter()、shuffle()等方法，支持并行处理大规模数据集，无需手动编写复杂的循环逻辑。
内存友好：支持流式加载和内存映射，即使处理GB级别的数据集也不会导致内存溢出。
跨框架兼容：可以轻松转换为pandas、NumPy、TensorFlow、PyTorch等格式，适配主流机器学习框架。

缺点

对非文本数据支持有限：虽然可以处理图像、音频类数据集，但功能不如文本数据集完善，预处理选项较少。
网络依赖较强：加载Hugging Face Hub上的数据集需要稳定的网络环境，离线使用时需要提前下载数据集。
新手学习成本：部分高级功能（如自定义数据集加载脚本）的使用需要熟悉库的底层逻辑，对纯小白不够友好。

1.3 License类型

二、Datasets库安装与环境配置

2.1 安装命令

Datasets库的安装非常简单，支持pip和conda两种方式，推荐使用pip安装，因为它的更新速度更快，适配性更强。

2.1.1 pip安装（推荐）

打开命令行终端，输入以下命令即可完成安装：

pip install datasets

如果需要处理特定格式的数据集（如音频、图像、Parquet文件），可以安装对应的依赖包：

pip install datasets[audio,vision,parquet]

2.1.2 conda安装

如果你的开发环境是基于conda管理的，可以使用以下命令安装：

conda install -c huggingface -c conda-forge datasets

2.2 环境验证

安装完成后，我们可以通过一段简单的Python代码验证是否安装成功。打开Python交互式环境或新建一个.py文件，输入以下代码：

import datasets
print(f"Datasets库版本：{datasets.__version__}")

运行代码后，如果终端输出了具体的版本号（如2.14.5），则说明安装成功；如果出现ModuleNotFoundError，则说明安装失败，需要重新检查安装命令或Python环境。

三、Datasets库核心功能与代码示例

3.1 加载开源数据集（Hugging Face Hub）

Datasets库最核心的功能之一就是一键加载Hugging Face Hub上的开源数据集。以经典的文本分类数据集imdb（电影评论情感分析数据集）为例，我们来演示如何加载并查看数据集的基本信息。

3.1.1 加载imdb数据集

from datasets import load_dataset

# 加载imdb数据集，该数据集包含train、test、unsupervised三个子集
dataset = load_dataset("imdb")

# 打印数据集的结构
print("数据集结构：", dataset)
# 打印训练集的第一条数据
print("训练集第一条数据：", dataset["train"][0])

代码说明

load_dataset()函数是加载数据集的核心函数，传入数据集名称即可自动从Hugging Face Hub下载并加载数据集。
imdb数据集是一个DatasetDict对象，包含train（训练集）、test（测试集）、unsupervised（无监督数据集）三个子集。
通过下标索引dataset["train"][0]可以查看训练集的第一条数据，数据格式为字典，包含text（电影评论文本）和label（情感标签，0为负面，1为正面）两个字段。

运行结果示例

数据集结构： DatasetDict({
    train: Dataset({
        features: ['text', 'label'],
        num_rows: 25000
    })
    test: Dataset({
        features: ['text', 'label'],
        num_rows: 25000
    })
    unsupervised: Dataset({
        features: ['text', 'label'],
        num_rows: 50000
    })
})
训练集第一条数据： {'text': 'I rented I AM CURIOUS-YELLOW...', 'label': 0}

3.1.2 加载指定子集与拆分数据

如果我们只需要加载imdb数据集的训练集和测试集，可以通过split参数指定；同时，还可以使用train_test_split()方法将数据集拆分为新的训练集和验证集。

from datasets import load_dataset

# 只加载train和test子集
dataset = load_dataset("imdb", split=["train", "test"])
train_dataset, test_dataset = dataset

# 将训练集拆分为训练集（80%）和验证集（20%）
train_val_dataset = train_dataset.train_test_split(test_size=0.2, seed=42)
train_new = train_val_dataset["train"]
val_new = train_val_dataset["test"]

print(f"新训练集大小：{len(train_new)}")
print(f"验证集大小：{len(val_new)}")
print(f"测试集大小：{len(test_dataset)}")

代码说明

split参数传入一个列表，可以指定加载的数据集子集，返回的结果是一个列表，需要手动解包为对应的数据集对象。
train_test_split()方法用于拆分数据集，test_size参数指定验证集的比例，seed参数用于固定随机种子，保证实验的可重复性。
len()函数可以查看数据集的样本数量，imdb原训练集有25000条数据，拆分后新训练集有20000条，验证集有5000条。

运行结果示例

新训练集大小：20000
验证集大小：5000
测试集大小：25000

3.2 加载本地数据集

除了加载Hugging Face Hub上的开源数据集，Datasets库还支持加载本地的CSV、JSON、文本等格式的数据集。下面我们以CSV格式为例，演示如何加载本地数据集。

3.2.1 准备本地CSV数据集

首先，我们创建一个名为local_data.csv的CSV文件，内容如下：

text,label
This movie is amazing,1
This movie is terrible,0
I love this film,1
The plot is boring,0
The acting is great,1

该文件包含两列，text列是电影评论文本，label列是情感标签。

3.2.2 加载本地CSV数据集

from datasets import load_dataset

# 加载本地CSV数据集
local_dataset = load_dataset("csv", data_files="local_data.csv")

# 打印数据集结构
print("本地数据集结构：", local_dataset)
# 打印所有数据
print("本地数据集所有数据：", local_dataset["train"])
# 访问单条数据的字段
for data in local_dataset["train"]:
    print(f"文本：{data['text']}，标签：{data['label']}")

代码说明

加载本地数据集时，load_dataset()函数的第一个参数需要指定数据格式（如csv、json、text），data_files参数指定本地数据文件的路径。
对于单文件的CSV数据集，加载后默认会生成一个名为train的子集。
可以通过遍历数据集对象，访问每条数据的具体字段。

运行结果示例

本地数据集结构： DatasetDict({
    train: Dataset({
        features: ['text', 'label'],
        num_rows: 5
    })
})
本地数据集所有数据： Dataset({
    features: ['text', 'label'],
    num_rows: 5
})
文本：This movie is amazing，标签：1
文本：This movie is terrible，标签：0
文本：I love this film，标签：1
文本：The plot is boring，标签：0
文本：The acting is great，标签：1

3.2.3 加载多个本地数据文件

如果本地有多个数据文件（如训练集和测试集分别存储在不同的CSV文件中），可以通过data_files参数传入一个字典来指定：

from datasets import load_dataset

# 定义多个数据文件的路径
data_files = {
    "train": "train_data.csv",
    "test": "test_data.csv"
}

# 加载多个CSV文件
multi_dataset = load_dataset("csv", data_files=data_files)

print("多文件数据集结构：", multi_dataset)
print(f"训练集样本数：{len(multi_dataset['train'])}")
print(f"测试集样本数：{len(multi_dataset['test'])}")

代码说明

data_files参数传入一个字典时，字典的键是数据集子集的名称（如train、test），值是对应的数据文件路径。
加载后会生成一个DatasetDict对象，包含指定的多个子集，方便后续分别处理训练集和测试集。

3.3 数据集预处理

加载数据集后，通常需要进行预处理（如文本分词、长度截断、特征提取等），Datasets库提供了map()方法，可以高效地对数据集进行批量预处理。下面我们以文本分词为例，演示如何使用map()方法。

3.3.1 安装分词器依赖

我们使用Hugging Face的transformers库中的分词器，需要先安装该库：

pip install transformers

3.3.2 对imdb数据集进行分词预处理

from datasets import load_dataset
from transformers import AutoTokenizer

# 加载imdb数据集的训练集
dataset = load_dataset("imdb", split="train")
# 加载预训练分词器（以bert-base-uncased为例）
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

# 定义预处理函数
def preprocess_function(examples):
    # 对文本进行分词，设置最大长度为128，超过截断，不足补齐
    return tokenizer(
        examples["text"],
        truncation=True,
        padding="max_length",
        max_length=128
    )

# 对数据集进行批量预处理，开启并行处理
tokenized_dataset = dataset.map(
    preprocess_function,
    batched=True,
    num_proc=4  # 使用4个进程并行处理
)

# 查看预处理后的数据集特征
print("预处理后的特征：", tokenized_dataset.features)
# 查看第一条预处理后的数据
print("第一条预处理后的数据：", tokenized_dataset[0])

代码说明

AutoTokenizer.from_pretrained()函数用于加载预训练分词器，bert-base-uncased是一个常用的英文预训练模型。
preprocess_function()是自定义的预处理函数，接收一个包含批量数据的字典examples，返回分词后的结果，包含input_ids、attention_mask等字段。
map()方法的batched=True参数表示批量处理数据，num_proc参数指定并行处理的进程数，能够大幅提升预处理速度。
预处理后的数据集会新增分词相关的字段，这些字段可以直接输入到预训练模型中进行训练。

运行结果示例

预处理后的特征： {'text': Value(dtype='string', id=None), 'label': ClassLabel(names=['neg', 'pos'], id=None), 'input_ids': Sequence(feature=Value(dtype='int32', id=None), length=-1, id=None), 'token_type_ids': Sequence(feature=Value(dtype='int32', id=None), length=-1, id=None), 'attention_mask': Sequence(feature=Value(dtype='int32', id=None), length=-1, id=None)}
第一条预处理后的数据： {'text': 'I rented I AM CURIOUS-YELLOW...', 'label': 0, 'input_ids': [101, 1045, 5975, ...], 'token_type_ids': [0, 0, 0, ...], 'attention_mask': [1, 1, 1, ...]}

3.4 数据集筛选与排序

Datasets库提供了filter()和sort()方法，用于对数据集进行筛选和排序。下面我们演示如何筛选出imdb数据集中文本长度大于100的样本，以及如何按文本长度对数据集进行排序。

3.4.1 数据集筛选

from datasets import load_dataset

# 加载imdb训练集
dataset = load_dataset("imdb", split="train")

# 定义筛选函数：保留文本长度大于100的样本
def filter_function(example):
    return len(example["text"]) > 100

# 对数据集进行筛选
filtered_dataset = dataset.filter(filter_function)

print(f"原数据集样本数：{len(dataset)}")
print(f"筛选后数据集样本数：{len(filtered_dataset)}")
# 查看筛选后第一条数据的文本长度
print(f"筛选后第一条数据文本长度：{len(filtered_dataset[0]['text'])}")

代码说明

filter_function()是自定义的筛选函数，接收单条数据example，返回True表示保留该样本，返回False表示过滤掉该样本。
filter()方法会遍历数据集的所有样本，根据筛选函数的结果保留符合条件的样本。

运行结果示例

原数据集样本数：25000
筛选后数据集样本数：24890
筛选后第一条数据文本长度：178

3.4.2 数据集排序

from datasets import load_dataset

# 加载imdb训练集
dataset = load_dataset("imdb", split="train")

# 为数据集添加文本长度字段
def add_length_field(example):
    example["text_length"] = len(example["text"])
    return example

dataset_with_length = dataset.map(add_length_field)

# 按文本长度升序排序
sorted_dataset_asc = dataset_with_length.sort("text_length")
# 按文本长度降序排序
sorted_dataset_desc = dataset_with_length.sort("text_length", reverse=True)

print(f"最短文本长度：{sorted_dataset_asc[0]['text_length']}")
print(f"最长文本长度：{sorted_dataset_desc[0]['text_length']}")

代码说明

首先通过map()方法为数据集添加一个text_length字段，存储每条数据的文本长度。
sort()方法接收一个字段名，按照该字段的值对数据集进行排序，reverse=True表示降序排序，默认是升序排序。

运行结果示例

最短文本长度：14
最长文本长度：13704

3.5 数据集格式转换

Datasets库的数据集对象可以轻松转换为pandas、NumPy、PyTorch、TensorFlow等格式，适配不同的机器学习框架。下面我们演示如何将数据集转换为这些格式。

3.5.1 转换为pandas DataFrame格式

from datasets import load_dataset

# 加载imdb训练集
dataset = load_dataset("imdb", split="train[:100]")  # 只取前100条数据

# 转换为pandas DataFrame
df = dataset.to_pandas()

# 查看DataFrame的前5行数据
print(df.head())
# 查看DataFrame的基本信息
print(df.info())

代码说明

to_pandas()方法可以将数据集转换为pandas的DataFrame对象，方便使用pandas进行数据分析和可视化。
为了避免数据量过大，我们通过split="train[:100]"只取训练集的前100条数据。

运行结果示例

                                                text  label
0  I rented I AM CURIOUS-YELLOW...                    0
1  "I Am Curious: Yellow" is a...                    0
2  If only to avoid making this...                    0
3  This film was probably intende...                  0
4  I saw this movie when I was ...                   0
<class 'pandas.core.frame.DataFrame'>
RangeIndex: 100 entries, 0 to 99
Data columns (total 2 columns):
 #   Column  Non-Null Count  Dtype 
    --  -- 
 0   text    100 non-null    object
 1   label   100 non-null    int64 
dtypes: int64(1), object(1)
memory usage: 1.7+ KB
None

3.5.2 转换为NumPy数组格式

from datasets import load_dataset

# 加载imdb训练集
dataset = load_dataset("imdb", split="train[:100]")

# 转换为NumPy数组
numpy_array = dataset.to_numpy()

# 查看NumPy数组的形状和第一条数据
print(f"NumPy数组形状：{numpy_array.shape}")
print(f"第一条数据：{numpy_array[0]}")

代码说明

to_numpy()方法将数据集转换为NumPy数组，数组中的每个元素是一个字典，包含数据集的所有字段。

运行结果示例

NumPy数组形状：(100,)
第一条数据： {'text': 'I rented I AM CURIOUS-YELLOW...', 'label': 0}

3.5.3 转换为PyTorch张量格式

from datasets import load_dataset
from transformers import AutoTokenizer

# 加载imdb训练集并进行分词预处理
dataset = load_dataset("imdb", split="train[:100]")
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

def preprocess_function(examples):
    return tokenizer(
        examples["text"],
        truncation=True,
        padding="max_length",
        max_length=128
    )

tokenized_dataset = dataset.map(preprocess_function, batched=True)

# 转换为PyTorch张量格式
torch_dataset = tokenized_dataset.with_format("torch")

# 查看张量格式的输入ids和标签
print(f"Input IDs张量形状：{torch_dataset['input_ids'].shape}")
print(f"标签张量形状：{torch_dataset['label'].shape}")

代码说明

with_format("torch")方法将数据集转换为PyTorch张量格式，转换后可以直接通过下标访问张量数据。
转换后的数据集可以直接输入到PyTorch模型中进行训练，无需手动转换数据类型。

运行结果示例

Input IDs张量形状： torch.Size([100, 128])
标签张量形状： torch.Size([100])

四、实际案例：基于Datasets库的文本分类任务

为了让大家更好地理解Datasets库的实际应用，我们结合transformers库搭建一个简单的文本分类模型，完成imdb电影评论情感分析任务。

4.1 案例流程

加载imdb数据集并进行预处理。
加载预训练模型和分词器。
构建训练参数并训练模型。
在测试集上评估模型性能。

4.2 完整代码实现

# 导入必要的库
from datasets import load_dataset, load_metric
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    TrainingArguments,
    Trainer
)
import numpy as np

# 步骤1：加载数据集并进行预处理
# 加载imdb的训练集和测试集
dataset = load_dataset("imdb")
train_dataset = dataset["train"]
test_dataset = dataset["test"]

# 加载预训练分词器
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

# 定义预处理函数
def preprocess_function(examples):
    return tokenizer(
        examples["text"],
        truncation=True,
        padding="max_length",
        max_length=128
    )

# 对训练集和测试集进行预处理
tokenized_train = train_dataset.map(preprocess_function, batched=True, num_proc=4)
tokenized_test = test_dataset.map(preprocess_function, batched=True, num_proc=4)

# 设置数据集格式为PyTorch张量
tokenized_train.set_format(
    type="torch",
    columns=["input_ids", "attention_mask", "label"]
)
tokenized_test.set_format(
    type="torch",
    columns=["input_ids", "attention_mask", "label"]
)

# 步骤2：加载预训练模型
model = AutoModelForSequenceClassification.from_pretrained(
    "bert-base-uncased",
    num_labels=2  # 二分类任务，标签数为2
)

# 步骤3：定义评估指标
metric = load_metric("accuracy")

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)

# 步骤4：设置训练参数
training_args = TrainingArguments(
    output_dir="./imdb_sentiment_model",  # 模型输出目录
    num_train_epochs=2,  # 训练轮数
    per_device_train_batch_size=16,  # 训练批次大小
    per_device_eval_batch_size=16,  # 评估批次大小
    evaluation_strategy="epoch",  # 每轮训练后进行评估
    save_strategy="epoch",  # 每轮训练后保存模型
    logging_dir="./logs",  # 日志目录
    logging_steps=100,  # 每100步记录一次日志
    learning_rate=2e-5,  # 学习率
    weight_decay=0.01,  # 权重衰减
)

# 步骤5：构建Trainer并开始训练
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_test,
    compute_metrics=compute_metrics,
)

# 开始训练
trainer.train()

# 步骤6：在测试集上评估模型
eval_results = trainer.evaluate()
print(f"测试集准确率：{eval_results['eval_accuracy']:.4f}")

4.3 代码说明

数据集预处理：我们使用bert-base-uncased分词器对文本进行分词，设置最大长度为128，并将数据集转换为PyTorch张量格式。
模型加载：AutoModelForSequenceClassification是一个用于序列分类任务的预训练模型，num_labels=2表示二分类任务。
评估指标：使用accuracy指标评估模型性能，compute_metrics函数用于计算预测结果的准确率。
训练参数设置：TrainingArguments包含了训练过程中的所有参数，如训练轮数、批次大小、学习率等。
模型训练与评估：Trainer是transformers库提供的训练工具，调用train()方法开始训练，evaluate()方法在测试集上评估模型性能。

4.4 预期结果

经过2轮训练后，模型在imdb测试集上的准确率可以达到90%以上，具体数值会因硬件环境和随机种子的不同而略有差异。

五、相关资源链接

Pypi地址：https://pypi.org/project/datasets
Github地址：https://github.com/huggingface/datasets
官方文档地址：https://huggingface.co/docs/datasets/index

关注我，每天分享一个实用的Python自动化工具。

实用工具

Python实用工具：Datasets库快速上手指南_数据处理与加载必备神器

一、Datasets库核心概述

1.1 用途与工作原理

1.2 优缺点分析

1.3 License类型

二、Datasets库安装与环境配置

2.1 安装命令

2.1.1 pip安装（推荐）

2.1.2 conda安装

2.2 环境验证

三、Datasets库核心功能与代码示例

3.1 加载开源数据集（Hugging Face Hub）

3.1.1 加载imdb数据集

3.1.2 加载指定子集与拆分数据

3.2 加载本地数据集

3.2.1 准备本地CSV数据集

3.2.2 加载本地CSV数据集

3.2.3 加载多个本地数据文件

3.3 数据集预处理

3.3.1 安装分词器依赖

3.3.2 对imdb数据集进行分词预处理

3.4 数据集筛选与排序

3.4.1 数据集筛选

3.4.2 数据集排序

3.5 数据集格式转换

3.5.1 转换为pandas DataFrame格式

3.5.2 转换为NumPy数组格式

3.5.3 转换为PyTorch张量格式

四、实际案例：基于Datasets库的文本分类任务

4.1 案例流程

4.2 完整代码实现

4.3 代码说明

4.4 预期结果

五、相关资源链接

更多文章

Python实用工具：Faker库完全指南——快速生成逼真测试数据

Python实用工具：Datasets库快速上手指南_数据处理与加载必备神器

Python实用工具：psycopg3 完全指南——PostgreSQL数据库交互最佳实践

Python实用工具：aioprometheus异步监控指标采集与上报教程