Implementation:NVIDIA NeMo Aligner PickScore Dataset
| Implementation Details | |
|---|---|
| Name | PickScore_Dataset |
| Type | API Doc |
| Module | nemo_aligner.data.mm |
| Repository | NeMo Aligner |
| Last Updated | 2026-02-08 00:00 GMT |
Overview
A PyTorch Dataset class for loading image-text preference pairs used in PickScore reward model training within the NeMo Aligner multimodal alignment pipeline.
Description
The PickScoreDataset class extends torch.utils.data.Dataset to load and serve image-text preference data for training PickScore-based reward models. Each sample consists of two images (jpg_0, jpg_1) paired with a text caption and preference labels indicating which image better matches the caption. The dataset reads from Arrow file shards on disk (produced by the HuggingFace datasets library), concatenates them, and provides shuffled access with an optional tie-filtering mechanism that removes ambiguous samples where neither image is preferred.
A companion build_train_valid_datasets factory function is also provided, which instantiates train, validation, and optionally test splits from a single model configuration object.
Usage
Used when training a multimodal reward model (e.g., PickScore) that learns human preferences over generated images. The dataset is typically constructed via build_train_valid_datasets and passed to a DataLoader for the training loop.
Code Reference
Source Location
- Repository: NeMo Aligner
- File:
nemo_aligner/data/mm/pickscore_dataset.py(L59-115 PickScoreDataset, L38-56 build_train_valid_datasets)
Signature
class PickScoreDataset(Dataset):
def __init__(
self,
model_cfg,
tokenizer=None,
stop_idx=None,
consumed_samples=0,
path=None,
seed: int = 42,
split: str = "train",
):
...
def __len__(self) -> int:
...
def __getitem__(self, i: int) -> dict[str, Any]:
...
def build_train_valid_datasets(
model_cfg, consumed_samples, tokenizer=None, seed=None, return_test_data=False,
):
...
Import
from nemo_aligner.data.mm.pickscore_dataset import PickScoreDataset
from nemo_aligner.data.mm.pickscore_dataset import build_train_valid_datasets
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| model_cfg | DictConfig / dict | Yes | Model configuration; must contain data.data_path pointing to the Arrow shard directory, and optional per-split filter_ties flag
|
| tokenizer | object | No | Tokenizer instance (reserved for future use; not consumed in current implementation) |
| stop_idx | int | No | Unused parameter (reserved) |
| consumed_samples | int | No | Number of samples already consumed (default 0; reserved for resumption) |
| path | str | No | Explicit data path; overrides model_cfg.data.data_path if provided
|
| seed | int | No | Random seed for index shuffling (default 42); set to None to disable shuffling
|
| split | str | No | One of "train", "val", or "test" (default "train")
|
Outputs
| Name | Type | Description |
|---|---|---|
| img_0 | torch.FloatTensor | First candidate image as a float tensor with shape (H, W, 3)
|
| img_1 | torch.FloatTensor | Second candidate image as a float tensor with shape (H, W, 3)
|
| label | torch.FloatTensor | Preference label tensor of shape (2,) containing [label_0, label_1]
|
| prompt | str | Text caption associated with the image pair |
| time_step | torch.Tensor | Scalar tensor [0.0] (placeholder for diffusion timestep compatibility)
|
Usage Examples
from omegaconf import OmegaConf
from nemo_aligner.data.mm.pickscore_dataset import PickScoreDataset, build_train_valid_datasets
# Using the factory function (recommended)
train_ds, val_ds = build_train_valid_datasets(
model_cfg=cfg.model,
consumed_samples=0,
tokenizer=tokenizer,
seed=42,
)
# Direct instantiation
dataset = PickScoreDataset(
model_cfg=cfg.model,
tokenizer=None,
split="val",
seed=42,
)
sample = dataset[0]
print(sample["prompt"]) # caption text
print(sample["img_0"].shape) # (H, W, 3)
print(sample["label"]) # tensor([label_0, label_1])