Dataset formats
This guide provides an overview of the dataset formats supported by each trainer in TRL. Since conversational datasets are very common, we also provide a guide on how to use them, and how to convert them into a standard dataset format for TRL trainers.
Overview of the dataset formats and types
The format of a dataset refers to how the data is structured, typically categorized as either standard or conversational. The type is associated with the specific task the dataset is designed for, such as prompt-only or preference. Each type is characterized by its columns, which vary according to the task, as shown in the table.
| Type \ Format | Standard | Conversational |
|---|---|---|
| Language modeling | | |
| Prompt-only | | |
| Prompt-completion | | |
| Preference | | |
| Unpaired preference | | |
Standard dataset format
The standard dataset format typically consists of plain text strings. The columns in the dataset vary depending on the task. This is the format expected by TRL trainers. Below are examples of standard dataset formats for different tasks:
# Language modeling
example = {"text": "The sky is blue."}
# Preference
example = {"chosen": "The sky is blue.", "rejected": "The sky is green."}Conversational dataset format
Conversational datasets are used for tasks involving dialogues or chat interactions between users and assistants. Unlike standard dataset formats, these contain sequences of messages where each message has a role (e.g., "user" or "assistant") and content (the message text).
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!"},
]Just like standard datasets, the columns in conversational datasets vary depending on the task. For instance, a preference dataset would include columns like "chosen" and "rejected" to compare responses:
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."},
],
}Conversational datasets are useful for training chat models, but must be converted into a standard format before being used with TRL trainers. This is typically done using chat templates specific to the model being used. For more information, refer to the Working with conversational datasets in TRL section.
Language modeling
A language modeling dataset consists of a column "text" (or "messages" for conversational datasets) containing a full sequence of text.
language_modeling_example = {"text": "The sky is blue."}Prompt-only
In a prompt-only dataset, only the initial prompt (the question or partial sentence) is provided under the key "prompt". The training typically involves generating the completion based on this prompt, where the model learns to continue or complete the given input.
prompt_only_example = {"prompt": "The sky is"}While both the prompt-only and language modeling formats are similar, they differ in how the input is handled. In the prompt-only format, the prompt represents a partial input that expects the model to complete or continue, while in the language modeling format, the input is treated as a complete sentence or sequence. These two formats are processed differently by TRL. Below is an example showing the difference in the output of the apply_chat_template function for each format:
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|>'}- The prompt-only output includes a
'<|assistant|>\n', indicating the beginning of the assistant’s turn and expecting the model to generate a completion. - In contrast, the language modeling output treats the input as a complete sequence and terminates it with
'<|endoftext|>', signaling the end of the text and not expecting any additional content.
Prompt-completion
A prompt-completion dataset includes a "prompt" and a "completion".
prompt_completion_example = {"prompt": "The sky is", "completion": " blue."}Preference
A preference dataset is used for tasks where the model is trained to choose between two or more possible completions to the same prompt. This dataset includes a "prompt", a "chosen" completion, and a "rejected" completion. The model is trained to select the "chosen" response over the "rejected" response.
Some dataset may not include the "prompt" column, in which case the prompt is implicit and directly included in the "chosen" and "rejected" completions. We recommend using explicit prompts whenever possible.
preference_example = {"prompt": "The sky is", "chosen": " blue.", "rejected": " green."} # recommended
# or,
preference_example = {"chosen": "The sky is blue.", "rejected": "The sky is green."}Unpaired preference
An unpaired preference dataset is similar to a preference dataset but instead of having "chosen" and "rejected" completions for the same prompt, it includes a single "completion" and a "label" indicating whether the completion is preferred or not.
unpaired_preference_example = {"prompt": "The sky is", "completion": " blue.", "label": True}Which dataset format to use?
Choosing the right dataset format depends on the task you are working on and the specific requirements of the TRL trainer you are using. Below is a brief overview of the dataset formats supported by each TRL trainer.
| Trainer | Expected dataset format |
|---|---|
| BCOTrainer | Unpaired preference |
| CPOTrainer | Preference (explicit prompt) |
| DPOTrainer | Preference (explicit prompt) |
| IterativeSFTTrainer | Unpaired preference |
| KTOTrainer | Unpaired preference |
| NashMDTrainer | Prompt-only |
| OnlineDPOTrainer | Prompt-only |
| ORPOTrainer | Preference (explicit prompt) |
| PPOv2Trainer | Tokenized language modeling |
| RewardTrainer | Preference (implicit prompt) |
| SFTTrainer | Language modeling |
| XPOTrainer | Prompt-only |
TRL trainers only support standard dataset formats, for now. If you have a conversational dataset, you must first convert it into a standard format. For more information on how to work with conversational datasets, refer to the Working with conversational datasets in TRL section.
Working with conversational datasets in TRL
Conversational datasets are increasingly common, especially for training chat models. However, TRL trainers (except SFTTrainer) don’t support conversational datasets in their raw format. These datasets must first be converted into a standard format. Fortunately, TRL offers tools to easily handle this conversion, which are detailed below.
Converting a conversational dataset into a standard dataset
TRL trainers do not support conversational datasets in their raw format. To use them, you need to convert them into a standard dataset format using a chat template. This template is provided by the tokenizer of the model you use.
For detailed instructions on using chat templating, refer to the Chat templating section in the transformers documentation.
In TRL, the method you apply to convert the dataset will vary depending on the task. Fortunately, TRL provides a helper function called apply_chat_template() to simplify this process. Here’s an example of how to use it:
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|>'}Alternatively, you can use the map method to apply the template across an entire dataset:
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|>']}We recommend using the apply_chat_template() function rather than directly calling tokenizer.apply_chat_template. Handling chat templates nonlanguage modeling datasets can be tricky and may lead to issues, such as inserting a system prompt in the middle of a conversation. For additional examples, see #1930 (comment). The apply_chat_template() is designed to handle these intricacies and ensure the correct application of chat templates for various tasks.
It’s important to note that chat templates are model-specific. For example, if you use the chat template from meta-llama/Meta-Llama-3.1-8B-Instruct with the above example, you get a different output:
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'}Always use the chat template associated with the model you’re working with. Using the wrong template can lead to inaccurate or unexpected results.
Using any dataset with TRL: preprocessing and conversion
Many datasets come in formats tailored to specific tasks, which might not be directly compatible with TRL. To use such datasets with TRL, you may need to preprocess and convert them into the required format.
To make this easier, we provide a set of example scripts that cover common dataset conversions.
Example: UltraFeedback dataset
Let’s take the UltraFeedback dataset as an example. Here’s a preview of the dataset:
As shown above, the dataset format does not match the expected structure. It’s not in a conversational format, the column names differ, and the results pertain to different models (e.g., Bard, GPT-4) and aspects (e.g., “helpfulness”, “honesty”).
By using the provided conversion script examples/datasets/ultrafeedback.py, you can transform this dataset into an unpaired preference format, and push it to the Hub:
python examples/datasets/ultrafeedback.py --push_to_hub --repo_id trl-lib/ultrafeedback-gpt-3.5-turbo-helpfulness
Once converted, the dataset will look like this:
Now, you can use this dataset with TRL!
By adapting the provided scripts or creating your own, you can convert any dataset into a format compatible with TRL.
Utilities for converting dataset types
This section provides example code to help you convert between different dataset types. While some conversions can be performed after applying the chat template (i.e., in the standard format), we recommend performing the conversion before applying the chat template to ensure it works consistently.
For simplicity, some of the examples below do not follow this recommendation and use the standard format. However, the conversions can be applied directly to the conversational format without modification.
| From \ To | Language modeling | Prompt-completion | Prompt-only | Preference with implicit prompt | Preference | Unpaired preference |
|---|---|---|---|---|---|---|
| Language modeling | N/A | N/A | N/A | N/A | N/A | N/A |
| Prompt-completion | 🔗 | N/A | 🔗 | N/A | N/A | N/A |
| Prompt-only | N/A | N/A | N/A | N/A | N/A | N/A |
| Preference with implicit prompt | 🔗 | 🔗 | 🔗 | N/A | 🔗 | 🔗 |
| Preference | 🔗 | 🔗 | 🔗 | 🔗 | N/A | 🔗 |
| Unpaired preference | 🔗 | 🔗 | 🔗 | N/A | N/A | N/A |
From prompt-completion to language modeling dataset
To convert a prompt-completion dataset into a language modeling dataset, concatenate the prompt and the completion.
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 prompt-completion to prompt-only dataset
To convert a prompt-completion dataset into a prompt-only dataset, remove the completion.
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'}From preference with implicit prompt to language modeling dataset
To convert a preference with implicit prompt dataset into a language modeling dataset, remove the rejected, and rename the column "chosen" to "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.'}From preference with implicit prompt to prompt-completion dataset
To convert a preference dataset with implicit prompt into a prompt-completion dataset, extract the prompt with extract_prompt(), remove the rejected, and rename the column "chosen" to "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.'}]}From preference with implicit prompt to prompt-only dataset
To convert a preference dataset with implicit prompt into a prompt-only dataset, extract the prompt with extract_prompt(), and remove the rejected and the chosen.
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?'}]}From implicit to explicit prompt preference dataset
To convert a preference dataset with implicit prompt into a preference dataset with explicit prompt, extract the prompt with 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.'}]}From preference with implicit prompt to unpaired preference dataset
To convert a preference dataset with implicit prompt into an unpaired preference dataset, extract the prompt with extract_prompt(), and unpair the dataset with 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}From preference to language modeling dataset
To convert a preference dataset into a language modeling dataset, remove the rejected, concatenate the prompt and the chosen into the "text" column.
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.'}From preference to prompt-completion dataset
To convert a preference dataset into a prompt-completion dataset, remove the rejected, and rename the column "chosen" to "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 preference to prompt-only dataset
To convert a preference dataset into a prompt-only dataset, remove the rejected and the chosen.
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 explicit to implicit prompt preference dataset
To convert a preference dataset with implicit prompt into a preference dataset with explicit prompt, concatenate the prompt to both chosen and rejected, and remove the prompt.
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.'}]}From preference to unpaired preference dataset
To convert dataset into an unpaired preference dataset, unpair the dataset with 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}From unpaired preference to language modeling dataset
To convert an unpaired preference dataset into a language modeling dataset, concatenate the prompt and the completion into the "text" column, and remove the prompt, completion and label columns.
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 unpaired preference to prompt-completion dataset
To convert an unpaired preference dataset into a prompt-completion dataset, remove the label columns.
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 unpaired preference to prompt-only dataset
To convert an unpaired preference dataset into a prompt-only dataset, remove the completion and the label columns.
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'}