lightrft.datasets.process_reward_dataset

class lightrft.datasets.process_reward_dataset.ProcessRewardDataset(*args: Any, **kwargs: Any)

Bases: Dataset

Dataset for process reward model

Args:

dataset: dataset for reward model self.tokenizer: self.tokenizer for reward model self.max_length: max length of input

collate_fn(item_list)

Collate function to batch process reward model samples.

Parameters:

item_list (list) – List of tuples (input_id, input_mask, label_id).

Returns:

Batched tensors (input_ids, input_masks, label_ids).

Return type:

tuple

packing_collate_fn(item_list)

Collate function for packing multiple process reward samples.

Parameters:

item_list (list) – List of tuples (input_id, input_mask, label_id).

Returns:

Packed tensors (packed_input_ids, packed_attention_masks, packed_label_ids, infos).

Return type:

tuple