Add datakit normalize steps for nsf_awards and nemotron v1/v2

Extend the convention established by common_corpus and starcoder2-extras
(#4626): expose a normalize_<dataset>_step factory in each datakit
download module and wire experiments through download -> normalize ->
tokenize.

Since normalize now processes a single directory (#4886), datasets with
multiple sub-datasets (nemotron v1/v2) get one normalize step per split:

- nsf_awards: one step, id_field="awd_id", file_extensions=(".parquet",)
- nemotron_v1: one step per quality/kind split (7 splits defined in
  NEMOTRON_V1_SPLITS), file_extensions=(".jsonl.gz",)
- nemotron_v2: one step per (family, subset) pair,
  file_extensions=(".parquet",)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: ravwojdyla

experiments/pretraining_datasets/nemotron_v2.py

@@ -5,23 +5,36 @@
 Nemotron v2 pre-training dataset tokenization.
 
 Download definitions live in marin.datakit.download.nemotron_v2.
-This file wires them into tokenization steps for experiment pipelines.
+This file wires them into normalize + tokenization steps for experiment
+pipelines.
 """
 
 import os.path
 
-from marin.datakit.download.nemotron_v2 import NEMOTRON_V2_DATASETS, download_nemotron_v2_step
+from marin.datakit.download.nemotron_v2 import (
+    NEMOTRON_V2_DATASETS,
+    download_nemotron_v2_step,
+    normalize_nemotron_v2_step,
+)
 from marin.execution.executor import ExecutorStep, this_output_path, versioned
 from marin.processing.tokenize import TokenizeConfig, tokenize
 from marin.processing.tokenize.data_configs import TokenizerStep
 
 # ============================================================================
-# RAW DATASET DOWNLOADS
+# RAW DATASET DOWNLOADS AND NORMALIZED OUTPUTS
 # ============================================================================
 
-downloads: dict[str, ExecutorStep] = {
-    family: download_nemotron_v2_step(family).as_executor_step() for family in NEMOTRON_V2_DATASETS
-}
+_downloads = {family: download_nemotron_v2_step(family) for family in NEMOTRON_V2_DATASETS}
+
+downloads: dict[str, ExecutorStep] = {family: step.as_executor_step() for family, step in _downloads.items()}
+
+# One normalize step per (family, subset) — normalize now processes a single directory.
+normalized: dict[str, dict[str, ExecutorStep]] = {}
+for _family, _dl in _downloads.items():
+    normalized[_family] = {
+        subset: normalize_nemotron_v2_step(_dl, family=_family, subset=subset).as_executor_step()
+        for subset in NEMOTRON_V2_DATASETS[_family].subsets
+    }
 
 
 # ============================================================================
@@ -34,23 +47,28 @@ def tokenize_nemotron_v2_family(
     *,
     tokenizer: str | None = None,
 ) -> dict[str, TokenizerStep]:
-    """Generate tokenization steps for all subsets of a Nemotron HF dataset family."""
+    """Generate tokenization steps for all subsets of a Nemotron HF dataset family.
+
+    Each subset has its own normalize step; tokenize reads from its
+    ``outputs/main/`` directory.
+    """
     if tokenizer is None:
         from experiments.llama import llama3_tokenizer
 
         tokenizer = llama3_tokenizer
 
     info = NEMOTRON_V2_DATASETS[family]
-    download_step = downloads[family]
+    family_normalized = normalized[family]
 
     steps: dict[str, ExecutorStep[TokenizeConfig]] = {}
-    for subset, glob_pattern in info.subsets.items():
+    for subset in info.subsets:
         output_name = os.path.join("tokenized", family, subset)
+        normalized_step = family_normalized[subset]
         step = ExecutorStep(
             name=output_name,
             fn=tokenize,
             config=TokenizeConfig(
-                train_paths=[download_step / glob_pattern],
+                train_paths=[normalized_step / "outputs/main/**/*.parquet"],
                 validation_paths=versioned([]),
                 cache_path=this_output_path(),
                 tokenizer=versioned(tokenizer),

experiments/pretraining_datasets/nsf_awards.py

@@ -1,20 +1,21 @@
 # Copyright The Marin Authors
 # SPDX-License-Identifier: Apache-2.0
 
-"""NSF awards dataset download and tokenization."""
+"""NSF awards dataset download, normalization, and tokenization."""
 
 from experiments.marin_models import marin_tokenizer
-from marin.datakit.download.nsf_awards import download_nsf_awards_step
+from marin.datakit.download.nsf_awards import download_nsf_awards_step, normalize_nsf_awards_step
 from marin.execution.executor import ExecutorStep, output_path_of, this_output_path, versioned
 from marin.processing.tokenize import TokenizeConfig, tokenize
 
 nsf_awards_download = download_nsf_awards_step().as_executor_step()
+nsf_awards_normalized = normalize_nsf_awards_step(nsf_awards_download).as_executor_step()
 
 nsf_awards_tokenized = ExecutorStep(
     name="tokenized/nsf_awards",
     fn=tokenize,
     config=TokenizeConfig(
-        train_paths=[output_path_of(nsf_awards_download, "*.parquet")],
+        train_paths=[output_path_of(nsf_awards_normalized, "outputs/main/*.parquet")],
         validation_paths=versioned([]),
         cache_path=this_output_path(),
         tokenizer=versioned(marin_tokenizer),

lib/marin/src/marin/datakit/download/nemotron_v1.py

@@ -11,6 +11,7 @@
 import requests
 import zstandard
 from rigging.filesystem import open_url
+from marin.datakit.normalize import normalize_step
 from marin.execution.step_spec import StepSpec
 from marin.utils import fsspec_exists
 from fray.cluster import ResourceConfig
@@ -122,3 +123,40 @@ def download_nemotron_v1_step() -> StepSpec:
         # NOTE: use the existing output to avoid re-downloading. Yes this is missing the `n`.
         override_output_path="raw/nemotro-cc-eeb783",
     )
+
+
+_NEMOTRON_V1_DATA_ROOT = "contrib/Nemotron/Nemotron-CC/data-jsonl"
+
+# Maps split name → relative path under data-jsonl/ that the normalize
+# step should point at. Each split gets its own normalize StepSpec because
+# normalize now processes a single directory (no subdirectory grouping).
+NEMOTRON_V1_SPLITS: dict[str, str] = {
+    "hq_actual": "quality=high/kind=actual",
+    "hq_synth": "quality=high/kind=synthetic",
+    "medium_high": "quality=medium-high",
+    "medium": "quality=medium",
+    "medium_low": "quality=medium-low",
+    "low_actual": "quality=low/kind=actual",
+    "low_synth": "quality=low/kind=synthetic",
+}
+
+
+def normalize_nemotron_v1_step(download: StepSpec, *, split: str) -> StepSpec:
+    """Normalize one Nemotron-CC v1 split.
+
+    The download writes dolma-format records ``{id, text, source, format,
+    metadata}`` as ``.jsonl.zst`` under nested ``quality=<x>/kind=<y>/``
+    directories. Each split gets its own normalize step pointing at the
+    corresponding subdirectory.
+    """
+    if split not in NEMOTRON_V1_SPLITS:
+        raise ValueError(f"Unknown split {split!r}. Choose from: {sorted(NEMOTRON_V1_SPLITS)}")
+    rel_path = NEMOTRON_V1_SPLITS[split]
+    return normalize_step(
+        name=f"normalized/nemotron_v1/{split}",
+        download=download,
+        text_field="text",
+        id_field="id",
+        file_extensions=(".jsonl.zst",),
+        input_path=f"{download.output_path}/{_NEMOTRON_V1_DATA_ROOT}/{rel_path}",
+    )

lib/marin/src/marin/datakit/download/nemotron_v2.py

@@ -13,6 +13,7 @@
 from dataclasses import dataclass, field
 
 from marin.datakit.download.huggingface import download_hf_step
+from marin.datakit.normalize import normalize_step
 from marin.execution.step_spec import StepSpec
 
 
@@ -131,3 +132,24 @@ def download_nemotron_v2_step(family: str) -> StepSpec:
         revision=info.revision,
         override_output_path=info.override_output_path,
     )
+
+
+def normalize_nemotron_v2_step(download: StepSpec, *, family: str, subset: str) -> StepSpec:
+    """Normalize one subset of a Nemotron v2 family.
+
+    Each subset gets its own normalize step because normalize now processes a
+    single directory. The subset's glob pattern (e.g. ``Diverse-QA/**/*.parquet``)
+    is used to derive the input subdirectory under the family download.
+    """
+    info = NEMOTRON_V2_DATASETS[family]
+    glob_pattern = info.subsets[subset]
+    # Extract the directory prefix from the glob (e.g. "Diverse-QA/**/*.parquet" → "Diverse-QA")
+    subset_dir = glob_pattern.split("/**")[0]
+    return normalize_step(
+        name=f"normalized/{family}/{subset}",
+        download=download,
+        text_field="text",
+        id_field="id",
+        file_extensions=(".parquet",),
+        input_path=f"{download.output_path}/{subset_dir}",
+    )

lib/marin/src/marin/datakit/download/nsf_awards.py

@@ -15,6 +15,7 @@
 from io import BytesIO
 
 import requests
+from marin.datakit.normalize import normalize_step
 from marin.execution.step_spec import StepSpec
 from zephyr import Dataset, ZephyrContext, counters
 from zephyr.writers import write_parquet_file
@@ -151,3 +152,14 @@ def _run(output_path: str) -> None:
         fn=_run,
         hash_attrs={"min_year": min_year, "max_year": max_year},
     )
+
+
+def normalize_nsf_awards_step(download: StepSpec) -> StepSpec:
+    """Normalize NSF awards: generate content-hash IDs, preserve awd_id as source_id."""
+    return normalize_step(
+        name="normalized/nsf-awards",
+        download=download,
+        text_field="text",
+        id_field="awd_id",
+        file_extensions=(".parquet",),
+    )