integrations.nemo_gym.dataset
integrations.nemo_gym.dataset
Dataset loading for NeMo Gym JSONL files.
Converts NeMo Gym JSONL format into HuggingFace Datasets compatible with TRL’s GRPOTrainer. Supports multi-environment routing via: 1. Per-dataset server_name (all rows in a file go to one server) 2. Per-row agent_ref.name (each row specifies its own server)
Functions
| Name | Description |
|---|---|
| load_nemo_gym_datasets | Load and merge NeMo Gym JSONL datasets with multi-environment support. |
load_nemo_gym_datasets
integrations.nemo_gym.dataset.load_nemo_gym_datasets(gym_dir, dataset_configs)Load and merge NeMo Gym JSONL datasets with multi-environment support.
Each dataset config should have
- path: JSONL file path (absolute, or relative to gym_dir)
- server_name: Default NeMo Gym server for this dataset. Can be overridden per-row if the JSONL has an “agent_ref” field.
- max_samples (optional): Max number of samples to use from this dataset
Per-row routing: If a JSONL row has an “agent_ref”: {“name”: “…”} field, that takes precedence over the dataset-level server_name. This allows mixing environments within a single dataset file (matching TRL’s pattern).
The output dataset has columns
- prompt: list[dict] chat format
- resources_server_ref: dict with {“name”: server_name}
- verify_extra: dict with original JSONL data for verify requests
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| gym_dir | str | Path to the NeMo Gym directory. | required |
| dataset_configs | list[dict] | List of dataset configuration dicts. | required |
Returns
| Name | Type | Description |
|---|---|---|
| Dataset | A HuggingFace Dataset ready for GRPOTrainer. |