手把手教你用 Hugging Face Datasets 清洗和整理数据集

在日常的机器学习项目中，我们经常会遇到这样的情况：下载好的数据集总是带着各种 “小毛病”—— 冗余的列、格式混乱的文本、参差不齐的样本…… 这些问题就像拦路虎，让数据无法直接用于模型训练。别担心，今天我们就来聊聊 Hugging Face Datasets 库，看看如何用它高效地清洗和整理数据集，为后续的模型训练铺好路。

一、数据加载：从 TSV 文件到 Dataset 对象

我们以加州大学欧文分校的药物审查数据集为例，这个数据集包含患者评论、病情、评分等信息，存储在 TSV 文件中。TSV 文件其实是 CSV 的 “近亲”，只是用制表符作为分隔符。我们可以轻松使用 Hugging Face Datasets 的load_dataset函数来加载它：

python

from datasets import load_dataset
data_files = {"train": "drugsComTrain_raw.tsv", "test": "drugsComTest_raw.tsv"}
drug_dataset = load_dataset("csv", data_files=data_files, delimiter="\t")

加载完数据后，为了快速了解数据特点，我们可以生成一个随机样本。通过链式调用shuffle和select函数，我们就能得到一个包含 1000 个示例的随机样本：

python

drug_sample = drug_dataset["train"].shuffle(seed=42).select(range(1000))

查看样本数据，我们发现了一些常见问题：有冗余列 “Unnamed: 0”，病情标签 “condition” 大小写混合，评论文本里包含 HTML 字符如'，而且评论长度不一。这些都是我们接下来要解决的 “小麻烦”。

二、数据清洗：让数据整整齐齐

列操作：重命名、过滤与统一格式

首先处理冗余列。我们猜测 “Unnamed: 0” 可能是患者 ID，通过unique函数验证后，使用rename_column函数将其重命名为 “patient_id”：

python

drug_dataset = drug_dataset.rename_column(original_column_name="Unnamed: 0", new_column_name="patient_id")

接着处理病情标签。我们打算将其全部转换为小写，但在转换过程中遇到了问题 —— 有些行的 “condition” 是 None 值，无法调用lower方法。这时，filter函数就派上用场了，我们用它过滤掉 “condition” 为 None 的行：

python

drug_dataset = drug_dataset.filter(lambda x: x["condition"] is not None)

过滤后，再使用map函数将 “condition” 列统一转换为小写：

python

def lowercase_condition(example):
    return {"condition": example["condition"].lower()}
drug_dataset = drug_dataset.map(lowercase_condition)

文本清洗：解码 HTML 字符与计算评论长度

评论文本中的 HTML 字符会影响后续处理，我们可以利用 Python 的html模块来解码。使用map函数对所有评论文本进行处理：

python

import html
drug_dataset = drug_dataset.map(lambda x: {"review": html.unescape(x["review"])})

评论长度也是一个重要特征。我们定义一个函数计算每条评论的字数，并通过map函数创建新列 “review_length”：

python

def compute_review_length(example):
    return {"review_length": len(example["review"].split())}
drug_dataset = drug_dataset.map(compute_review_length)

对于过短的评论，比如字数少于 30 的，可能包含的信息不足，我们可以用filter函数将其过滤掉：

python

drug_dataset = drug_dataset.filter(lambda x: x["review_length"] > 30)

三、数据处理高级技巧：提升效率与灵活性

Dataset.map () 的优化技巧

map函数是数据处理的得力助手，它有两个实用的参数batched和num_proc。设置batched=True可以批量处理数据，显著提升处理速度。比如处理 HTML 解码时，批处理比单例处理快数倍。当使用慢速分词器时，设置num_proc参数进行多进程处理，能进一步加速数据处理：

python

slow_tokenizer = AutoTokenizer.from_pretrained("bert-base-cased", use_fast=False)
def slow_tokenize_function(examples):
    return slow_tokenizer(examples["review"], truncation=True)
tokenized_dataset = drug_dataset.map(slow_tokenize_function, batched=True, num_proc=8)

当分词生成多个样本导致数据长度不匹配时，我们可以通过remove_columns参数删除旧列，或者利用overflow_to_sample_mapping字段映射原始样本索引来保留旧字段，灵活处理数据结构变化。

与 Pandas 交互：取长补短

Hugging Face Datasets 与 Pandas 的交互非常方便。通过set_format("pandas")函数，我们可以将数据集转为 Pandas DataFrame，利用 Pandas 的groupby、value_counts等高级功能进行数据分析：

python

drug_dataset.set_format("pandas")
train_df = drug_dataset["train"][:]
frequencies = (
    train_df["condition"]
    .value_counts()
    .to_frame()
    .reset_index()
    .rename(columns={"index": "condition", "count": "frequency"})
)

分析完成后，再通过Dataset.from_pandas()函数将 DataFrame 转回 Dataset 对象，无缝衔接后续操作。

四、数据集划分与保存：为训练做好准备

创建验证集：避免过拟合

为了更好地评估模型性能，我们需要创建验证集。使用train_test_split函数将训练集划分为训练集和验证集，设置固定的随机种子确保结果可复现：

python

drug_dataset_clean = drug_dataset["train"].train_test_split(train_size=0.8, seed=42)
drug_dataset_clean["validation"] = drug_dataset_clean.pop("test")
drug_dataset_clean["test"] = drug_dataset["test"]

保存数据集：多种格式任选

Hugging Face Datasets 支持多种格式保存数据集。Arrow 格式适合存储大型数据集，具有高效加载和分片存储的优点：

python

drug_dataset_clean.save_to_disk("drug-reviews")

如果需要通用格式，也可以保存为 CSV 或 JSON 格式。保存为 JSON Lines 格式时，每个样本占一行 JSON，方便后续加载和处理：

python

for split, dataset in drug_dataset_clean.items():
    dataset.to_json(f"drug-reviews-{split}.jsonl")

五、核心工具与最佳实践

在整个数据处理过程中，map、filter、shuffle、select、sort等函数是我们的核心工具。性能优化方面，优先使用批处理和快速分词器，合理设置多进程参数。同时，固定随机种子是保证结果可复现的关键，这在团队协作和模型调试中非常重要。

通过以上步骤，我们已经将一个 “脏乱差” 的数据集清洗整理成适合模型训练的格式。Hugging Face Datasets 库的强大功能让数据处理变得高效又灵活，大大减少了我们在数据预处理上的时间和精力。

希望本文能为你在数据集处理上提供实用的帮助。如果你觉得有用，欢迎点赞收藏，后续我们会分享更多 Hugging Face 库的实用技巧，比如模型训练、结果评估等。如果你在实际操作中遇到问题，或者有想了解的技术点，欢迎在评论区留言，我们一起探讨！