Add Huggingface support for data preparation

docs/how-to-faq.md

@@ -259,13 +259,23 @@ Gitlab model names are case sensitive. There can be models named 'My_Model' and
 :::
 
 Downloading a dataset from Huggingface is straightforward:
+
+
+**For JSONL files** (specify `artifact_fpath`):
 ```bash
 ng_download_dataset_from_hf \
-    +repo_id=NVIDIA/NeMo-Gym-Instruction_Following-multineedle-{your dataset name} \
-    +artifact_fpath=multineedle_benchmark.jsonl \
-    +output_fpath=data/multineedle_benchmark_hf.jsonl
+    +repo_id=nvidia/Nemotron-RL-instruction_following \
+    +output_fpath=data/train.jsonl \
+    +artifact_fpath=train.jsonl
 ```
 
+**For parquet datasets** (omit `artifact_fpath`, specify `split`):
+```bash
+ng_download_dataset_from_hf \
+    +repo_id=nvidia/Nemotron-RL-knowledge-mcqa \
+    +output_fpath=data/train.jsonl \
+    +split=train
+```
 
 # How To: Prepare and validate data for PR submission or RL training
 When you use `ng_init_resources_server +entrypoint=resources_servers/example_multi_step` to initialize a resources server, you will get a config.yaml that looks like the below code block. The dataset information for training, validation, and example will be inside the scope of your agent config (e.g. under simple_agent) and is a list of dataset objects.
@@ -317,7 +327,8 @@ A dataset object consists of:
 - Type: train, validation, or example. Train and validation are as used in NeMo RL or other train frameworks. More information about the example type is in the next section.
 - Jsonl fpath: the local file path to your jsonl file for this dataset.
 - Num repeats: optionally repeat each row when preparing or collating data. Defaults to 1 if unspecified.
-- Gitlab identifier: The remote path to the dataset as held in the Gitlab dataset registry. This field is required for train and validation datasets. (Not required for example datasets since those are required to be committed to Git).
+- Gitlab identifier: (NVIDIA internal) The remote path to the dataset as held in the Gitlab dataset registry. This field is required for train and validation datasets. (Not required for example datasets since those are required to be committed to Git).
+- HuggingFace identifier: (Public) The remote path to the dataset on HuggingFace. A `artifact_fpath` or `split` must be provided for downloads.
 - License: The license of that dataset. Required for train and validation datasets and not required for example datasets, similar in principle to the Gitlab identifier.
 - Start idx, end idx: used for slicing your dataset.
 ```yaml

nemo_gym/config_types.py

@@ -146,6 +146,13 @@ class DeleteJsonlDatasetGitlabConfig(BaseNeMoGymCLIConfig):
     dataset_name: str
 
 
+class JsonlDatasetHuggingFaceIdentifer(BaseModel):
+    repo_id: str = Field(description="The repo id.")
+    artifact_fpath: Optional[str] = Field(
+        default=None, description="The filepath to the artifact in the repo (for JSONL). Omit for parquet."
+    )
+
+
 class BaseUploadJsonlDatasetHuggingFaceConfig(BaseNeMoGymCLIConfig):
     hf_token: str
     hf_organization: str
@@ -173,11 +180,14 @@ class UploadJsonlDatasetHuggingFaceMaybeDeleteConfig(BaseUploadJsonlDatasetHuggi
     delete_from_gitlab: Optional[bool] = False
 
 
-class DownloadJsonlDatasetHuggingFaceConfig(BaseNeMoGymCLIConfig):
-    output_fpath: str
-    hf_token: str
-    artifact_fpath: str
-    repo_id: str
+class DownloadJsonlDatasetHuggingFaceConfig(JsonlDatasetHuggingFaceIdentifer, BaseNeMoGymCLIConfig):
+    output_fpath: str = Field(description="Where to save the downloaded dataset (as JSONL).")
+    hf_token: Optional[str] = Field(default=None, description="The huggingface token.")
+    split: str = Field(default="train", description="Dataset split (only used for parquet).")
+
+
+class PrepareDataConfig(BaseNeMoGymCLIConfig):
+    data_source: Literal["gitlab", "huggingface"] = "huggingface"
 
 
 DatasetType = Union[Literal["train"], Literal["validation"], Literal["example"]]
@@ -190,6 +200,7 @@ class DatasetConfig(BaseModel):
 
     num_repeats: int = Field(default=1, ge=1)
     gitlab_identifier: Optional[JsonlDatasetGitlabIdentifer] = None
+    huggingface_identifier: Optional[JsonlDatasetHuggingFaceIdentifer] = None
     license: Optional[
         Union[
             Literal["Apache 2.0"],
@@ -204,7 +215,6 @@ class DatasetConfig(BaseModel):
     @model_validator(mode="after")
     def check_train_validation_sets(self) -> "DatasetConfig":
         if self.type in ["train", "validation"]:
-            assert self.gitlab_identifier is not None, f"A Gitlab path is required for {self.name}"
             assert self.license is not None, f"A license is required for {self.name}"
 
         return self

nemo_gym/dataset_orchestrator.py

@@ -21,6 +21,7 @@
 )
 from nemo_gym.gitlab_utils import delete_model_from_gitlab, is_model_in_gitlab
 from nemo_gym.hf_utils import download_jsonl_dataset as download_jsonl_dataset_from_hf
+from nemo_gym.hf_utils import download_parquet_dataset_as_jsonl
 from nemo_gym.hf_utils import upload_jsonl_dataset as upload_jsonl_dataset_to_hf
 from nemo_gym.server_utils import get_global_config_dict
 
@@ -73,7 +74,18 @@ def upload_jsonl_dataset_to_hf_and_delete_gitlab_cli() -> None:  # pragma: no co
 def download_jsonl_dataset_from_hf_cli() -> None:  # pragma: no cover
     global_config = get_global_config_dict()
     config = DownloadJsonlDatasetHuggingFaceConfig.model_validate(global_config)
-    download_jsonl_dataset_from_hf(config)
+
+    if config.artifact_fpath:
+        download_jsonl_dataset_from_hf(config)
+    else:
+        print(
+            f"""\
+[Nemo-Gym] - No artifact_fpath provided.
+Downloading the entire '{config.split}' split from '{config.repo_id}'
+and saving as JSONL at '{config.output_fpath}'.
+"""
+        )
+        download_parquet_dataset_as_jsonl(config)
 
 
 def delete_jsonl_dataset_from_gitlab_cli() -> None:  # pragma: no cover

nemo_gym/hf_utils.py

@@ -16,6 +16,7 @@
 from pathlib import Path
 
 import yaml
+from datasets import load_dataset
 from huggingface_hub import HfApi, hf_hub_download
 from huggingface_hub.utils import HfHubHTTPError
 from scripts.update_resource_servers import visit_resource_server
@@ -54,6 +55,19 @@ def check_jsonl_format(file_path: str) -> bool:  # pragma: no cover
     return True
 
 
+def download_parquet_dataset_as_jsonl(
+    config: DownloadJsonlDatasetHuggingFaceConfig,
+) -> None:  # pragma: no cover
+    """Download a HF dataset and save as JSONL"""
+    try:
+        ds = load_dataset(config.repo_id, split=config.split, token=config.hf_token)
+        ds.to_json(config.output_fpath)
+        print(f"[Nemo-Gym] - Downloaded and converted dataset to: {config.output_fpath}")
+    except Exception as e:
+        print(f"[Nemo-Gym] - Error downloading/converting dataset: {e}")
+        raise
+
+
 def upload_jsonl_dataset(
     config: UploadJsonlDatasetHuggingFaceConfig,
 ) -> None:  # pragma: no cover

nemo_gym/train_data_utils.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import json
+import sys
 from abc import abstractmethod
 from collections import Counter, defaultdict
 from math import sqrt
@@ -34,6 +35,7 @@
     DatasetConfig,
     DatasetType,
     DownloadJsonlDatasetGitlabConfig,
+    DownloadJsonlDatasetHuggingFaceConfig,
     ServerInstanceConfig,
 )
 from nemo_gym.gitlab_utils import download_jsonl_dataset
@@ -42,6 +44,12 @@
     GlobalConfigDictParserConfig,
     get_global_config_dict,
 )
+from nemo_gym.hf_utils import (
+    download_jsonl_dataset as download_jsonl_dataset_from_hf,
+)
+from nemo_gym.hf_utils import (
+    download_parquet_dataset_as_jsonl,
+)
 
 
 class TrainDataProcessorConfig(BaseNeMoGymCLIConfig):
@@ -53,6 +61,10 @@ class TrainDataProcessorConfig(BaseNeMoGymCLIConfig):
         default=False,
         description="Whether or not to download missing datasets. By default, no datasets will be downloaded.",
     )
+    data_source: Literal["gitlab", "huggingface"] = Field(
+        default="huggingface",
+        description="Where to download missing datasets from: 'gitlab' (NVIDIA internal) or 'huggingface' (external).",
+    )
 
     @property
     def in_scope_dataset_types(self) -> List[DatasetType]:
@@ -437,16 +449,67 @@ def load_datasets(
                 "Missing local datasets. You must provide local datasets since download is disabled. Run with `+should_download=true` to enable downloading."
             )
 
-        for (
-            server_name,
-            datasets,
-        ) in local_datasets_not_found.items():  # pragma: no cover
-            for d in datasets:
-                download_config = DownloadJsonlDatasetGitlabConfig.model_validate(
-                    d.gitlab_identifier.model_dump() | {"output_fpath": d.jsonl_fpath}
-                )
-                print(f"Downloading dataset `{d.name}` from `{server_name}` using {download_config}")
-                download_jsonl_dataset(download_config)
+        if local_datasets_not_found:
+            backend = config.data_source
+            is_valid, error_msg = validate_backend_credentials(backend)
+            global_config = get_global_config_dict()
+            if not is_valid:
+                print(f"Cannot download datasets: {error_msg}")
+                sys.exit(1)
+
+            for (
+                server_name,
+                datasets,
+            ) in local_datasets_not_found.items():  # pragma: no cover
+                for d in datasets:
+                    try:
+                        if backend == "gitlab":
+                            if d.gitlab_identifier is None:
+                                print(f"Dataset `{d.name}` missing gitlab_identifier for GitLab backend")
+                                continue
+
+                            download_config = DownloadJsonlDatasetGitlabConfig.model_validate(
+                                d.gitlab_identifier.model_dump() | {"output_fpath": d.jsonl_fpath}
+                            )
+                            print(
+                                f"Downloading dataset `{d.name}` for `{server_name}` from {backend} using {download_config}"
+                            )
+                            download_jsonl_dataset(download_config)
+
+                        elif backend == "huggingface":
+                            if d.huggingface_identifier is None:
+                                print(f"Dataset `{d.name}` missing huggingface_identifier for HuggingFace backend")
+                                continue
+
+                            if d.huggingface_identifier.artifact_fpath is None:
+                                download_config = DownloadJsonlDatasetHuggingFaceConfig.model_validate(
+                                    {
+                                        "repo_id": d.huggingface_identifier.repo_id,
+                                        "output_fpath": d.jsonl_fpath,
+                                        "split": d.type,
+                                        "hf_token": global_config.get("hf_token"),
+                                    }
+                                )
+                                print(
+                                    f"Downloading parquet dataset `{d.name}` from {d.huggingface_identifier.repo_id}..."
+                                )
+                                download_parquet_dataset_as_jsonl(download_config)
+                            else:
+                                download_config = DownloadJsonlDatasetHuggingFaceConfig.model_validate(
+                                    {
+                                        "repo_id": d.huggingface_identifier.repo_id,
+                                        "artifact_fpath": d.huggingface_identifier.artifact_fpath,
+                                        "output_fpath": d.jsonl_fpath,
+                                        "hf_token": global_config["hf_token"],
+                                    }
+                                )
+                                print(
+                                    f"Downloading dataset `{d.name}` for `{server_name}` from {backend}: {d.huggingface_identifier.repo_id}"
+                                )
+                                download_jsonl_dataset_from_hf(download_config)
+
+                    except Exception as e:
+                        print(f"Failed to download dataset `{d.name}` from {backend}: {e}")
 
     ########################################
     # Validate samples and aggregate metrics
@@ -715,6 +778,36 @@ def collate_samples(
         print(f"View your final data!{final_fpaths_str}")
 
 
+def validate_backend_credentials(backend: str) -> tuple[bool, str]:
+    """Check if required env variables are present for the chosen backend"""
+    global_config = get_global_config_dict()
+
+    if backend == "gitlab":
+        required = ["mlflow_tracking_uri", "mlflow_tracking_token"]
+        missing = [k for k in required if k not in global_config or not global_config[k]]
+        if missing:
+            return False, (
+                f"GitLab backend selected but missing credentials: {missing}\n"
+                f"Add to env.yaml:\n"
+                f"  mlflow_tracking_uri: <your_gitlab_uri>\n"
+                f"  mlflow_tracking_token: <your_gitlab_token>"
+            )
+
+    elif backend == "huggingface":
+        required = ["hf_token", "hf_organization", "hf_collection_slug"]
+        missing = [k for k in required if k not in global_config or not global_config[k]]
+        if missing:
+            return False, (
+                f"HuggingFace backend selected but missing credentials: {missing}\n"
+                f"Add to env.yaml:\n"
+                f"  hf_token: <your_hf_token>\n"
+                f"  hf_organization: <your_org>\n"
+                f"  hf_collection_slug: <your_collection_slug>"
+            )
+
+    return True, ""
+
+
 def prepare_data():  # pragma: no cover
     global_config_dict = get_global_config_dict(
         global_config_dict_parser_config=GlobalConfigDictParserConfig(

pyproject.toml

@@ -153,6 +153,10 @@ dependencies = [
     # Updated: Fri Nov 07, 2025 with psutil==6.1.1
     # License: BSD 3-Clause https://github.com/giampaolo/psutil/blob/master/LICENSE
     "psutil",
+
+    # HuggingFace datasets: for loading and converting parquet datasets
+    # License: Apache 2.0 https://github.com/huggingface/datasets/blob/main/LICENSE
+    "datasets",
 ]
 
 [dependency-groups]

resources_servers/equivalence_llm_judge/configs/equivalence_llm_judge.yaml

@@ -138,4 +138,5 @@ equivalence_llm_judge_simple_agent:
         type: train
         license: Apache 2.0
         jsonl_fpath: resources_servers/equivalence_llm_judge/data/train.jsonl
-        dataset_url: https://huggingface.co/datasets/nvidia/Nemotron-RL-knowledge-openQA
+        huggingface_identifier:
+          repo_id: nvidia/Nemotron-RL-knowledge-openQA

resources_servers/example_multi_step/data/example_metrics.json

@@ -70,5 +70,4 @@
         "Median": 299.0,
         "Standard deviation": 0.0
     }
-}
-
+}

resources_servers/google_search/configs/google_search.yaml

@@ -25,7 +25,9 @@ google_search_simple_agent:
           version: 0.0.1
           artifact_fpath: MCQA_syn_gpqa_1_2_difficulty_filtered_responses_api.jsonl
         license: Apache 2.0
-        dataset_url: https://huggingface.co/datasets/nvidia/Nemotron-RL-knowledge-web_search-mcqa
+        huggingface_identifier:
+          repo_id: nvidia/Nemotron-RL-knowledge-web_search-mcqa
+          artifact_fpath: mcqa_search.jsonl
       - name: example
         type: example
         jsonl_fpath: resources_servers/google_search/data/example.jsonl

resources_servers/google_search/data/_train_metrics.json

@@ -0,0 +1,58 @@
+{
+    "name": "train",
+    "type": "train",
+    "jsonl_fpath": "resources_servers/google_search/data/train.jsonl",
+    "num_repeats": 1,
+    "gitlab_identifier": {
+        "dataset_name": "search_STEM_syn_gpqa_v1_2_difficulty_filtered",
+        "version": "0.0.1",
+        "artifact_fpath": "MCQA_syn_gpqa_1_2_difficulty_filtered_responses_api.jsonl"
+    },
+    "dataset_url": "https://huggingface.co/datasets/nvidia/Nemotron-RL-knowledge-web_search-mcqa",
+    "license": "Apache 2.0",
+    "Number of examples": 2907,
+    "Number of tools": {
+        "Total # non-null values": 2907,
+        "Average": 2.0,
+        "Min": 2.0,
+        "Max": 2.0,
+        "Median": 2.0,
+        "Standard deviation": 0.0
+    },
+    "Json-dumped number of words (proxy for token count)": {
+        "Total # non-null values": 2907,
+        "Average": 282.98,
+        "Min": 207.0,
+        "Max": 439.0,
+        "Median": 280.91,
+        "Standard deviation": 35.33
+    },
+    "Number of turns": {
+        "Total # non-null values": 2907,
+        "Average": 1.0,
+        "Min": 1.0,
+        "Max": 1.0,
+        "Median": 1.0,
+        "Standard deviation": 0.0
+    },
+    "Temperature": {
+        "Total # non-null values": 2907,
+        "Average": 0.6,
+        "Min": 0.6,
+        "Max": 0.6,
+        "Median": 0.6,
+        "Standard deviation": 0.0
+    },
+    "expected_answer": {
+        "unique_count": 4,
+        "total_count": 2907
+    },
+    "task_difficulty_qwen3_32b_avg_8": {
+        "Total # non-null values": 2907,
+        "Average": 0.374,
+        "Min": 0.25,
+        "Max": 0.5,
+        "Median": 0.374,
+        "Standard deviation": 0.103
+    }
+}