Add NeMo Curator PII redaction for pipeline-scale detection

Introduction

When you push Presidio past a few thousand documents, two walls show up: regex misses contextual person names and organization-linked addresses, and single-process detection chokes on multi-gigabyte corpora. NeMo Curator addresses both — a transformer-based NER model that reads surrounding context, wrapped in a Dask-distributed pipeline that scales horizontally across cores or workers. Get this wrong and your training corpus ships with names, account numbers, and patient identifiers embedded in prose, which is the kind of leak that triggers data-deletion requests across an entire fine-tune. By the end you'll be able to configure a NeMo Curator PII redaction pipeline on a DocumentDataset, and layer it behind Presidio so the expensive transformer pass runs only on documents the cheap regex pass left ambiguous.

Key Terminology

• NeMo Curator — NVIDIA's open-source data curation toolkit that wraps transformer-based NER models in Dask-distributed pipelines; it supplies the high-recall second pass that catches PII Presidio's regex+spaCy first pass leaves behind.
• PiiModifier — NeMo Curator's modifier class that wraps a transformer-based NER model and applies token-level PII detection and redaction across a DocumentDataset inside a pipeline.
• DocumentDataset — NeMo Curator's batch-oriented wrapper over a Dask DataFrame; PII modules consume and emit this type so detection scales across partitions and workers.
• Layered detection pipeline — a two-stage architecture in which Presidio runs a cheap regex+spaCy first pass and only documents with low-confidence or empty results are routed to NeMo Curator's transformer-based second pass.
• Anonymize action — the post-detection transformation PiiModifier applies per detected span (replace writes [PERSON]-style placeholders, mask overwrites characters, remove deletes the span); chosen per downstream consumer's needs.

Concepts

Why NeMo Curator complements Presidio

NeMo Curator is NVIDIA's data curation toolkit for cleaning and preparing large-scale datasets for language-model training and retrieval workloads. Its PII module applies transformer-based NER that detects sensitive information missed by regex and rule-based systems. Presidio's strength is structured identifiers — emails, phone numbers, credit cards — where regex anchored by spaCy NER hits high precision cheaply. NeMo Curator's strength is contextual entities: a person's name buried in a chat transcript, an organization implied by an address, financial references phrased in prose. The right deployment is layered, not either-or: Presidio first, then NeMo Curator only on what Presidio couldn't confidently resolve (see Code Walkthrough).

Pipeline-scale architecture

The DocumentDataset wraps a Dask DataFrame, and every Sequential pipeline operates on partitions in parallel. Scale-out is therefore a partition-count knob, not a code change — the same PiiModifier you configure on a laptop runs unchanged on a cluster. batch_size controls transformer forward-pass size on each worker; npartitions controls how the workload spreads across cores. Tune both together: raise batch_size until per-worker memory is saturated, then raise npartitions to spread load rather than overloading a single partition. Computing per-entity-type precision/recall against a labeled test set before promoting the layered pipeline is non-negotiable — Presidio typically wins on EMAIL_ADDRESS while NeMo Curator wins on PERSON, and the routing threshold should reflect that.

Anonymize actions and entity coverage

PiiModifier exposes three redaction strategies via anonymize_action: replace substitutes type-tagged placeholders like [PERSON] (preserves sentence structure for downstream LM training), mask overwrites characters (preserves length for log alignment), and remove deletes the span entirely (smallest output but breaks token offsets). DEFAULT_ENTITIES covers the common PII set; pass an explicit list to supported_entities when you need a narrower scope — pruning unused types keeps the transformer from spending cycles scoring entities your downstream policy will ignore.

Code Walkthrough

Now that you have seen the concepts above, the walkthrough below turns them into working code.

The first snippet configures PiiModifier and runs it over a DocumentDataset — the minimum viable NeMo Curator redaction job. The second wires it behind Presidio, routing only low-confidence documents to the expensive transformer pass.

1import pandas as pd
2import dask.dataframe as dd
3from nemo_curator import Sequential
4from nemo_curator.modules.modify import PiiModifier
5from nemo_curator.pii.constants import DEFAULT_ENTITIES
6from nemo_curator.datasets import DocumentDataset
7
8pii_modifier = PiiModifier(
9    language="en",
10    supported_entities=DEFAULT_ENTITIES,
11    anonymize_action="replace",
12    device="cpu",
13    batch_size=16,
14)
15
16documents_df = dd.from_pandas(
17    pd.DataFrame({
18        "text": [
19            "Dr. Sarah Chen at MIT prescribed medication for patient ID 4829103.",
20            "Contact support@acme.com or call 555-867-5309 for billing.",
21        ]
22    }),
23    npartitions=2,
24)
25
26dataset = DocumentDataset(documents_df)
27redacted_dataset = Sequential([pii_modifier])(dataset)
28print(redacted_dataset.df.compute())

1def layered_pii_pipeline(documents, analyzer, pii_modifier, threshold=0.7):
2    first_pass = []
3    needs_second = []
4
5    for doc in documents:
6        results = analyzer.analyze(text=doc["text"], language="en")
7        high = [r for r in results if r.score >= threshold]
8        low = [r for r in results if r.score < threshold]
9
10        first_pass.append({
11            "doc_id": doc["id"],
12            "detections": high,
13            "text": doc["text"],
14        })
15
16        if low or not high:
17            needs_second.append(doc)
18
19    second_df = dd.from_pandas(pd.DataFrame(needs_second), npartitions=4)
20    second_dataset = DocumentDataset(second_df)
21    second_results = Sequential([pii_modifier])(second_dataset).df.compute()
22    return first_pass, second_results

You'll know it works when the redacted output for the sample text shows [PERSON], [EMAIL_ADDRESS], and [PHONE_NUMBER] placeholders in place of the originals, and the layered pipeline routes only a fraction of documents to the second pass — measure that fraction; if it's near 100%, your threshold is too high or Presidio is misconfigured.

Do's and Don'ts

Now that you have worked through the implementation, the practices below separate a durable approach from a fragile one.

Do's

1. ✓Do configure supported_entities explicitly — start from DEFAULT_ENTITIES and prune so the transformer does not waste cycles scoring types your redaction policy will ignore.
2. ✓Do tune batch_size and npartitions together — raise batch_size until per-worker memory is saturated, then increase npartitions to spread load rather than overload a single partition.
3. ✓Do route only low-confidence or empty Presidio results to the second pass — treat the threshold (e.g. 0.7) as the cost knob; higher thresholds send more documents to the transformer and raise compute spend.

Don'ts

1. ✗Don't call PiiModifier on every document by default — the transformer's per-document cost makes it a second-pass tool; sending the whole corpus through negates the layered design.
2. ✗Don't run device="cpu" on a GPU-equipped worker for production-scale corpora — leaving the GPU idle wastes the most expensive resource in the cluster.
3. ✗Don't mix raw Pandas DataFrames with DocumentDataset mid-pipeline — convert through dd.from_pandas(..., npartitions=...) so partitioning is explicit and Dask can schedule transformer batches in parallel.