🏷️ 使用 ⚗️ distilabel 为 LLM 对齐构建 AI 反馈 (AIF) 数据集

社区文章发布于 2023 年 12 月 5 日

摘要

引言

什么是 distilabel？

安装

构建 HelpSteer-like AIF 数据集
1. 定义 HelpSteerTask

2. 从 Hub 加载 HelpSteer 数据集

3. 运行管道

4. 导出到 Argilla（可选）

Colab Notebook

下一步？

参考文献

概述

distilabel 是一个 AI 反馈 (AIF) 框架，用于使用 LLM 生成和标注数据集，得益于其可扩展性，您可以使用它通过流行的 LLM 引擎（如 🤗 transformers、🤗 Inference Endpoints、vLLM、llama.cpp 等）生成任何合成数据集。HelpSteer 是 NVIDIA 发布的一个有用性数据集，它使用人工标注员评估提示-响应对数据集，除了有用性之外，还包括：正确性、冗长性、连贯性和复杂性。这篇博文解释了如何使用 distilabel 生成 HelpSteer-like 数据集，用于 LLM 对齐，使用 AIF 代替人工标注，同时还展示了如何通过 Argilla 集成来整合人工反馈。

引言

最近的模型，例如 HuggingFaceH4/zephyr-7b-beta 和 allenai/tulu-2-7b，已经证明可以通过使用 AI 反馈 (AIF) 数据集进行直接偏好优化 (DPO) 来微调强大的开源 LLM。他们的方法包括使用监督微调 (SFT) 对一个强大的基础 LLM 进行微调，然后使用 DPO 进行意图对齐，因为大多数情况下，使用 SFT 生产的 LLM 在基准测试中表现良好，但在真实世界场景中无法泛化，即生成的文本的意图与自然提示的意图不一致。

DPO 微调在短时间内改善了 LLM 对齐，相比于 SFT 微调，在这种情况下使用合成生成的数据集进行 AIF；并且在 AIF 领域有一些令人兴奋的研究，例如 UltraFeedback、JudgeLM 或 Prometheus。

然而，超越研究工作并大规模应用 AIF 是不同的，这就是 distilabel 最初创建的原因，旨在以稳健、高效和可扩展的方式实施 AIF 方法，允许任何人针对自己的用例大规模构建自定义合成数据集。唯一的不足之处在于 LLM 有时无法产生我们预期的结果，因此人类参与（human-in-the-loop）对于提高数据集质量是开源 LLM 模型下一步的重大飞跃。

distilabel 旨在通过其与 Argilla 的集成来弥合这一差距，Argilla 是一个人类参与数据标注工具，允许任何人将人工反馈整合到他们的数据集中。

什么是 distilabel？

distilabel 是一个 AIF 框架，用于使用 LLM 生成和标注数据集，得益于其可扩展性，您可以使用它通过流行的 LLM 引擎（如 🤗 transformers、🤗 Inference Endpoints、vLLM、llama.cpp 等）生成任何合成数据集。

从上述工作流程中，让我们分解不同的步骤

任务：任务是一个类，定义 LLM 的输入和输出参数，以及用于生成数据集的提示。它还负责解析 LLM 输出并返回一个包含输出参数及其相应值的字典。
LLMs：LLM 是将用于生成数据集的模型。它们被定义为实现 generate 方法的类，该方法接收一个提示并返回生成的输出。LLM 为一些流行的框架/引擎（如 🤗 transformers、🤗 Inference Endpoints、vLLM、llama.cpp 等）提供了包装器。
管道：管道是一个类，它负责编排所提供的 🤗 数据集的生成、标注或两者结合。它负责生成提示、生成数据集、标注数据集，最后返回已标注的数据集。它实现了一些优化机制，以避免在任何步骤中中断，从而提高健壮性，因为期望 LLM 始终产生预期输出并非总是如此。

下面您将找到一个端到端示例，说明如何生成一个类似于 nvidia/HelpSteer 的数据集，但使用 AIF 而不是人工标注。

安装

您可以通过 pip 安装它，如下所示，它还将安装 openai 和 argilla 附加组件，它们分别用于 OpenAI 集成和 Argilla 集成，分别用于收集 AIF 和导出到 Argilla。请注意，🤗 datasets 是一个核心依赖项，因此不需要通过任何附加组件安装，而其他一些包可能需要。

pip install distilabel[openai,argilla]

构建 HelpSteer-like AIF 数据集

让我们简要介绍一下 HelpSteer 数据集，这是 NVIDIA 发布的一个有用性数据集，它使用人工标注员评估提示-响应对数据集，除了有用性之外，还包括：正确性、冗长性、连贯性和复杂性。提示-响应数据集包含一组提示，每个提示最多可以包含 4 个响应，这些响应是使用 NVIDIA 内部的 43B 参数 LLM 生成的。

有关数据收集和标注过程的更多信息，请参阅 nvidia/HelpSteer。

1. 定义 HelpSteerTask

如前所述，我们从 HelpSteerTask 的定义开始，所以我们需要继承 distilabel.tasks.Task 并实现 input_args_names、output_args_names、generate_prompt 和 parse_output 方法。可选地，我们还可以决定实现 to_argilla_dataset 和 to_argilla_record 方法，以便稍后能够轻松地将标注数据集导出到 Argilla。

理想情况下，由于提示将定义稍后使用的 LLM 如何处理信息，我们应该在实际定义提示之前进行一些实验和提示工程。在这种情况下，我们使用一个提示，要求 LLM 根据 HelpSteer 中呈现的不同区域评估提示-响应对，并使用 XLM 格式，因为这是通过 OpenAI 的 Playground 进行提示工程后被证明是最佳的，也遵循了 OpenAI - 提示工程中描述的注意事项。

import re
from typing import Any, Dict, List

import argilla as rg
from distilabel.tasks import Task


class HelpSteerTask(Task):
    system_prompt: str = (
        "You are a helpful assistant and you are asked to evaluate a generated"
        " response for a given prompt.The aspects that you need to evaluate are:"
        " correctness, coherence, complexity, verbosity, and helpfulness; the"
        " ratings are integers that go from 0 to 4, meaning the higher the"
        " better. You should expect and produce the following:\n\n<prompt>..."
        "</prompt>\n<response>...</response>\n<correctness>X</correctness>\n"
        "<coherence>X</coherence>\n<complexity>X</complexity>\n<verbosity>X"
        "</verbosity>\n<helpfulness>X</helpfulness>\n"
    )

    @property
    def input_args_names(self) -> List[str]:
        return ["prompt", "response"]

    @property
    def output_args_names(self) -> List[str]:
        return ["correctness", "coherence", "complexity", "verbosity", "helpfulness"]

    def generate_prompt(self, prompt: str, response: str) -> List[Dict[str, str]]:
        return [
            {
                "role": "system",
                "content": self.system_prompt,
            },
            {
                "role": "user",
                "content": f"<prompt>{prompt}</prompt>/n<response>{response}</response>"
            },
        ]

    def parse_output(self, output: str) -> Dict[str, int]:
        matches = re.findall(r"<(\w+)>(\d+)</\1>", output)
        return dict((key, int(value)) for key, value in matches)

    # Optional methods to export to Argilla later-on
    def to_argilla_dataset(self, dataset_row: Dict[str, Any]) -> rg.FeedbackDataset:
        return rg.FeedbackDataset(
            fields=[
                rg.TextField(name=input_arg_name) for input_arg_name in self.input_args_names
            ],
            questions=[
                # We need to shift the ratings 1 to the right, as Argilla won't allow 0, so the Likert 5 scale has to be 1-5 instead of 0-4
                rg.RatingQuestion(name=output_arg_name, values=[1, 2, 3, 4, 5]) for output_arg_name in self.output_args_names
            ],
            guidelines="https://huggingface.co/datasets/nvidia/HelpSteer",
        )

    # Optional methods to export to Argilla later-on
    def to_argilla_record(self, dataset_row: Dict[str, Any]) -> rg.FeedbackRecord:
        return rg.FeedbackRecord(
            fields=dict((input_arg_name, dataset_row[input_arg_name]) for input_arg_name in self.input_args_names),
            suggestions=[
                # We need to shift the ratings 1 to the right, as Argilla won't allow 0, so the Likert 5 scale has to be 1-5 instead of 0-4
                rg.SuggestionSchema(question_name=output_arg_name, value=dataset_row[output_arg_name] + 1) for output_arg_name in self.output_args_names
            ],
        )

2. 从 Hub 加载 HelpSteer 数据集

然后我们从 HuggingFace Hub 加载数据集，并只保留前 1000 行，因为这只是为了在一个简短且可重现的子集上展示 distilabel。除此之外，我们还删除了 Scale AI 在人工标注过程中生成的列，因为这些是我们希望使用 AIF 生成的列。

from datasets import load_dataset

dataset = load_dataset("nvidia/HelpSteer", split="train[:1000]")

dataset = dataset.remove_columns(column_names=["helpfulness", "correctness", "coherence", "complexity", "verbosity"])

3. 运行管道

然后我们需要定义 LLM 和 Pipeline 来运行标注数据集的生成。

我们将使用 distilabel.llm.OpenAILLM 定义 LLM，它要求我们提供要使用的 task（已在上面定义），以及一些生成 kwargs，例如要使用的 model、要生成的 max_new_tokens、要使用的 temperature 和用于并行生成数据集的 num_threads。

由于 OpenAI 的成本高于其他 OSS 解决方案，我们的计划是使用 OSS 模型进行 AIF 并训练我们自己的模型，以便在定义用于 AIF 的 LLM 时为用户提供更大的灵活性和替代方案；但至少在最初的迭代中，我们将依赖 OpenAI，因为它在标注（即生成结构化输出）方面表现出最强的实力。一个用于收集 AIF 的开源 LLM 的不错示例是 kaist-ai/prometheus-7b-v1.0，但它不适合此任务。

然后我们只需将 OpenAILLM 作为 labeller 提供给 Pipeline，然后我们调用 generate 方法，将预加载的 🤗 Dataset 作为 dataset 参数，我们将获得标注的数据集作为结果，这是一个继承自 🤗 Dataset 的类，包含一些预定义的方法，可以方便地将 AIF 数据集导出到 Argilla，以防用户愿意添加人工参与以获得更高质量的数据。

import os
from distilabel.llm import OpenAILLM
from distilabel.pipeline import Pipeline

pipeline = Pipeline(
    labeller=OpenAILLM(
        task=HelpSteerTask(),
        model="gpt-4",
        max_new_tokens=128,  # We just need 43
        temperature=0.0,
        num_threads=4,
        openai_api_key=os.getenv("OPENAI_API_KEY") or "sk-...",
    ),
)

dataset = pipeline.generate(dataset=dataset, display_progress_bar=True, verbose=False)

4. 导出到 Argilla（可选）

此步骤是可选的，但如果您还没有 Argilla 实例运行，请随时开始使用 HuggingFace Spaces 中的 Argilla 模板。

此外，如果我们在 HelpSteerTask 类中实现了 to_argilla_dataset 和 to_argilla_record 方法，我们可以通过 to_argilla 方法将标注数据集导出到 Argilla。因此，to_argilla 方法将返回格式为 rg.FeedbackDataset 的数据集和格式为 rg.FeedbackRecord 的记录，这些记录可以用于一行代码将数据集推送到 Argilla。

import argilla as rg

rg.init(api_url="<ARGILLA_API_URL>", api_key="<ARGILLA_API_KEY>")

rg_dataset = dataset.to_argilla()
rg_dataset.push_to_argilla(name="HelpSteer-AIF", workspace="admin")

Colab Notebook

为了使用 AIF 重现 HelpSteer 数据集标注，您可以使用上面的代码片段，或者直接使用以下 Google Colab Notebook。

接下来是什么？

我们很乐意收到您对 distilabel 的反馈，以及对新示例、新功能等的需求，以继续推动我们的 AIF 框架，使使用 AIF 生成和标注数据集比以往任何时候都更容易、更简单！

我们接下来计划做什么

添加模型池化机制，能够一次使用多个 LLM 生成响应，从而获得更多样化的响应进行评估。
探索一些开源 LLM 作为标注器，而不是我们使用的默认 OpenAI 模型（实际上您可以在 distilabel 中使用任何开源模型作为标注器，但我们尚未测试过）
包含更多示例和用例，这也意味着扩展我们的文档。
寻求性能改进，使其在生成/标注方面达到最佳。

除此之外，我们可能会处理一些相关的小功能和/或错误修复，但请随时提出问题或发起讨论与我们交流！

参考文献

GitHub 仓库：argilla-io/distilabel
HelpSteer 🤗 数据集：nvidia/HelpSteer
HelpSteer 论文：https://arxiv.org/abs/2311.09528
带 AIF 的 HelpSteer 🤗 数据集：alvarobartt/HelpSteer-AIF
使用 distilabel 生成的 AIF 数据集

社区

通过拖放到文本输入框、粘贴或点击此处上传图片、音频和视频。

点击或粘贴此处以上传图片

· 注册或登录评论