TRL documentation

数据集格式和类型

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.17.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

数据集格式和类型

本指南概述了 TRL 中每个 trainer 支持的数据集格式和类型。

数据集格式和类型概述

  • 数据集的 format 指的是数据的结构方式,通常分为 standard 或 conversational 两种类型。
  • type 与数据集所针对的具体任务相关联,例如 prompt-only 或 preference。每种类型都由其列来表征,这些列会根据任务的不同而变化,如表所示。
类型 \ 格式标准对话式
语言建模
{"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
language_modeling_example = {"text": "The sky is blue."}
# Preference
preference_example = {"prompt": "The sky is", "chosen": " blue.", "rejected": " green."}
# Unpaired preference
unpaired_preference_example = {"prompt": "The sky is", "completion": " blue.", "label": True}

对话式

对话数据集用于涉及用户与助手之间对话或聊天交互的任务。与标准数据集格式不同,这些数据集包含消息序列,其中每条消息都有一个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!"},
]

与标准数据集一样,对话式数据集的列会根据任务的不同而变化。以下是不同任务的对话式数据集格式示例:

# Prompt-completion
prompt_completion_example = {"prompt": [{"role": "user", "content": "What color is the sky?"}],
                             "completion": [{"role": "assistant", "content": "It is blue."}]}
# Preference
preference_example = {
    "prompt": [{"role": "user", "content": "What color is the sky?"}],
    "chosen": [{"role": "assistant", "content": "It is blue."}],
    "rejected": [{"role": "assistant", "content": "It is green."}],
}

工具调用

某些聊天模板支持工具调用,允许模型在生成过程中与外部函数(即工具)进行交互。这扩展了模型的对话能力,使其在决定调用工具时,能够输出"tool_calls"字段,而不是标准的"content"消息。

助手发起工具调用后,该工具会执行并返回其输出。随后,助手可以处理此输出,并据此继续对话。

以下是一个工具调用交互的简单示例:

messages = [
    {"role": "user", "content": "Turn on the living room lights."},
    {"role": "assistant", "tool_calls": [
        {"type": "function", "function": {
            "name": "control_light",
            "arguments": {"room": "living room", "state": "on"}
        }}]
    },
    {"role": "tool", "name": "control_light", "content": "The lights in the living room are now on."},
    {"role": "assistant", "content": "Done!"}
]

在为工具调用准备监督微调(SFT)数据集时,确保数据集中包含一个名为 tools 的额外列非常重要。该列包含模型可用的工具列表,通常由聊天模板用于构建系统提示。

工具必须以编码化的 JSON 模式格式进行指定。您可以使用 get_json_schema 实用程序从 Python 函数签名自动生成此模式:

import json
from transformers.utils import get_json_schema

def control_light(room: str, state: str) -> str:
    """
    Controls the lights in a room.

    Args:
        room: The name of the room.
        state: The desired state of the light ("on" or "off").

    Returns:
        str: A message indicating the new state of the lights.
    """
    return f"The lights in {room} are now {state}."

# Generate JSON schema
json_schema = get_json_schema(control_light)

生成的模式如下所示:

{"type": "function", "function": {"name": "control_light", "description": "Controls the lights in a room.", "parameters": {"type": "object", "properties": {"room": {"type": "string", "description": "The name of the room."}, "state": {"type": "string", "description": "The desired state of the light (\"on\" or \"off\")."}}, "required": ["room", "state"]}, "return": {"type": "string", "description": "str: A message indicating the new state of the lights."}}}

SFT 的完整数据集条目可能如下所示:

{"messages": messages, "tools": [json_schema]}

要获取 Dataset,您需要使用 Json() 类型作为工具参数,因为它们是任意 JSON 对象,而不是具有固定字段和类型的字典:

from datasets import Dataset

data = [
    {"messages": messages1, "tools": [json_schema1]},
    {"messages": messages2, "tools": [json_schema2]},
]
# auto-apply the Json() type
dataset = Dataset.from_list(data, on_mixed_types="use_json")

# or specify the features manually
from datasets import Features, Json, List, Value

features = Features(
    {
        "messages": List({"role": Value("string"), "content": Value("string"), "tool_calls": List(Json())}),
        "tools": List(Json()),
    }
)
dataset = Dataset.from_list(data, features=features)

在较旧版本的 datasets(<4.7.0)中,由于没有 Json() 类型,您应将 tools 存储为 JSON str(使用 json.dumps([...])):

dataset = Dataset.from_list(
    [{"messages": messages1, "tools": json.dumps([json_schema1])},
     {"messages": messages2, "tools": json.dumps([json_schema2])}]
)

有关工具调用的更详细信息,请参阅 transformers 文档中的 Tool Calling 部分 以及博客文章 Tool Use, Unified。

Harmony

Harmony 响应格式随OpenAI GPT OSS 模型一同推出。它在对话格式的基础上进行了扩展,为推理、函数调用以及关于模型行为的元数据增加了更丰富的结构。主要特性包括:

  • 开发者角色 – 提供高层级指令(类似于系统提示词),并列出可用的工具。

  • 通道 – 将不同类型的助手输出分离到不同的流中:

    • analysis – 用于内部推理,来自键 "thinking"
    • final – 面向用户的回答,来自键 "content"
    • commentary – 用于工具调用或元注释
  • 推理力度 – 指示模型应展示多少思考过程(例如,"low"、"medium"、"high")。

  • 模型身份 – 明确定义助手的角色。

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("openai/gpt-oss-20b")

messages = [
    {"role": "developer", "content": "Use a friendly tone."},
    {"role": "user", "content": "What is the meaning of life?"},
    {"role": "assistant", "thinking": "Deep reflection...", "content": "The final answer is..."},
]

print(
    tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        reasoning_effort="low",
        model_identity="You are HuggingGPT, a large language model trained by Hugging Face.",
    )
)

这将生成:

<|start|>system<|message|>You are HuggingGPT, a large language model trained by Hugging Face.
Knowledge cutoff: 2024-06
Current date: 2025-08-03

Reasoning: low

# Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|><|start|>developer<|message|># Instructions

Use a friendly tone.<|end|><|start|>user<|message|>What is the meaning of life?<|end|><|start|>assistant<|channel|>analysis<|message|>Deep reflection...<|end|><|start|>assistant<|channel|>final<|message|>The final answer is...<|return|>

有关消息结构、支持的字段以及高级用法的完整详细信息,请参阅 Harmony 文档。

类型

语言建模

语言建模数据集包含一列 "text"(对于对话数据集则为 "messages"),其中包含完整的文本序列。

# Standard format
language_modeling_example = {"text": "The sky is blue."}
# Conversational format
language_modeling_example = {"messages": [
    {"role": "user", "content": "What color is the sky?"},
    {"role": "assistant", "content": "It is blue."}
]}

仅提示

在仅包含提示的数据集中,初始提示(问题或不完整的句子)仅通过键 "prompt" 提供。训练通常涉及基于该提示生成补全内容,模型在此过程中学习如何续写或补全给定的输入。

# Standard format
prompt_only_example = {"prompt": "The sky is"}
# Conversational format
prompt_only_example = {"prompt": [{"role": "user", "content": "What color is the sky?"}]}

有关仅包含提示的数据集的示例,请参阅Prompt-only datasets collection。

虽然仅提示(prompt-only)类型和语言建模(language modeling)类型相似,但它们在输入的处理方式上有所不同。在仅提示类型中,提示代表一个部分输入,期望模型对其进行补全或续写;而在语言建模类型中,输入被视为一个完整的句子或序列。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 type
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 type
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"。

# Standard format
prompt_completion_example = {"prompt": "The sky is", "completion": " blue."}
# Conversational format
prompt_completion_example = {"prompt": [{"role": "user", "content": "What color is the sky?"}],
                             "completion": [{"role": "assistant", "content": "It is blue."}]}

有关提示-补全数据集的示例,请参阅 Prompt-completion datasets collection。

偏好

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

# Standard format
## Explicit prompt (recommended)
preference_example = {"prompt": "The sky is", "chosen": " blue.", "rejected": " green."}
# Implicit prompt
preference_example = {"chosen": "The sky is blue.", "rejected": "The sky is green."}

# Conversational format
## Explicit prompt (recommended)
preference_example = {"prompt": [{"role": "user", "content": "What color is the sky?"}],
                      "chosen": [{"role": "assistant", "content": "It is blue."}],
                      "rejected": [{"role": "assistant", "content": "It is green."}]}
## Implicit prompt
preference_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."}]}

有关偏好数据集的示例,请参阅偏好数据集集合。

一些偏好数据集可以在 Hugging Face Hub 上通过 dpo 标签找到。您还可以浏览 librarian-bots 的 DPO 合集 来查找偏好数据集。

非配对偏好

非配对偏好数据集与偏好数据集类似,但区别在于,它并非为同一提示词提供"chosen"和"rejected"补全,而是包含单个"completion"以及一个"label",用于指示该补全是否受到偏好。

# Standard format
unpaired_preference_example = {"prompt": "The sky is", "completion": " blue.", "label": True}
# Conversational format
unpaired_preference_example = {"prompt": [{"role": "user", "content": "What color is the sky?"}],
                               "completion": [{"role": "assistant", "content": "It is blue."}],
                               "label": True}

有关非配对偏好数据集的示例,请参阅非配对偏好数据集集合。

应该使用哪种数据集类型?

选择合适的数据集类型取决于您正在处理的任务以及所使用的 TRL 训练器的具体要求。以下是每个 TRL 训练器所支持的数据集类型的简要概述。

Trainer预期的数据集类型
DPOTrainer偏好(建议显式提示)
GRPOTrainer仅提示词
KTOTrainer非配对偏好或偏好(建议显式提示)
RewardTrainer偏好(建议隐式提示)
RLOOTrainer仅提示词
SFTTrainer语言建模或提示补全
experimental.cpo.CPOTrainer偏好(建议显式提示)
experimental.gkd.GKDTrainer提示-补全
experimental.online_dpo.OnlineDPOTrainer仅提示词
experimental.orpo.ORPOTrainer偏好(建议显式提示)

使用任意数据集配合 TRL:预处理与转换

许多数据集采用针对特定任务定制的数据格式,可能无法直接与 TRL 兼容。若要在 TRL 中使用此类数据集,您可能需要先对其进行预处理,并将其转换为所需的格式。

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

示例:UltraFeedback 数据集

我们以 UltraFeedback 数据集 为例。以下是该数据集的预览:

如上所示,数据集格式与预期结构不匹配。它并非对话格式,列名也不同,且结果涉及不同的模型(例如 Bard、GPT-4)和不同的方面(例如“helpfulness”、“honesty”)。

通过使用提供的转换脚本 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 兼容的格式。

用于转换数据集类型的实用工具

本节提供示例代码,帮助您转换不同的数据集类型。虽然某些转换可以在应用聊天模板之后(即标准格式下)执行,但我们建议在应用聊天模板之前执行转换,以确保其一致有效。

为简洁起见,下文中的部分示例未遵循此建议,而是使用了标准格式。不过,这些转换可以直接应用于对话格式,无需任何修改。

从 \ 到语言建模提示-补全仅提示带有隐式提示的偏好偏好非配对偏好
语言建模N/AN/AN/AN/AN/AN/A
提示-补全🔗N/A🔗N/AN/AN/A
仅提示N/AN/AN/AN/AN/AN/A
带有隐式提示的偏好🔗🔗🔗N/A🔗🔗
偏好🔗🔗🔗🔗N/A🔗
非配对偏好🔗🔗🔗N/AN/AN/A

从提示补全到语言建模数据集

要将提示-补全数据集转换为语言建模数据集,请将提示和补全拼接在一起。

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}

请注意,偏好数据集中的"chosen"和"rejected"补全可能都是好的或都是坏的。 在应用unpair_preference_dataset()之前,请确保所有"chosen"补全都可以被标记为好的,所有"rejected"补全都可以被标记为坏的。 这可以通过检查每个补全的绝对评分来确保,例如来自奖励模型的评分。

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

要将偏好数据集转换为语言建模数据集,请移除 rejected,并将 prompt 和 chosen 拼接起来,放入 "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'}

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

要将带有显式提示的偏好数据集转换为带有隐式提示的偏好数据集,请将提示拼接到 chosen 和 rejected 中,然后移除提示。

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}

请注意,偏好数据集中的"chosen"和"rejected"补全可能都是好的或都是坏的。 在应用unpair_preference_dataset()之前,请确保所有"chosen"补全都可以被标记为好的,所有"rejected"补全都可以被标记为坏的。 这可以通过检查每个补全的绝对评分来确保,例如来自奖励模型的评分。

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

要将非配对偏好数据集转换为语言建模数据集,请将提示与优质补全拼接后放入 "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.filter(lambda x: x["label"]).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.filter(lambda x: x["label"]).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'}

视觉数据集

一些训练器还支持使用图文对来微调视觉-语言模型(VLM)。在此场景下,建议使用对话格式,因为每个模型对文本中图像占位符的处理方式各不相同。

对话式视觉数据集与标准对话式数据集在两个关键方面有所不同:

  1. 数据集必须包含键 images(其中包含图像数据,格式为 PIL 图像列表)或 image(其中包含单个 PIL 图像)。
  2. 消息中的 "content" 字段必须是一个字典列表,其中每个字典指定数据类型:"image" 或 "text"。

示例:

# Textual dataset:
"content": "What color is the sky?"

# Vision dataset:
"content": [
    {"type": "image"},
    {"type": "text", "text": "What color is the sky in the image?"}
]

一个对话式视觉数据集的示例是 openbmb/RLAIF-V-Dataset。以下是该数据集训练数据的嵌入视图,您可以直接进行探索:

在数据集中混合纯文本数据和视觉语言数据是可行的,但需要 transformers 4.57.0 或更高版本。示例:

dataset = Dataset.from_dict({
    "prompt": [
        [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "What color is the sky in the image?"}]}],
        [{"role": "user", "content": [{"type": "text", "text": "What is the capital of France?"}]}],
    ],
    "completion": [
        [{"role": "assistant", "content": [{"type": "text", "text": "It is blue."}]}],
        [{"role": "assistant", "content": [{"type": "text", "text": "Paris."}]}],
    ],
    "images": [
        [PIL.Image.open("path/to/sky_image1.png")],
        [],
    ],
})
Update on GitHub