nemo_rl.data_plane.grouping

Prompt-group identity shared by TQ advantage consumers.

Module Contents

Functions

row_group_ids	Return each row’s explicit prompt-group id in metadata order.
group_index_column	Factorize prompt groups into a [N, 1] integer equality key.

API

nemo_rl.data_plane.grouping.row_group_ids(

meta: nemo_rl.data_plane.interfaces.KVBatchMeta,

) → list[str]

Return each row’s explicit prompt-group id in metadata order.

Missing tags raise instead of falling back to tokens, which can be equal for distinct prompts or change during a rollout.

nemo_rl.data_plane.grouping.group_index_column(group_ids: list[str]) → torch.Tensor

Factorize prompt groups into a [N, 1] integer equality key.

Indices are local to this call. Preserve the original group ids across filtering and concatenation, then factorize the final selected batch; concatenating independently factorized columns would merge distinct groups.