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

概述

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 等）生成任何合成数据集。

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

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