Hugging Face 的徽标

TRL 文档

数据集格式

Hugging Face's logo
加入 Hugging Face 社区

并获得增强型文档体验

以开始使用

数据集格式

本指南概述了 TRL 中每个训练器支持的数据集格式。由于对话数据集非常常见，我们还提供有关如何使用它们以及如何将它们转换为 TRL 训练器的标准数据集格式的指南。

数据集格式和类型的概述

数据集的格式是指数据的结构方式，通常分为标准或对话两种。类型与数据集设计用于的特定任务相关联，例如仅提示或偏好。每种类型都以其列为特征，这些列根据任务而异，如表所示。

类型 \ 格式 标准 对话
语言建模 
{"text": "The sky is blue."}
 
{"messages": [{"role": "user", "content": "What color is the sky?"},
              {"role": "assistant", "content": "It is blue."}]}
仅提示 
{"prompt": "The sky is"}
 
{"prompt": [{"role": "user", "content": "What color is the sky?"}]}
提示-完成 
{"prompt": "The sky is",
 "completion": " blue."}
 
{"prompt": [{"role": "user", "content": "What color is the sky?"}],
 "completion": [{"role": "assistant", "content": "It is blue."}]}
偏好 
{"prompt": "The sky is",
 "chosen": " blue.",
 "rejected": " green."}
或者，带有隐式提示
{"chosen": "The sky is blue.",
 "rejected": "The sky is green."}
 
{"prompt": [{"role": "user", "content": "What color is the sky?"}],
 "chosen": [{"role": "assistant", "content": "It is blue."}],
 "rejected": [{"role": "assistant", "content": "It is green."}]}
或者，带有隐式提示
{"chosen": [{"role": "user", "content": "What color is the sky?"},
              {"role": "assistant", "content": "It is blue."}],
  "rejected": [{"role": "user", "content": "What color is the sky?"},
                {"role": "assistant", "content": "It is green."}]}
非配对偏好 
{"prompt": "The sky is",
 "completion": " blue.",
 "label": True}
 
{"prompt": [{"role": "user", "content": "What color is the sky?"}],
 "completion": [{"role": "assistant", "content": "It is green."}],
 "label": False}

标准数据集格式

标准数据集格式通常由纯文本字符串组成。数据集中的列根据任务而异。这是 TRL 训练器所需的格式。以下是不同任务的标准数据集格式示例

# Language modeling
example = {"text": "The sky is blue."}
# Preference
example = {"chosen": "The sky is blue.", "rejected": "The sky is green."}

对话数据集格式

对话数据集用于涉及用户和助手之间对话或聊天交互的任务。与标准数据集格式不同，这些数据集包含一系列消息，其中每条消息都具有role（例如，"user"或"assistant"）和content（消息文本）。

messages = [
    {"role": "user", "content": "Hello, how are you?"},
    {"role": "assistant", "content": "I'm doing great. How can I help you today?"},
    {"role": "user", "content": "I'd like to show off how chat templating works!"},
]

与标准数据集一样，对话数据集中的列会因任务而异。例如，偏好数据集将包含诸如"chosen"和"rejected"之类的列以比较响应

example = {
    "chosen": [
        {"role": "user", "content": "What color is the sky?"},
        {"role": "assistant", "content": "It is blue."},
    ],
    "rejected": [
        {"role": "user", "content": "What color is the sky?"},
        {"role": "assistant", "content": "It is green."},
    ],
}

对话数据集可用于训练聊天模型，但在使用 TRL 训练器之前，必须将其转换为标准格式。这通常使用特定于所用模型的聊天模板来完成。有关更多信息，请参阅在 TRL 中使用对话数据集部分。

语言建模

语言建模数据集包含一个名为"text"（或对话数据集中的"messages"）的列，其中包含完整的文本序列。

language_modeling_example = {"text": "The sky is blue."}

仅提示

在仅提示数据集，仅提供"prompt"键下的初始提示（问题或部分句子）。训练通常涉及根据此提示生成完成，其中模型学习继续或完成给定的输入。

prompt_only_example = {"prompt": "The sky is"}

虽然仅提示和语言建模格式类似，但它们在处理输入的方式上有所不同。在仅提示格式中，提示表示一个部分输入，期望模型完成或继续，而在语言建模格式中，输入被视为完整的句子或序列。TRL 对这两种格式的处理方式不同。以下是一个示例，显示了每种格式的apply_chat_template函数输出的差异

from transformers import AutoTokenizer
from trl import apply_chat_template

tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-128k-instruct")

# Example for prompt-only format
prompt_only_example = {"prompt": [{"role": "user", "content": "What color is the sky?"}]}
apply_chat_template(prompt_only_example, tokenizer)
# Output: {'prompt': '<|user|>\nWhat color is the sky?<|end|>\n<|assistant|>\n'}

# Example for language modeling format
lm_example = {"messages": [{"role": "user", "content": "What color is the sky?"}]}
apply_chat_template(lm_example, tokenizer)
# Output: {'text': '<|user|>\nWhat color is the sky?<|end|>\n<|endoftext|>'}
  • 仅提示输出包括一个'<|assistant|>\n'，指示助手轮次的开始，并期望模型生成完成。
  • 相反，语言建模输出将输入视为完整的序列，并以'<|endoftext|>'结束，表示文本的结束，并且不期望任何其他内容。

提示完成

提示完成数据集包含一个"prompt"和一个"completion"。

prompt_completion_example = {"prompt": "The sky is", "completion": " blue."}

偏好

偏好数据集用于模型被训练以在对同一个提示的两个或多个可能的完成之间进行选择的任务。该数据集包含一个"prompt"、一个"chosen"完成和一个"rejected"完成。模型被训练为选择"chosen"响应而不是"rejected"响应。某些数据集可能不包含"prompt"列，在这种情况下，提示是隐式的，直接包含在"chosen"和"rejected"完成中。我们建议尽可能使用显式提示。

preference_example = {"prompt": "The sky is", "chosen": " blue.", "rejected": " green."}  # recommended
# or,
preference_example = {"chosen": "The sky is blue.", "rejected": "The sky is green."}

非配对偏好

非配对偏好数据集类似于偏好数据集，但不是为同一个提示提供"chosen"和"rejected"完成，而是包含单个"completion"和一个"label"，指示完成是否为偏好。

unpaired_preference_example = {"prompt": "The sky is", "completion": " blue.", "label": True}

使用哪种数据集格式？

选择正确的数据集格式取决于您正在执行的任务以及您使用的 TRL 训练器的具体要求。以下是对每个 TRL 训练器支持的数据集格式的简要概述。

训练器 预期数据集格式
BCOTrainer 非配对偏好
CPOTrainer 偏好（显式提示）
DPOTrainer 偏好（显式提示）
IterativeSFTTrainer 非配对偏好
KTOTrainer 非配对偏好
NashMDTrainer 仅提示
OnlineDPOTrainer 仅提示
ORPOTrainer 偏好（显式提示）
PPOv2Trainer 标记化的语言建模
RewardTrainer 偏好（隐式提示）
SFTTrainer 语言建模
XPOTrainer 仅提示

TRL 训练器目前只支持标准数据集格式，详情请见此处。如果您有对话数据集，您必须先将其转换为标准格式。有关如何使用对话数据集的更多信息，请参阅在 TRL 中使用对话数据集部分。

在 TRL 中使用对话数据集

对话数据集越来越普遍，特别是在训练聊天模型方面。但是，TRL 训练器（除了SFTTrainer）不支持其原始格式的对话数据集。这些数据集必须首先转换为标准格式。幸运的是，TRL 提供了易于处理此转换的工具，这些工具将在下面详细介绍。

将对话数据集转换为标准数据集

TRL 训练器不支持其原始格式的对话数据集。要使用它们，您需要使用聊天模板将其转换为标准数据集格式。此模板由您使用的模型的标记器提供。

有关使用聊天模板的详细说明，请参阅transformers 文档中的聊天模板部分.

在 TRL 中，您用于转换数据集的方法会因任务而异。幸运的是，TRL 提供了一个名为apply_chat_template()的辅助函数来简化此过程。以下是如何使用它的示例

from transformers import AutoTokenizer
from trl import apply_chat_template

tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-128k-instruct")

example = {
    "prompt": [{"role": "user", "content": "What color is the sky?"}],
    "completion": [{"role": "assistant", "content": "It is blue."}]
}

apply_chat_template(example, tokenizer)
# Output:
# {'prompt': '<|user|>\nWhat color is the sky?<|end|>\n<|assistant|>\n', 'completion': 'It is blue.<|end|>\n<|endoftext|>'}

或者，您可以使用map方法在整个数据集上应用模板

from datasets import Dataset
from trl import apply_chat_template

dataset_dict = {
    "prompt": [[{"role": "user", "content": "What color is the sky?"}],
               [{"role": "user", "content": "Where is the sun?"}]],
    "completion": [[{"role": "assistant", "content": "It is blue."}],
                   [{"role": "assistant", "content": "In the sky."}]]
}

dataset = Dataset.from_dict(dataset_dict)
dataset = dataset.map(apply_chat_template, fn_kwargs={"tokenizer": tokenizer})
# Output:
# {'prompt': ['<|user|>\nWhat color is the sky?<|end|>\n<|assistant|>\n',
#             '<|user|>\nWhere is the sun?<|end|>\n<|assistant|>\n'],
#  'completion': ['It is blue.<|end|>\n<|endoftext|>', 'In the sky.<|end|>\n<|endoftext|>']}

我们建议使用apply_chat_template()函数，而不是直接调用tokenizer.apply_chat_template。处理聊天模板非语言建模数据集可能很棘手，并可能导致问题，例如在对话中间插入系统提示。有关其他示例，请参阅#1930 (comment)。apply_chat_template()旨在处理这些复杂性，并确保为各种任务正确应用聊天模板。

需要注意的是，聊天模板是特定于模型的。例如，如果您使用来自 meta-llama/Meta-Llama-3.1-8B-Instruct 的聊天模板与上面的示例一起使用，您将获得不同的输出。

apply_chat_template(example, AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B-Instruct"))
# Output:
# {'prompt': '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\nWhat color is the sky?<|im_end|>\n<|im_start|>assistant\n',
#  'completion': 'It is blue.<|im_end|>\n'}

始终使用与您正在使用的模型相关的聊天模板。使用错误的模板会导致不准确或意外的结果。

使用任何数据集与 TRL：预处理和转换

许多数据集以适合特定任务的格式提供，这些格式可能与 TRL 不直接兼容。为了使用 TRL 使用此类数据集，您可能需要对它们进行预处理并将其转换为所需的格式。

为了简化此过程，我们提供了一组 示例脚本，这些脚本涵盖了常见的数据集转换。

示例：UltraFeedback 数据集

让我们以 UltraFeedback 数据集 作为示例。以下是数据集的预览。

如上所示，数据集格式与预期结构不匹配。它不是对话格式，列名也不同，结果与不同的模型（例如 Bard、GPT-4）和方面（例如，“帮助性”、“诚实性”）相关。

通过使用提供的转换脚本 examples/datasets/ultrafeedback.py，您可以将此数据集转换为非配对的偏好格式，并将其推送到 Hub。

python examples/datasets/ultrafeedback.py --push_to_hub --repo_id trl-lib/ultrafeedback-gpt-3.5-turbo-helpfulness

转换后，数据集将如下所示。

现在，您可以使用 TRL 使用此数据集！

通过调整提供的脚本或创建自己的脚本，您可以将任何数据集转换为与 TRL 兼容的格式。

用于转换数据集类型的实用程序

本节提供示例代码，以帮助您在不同数据集类型之间进行转换。虽然某些转换可以在应用聊天模板后执行（即以标准格式），但我们建议在应用聊天模板之前执行转换，以确保其始终如一地工作。

为了简单起见，以下示例中的一些示例未遵循此建议，并使用标准格式。但是，转换可以直接应用于对话格式，无需修改。

从 \ 到 语言建模 提示-完成 仅提示 带隐式提示的偏好 偏好 非配对偏好
语言建模 不适用 不适用 不适用 不适用 不适用 不适用
提示-完成 🔗 不适用 🔗 不适用 不适用 不适用
仅提示 不适用 不适用 不适用 不适用 不适用 不适用
带隐式提示的偏好 🔗 🔗 🔗 不适用 🔗 🔗
偏好 🔗 🔗 🔗 🔗 不适用 🔗
非配对偏好 🔗 🔗 🔗 不适用 不适用 不适用

从提示-完成到语言建模数据集

要将提示-完成数据集转换为语言建模数据集，请将提示和完成连接起来。

from datasets import Dataset

dataset = Dataset.from_dict({
    "prompt": ["The sky is", "The sun is"],
    "completion": [" blue.", " in the sky."],
})

def concat_prompt_completion(example):
    return {"text": example["prompt"] + example["completion"]}

dataset = dataset.map(concat_prompt_completion, remove_columns=["prompt", "completion"])
>>> dataset[0]
{'text': 'The sky is blue.'}

从提示-完成到仅提示数据集

要将提示-完成数据集转换为仅提示数据集，请删除完成。

from datasets import Dataset

dataset = Dataset.from_dict({
    "prompt": ["The sky is", "The sun is"],
    "completion": [" blue.", " in the sky."],
})

dataset = dataset.remove_columns("completion")
>>> dataset[0]
{'prompt': 'The sky is'}

从带隐式提示的偏好到语言建模数据集

要将带隐式提示的偏好数据集转换为语言建模数据集，请删除被拒绝的，并将列 "chosen" 重命名为 "text"。

from datasets import Dataset

dataset = Dataset.from_dict({
    "chosen": ["The sky is blue.", "The sun is in the sky."],
    "rejected": ["The sky is green.", "The sun is in the sea."],
})

dataset = dataset.rename_column("chosen", "text").remove_columns("rejected")
>>> dataset[0]
{'text': 'The sky is blue.'}

从带隐式提示的偏好到提示-完成数据集

要将带隐式提示的偏好数据集转换为提示-完成数据集，请使用 extract_prompt() 提取提示，删除被拒绝的，并将列 "chosen" 重命名为 "completion"。

from datasets import Dataset
from trl import extract_prompt

dataset = Dataset.from_dict({
    "chosen": [
        [{"role": "user", "content": "What color is the sky?"}, {"role": "assistant", "content": "It is blue."}],
        [{"role": "user", "content": "Where is the sun?"}, {"role": "assistant", "content": "In the sky."}],
    ],
    "rejected": [
        [{"role": "user", "content": "What color is the sky?"}, {"role": "assistant", "content": "It is green."}],
        [{"role": "user", "content": "Where is the sun?"}, {"role": "assistant", "content": "In the sea."}],
    ],
})
dataset = dataset.map(extract_prompt).remove_columns("rejected").rename_column("chosen", "completion")
>>> dataset[0]
{'prompt': [{'role': 'user', 'content': 'What color is the sky?'}], 'completion': [{'role': 'assistant', 'content': 'It is blue.'}]}

从带隐式提示的偏好到仅提示数据集

要将带隐式提示的偏好数据集转换为仅提示数据集，请使用 extract_prompt() 提取提示，并删除被拒绝的和选择的。

from datasets import Dataset
from trl import extract_prompt

dataset = Dataset.from_dict({
    "chosen": [
        [{"role": "user", "content": "What color is the sky?"}, {"role": "assistant", "content": "It is blue."}],
        [{"role": "user", "content": "Where is the sun?"}, {"role": "assistant", "content": "In the sky."}],
    ],
    "rejected": [
        [{"role": "user", "content": "What color is the sky?"}, {"role": "assistant", "content": "It is green."}],
        [{"role": "user", "content": "Where is the sun?"}, {"role": "assistant", "content": "In the sea."}],
    ],
})
dataset = dataset.map(extract_prompt).remove_columns(["chosen", "rejected"])
>>> dataset[0]
{'prompt': [{'role': 'user', 'content': 'What color is the sky?'}]}

从隐式到显式提示偏好数据集

要将带隐式提示的偏好数据集转换为带显式提示的偏好数据集，请使用 extract_prompt() 提取提示。

from datasets import Dataset
from trl import extract_prompt

dataset = Dataset.from_dict({
    "chosen": [
        [{"role": "user", "content": "What color is the sky?"}, {"role": "assistant", "content": "It is blue."}],
        [{"role": "user", "content": "Where is the sun?"}, {"role": "assistant", "content": "In the sky."}],
    ],
    "rejected": [
        [{"role": "user", "content": "What color is the sky?"}, {"role": "assistant", "content": "It is green."}],
        [{"role": "user", "content": "Where is the sun?"}, {"role": "assistant", "content": "In the sea."}],
    ],
})

dataset = dataset.map(extract_prompt)
>>> dataset[0]
{'prompt': [{'role': 'user', 'content': 'What color is the sky?'}],
 'chosen': [{'role': 'assistant', 'content': 'It is blue.'}],
 'rejected': [{'role': 'assistant', 'content': 'It is green.'}]}

从带有隐式提示的偏好数据集到无配对偏好数据集

要将带有隐式提示的偏好数据集转换为无配对偏好数据集，请使用 extract_prompt() 提取提示，并使用 unpair_preference_dataset() 解除数据集的配对。

from datasets import Dataset
from trl import extract_prompt, unpair_preference_dataset

dataset = Dataset.from_dict({
    "chosen": [
        [{"role": "user", "content": "What color is the sky?"}, {"role": "assistant", "content": "It is blue."}],
        [{"role": "user", "content": "Where is the sun?"}, {"role": "assistant", "content": "In the sky."}],
    ],
    "rejected": [
        [{"role": "user", "content": "What color is the sky?"}, {"role": "assistant", "content": "It is green."}],
        [{"role": "user", "content": "Where is the sun?"}, {"role": "assistant", "content": "In the sea."}],
    ],
})

dataset = dataset.map(extract_prompt)
dataset = unpair_preference_dataset(dataset)
>>> dataset[0]
{'prompt': [{'role': 'user', 'content': 'What color is the sky?'}],
 'completion': [{'role': 'assistant', 'content': 'It is blue.'}],
 'label': True}

从偏好数据集到语言建模数据集

要将偏好数据集转换为语言建模数据集，请删除拒绝的样本，并将提示和选定的样本连接到 "text" 列中。

from datasets import Dataset

dataset = Dataset.from_dict({
    "prompt": ["The sky is", "The sun is"],
    "chosen": [" blue.", " in the sky."],
    "rejected": [" green.", " in the sea."],
})

def concat_prompt_chosen(example):
    return {"text": example["prompt"] + example["chosen"]}

dataset = dataset.map(concat_prompt_chosen, remove_columns=["prompt", "chosen", "rejected"])
>>> dataset[0]
{'text': 'The sky is blue.'}

从偏好数据集到提示-完成数据集

要将偏好数据集转换为提示-完成数据集，请删除拒绝的样本，并将 "chosen" 列重命名为 "completion"。

from datasets import Dataset

dataset = Dataset.from_dict({
    "prompt": ["The sky is", "The sun is"],
    "chosen": [" blue.", " in the sky."],
    "rejected": [" green.", " in the sea."],
})

dataset = dataset.remove_columns("rejected").rename_column("chosen", "completion")
>>> dataset[0]
{'prompt': 'The sky is', 'completion': ' blue.'}

从偏好数据集到仅提示数据集

要将偏好数据集转换为仅提示数据集，请删除拒绝的样本和选定的样本。

from datasets import Dataset

dataset = Dataset.from_dict({
    "prompt": ["The sky is", "The sun is"],
    "chosen": [" blue.", " in the sky."],
    "rejected": [" green.", " in the sea."],
})

dataset = dataset.remove_columns(["chosen", "rejected"])
>>> dataset[0]
{'prompt': 'The sky is'}

从显式提示到隐式提示偏好数据集

要将带有隐式提示的偏好数据集转换为带有显式提示的偏好数据集，请将提示连接到选定的样本和拒绝的样本，然后删除提示。

from datasets import Dataset

dataset = Dataset.from_dict({
    "prompt": [
        [{"role": "user", "content": "What color is the sky?"}],
        [{"role": "user", "content": "Where is the sun?"}],
    ],
    "chosen": [
        [{"role": "assistant", "content": "It is blue."}],
        [{"role": "assistant", "content": "In the sky."}],
    ],
    "rejected": [
        [{"role": "assistant", "content": "It is green."}],
        [{"role": "assistant", "content": "In the sea."}],
    ],
})

def concat_prompt_to_completions(example):
    return {"chosen": example["prompt"] + example["chosen"], "rejected": example["prompt"] + example["rejected"]}

dataset = dataset.map(concat_prompt_to_completions, remove_columns="prompt")
>>> dataset[0]
{'chosen': [{'role': 'user', 'content': 'What color is the sky?'}, {'role': 'assistant', 'content': 'It is blue.'}],
 'rejected': [{'role': 'user', 'content': 'What color is the sky?'}, {'role': 'assistant', 'content': 'It is green.'}]}

从偏好数据集到无配对偏好数据集

要将数据集转换为无配对偏好数据集，请使用 unpair_preference_dataset() 解除数据集的配对。

from datasets import Dataset
from trl import unpair_preference_dataset

dataset = Dataset.from_dict({
    "prompt": [
        [{"role": "user", "content": "What color is the sky?"}],
        [{"role": "user", "content": "Where is the sun?"}],
    ],
    "chosen": [
        [{"role": "assistant", "content": "It is blue."}],
        [{"role": "assistant", "content": "In the sky."}],
    ],
    "rejected": [
        [{"role": "assistant", "content": "It is green."}],
        [{"role": "assistant", "content": "In the sea."}],
    ],
})

dataset = unpair_preference_dataset(dataset)
>>> dataset[0]
{'prompt': [{'role': 'user', 'content': 'What color is the sky?'}],
 'completion': [{'role': 'assistant', 'content': 'It is blue.'}],
 'label': True}

从无配对偏好数据集到语言建模数据集

要将无配对偏好数据集转换为语言建模数据集，请将提示和完成连接到 "text" 列中，并删除提示、完成和标签列。

from datasets import Dataset

dataset = Dataset.from_dict({
    "prompt": ["The sky is", "The sun is", "The sky is", "The sun is"],
    "completion": [" blue.", " in the sky.", " green.", " in the sea."],
    "label": [True, True, False, False],
})

def concatenate_prompt_completion(example):
    return {"text": example["prompt"] + example["completion"]}

dataset = dataset.map(concatenate_prompt_completion).remove_columns(["prompt", "completion", "label"])
>>> dataset[0]
{'text': 'The sky is blue.'}

从无配对偏好数据集到提示-完成数据集

要将无配对偏好数据集转换为提示-完成数据集，请删除标签列。

from datasets import Dataset

dataset = Dataset.from_dict({
    "prompt": ["The sky is", "The sun is", "The sky is", "The sun is"],
    "completion": [" blue.", " in the sky.", " green.", " in the sea."],
    "label": [True, True, False, False],
})

dataset = dataset.remove_columns(["label"])
>>> dataset[0]
{'prompt': 'The sky is', 'completion': ' blue.'}

从无配对偏好数据集到仅提示数据集

要将无配对偏好数据集转换为仅提示数据集，请删除完成和标签列。

from datasets import Dataset

dataset = Dataset.from_dict({
    "prompt": ["The sky is", "The sun is", "The sky is", "The sun is"],
    "completion": [" blue.", " in the sky.", " green.", " in the sea."],
    "label": [True, True, False, False],
})

dataset = dataset.remove_columns(["completion", "label"])
>>> dataset[0]
{'prompt': 'The sky is'}
< > 在 GitHub 上更新

←命令行界面 (CLI) 入门 PPO 训练常见问题解答→