Hugging Face load_dataset 方法使用

load_dataset 是 Hugging Face datasets 库中最常用的函数之一，它支持从 Hugging Face Hub 或本地加载各种格式的数据集（如 CSV、JSON、Parquet、SQL、文本等）。以下是 load_dataset 的常见用法：

1. 从 Hugging Face Hub 加载公开数据集

from datasets import load_dataset

# 加载 IMDb 电影评论数据集
dataset = load_dataset("imdb")

# 查看数据集的结构
print(dataset)

说明：

• "imdb" 是数据集名称，Hugging Face Hub 提供大量公开数据集，可以直接通过名称加载。
• 该数据集包含 "train", "test", "unsupervised" 三个子集，print(dataset) 会显示数据集的结构。

2. 加载特定子集（Split）

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

# 只加载测试集
test_dataset = load_dataset("imdb", split="test")

# 加载部分数据（切片）
subset = load_dataset("imdb", split="train[:1000]")  # 训练集前 1000 条数据

说明：

• split="train" 仅加载训练集。
• split="train[:1000]" 仅加载训练集的前 1000 条数据，这在调试或快速实验时很有用。

3. 加载多语言数据集

# 加载 XNLI 数据集的英文部分
dataset = load_dataset("xnli", "en")

# 加载 XNLI 数据集的中文部分
dataset = load_dataset("xnli", "zh")

说明：

• "xnli" 是一个跨语言自然语言推理数据集。
• 第二个参数 "en" 或 "zh" 指定语言版本。

4. 从本地文件加载数据

4.1. 加载 CSV 文件

dataset = load_dataset("csv", data_files="data.csv")

4.2. 加载 JSON 文件

dataset = load_dataset("json", data_files="data.json")

4.3. 加载多个文件

dataset = load_dataset("csv", data_files={"train": "train.csv", "test": "test.csv"})

说明：

• data_files="file.csv" 直接加载单个 CSV 文件。
• data_files={"train": "train.csv", "test": "test.csv"} 指定不同文件作为不同的子集。

5. 加载压缩格式数据

dataset = load_dataset("csv", data_files="data.zip")

说明：

• load_dataset 支持加载 .zip、.tar.gz 等格式的压缩文件，前提是其中包含可以解析的数据（如 CSV、JSON）。

6. 从本地目录加载 Hugging Face 格式数据集

dataset = load_dataset("/root/DS_ft/medical-o1-reasoning-SFT", "en")

说明：

• /root/DS_ft/medical-o1-reasoning-SFT 是本地存储的数据集路径。
• "en" 可能是数据集的子配置。

7. 远程加载（支持 trust_remote_code）

dataset = load_dataset("custom_dataset_name", trust_remote_code=True)

说明：

• 如果数据集脚本包含远程代码（如数据转换或解析逻辑），需要 trust_remote_code=True 允许执行远程代码。

8. 处理大规模数据（Streaming 模式）

dataset = load_dataset("imdb", split="train", streaming=True)

for example in dataset:
    print(example)

说明：

• streaming=True 使数据集以流式方式加载，适用于超大规模数据集，可避免内存占用过高。

9. 指定缓存目录

dataset = load_dataset("imdb", cache_dir="/custom/cache/path")

说明：

• cache_dir 允许指定数据集缓存目录，避免重复下载，提高加载速度。

10. 加载特定列

dataset = load_dataset("csv", data_files="data.csv", keep_in_memory=True, features={"text": "string", "label": "int32"})

说明：

• keep_in_memory=True 让数据存入内存，适用于小型数据集。
• features 指定数据类型，如 "text" 是字符串，"label" 是整数。

总结

load_dataset 功能非常强大，适用于不同场景的数据加载需求。