lightrft.datasets.image_reward_db

class lightrft.datasets.image_reward_db.ImageRewardDBHandler

Bases: BaseDataHandler

Data Handler for ImageRewardDB dataset.

Paper: https://arxiv.org/abs/2304.05977 Dataset Repo: https://huggingface.co/datasets/zai-org/ImageRewardDB

get_media_info(item: Dict[str, Any]) → Dict[str, Dict[str, str]]

Extract path info for chosen and rejected images.

Parameters:

item (Dict[str, Any]) – A data item from load_data

Returns:

Dict containing local paths for ‘preferred_image’ and ‘rejected_image’

Return type:

Dict[str, Dict[str, str]]

Example:

info = handler.get_media_info(item)

load_data(path: str) → List[Dict[str, Any]]

Load ImageRewardDB shards and build preference pairs.

This method scans the given dataset root for ImageRewardDB JSON shards and aggregates image entries by prompt_id. For each prompt group, it constructs all unordered pairs of images and determines the preferred image based on the rank field (smaller is better, i.e., 1 is best).

Parameters:

path (str) – Path to the dataset root directory of ImageRewardDB.

Returns:

List of preference pair dictionaries.

Return type:

List[Dict[str, Any]]

Example:

data = handler.load_data("path/to/ImageRewardDB")

parse_item(item: Dict[str, Any], media_content: Dict[str, Any], config: Dict[str, Any]) → Tuple[List[Dict], List[Dict], Dict]

Parse a single ImageRewardDB item into message pairs for ranking.

Parameters:

• item (Dict[str, Any]) – Raw data item from ImageRewardDB dataset.

• media_content (Dict[str, Any]) – Loaded media content with ‘preferred_image’ and ‘rejected_image’ keys.

• config (Dict[str, Any]) – Configuration dict with task instructions and max_pixels

Returns:

A tuple of (messages0, messages1, metadata)

Return type:

Tuple[List[Dict], List[Dict], Dict]

Example:

msg0, msg1, other = handler.parse_item(item, media_content, config)

task_type = 'text-to-image'