Implementation:OpenRLHF OpenRLHF PromptDataset init

Knowledge Sources	OpenRLHF
Domains	Data_Processing, Reinforcement_Learning
Last Updated	2026-02-07 00:00 GMT

Overview

Concrete tool for constructing prompt-only datasets for RL generation provided by OpenRLHF.

Description

The PromptDataset class processes raw datasets into prompt-only format with optional labels and data source tracking. It applies chat templates, stores prompts as strings (not tokenized), and tracks data source metadata. The collate function returns (datasources, prompts, labels) tuples for batch processing.

Usage

Instantiate with a blended dataset for PPO, GRPO, rejection sampling, or iterative DPO training.

Code Reference

Source Location

• Repository: OpenRLHF
• File: openrlhf/datasets/prompts_dataset.py
• Lines: L21-76 (class), L31-58 (__init__)

Signature

class PromptDataset(Dataset):
    def __init__(
        self,
        dataset,               # datasets.Dataset: raw data
        tokenizer,             # tokenizer (for chat template)
        strategy,              # DeepspeedStrategy
        input_template=None,   # str: prompt formatting template
    ) -> None:

Import

from openrlhf.datasets import PromptDataset

I/O Contract

Inputs

Name	Type	Required	Description
dataset	datasets.Dataset	Yes	Raw dataset with prompt data
tokenizer	PreTrainedTokenizer	Yes	Tokenizer for chat template
strategy	DeepspeedStrategy	Yes	Training strategy

Outputs

Name	Type	Description
__getitem__ returns	Tuple	(datasource, prompt_str, label_str)
collate_fn returns	Tuple	(datasources_list, prompts_list, labels_list)

Usage Examples

from openrlhf.datasets import PromptDataset
from openrlhf.datasets.utils import blending_datasets

raw_data = blending_datasets(args.prompt_data, strategy=strategy)
prompt_dataset = PromptDataset(
    raw_data, tokenizer, strategy,
    input_template=args.input_template,
)

Related Pages

Implements Principle

• Principle:OpenRLHF_OpenRLHF_Prompt_Dataset_Construction