NeMo Curator

NeMo Curator is a Python library specifically designed for scalable and efficient dataset preparation. It greatly accelerates data curation by leveraging GPUs with Dask and RAPIDS, resulting in significant time savings. The library provides a customizable and modular interface, simplifying pipeline expansion and accelerating model convergence through the preparation of high-quality tokens.

At the core of the NeMo Curator is the DocumentDataset which serves as the the main dataset class. It acts as a straightforward wrapper around a Dask DataFrame. The Python library offers easy-to-use methods for expanding the functionality of your curation pipeline while eliminating scalability concerns.

Key Features

NeMo Curator provides a collection of scalable data-mining modules. Some of the key features include:

Data download and text extraction

• Default implementations for downloading and extracting Common Crawl, Wikipedia, and ArXiv data
• Easily customize the download and extraction and extend to other datasets

Language identification and separation

• Language identification with fastText and pycld2

Text reformatting and cleaning

• Fix unicode decoding errors via ftfy

Quality filtering

• Multilingual heuristic-based filtering
• Classifier-based filtering via fastText

Document-level deduplication

• Both exact and fuzzy deduplication are accelerated using cuDF and Dask
• For fuzzy deduplication, our implementation follows the method described in Microsoft Turing NLG 530B

Multilingual downstream-task decontamination

• Our implementation follows the approach of OpenAI GPT3 and Microsoft Turing NLG 530B

Distributed data classification

• Multi-node, multi-GPU classifier inference
• Provides sophisticated domain and quality classification
• Flexible interface for extending to your own classifier network

Personal identifiable information (PII) redaction

• Identification tools for removing addresses, credit card numbers, social security numbers, and more

These modules offer flexibility and permit reordering, with only a few exceptions. In addition, the NeMo Framework Launcher provides pre-built pipelines that can serve as a foundation for your customization use cases.

Resources

• Documentation
• Examples
• Tutorials

Get Started

This section explains how to install NeMo Curator and use the Python library, Python modules, and CLI scripts. It also includes a list of tutorials to help you get started right away. Finally, this section explains how to use the NeMo Framework Launcher as an alternative method for interfacing with NeMo Curator.

Requirements

Before installing NeMo Curator, ensure that the following requirements are met:

• Python 3.10
• Ubuntu 22.04/20.04
• NVIDIA GPU (optional)
  • Volta™ or higher (compute capability 7.0+)
  • CUDA 12 (or above)

Install NeMo Curator

You can install NeMo-Curator from PyPi, from source or get it through the NeMo Framework container.

PyPi

NeMo Curator can be installed via PyPi as follows -

To install the CPU-only modules:

pip install nemo-curator

To install the CPU and CUDA-accelerated modules:

pip install --extra-index-url https://pypi.nvidia.com nemo-curator[cuda12x]

From Source

1. Clone the NeMo Curator repository in GitHub.

   git clone https://github.com/NVIDIA/NeMo-Curator.git
   cd NeMo-Curator

2. Install the modules that you need.

   To install the CPU-only modules:

   pip install .

   To install the CPU and CUDA-accelerated modules:

   pip install --extra-index-url https://pypi.nvidia.com ".[cuda12x]"

Install from the NeMo Framework Container

NeMo Curator is available in the NeMo Framework Container. The latest release of NeMo Curator comes preinstalled in the container.

If you want the latest commit inside the container, uninstall the existing version using:

pip uninstall nemo-curator

And follow the instructions for installing from source from above.

Use the Python Library

The following snippet demonstrates how to create a small data curation pipeline that downloads and curates a small subset of the Common Crawl dataset.

# Download your dataset
dataset = download_common_crawl("/datasets/common_crawl/", "2021-04", "2021-10", url_limit=10)
# Build your pipeline
curation_pipeline = Sequential([
  Modify(UnicodeReformatter()),
  ScoreFilter(WordCountFilter(min_words=80)),
  ScoreFilter(FastTextQualityFilter(model_path="model.bin")),
  TaskDecontamination([Winogrande(), Squad(), TriviaQA()])
])
# Curate your dataset
curated_dataset = curation_pipeline(dataset)

Explore NeMo Curator Tutorials

To get started with NeMo Curator, you can follow the tutorials available here: [Tutorials] (https://github.com/NVIDIA/NeMo-Curator/tree/main/tutorials). These tutorials include:

• tinystories which focuses on data curation for training from scratch.
• peft-curation which focuses on data curation for parameter-efficient fine-tuning use-cases.

Access Python Modules

The Data Curation section of the NeMo Framework User Guide provides in-depth information about how the Python modules work. The examples directory in the GitHub repository provides scripts that showcase these modules.

Use CLI Scripts

NeMo Curator also offers CLI scripts for you to use. The scripts in nemo_curator/scripts map closely to the supplied Python modules. Refer to the NeMo Framework User Guide for more information about the Python modules and scripts.

Use NeMo Framework Launcher

As an alternative method for interfacing with NeMo Curator, you can use the NeMo Framework Launcher. The launcher enables you to easily configure the parameters and cluster. It can also automatically generate the SLURM batch scripts that wrap around the CLI scripts required to run your pipeline.

In addition, other methods are available to run NeMo Curator on SLURM. For example, refer to the example scripts in examples/slurm for information on how to run NeMo Curator on SLURM without the NeMo Framework Launcher.

Module Ablation and Compute Performance

The modules within NeMo Curator were primarily designed to curate high-quality documents from Common Crawl snapshots in a scalable manner. To evaluate the quality of the curated Common Crawl documents, we conducted a series of ablation experiments. In these experiments, we trained a 357M-parameter GPT-style model using datasets generated at various stages of our data curation pipeline, which was implemented in NeMo Curator.

The following figure shows that the use of different data curation modules implemented in NeMo Curator led to improved model zero-shot downstream task performance.

In terms of scalability and compute performance, using the combination of RAPIDS and Dask fuzzy deduplication enabled us to deduplicate the 1.1 Trillion token Red Pajama dataset in 1.8 hours with 64 NVIDIA A100 Tensor Core GPUs.

Additionally, using the CPU-based modules, the following table shows the time required and resulting data size reduction for each processing step Common Crawl snapshot from November/December of 2020 using 30 CPU nodes (with hardware similar to the c5.24xlarge Amazon AWS C5 instance).

Dataset	Download and text extraction		Text cleaning		Quality filtering
	Time	Output Size	Time	Output Size	Time	Output Size
Common Crawl 2020-50	36 hrs	2.8 TB	1 hr	2.8 TB	0.2 hr	0.52 TB

Contribute to NeMo Curator

We welcome community contributions! Please refer to CONTRIBUTING.md for the process.

To contribute an article to the collection, please submit a pull request to the gh-pages-src branch of this repository. For detailed information, please consult the README located at the gh-pages-src branch.