Merge branch 'main' into video_nvdec_fix

Author: ayushdg

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",
 ]

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
 
 

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: