准备数据以微调 Hugging Face 模型

项目
2025/04/23

本文演示了如何使用 Hugging Face Transformers 和 Hugging Face Datasets 准备数据以微调开源大型语言模型。

要求

用于机器学习的 Databricks Runtime 13.0 及更高版本。本指南中的示例使用 Databricks Runtime 13.0 ML 及更高版本中包含的 Hugging Face 数据集。
一个启用了 Unity Catalog 的工作区。还需要具有以下权限才能将数据写入 Unity Catalog 卷：
- 要将文件上传到的卷的 WRITE VOLUME 特权。
- 父架构上的 USE SCHEMA 特权。
- 父目录上的 USE CATALOG 特权。
用于下载大型数据集的重要计算资源。提供的示例笔记本中使用的大型数据集需要超过一天的下载时间。

从 Hugging Face 加载数据

Hugging Face Datasets 是一个 Hugging Face 库，用于访问和共享音频、计算机视觉和自然语言处理 (NLP) 任务的数据集。借助 Hugging Face datasets，你可以从不同位置加载数据。 datasets 库包含用于从 Hugging Face Hub 读取数据集的实用工具。使用 load_dataset 函数可以从 Hugging Face Hub 下载和读取许多数据集。在 Hugging Face 文档中详细了解如何使用 Hugging Face Datasets 加载数据。

from datasets import load_dataset
dataset = load_dataset("imdb")

Hugging Face Hub 中的某些数据集提供了调用 load_dataset 时下载和生成的数据的大小。在使用 load_dataset_builder 下载数据集之前，可以使用 load_dataset 来了解大小。

from datasets import load_dataset_builder
from psutil._common import bytes2human

def print_dataset_size_if_provided(*args, **kwargs):
  dataset_builder = load_dataset_builder(*args, **kwargs)

  if dataset_builder.info.download_size and dataset_builder.info.dataset_size:
    print(f'download_size={bytes2human(dataset_builder.info.download_size)}, dataset_size={bytes2human(dataset_builder.info.dataset_size)}')
  else:
    print('Dataset size is not provided by uploader')

print_dataset_size_if_provided("imdb")

有关如何在 Azure Databricks 中为不同大小的数据下载和准备数据集的指导，请参阅从 Hugging Face 最佳做法笔记本下载数据集。

格式化你的训练和评估数据

若要使用你自己的数据进行模型微调，必须首先将训练和评估数据的格式设置为 Spark 数据帧。然后，使用 Hugging Face datasets 库加载数据帧。

首先将训练数据的格式设置为符合训练器预期的表。对于文本分类，此表包含两列：文本列和标签列。

若要执行微调，需要提供一个模型。使用 Hugging Face Transformers AutoClasses 库可用于轻松加载模型和配置设置，包括用于自然语言处理的各种 Auto Models。

例如，Hugging Face transformers 提供 AutoModelForSequenceClassification 作为文本分类的模型加载器，它要求提供整数 ID 作为类别标签。但是，如果你有带字符串标签的数据帧，则还必须在创建模型时指定整数标签和字符串标签之间的映射。可按如下所示方式收集此信息：

labels = df.select(df.label).groupBy(df.label).count().collect()
id2label = {index: row.label for (index, row) in enumerate(labels)}
label2id = {row.label: index for (index, row) in enumerate(labels)}

然后使用 Pandas UDF 创建整数 ID 作为标签列：

from pyspark.sql.functions import pandas_udf
import pandas as pd
@pandas_udf('integer')
def replace_labels_with_ids(labels: pd.Series) -> pd.Series:
  return labels.apply(lambda x: label2id[x])

df_id_labels = df.select(replace_labels_with_ids(df.label).alias('label'), df.text)

在 Spark DataFrame 中加载 Hugging Face 数据集

Hugging Face datasets 支持使用 datasets.Dataset.from_spark 从 Spark 数据帧进行加载。若要详细了解 from_spark() 方法，请参阅 Hugging Face 文档。

例如，如果你有 train_df 和 test_df 数据帧，则可以使用以下代码为每个数据帧创建数据集：

import datasets
train_dataset = datasets.Dataset.from_spark(train_df, cache_dir="/Volumes/main/default/my-volume/train")
test_dataset = datasets.Dataset.from_spark(test_df, cache_dir="/Volumes/main/default/my-volume/test")

Dataset.from_spark 会缓存数据集。此示例描述了驱动程序上的模型训练，因此必须向驱动程序提供数据。此外，由于缓存具体化是使用 Spark 并行化的，因此，提供的 cache_dir 必须可供所有工作器访问。要满足这些约束，cache_dir 应为 Unity Catalog 卷路径。

可使用 Unity Catalog 管理对卷的访问权限。

如果数据集很大，将其写入 Unity Catalog 可能需要很长时间。为了加快该过程，可以使用 working_dir 参数让 Hugging Face datasets 将数据集写入磁盘上的临时位置，然后将其移动到 Unity Catalog。例如，若要使用 SSD 作为临时位置，请运行以下命令：

import datasets
dataset = datasets.Dataset.from_spark(
  train_df,
  cache_dir="/Volumes/main/default/my-volume/train",
  working_dir="/local_disk0/tmp/train",
)

数据集的缓存

缓存是 datasets 提高效率的方法之一。缓存会存储所有已下载和处理的数据集，因此，当用户需要使用中间数据集时，系统会直接从缓存中重新加载这些数据集。

数据集的默认缓存目录是 ~/.cache/huggingface/datasets。终止群集时，缓存数据也会丢失。若要在群集终止时保留缓存文件，Databricks 建议通过设置环境变量 HF_DATASETS_CACHE 将缓存位置更改为 Unity Catalog 卷路径：

import os
os.environ["HF_DATASETS_CACHE"] = "/Volumes/main/default/my-volume/"

微调模型

数据准备就绪后，可以使用它来微调 Hugging Face 模型。

笔记本：从 Hugging Face 下载数据集

此示例笔记本提供了使用 Hugging Face load_dataset 函数在 Azure Databricks 中为不同大小的数据下载和准备数据集的建议最佳做法。

从 Hugging Face 最佳做法笔记本下载数据集

获取笔记本

通过