Implementation:Huggingface Trl Get Dataset Apply Chat Template

Knowledge Sources	TRL TRL Docs
Domains	NLP, Training
Last Updated	2026-02-06 17:00 GMT

Overview

Concrete tools for loading dataset mixtures and applying chat templates to conversational data for SFT, provided by the TRL library.

Description

The TRL library provides two key functions for dataset preparation:

get_dataset() loads one or more datasets specified in a DatasetMixtureConfig, concatenates them, and optionally splits into train/test sets. Each dataset in the mixture can specify its own path, configuration name, data directory, data files, split, and column selection.

apply_chat_template() takes a single conversational example (a dictionary with keys like "messages", "prompt", "completion", "chosen", "rejected") and renders it through the tokenizer's chat template. For prompt-completion data, it carefully separates the prompt from the completion by rendering the full conversation and then extracting the suffix that corresponds to the completion.

In the SFT script, dataset loading follows a priority order: if DatasetMixtureConfig.datasets is provided, it is used via get_dataset(); otherwise the script falls back to datasets.load_dataset() with ScriptArguments.dataset_name.

Usage

Use get_dataset() when mixing multiple datasets from the Hub into a single training set. Use apply_chat_template() when converting conversational examples into plain-text format for tokenization.

Code Reference

Source Location

• Repository: TRL
• File: trl/scripts/utils.py (lines 414-474, get_dataset)
• File: trl/data_utils.py (lines 186-316, apply_chat_template)
• File: trl/scripts/utils.py (lines 90-152, DatasetMixtureConfig; lines 56-88, DatasetConfig)

Signature

def get_dataset(mixture_config: DatasetMixtureConfig) -> DatasetDict:
    """
    Load a mixture of datasets based on the configuration.
    Returns a DatasetDict with a 'train' split (and optionally 'test').
    """
    ...


def apply_chat_template(
    example: dict[str, list[dict[str, str]]],
    tokenizer: PreTrainedTokenizerBase | ProcessorMixin,
    tools: list[dict | Callable] | None = None,
    **template_kwargs,
) -> dict[str, str]:
    """
    Apply a chat template to a conversational example.
    Supported key sets:
      - {"messages"}               -> language modeling
      - {"prompt", "completion"}   -> prompt-completion
      - {"prompt", "chosen", "rejected"} -> preference
      - {"chosen", "rejected"}     -> preference with implicit prompt
      - {"prompt"}                 -> prompt-only
    """
    ...


@dataclass
class DatasetMixtureConfig:
    datasets: list[DatasetConfig] = field(default_factory=list)
    streaming: bool = False
    test_split_size: float | None = None


@dataclass
class DatasetConfig:
    path: str
    name: str | None = None
    data_dir: str | None = None
    data_files: str | list[str] | dict[str, str] | None = None
    split: str = "train"
    columns: list[str] | None = None

Import

from trl import get_dataset, DatasetMixtureConfig
from trl.data_utils import apply_chat_template
from trl.scripts.utils import DatasetConfig

I/O Contract

Inputs (get_dataset)

Name	Type	Required	Description
mixture_config	DatasetMixtureConfig	Yes	Configuration specifying which datasets to load and how to combine them
mixture_config.datasets	list[DatasetConfig]	Yes	List of individual dataset configurations (path, name, split, etc.)
mixture_config.streaming	bool	No	Whether to load datasets in streaming mode; default: False
mixture_config.test_split_size	None	No	Fraction for train/test split; None means no split

Inputs (apply_chat_template)

Name	Type	Required	Description
example	dict[str, list[dict[str, str]]]	Yes	Single conversational example with supported keys (messages, prompt, completion, etc.)
tokenizer	ProcessorMixin	Yes	Tokenizer with a chat template defined
tools	None	No	Tool/function definitions for function-calling templates

Outputs

Name	Type	Description
dataset (from get_dataset)	DatasetDict	Combined dataset with "train" split and optionally "test" split
formatted_example (from apply_chat_template)	dict[str, str]	Example with conversational messages rendered to plain text (e.g., "text" for language modeling, "prompt"/"completion" for prompt-completion)

Usage Examples

Loading a Dataset Mixture

from trl import DatasetMixtureConfig, get_dataset
from trl.scripts.utils import DatasetConfig

mixture_config = DatasetMixtureConfig(
    datasets=[
        DatasetConfig(path="trl-lib/Capybara", split="train"),
        DatasetConfig(path="trl-lib/tldr", split="train"),
    ],
    test_split_size=0.05,
)

dataset = get_dataset(mixture_config)
print(dataset)
# DatasetDict({
#     train: Dataset({features: [...], num_rows: ...})
#     test: Dataset({features: [...], num_rows: ...})
# })

Applying Chat Template

from transformers import AutoTokenizer
from trl.data_utils import apply_chat_template

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")

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

formatted = apply_chat_template(example, tokenizer)
print(formatted["prompt"])
# '<|im_start|>user\nWhat color is the sky?<|im_end|>\n<|im_start|>assistant\n'
print(formatted["completion"])
# 'The sky is blue.<|im_end|>\n'

Using in the SFT Script

# config.yaml
datasets:
  - path: trl-lib/Capybara
    split: train
  - path: trl-lib/tldr
    split: train
test_split_size: 0.05

python trl/scripts/sft.py \
    --config config.yaml \
    --model_name_or_path Qwen/Qwen2-0.5B \
    --output_dir ./output

Related Pages

Implements Principle

• Principle:Huggingface_Trl_SFT_Dataset_Preparation