准备数据以微调 Hugging Face 模型

项目
03/01/2024

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

要求

用于机器学习的 Databricks Runtime 13.0 及更高版本。本指南中的示例使用 Databricks Runtime 13.0 ML 及更高版本中包含的 Hugging Face 数据集。

从 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 下载数据集之前，可以使用 load_dataset_builder 来了解大小。

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 数据帧加载 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="/dbfs/cache/train")
test_dataset = datasets.Dataset.from_spark(test_df, cache_dir="/dbfs/cache/test")

Dataset.from_spark 会缓存数据集。此示例描述了驱动程序上的模型训练，因此必须向驱动程序提供数据。此外，由于缓存具体化是使用 Spark 并行化的，因此，提供的 cache_dir 必须可供所有工作器访问。若要满足这些约束，cache_dir 应当是 Databricks 文件系统 (DBFS) 根卷或装入点。

DBFS 根卷可供工作区的所有用户访问，但只能用于没有访问限制的数据。如果数据需要访问控制，请使用装入点而不是 DBFS 根。

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

import datasets
dataset = datasets.Dataset.from_spark(
  train_df,
  cache_dir="/dbfs/cache/train",
  working_dir="/local_disk0/tmp/train",
)

数据集的缓存

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

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

import os
os.environ["HF_DATASETS_CACHE"] = "/dbfs/place/you/want/to/save"

微调模型

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

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

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

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

获取笔记本