lightrft.datasets.prompts_dataset

class lightrft.datasets.prompts_dataset.PromptDataset(*args: Any, **kwargs: Any)

Bases: Dataset

Dataset for PPO (Proximal Policy Optimization) model training.

This dataset processes and stores prompts, labels, and references by applying templates and tokenization as specified by the strategy and tokenizer.

Parameters:

• dataset (Any) – The raw dataset used for training

• tokenizer (Any) – The tokenizer used for processing the prompts

• strategy (Any) – A strategy object containing configuration in its args

• input_template (str | None) – Template for formatting input text

collate_fn(item_list: list[tuple[str, str, str]]) → tuple[list[str], list[str], list[str]]

Collate a list of samples into separate lists for prompts, labels, and references.

Parameters:

item_list (list[tuple[str, str, str]]) – A list of (prompt, label, reference) tuples

Returns:

Three lists containing prompts, labels, and references, respectively

Return type:

tuple[list[str], list[str], list[str]]

lightrft.datasets.prompts_dataset.preprocess_data(example: dict[str, Any], input_template: str | None = None, input_key: str | None = None, label_key: str | None = None, apply_chat_template: Callable | None = None) → Tuple[str, str, str]

Process a single example into a (prompt, label, reference) tuple.

It supports the following input formats:

• Plain text: example[“input”] or example[“prompt”]

• Chat list: a list of dictionaries with roles and content.

• Chat string: the string is treated as a user message.

Parameters:

• example (dict[str, Any]) – The data sample

• input_template (str | None) – Template to format the prompt

• input_key (str | None) – User-specified key for input extraction

• label_key (str | None) – User-specified key for label extraction

• apply_chat_template (Callable | None) – Function to apply a chat template

Returns:

The processed (prompt, label, reference) tuple

Return type:

Tuple[str, str, str]