Overhaul Iterate/Extract abstraction as concepts in favor of a simpler Stage that yields a Task

Open
#1,459 0 comments 0 reactions 1 assignee View on GitHub

@sarahyurick is already working on this.

Since Feb 4, 2026.

Assessment

This issue has not been assessed yet.

Description

enhancement

Description

PR #1458 (https://github.com/NVIDIA-NeMo/Curator/pull/1458) improves performance and stability by fusing document iterate + extract stages (benchmarked faster and less OOM-prone). During review, we identified a broader follow-up: the existing Iterate/Extract vocabulary and abstract classes are confusing and may overly constrain users to record-level extract operations.

Goal
Move toward a simpler, unified abstraction whose generic purpose is: “produce DocumentBatch from a list/group of files” (e.g., a ProcessingStage[FileGroupTask, DocumentBatch] style stage), enabling implementers to apply transformations in a flexible way (including vectorized operations where applicable), without forcing an Iterate/Extract split.

Why not in PR #1458
Iterate/Extract are used across current pipelines (cc/wiki/arxiv) and tutorials. A full refactor would touch many scripts, require updating tutorials, and needs careful testing—too large in scope for #1458.

Scope / Tasks

  • Design a unified stage API for “file group -> document batch” production.
  • Determine migration path from current Iterate/Extract:
    • Identify impacted pipelines and tutorials.
    • Provide an adapter/shim or staged deprecation plan.
  • Refactor download+extract implementations to use the new unified approach.
  • Update tutorials/docs and add/expand test coverage to ensure behavior remains correct.
  • Decide whether/when to remove/deprecate old Iterate/Extract abstract classes.

Acceptance criteria

  • A new unified stage exists and is usable for at least one end-to-end pipeline.
  • Tutorials and core pipelines remain working (or are updated) with test coverage.
  • Clear deprecation/migration guidance is documented for users of Iterate/Extract.
Dominant language
Python
Stars
1.8k
Forks
328
Avg merge
4d 5h
Merged PRs (30d)
30

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from NVIDIA-NeMo/Curator

All issues in NVIDIA-NeMo/Curator

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.