nemo_gym.data.task_selection

Select input rows by the task id each row carries.

Module Contents

Functions

Name	Description
`_content_without_bookkeeping`	The row as canonical JSON without Gym’s per-rollout bookkeeping (`_ng_*`), which differs between repeats.
`_first_row_of_each_task`	-
`_require_repeated_rows_to_be_identical`	-
`_require_task_id_in_every_row`	-
`_require_task_ids_in_input`	-
`row_task_id`	The id a row carries itself: a materialized row’s TaskId, else its first of `TASK_ID_FIELDS`.
`select_tasks`	Indices of the rows to run: the first row of each selected task, in input order.

API

nemo_gym.data.task_selection._content_without_bookkeeping(
    row: collections.abc.Mapping[str, pydantic.JsonValue]
) -> str

The row as canonical JSON without Gym’s per-rollout bookkeeping (_ng_*), which differs between repeats.

nemo_gym.data.task_selection._first_row_of_each_task(
    rows: collections.abc.Sequence[collections.abc.Mapping[str, pydantic.JsonValue]]
) -> dict[str, int]

nemo_gym.data.task_selection._require_repeated_rows_to_be_identical(
    rows: collections.abc.Sequence[collections.abc.Mapping[str, pydantic.JsonValue]],
    first_row_of_task: collections.abc.Mapping[str, int],
    task_ids: set[str]
) -> None

nemo_gym.data.task_selection._require_task_id_in_every_row(
    rows: collections.abc.Sequence[collections.abc.Mapping[str, pydantic.JsonValue]]
) -> None

nemo_gym.data.task_selection._require_task_ids_in_input(
    task_ids: collections.abc.Sequence[str],
    first_row_of_task: collections.abc.Mapping[str, int]
) -> None

nemo_gym.data.task_selection.row_task_id(
    row: collections.abc.Mapping[str, pydantic.JsonValue]
) -> str | None

The id a row carries itself: a materialized row’s TaskId, else its first of TASK_ID_FIELDS.

nemo_gym.data.task_selection.select_tasks(
    rows: collections.abc.Sequence[collections.abc.Mapping[str, pydantic.JsonValue]],
    task_ids: collections.abc.Sequence[str]
) -> list[int]

Indices of the rows to run: the first row of each selected task, in input order.

Each task is returned once even when the input repeats its row, so the caller’s repeat count alone decides the number of samples. Rows are selected only by an id stored in the row, never by position. Every row must carry an id; only the selected tasks’ rows must agree with each other, so a caller that selects one task per call does not re-check the whole input each time.

For example, a collated input that already repeats each task four times (a, a, a, a, b, b, b, b) with task_ids=["a"] gives one index, [0]. Run with num_repeats=2, that is two rollouts of a, not eight.