当前位置：首页 > news >正文

huggingface datasets库中的load_dataset方法-------deepseek问答记录

news 来源：原创 2025/8/2 16:04:55

1. 基本介绍

Hugging Face 的 datasets 库中的 load_dataset 方法是用于加载数据集的核心工具，它支持从多种来源（如本地文件、Hugging Face Hub、内存数据等）加载数据集，并返回标准的 Dataset 或 DatasetDict 对象，方便进行高效的数据处理和训练。以下是详细介绍：

基本用法

from datasets import load_datasetdataset = load_dataset(path, name=None, split=None, **kwargs)

主要参数说明

path (必填):
• Hugging Face Hub 数据集: 直接传入 Hub 上的数据集名称，例如 "glue"、"squad"、"imdb" 等。
• 本地文件/目录: 传入本地文件路径（支持 CSV/JSON/TXT 等格式），例如 "path/to/data.csv"。
• 自定义脚本: 传入本地数据集生成脚本的路径（需符合 datasets 库格式）。
name (可选):
• 指定数据集的子配置（例如 "glue" 数据集下有 "cola"、"sst2" 等子任务）。
• 示例: load_dataset("glue", name="sst2") 加载 GLUE 的 SST-2 子任务。
split (可选):
• 指定加载的数据集划分，如 "train"、"test"、"validation"，或组合（如 "train+test"）。
• 示例: split="train[:10%]" 加载训练集的前 10%。
其他常用参数:
• data_dir: 数据集文件的存储目录（适用于需要额外数据的场景）。
• data_files: 直接指定文件路径（支持通配符 *），例如 data_files={"train": "train.csv", "test": "test.csv"}。
• cache_dir: 自定义缓存目录（默认在 ~/.cache/huggingface/datasets）。
• streaming: 设为 True 时启用流式加载（适用于超大数据集，无需全量加载到内存）。

常见使用场景

1. 从 Hugging Face Hub 加载数据集

# 加载 GLUE 数据集的 SST-2 子任务，训练集
dataset = load_dataset("glue", "sst2", split="train")# 加载 SQuAD 问答数据集（返回 DatasetDict 包含 train 和 validation）
dataset_dict = load_dataset("squad")

2. 加载本地文件

# 加载 CSV 文件（自动推断格式）
dataset = load_dataset("csv", data_files="data.csv")# 加载多个 JSON 文件
dataset = load_dataset("json", data_files={"train": "train.json", "test": "test.json"})

3. 流式模式（处理大型数据集）

# 逐样本加载，避免内存不足
streaming_dataset = load_dataset("big_dataset", streaming=True)
for example in iter(streaming_dataset["train"]):process(example)

4. 自定义数据集脚本

若数据集未在 Hub 上，可编写脚本定义数据加载逻辑（需符合 Dataset Script 规范）：

dataset = load_dataset("path/to/dataset_script.py", name="my_config")

返回对象

• Dataset: 单一切片的数据集（如 split="train"），支持类似 Pandas 的索引和操作。
• DatasetDict: 包含多个切片的字典（例如 {"train": Dataset, "test": Dataset}）。

数据处理特性

内存高效：基于 Apache Arrow 格式，支持零拷贝读取。
预处理：可通过 .map() 方法快速应用预处理函数。
兼容性：可轻松转换为 Pandas DataFrame（.to_pandas()）或 NumPy 数组。

示例代码

# 加载 IMDb 电影评论数据集
dataset = load_dataset("imdb")# 查看训练集前 3 条样本
print(dataset["train"][:3])# 预处理：分词
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")def tokenize_function(examples):return tokenizer(examples["text"], truncation=True)tokenized_dataset = dataset.map(tokenize_function, batched=True)

注意事项

• 缓存机制：首次加载数据集时会下载或处理数据，后续调用直接读取缓存。
• 依赖安装：加载某些格式（如 Parquet）需安装额外依赖（pip install datasets[parquet]）。
• 版本控制：通过 revision 参数指定数据集版本（如 Git 分支、commit hash）。

通过 load_dataset，Hugging Face 提供了统一且高效的接口，极大简化了 NLP 任务中的数据加载流程。

2. load_datase返回值

load_dataset 方法返回的数据类型取决于传入的参数，通常是 Dataset 或 DatasetDict 对象。以下是具体规则：

1. 默认情况：返回 `DatasetDict`

• 当不指定 split 参数，且数据集包含预定义的多个划分（如 train/test/validation）时，返回一个 DatasetDict 对象。
• 示例：

from datasets import load_dataset# 加载 SQuAD 数据集（包含 train 和 validation）
dataset_dict = load_dataset("squad")
print(type(dataset_dict))  # 输出: <class 'datasets.dataset_dict.DatasetDict'># 访问训练集
train_data = dataset_dict["train"]

2. 指定 `split` 参数：返回 `Dataset`

• 当明确指定 split（如 split="train"），返回单个 Dataset 对象。
• 示例：

# 仅加载 IMDb 数据集的测试集
dataset = load_dataset("imdb", split="test")
print(type(dataset))  # 输出: <class 'datasets.arrow_dataset.Dataset'># 访问第一条数据
print(dataset[0])  # 输出: {"text": "Great movie!", "label": 1, ...}

3. 自定义数据文件：返回类型灵活

• 通过 data_files 参数加载本地文件时：
• 如果指定单个文件（如 data_files="data.csv"），返回 Dataset。
• 如果指定多个文件（如 data_files={"train": "train.csv", "test": "test.csv"}），返回 DatasetDict。

示例：

# 加载单个 CSV 文件
dataset = load_dataset("csv", data_files="data.csv", split="train")# 加载多个 JSON 文件
dataset_dict = load_dataset("json", data_files={"train": "train.json", "test": "test.json"})

4. 流式模式 (`streaming=True`)：返回迭代器

• 当启用流式加载（streaming=True）时，返回的是 IterableDataset 或 IterableDatasetDict，适合逐样本处理超大数据集。
• 示例：

# 流式加载维基百科数据集
dataset = load_dataset("wikipedia", "20220301.en", split="train", streaming=True)
print(type(dataset))  # 输出: <class 'datasets.iterable_dataset.IterableDataset'># 逐样本读取（无需全量加载到内存）
for example in dataset:print(example["text"])break  # 仅读取第一条

关键区别

返回类型	适用场景	数据访问方式
`DatasetDict`	多划分数据集（如 train/test）	通过字典键访问：`dataset_dict["train"]`
`Dataset`	单划分数据集	直接索引或切片：`dataset[0]`
`IterableDataset(Dict)`	流式处理超大数据集（避免内存爆炸）	通过迭代器逐条访问：`for example in dataset: ...`

常用操作

转换为其他格式：

# 转 Pandas DataFrame
df = dataset.to_pandas()# 转 Python 字典列表
data_list = dataset.to_list()

数据预处理：

# 使用 .map() 方法批量处理
def preprocess(example):example["text_length"] = len(example["text"])return exampleprocessed_dataset = dataset.map(preprocess)

过滤数据：

# 过滤短文本
filtered_dataset = dataset.filter(lambda x: len(x["text"]) > 100)

注意事项

缓存机制：首次加载数据集会下载或处理数据，之后直接从缓存读取（路径默认在 ~/.cache/huggingface/datasets）。
内存优化：Dataset 基于 Apache Arrow 格式，内存占用低，支持快速随机访问。
版本控制：通过 revision 参数加载特定版本的数据集（如 Git 分支、commit hash）。

简单总结

• 需要处理多个划分（如训练集、测试集） → DatasetDict。
• 只需单个划分 → Dataset。
• 处理超大数据集 → IterableDataset + 流式模式。

1. 基本介绍

基本用法

主要参数说明

常见使用场景

1. 从 Hugging Face Hub 加载数据集

2. 加载本地文件

3. 流式模式（处理大型数据集）

4. 自定义数据集脚本

返回对象

数据处理特性

示例代码

注意事项

2. load_datase返回值

1. 默认情况：返回 DatasetDict

2. 指定 split 参数：返回 Dataset

3. 自定义数据文件：返回类型灵活

4. 流式模式 (streaming=True)：返回迭代器

关键区别

常用操作

转换为其他格式：

数据预处理：

过滤数据：

注意事项

简单总结

相关文章：

1. 默认情况：返回 `DatasetDict`

2. 指定 `split` 参数：返回 `Dataset`

4. 流式模式 (`streaming=True`)：返回迭代器