nemo_gym.data.dataset_preparation

Module Contents

Classes

Name	Description
`Accumulator`	-
`AvgMinMax`	-
`CategoricalMetrics`	Accumulate a bounded categorical distribution.
`DatasetMetrics`	-
`DatasetValidatorState`	-
`StringMetrics`	-
`TrainDataProcessor`	-
`TrainDataProcessorConfig`	Prepare and validate training data, generating metrics and statistics for datasets.

Functions

Name	Description
`_declares_prepared_dataset`	-
`_materialized_task_parts`	-
`_observe_task_metric`	-
`_own_task_source`	Whether `row` already carries a non-empty `task_source` of its own.
`_prepare_command`	-
`_ref_name`	The `name` of a server reference given as a mapping, a model or nothing.
`_sidecar_excluded_dataset_fields`	The entry fields left out of the metrics sidecar’s record of `dataset`.
`aggregate_other_metrics`	Combines misc items (those other than response/response create params) into current metrics
`compute_sample_metrics`	-
`postprocess_other_metrics`	Aggregates metrics and merges current metrics (containing only AvgMinMax) with StringMetrics
`validate_backend_credentials`	Check if required env variables are present for the chosen backend

Data

MAX_CATEGORICAL_LABELS

_CATALOG_DATASET_FIELDS

__getattr__

API

class nemo_gym.data.dataset_preparation.Accumulator()

Bases: BaseModel

is_aggregated

bool = Field(default=False, exclude=True)

nemo_gym.data.dataset_preparation.Accumulator._add(
    other: typing.Self
) -> None

abstract

nemo_gym.data.dataset_preparation.Accumulator._aggregate() -> typing.Self

abstract

nemo_gym.data.dataset_preparation.Accumulator.add(
    other: typing.Self
) -> None

nemo_gym.data.dataset_preparation.Accumulator.aggregate() -> typing.Self

class nemo_gym.data.dataset_preparation.AvgMinMax()

Bases: Accumulator

float = Field(default=0, exclude=True)

average

float = Field(serialization_alias='Average', default=0)

max

float

mean

float = Field(default=0, exclude=True)

min

float

model_config

= ConfigDict(arbitrary_types_allowed=True)

stddev

float

total

int

nemo_gym.data.dataset_preparation.AvgMinMax._add(
    other: typing.Self
) -> None

nemo_gym.data.dataset_preparation.AvgMinMax._aggregate() -> typing.Self

nemo_gym.data.dataset_preparation.AvgMinMax.observe(
    x: float
) -> None

class nemo_gym.data.dataset_preparation.CategoricalMetrics()

Bases: Accumulator

Accumulate a bounded categorical distribution.

counts

Dict[str, int] = Field(default_factory=dict)

other_count

int = Field(default=0, exclude=True)

total_count

int = 0

unique_count

int = 0

nemo_gym.data.dataset_preparation.CategoricalMetrics._add(
    other: typing.Self
) -> None

nemo_gym.data.dataset_preparation.CategoricalMetrics._aggregate() -> typing.Self

nemo_gym.data.dataset_preparation.CategoricalMetrics.observe(
    value: str
) -> None

class nemo_gym.data.dataset_preparation.DatasetMetrics()

Bases: Accumulator

json_dumped_number_of_words

AvgMinMax

model_config

= ConfigDict(extra='allow')

number_of_examples

int

number_of_tasks

int

number_of_tools

AvgMinMax

number_of_turns

AvgMinMax

task_input_json_dumped_number_of_words

AvgMinMax

temperature

AvgMinMax

nemo_gym.data.dataset_preparation.DatasetMetrics._add(
    other: typing.Self
) -> None

nemo_gym.data.dataset_preparation.DatasetMetrics._aggregate() -> typing.Self

nemo_gym.data.dataset_preparation.DatasetMetrics.model_dump_for_output() -> typing.Dict[str, typing.Any]

Serialize existing metrics, adding task metrics only for materialized rows.

class nemo_gym.data.dataset_preparation.DatasetValidatorState()

Bases: BaseModel

key_counts

Counter = Field(default_factory=Counter)

metrics

DatasetMetrics = Field(default_factory=DatasetMetrics)

model_config

= ConfigDict(arbitrary_types_allowed=True)

offending_example_idxs

List[int] = Field(default_factory=list)

other_metrics

Dict[str, Any] = Field(default_factory=dict)

class nemo_gym.data.dataset_preparation.StringMetrics()

Bases: BaseModel

total_count

int

unique_count

int

class nemo_gym.data.dataset_preparation.TrainDataProcessor()

Bases: BaseModel

_all_server_instance_configs

List[ServerInstanceConfig] = PrivateAttr(default_factory=list)

nemo_gym.data.dataset_preparation.TrainDataProcessor._collate_samples_single_type(
    type: nemo_gym.config.config_types.DatasetType,
    server_instance_configs: typing.List[nemo_gym.config.config_types.ServerInstanceConfig],
    task_data_validation: str = 'warn',
    num_repeats: typing.Optional[typing.Union[int, typing.Dict[str, int]]] = None
) -> typing.List[pathlib.Path]

nemo_gym.data.dataset_preparation.TrainDataProcessor._dataset_metrics_hook_for(
    c: nemo_gym.config.config_types.ServerInstanceConfig,
    server_instance_configs: typing.List[nemo_gym.config.config_types.ServerInstanceConfig]
) -> nemo_gym.data.dataset_metrics.DatasetMetricHook | None

classmethod

nemo_gym.data.dataset_preparation.TrainDataProcessor._environment_servers_naming(
    c: nemo_gym.config.config_types.ServerInstanceConfig,
    server_instance_configs: typing.List[nemo_gym.config.config_types.ServerInstanceConfig]
) -> typing.List[nemo_gym.config.config_types.ServerInstanceConfig]

staticmethod

Every environment-server instance whose agent_server is the agent c, in config order.

nemo_gym.data.dataset_preparation.TrainDataProcessor._iter_dataset_lines(
    dataset_config: nemo_gym.config.config_types.DatasetConfig
)

nemo_gym.data.dataset_preparation.TrainDataProcessor._owning_environment_server(
    c: nemo_gym.config.config_types.ServerInstanceConfig,
    server_instance_configs: typing.List[nemo_gym.config.config_types.ServerInstanceConfig]
) -> typing.Optional[nemo_gym.config.config_types.ServerInstanceConfig]

classmethod

The environment-server instance that drives a self-contained agent’s episodes.

An environment server that names the agent as its agent_server owns the rows the agent declares (the Harbor layout: the agent block carries datasets so rollout collection finds them, the environment server seeds and verifies them), so its task_data.py, when it ships one, is the schema to validate against. Several environment servers may name one agent (one per taskset, each routing its own rows to one harness instance); they share the implementation the first one stands for here, and _taskset_folder tells them apart by the dataset’s taskset.

nemo_gym.data.dataset_preparation.TrainDataProcessor._owning_environment_server_impl(
    c: nemo_gym.config.config_types.ServerInstanceConfig,
    server_instance_configs: typing.List[nemo_gym.config.config_types.ServerInstanceConfig]
) -> typing.Optional[str]

classmethod

The implementation key of _owning_environment_server, or None.

nemo_gym.data.dataset_preparation.TrainDataProcessor._owning_resources_server_impl(
    c: nemo_gym.config.config_types.ServerInstanceConfig,
    server_instance_configs: typing.List[nemo_gym.config.config_types.ServerInstanceConfig]
) -> typing.Optional[str]

staticmethod

The resources-server implementation (directory) that owns a declaring instance’s data.

Datasets declared by a resources server belong to its own implementation; datasets declared by an agent belong to the resources server the agent references. Self-contained agents (no resources_server reference) return None here; the validator lookup then falls back to the environment server that drives the agent, then to the agent’s own directory.

nemo_gym.data.dataset_preparation.TrainDataProcessor._print_title(
    title: str
) -> None

nemo_gym.data.dataset_preparation.TrainDataProcessor._schema_home_for(
    c: nemo_gym.config.config_types.ServerInstanceConfig,
    server_instance_configs: typing.List[nemo_gym.config.config_types.ServerInstanceConfig]
) -> typing.Optional[tuple[str, str]]

classmethod

(implementation, base_folder) whose task_data.py validates c’s datasets, or None.

A resources server owns its own data and the data of agents that reference it. An agent without a reference is owned by the environment server that names it when that server ships a task_data.py (the Harbor layout), else by itself: a self-contained agent wrapped by a schema-less environment server (the legacy_agent layout) validates against its own task_data.py. An agent whose reference is dangling is skipped: its schema home is the (missing) server.

nemo_gym.data.dataset_preparation.TrainDataProcessor._schema_lookup_configs(
    server_instance_configs: typing.List[nemo_gym.config.config_types.ServerInstanceConfig]
) -> typing.List[nemo_gym.config.config_types.ServerInstanceConfig]

The in-scope instances plus every other instance of the run config, for schema-home lookups.

nemo_gym.data.dataset_preparation.TrainDataProcessor._task_data_validator_for(
    c: nemo_gym.config.config_types.ServerInstanceConfig,
    d: typing.Union[nemo_gym.config.config_types.DatasetConfig, nemo_gym.config.config_types.BenchmarkDatasetConfig],
    server_instance_configs: typing.List[nemo_gym.config.config_types.ServerInstanceConfig],
    task_data_validation: str = 'warn'
) -> typing.Optional[nemo_gym.data.task_data.TaskDataValidator]

classmethod

The validator for c’s dataset d, or None when no schema applies.

task_data_validation is the collation’s mode. A server schema that does not load is skipped with a warning in every mode, as before. A Harbor dataset’s own task_data.py that does not load is a warning in warn mode too, and the rows are then checked for the server’s reserved keys alone, as an untyped dataset’s are (the seed refuses every episode of that taskset with the same error); in error mode it fails the collation, as a row the schema rejects would.

nemo_gym.data.dataset_preparation.TrainDataProcessor._taskset_folder(
    c: nemo_gym.config.config_types.ServerInstanceConfig,
    d: typing.Union[nemo_gym.config.config_types.DatasetConfig, nemo_gym.config.config_types.BenchmarkDatasetConfig],
    server_instance_configs: typing.List[nemo_gym.config.config_types.ServerInstanceConfig]
) -> typing.Optional[pathlib.Path]

classmethod

The task folder of the Harbor taskset the agent c’s dataset d was materialized from, or None.

Rows route to the environment server whose tasksets carries their taskset (the generated run config names the agent’s dataset after the taskset and routes by taskset), so that server’s tasksets.<taskset>.folder is where the loader read the tasks from, whichever of the environment servers naming the agent comes first in the config. None for a dataset no such server carries.

nemo_gym.data.dataset_preparation.TrainDataProcessor._validate_aggregate_metrics(
    aggregate_metrics_dict: typing.Dict,
    metrics_fpath: pathlib.Path,
    source_metrics: bool = False
) -> typing.Optional[pathlib.Path]

Returns the conflicting metrics fpath if invalid. Else returns None

nemo_gym.data.dataset_preparation.TrainDataProcessor._validate_samples_and_aggregate_metrics_single_dataset(
    dataset_config: nemo_gym.config.config_types.DatasetConfig,
    dataset_metrics_hook: nemo_gym.data.dataset_metrics.DatasetMetricHook | None = None
) -> nemo_gym.data.dataset_preparation.DatasetValidatorState

nemo_gym.data.dataset_preparation.TrainDataProcessor._validate_samples_and_aggregate_metrics_single_sample(
    state: nemo_gym.data.dataset_preparation.DatasetValidatorState,
    sample_idx: int,
    sample_dict_str: str,
    require_responses: bool = True,
    dataset_metrics_hook: nemo_gym.data.dataset_metrics.DatasetMetricHook | None = None
) -> None

nemo_gym.data.dataset_preparation.TrainDataProcessor.collate_samples(
    config: nemo_gym.data.dataset_preparation.TrainDataProcessorConfig,
    server_instance_configs: typing.List[nemo_gym.config.config_types.ServerInstanceConfig],
    dataset_type_to_aggregate_metrics: typing.Dict[str, nemo_gym.data.dataset_preparation.DatasetMetrics]
) -> None

nemo_gym.data.dataset_preparation.TrainDataProcessor.load_and_validate_server_instance_configs(
    config: nemo_gym.data.dataset_preparation.TrainDataProcessorConfig,
    global_config_dict: omegaconf.DictConfig
) -> typing.List[nemo_gym.config.config_types.ServerInstanceConfig]

nemo_gym.data.dataset_preparation.TrainDataProcessor.load_datasets(
    config: nemo_gym.data.dataset_preparation.TrainDataProcessorConfig,
    server_instance_configs: typing.List[nemo_gym.config.config_types.ServerInstanceConfig]
) -> None

nemo_gym.data.dataset_preparation.TrainDataProcessor.run(
    global_config_dict: omegaconf.DictConfig
)

See the README section “How To: Prepare and validate data for PR submission or RL training”

nemo_gym.data.dataset_preparation.TrainDataProcessor.validate_samples_and_aggregate_metrics(
    server_instance_configs: typing.List[nemo_gym.config.config_types.ServerInstanceConfig],
    overwrite_metrics_conflicts: bool
) -> typing.Dict[str, nemo_gym.data.dataset_preparation.DatasetMetrics]

class nemo_gym.data.dataset_preparation.TrainDataProcessorConfig()

Bases: BaseNeMoGymCLIConfig

Prepare and validate training data, generating metrics and statistics for datasets.

Examples:

config_paths="resources_servers/example_multi_step/configs/example_multi_step.yaml,\
responses_api_models/openai_model/configs/openai_model.yaml"
gym dataset collate "+config_paths=[${config_paths}]"         +output_dirpath=data/example_multi_step         +mode=example_validation

data_source

Literal['gitlab', 'huggingface']

effective_task_data_validation

str

in_scope_dataset_types

List[DatasetType]

mode

Union[Literal['train_preparation'], Literal['example_validation']]

num_repeats

Optional[Union[int, Dict[str, int]]]

output_dirpath

str

overwrite_metrics_conflicts

bool

should_download

bool

task_data_validation

Literal['off', 'warn', 'error', 'auto']

nemo_gym.data.dataset_preparation._declares_prepared_dataset(
    value: object,
    dataset: nemo_gym.config.config_types.DatasetConfig
) -> bool

nemo_gym.data.dataset_preparation._materialized_task_parts(
    sample: collections.abc.Mapping[str, typing.Any]
) -> tuple[nemo_gym.core.episode_types.TaskId, collections.abc.Mapping[str, typing.Any]] | None

nemo_gym.data.dataset_preparation._observe_task_metric(
    metrics: nemo_gym.data.dataset_preparation.DatasetMetrics,
    name: str,
    value: nemo_gym.data.dataset_metrics.DatasetMetricValue
) -> None

nemo_gym.data.dataset_preparation._own_task_source(
    row: collections.abc.Mapping[str, typing.Any]
) -> bool

Whether row already carries a non-empty task_source of its own.

nemo_gym.data.dataset_preparation._prepare_command(
    global_config: omegaconf.DictConfig,
    dataset: nemo_gym.config.config_types.DatasetConfig
) -> str

nemo_gym.data.dataset_preparation._ref_name(
    ref: typing.Any
) -> typing.Optional[str]

The name of a server reference given as a mapping, a model or nothing.

nemo_gym.data.dataset_preparation._sidecar_excluded_dataset_fields(
    dataset: nemo_gym.config.config_types.DatasetConfig | nemo_gym.config.config_types.BenchmarkDatasetConfig,
    preparation: bool = True
) -> set[str]

The entry fields left out of the metrics sidecar’s record of dataset.

Agent and taskset select routing, not source data, so taskset and flat declarations of one file share its sidecar. Preparation (prepare_script, prepare_dependencies) controls how a local source file is produced and is not a property of the rows, and defaults must not churn existing sidecars. The catalog declarations (_CATALOG_DATASET_FIELDS) are left out for the same reason.

preparation False keeps the preparation fields: the per-split sidecar (<split>_metrics.json in the output folder, TrainDataProcessor.collate_samples) has always recorded them, and TrainDataProcessor._validate_aggregate_metrics requires every key of an existing sidecar, so leaving them out there would reject every prepared dataset whose sidecar names them.

nemo_gym.data.dataset_preparation.aggregate_other_metrics(
    metrics: typing.Dict[str, typing.Any],
    sample: typing.Dict[str, typing.Any]
) -> None

Combines misc items (those other than response/response create params) into current metrics

nemo_gym.data.dataset_preparation.compute_sample_metrics(
    sample_dict_str: str,
    require_responses: bool = True
) -> typing.Tuple[nemo_gym.data.dataset_preparation.DatasetMetrics, bool]

nemo_gym.data.dataset_preparation.postprocess_other_metrics(
    metrics: nemo_gym.data.dataset_preparation.DatasetMetrics,
    other_metrics: typing.Dict[str, typing.Any]
) -> None

Aggregates metrics and merges current metrics (containing only AvgMinMax) with StringMetrics

nemo_gym.data.dataset_preparation.validate_backend_credentials(
    backend: str
) -> tuple[bool, str]

Check if required env variables are present for the chosen backend

nemo_gym.data.dataset_preparation.MAX_CATEGORICAL_LABELS = 100

nemo_gym.data.dataset_preparation._CATALOG_DATASET_FIELDS = frozenset({'domain', 'verified', 'allowed_agents'})

nemo_gym.data.dataset_preparation.__getattr__ = moved_attr_getter(__name__, {'prepare_data': 'nemo_gym.cli.dataset'})