Remove MegatronTokenizerWriter from writer/__init__.py

[ArivunidhiA]

Description

Closes #1418

Remove MegatronTokenizerWriter re-export from writer/__init__.py to avoid transitively importing transformers (~150 MB peak memory) for users who only need JsonlWriter or ParquetWriter.

Updated internal references in tests and tutorials to use the direct import path.

Usage

# MegatronTokenizerWriter is now imported directly:
from nemo_curator.stages.text.io.writer.megatron_tokenizer import MegatronTokenizerWriter

Checklist

• I am familiar with the Contributing Guide.
• New or Existing tests cover these changes.
• The documentation is up to date with these changes.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR removes MegatronTokenizerWriter from writer/__init__.py to avoid transitively importing transformers (~150 MB peak memory) when users only need JsonlWriter or ParquetWriter. Internal references in tests and tutorials are updated to use the direct import path (from nemo_curator.stages.text.io.writer.megatron_tokenizer import MegatronTokenizerWriter).

Key changes:

• writer/__init__.py: Removes the MegatronTokenizerWriter import and its __all__ entry; JsonlWriter and ParquetWriter remain unaffected.
• tests/stages/text/io/writer/test_megatron_tokenizer.py: Consolidates the two separate imports into a single direct import from megatron_tokenizer module.
• tutorials/text/megatron-tokenizer/main.py: Updates import to the direct module path.

Notable observation: This is a breaking change for any external users who were importing MegatronTokenizerWriter from the package-level nemo_curator.stages.text.io.writer. The Breaking Changes section of docs/about/release-notes/index.md does not mention this import path change, which could surprise users who upgrade without reading commit history. Adding a note there would improve discoverability of the migration path.

Confidence Score: 4/5

• This PR is safe to merge; all internal references are correctly updated and the core logic is unchanged.
• The change is minimal and well-scoped: removing one import and updating two call sites. No logic is altered, existing tests cover the affected class, and all internal usages of the old import path have been updated. The one point deduction is because the Breaking Changes section in the release notes omits the import-path change, which is a public API break that external users will hit as an ImportError on upgrade.
• docs/about/release-notes/index.md — the Breaking Changes section should document the new import path for MegatronTokenizerWriter.

Important Files Changed

Filename	Overview
nemo_curator/stages/text/io/writer/init.py	Removes MegatronTokenizerWriter re-export and its all entry. Clean change with no issues; JsonlWriter and ParquetWriter remain correctly exported.
tests/stages/text/io/writer/test_megatron_tokenizer.py	Updates import to use the direct module path; consolidates two separate imports from megatron_tokenizer into one line. No logic changes, tests unchanged.
tutorials/text/megatron-tokenizer/main.py	Updates import to use the direct module path instead of the package-level re-export. No functional changes.

Flowchart

flowchart TD
    A["User imports MegatronTokenizerWriter"] --> B{Before PR}
    A --> C{After PR}

    B --> D["from nemo_curator.stages.text.io.writer import MegatronTokenizerWriter"]
    D --> E["writer/__init__.py"]
    E --> F["megatron_tokenizer.py\n(imports transformers ~150MB)"]

    C --> G["from nemo_curator.stages.text.io.writer.megatron_tokenizer import MegatronTokenizerWriter"]
    G --> F

    H["User imports JsonlWriter or ParquetWriter"] --> I["from nemo_curator.stages.text.io.writer import JsonlWriter"]
    I --> J["writer/__init__.py (Before PR)\nLoaded transformers transitively"]
    I --> K["writer/__init__.py (After PR)\nNo transformers import"]

    style J fill:#f88,stroke:#c33
    style K fill:#8f8,stroke:#3a3

Loading

_{Last reviewed commit: 7b4fcff}

[greptile-apps]

_{3 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[sarahyurick]

Thank you @ArivunidhiA !

[sarahyurick]

/ok to test 9a75edb

[greptile-apps]

_{3 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Additional Comments (1)

docs/about/release-notes/index.md
Missing breaking change entry for MegatronTokenizerWriter import path

The Breaking Changes section does not mention that MegatronTokenizerWriter is no longer re-exported from nemo_curator.stages.text.io.writer. Any external user who had:

from nemo_curator.stages.text.io.writer import MegatronTokenizerWriter

will receive an ImportError after upgrading. Without a release-notes entry, they have no obvious migration path. Consider adding:

- **MegatronTokenizerWriter Import Path Changed**: `MegatronTokenizerWriter` is no longer re-exported from `nemo_curator.stages.text.io.writer`. Use the direct import instead:
  ```python
  from nemo_curator.stages.text.io.writer.megatron_tokenizer import MegatronTokenizerWriter

</details>