From 9a75edbf801555219d230d9a0884f81c8157f41d Mon Sep 17 00:00:00 2001 From: Arivunidhi A Date: Sat, 14 Feb 2026 22:55:45 -0500 Subject: [PATCH] Remove MegatronTokenizerWriter from writer/__init__.py (#1418) Signed-off-by: Arivunidhi A --- nemo_curator/stages/text/io/writer/__init__.py | 2 -- tests/stages/text/io/writer/test_megatron_tokenizer.py | 3 +-- tutorials/text/megatron-tokenizer/main.py | 2 +- 3 files changed, 2 insertions(+), 5 deletions(-) diff --git a/nemo_curator/stages/text/io/writer/__init__.py b/nemo_curator/stages/text/io/writer/__init__.py index 828c37ff63..4b0362b21d 100644 --- a/nemo_curator/stages/text/io/writer/__init__.py +++ b/nemo_curator/stages/text/io/writer/__init__.py @@ -13,11 +13,9 @@ # limitations under the License. from nemo_curator.stages.text.io.writer.jsonl import JsonlWriter -from nemo_curator.stages.text.io.writer.megatron_tokenizer import MegatronTokenizerWriter from nemo_curator.stages.text.io.writer.parquet import ParquetWriter __all__ = [ "JsonlWriter", - "MegatronTokenizerWriter", "ParquetWriter", ] diff --git a/tests/stages/text/io/writer/test_megatron_tokenizer.py b/tests/stages/text/io/writer/test_megatron_tokenizer.py index 0670261a71..31e5bd2d6b 100644 --- a/tests/stages/text/io/writer/test_megatron_tokenizer.py +++ b/tests/stages/text/io/writer/test_megatron_tokenizer.py @@ -24,8 +24,7 @@ import pytest import nemo_curator.stages.text.io.writer.utils as writer_utils -from nemo_curator.stages.text.io.writer import MegatronTokenizerWriter -from nemo_curator.stages.text.io.writer.megatron_tokenizer import _INDEX_HEADER +from nemo_curator.stages.text.io.writer.megatron_tokenizer import _INDEX_HEADER, MegatronTokenizerWriter from nemo_curator.tasks import DocumentBatch diff --git a/tutorials/text/megatron-tokenizer/main.py b/tutorials/text/megatron-tokenizer/main.py index 714643549f..4e13a27591 100644 --- a/tutorials/text/megatron-tokenizer/main.py +++ b/tutorials/text/megatron-tokenizer/main.py @@ -22,7 +22,7 @@ from nemo_curator.core.client import RayClient from nemo_curator.pipeline import Pipeline from nemo_curator.stages.text.io.reader import ParquetReader -from nemo_curator.stages.text.io.writer import MegatronTokenizerWriter +from nemo_curator.stages.text.io.writer.megatron_tokenizer import MegatronTokenizerWriter def main(args: argparse.Namespace) -> None: