lightrft.datasets.imagegen_cot_reward

class lightrft.datasets.imagegen_cot_reward.ImageGenCoTRewardGRMHandler

Bases: BaseDataHandler

Data handler for ImageGen-CoT-Reward-5K dataset. For Text-to-Image generation task.

Paper: https://arxiv.org/pdf/2505.03318 Dataset Repo: https://huggingface.co/datasets/CodeGoat24/ImageGen-CoT-Reward-5K

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

Extract path info for the two images.

Parameters:

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

Returns:

Dict containing local paths for ‘image0’ and ‘image1’

Return type:

Dict[str, Dict[str, str]]

Example:

info = handler.get_media_info(item)

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

Loads data from json file.

Parameters:

path (str) – Path to the dataset JSON file

Returns:

List of samples with ‘data_root’ attached

Return type:

List[Dict[str, Any]]

Example:

handler = ImageGenCoTRewardHandler()
data = handler.load_data("path/to/ImageGen-CoT-Reward.json")

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

Parse a single ImageGen-CoT-Reward item into message pairs.

Parameters:

• item (Dict[str, Any]) – Raw data item from ImageGen-CoT-Reward dataset.

• media_content (Dict[str, Any]) – Loaded image content (PIL images/bytes)

• config (Dict[str, Any]) – Configuration for max_pixels

Returns:

A tuple of (messages, metadata)

Return type:

Tuple[List[Dict], Dict]

Example:

messages, other = handler.parse_item(item, media_content, config)

task_type = 'text-to-image'