feat: PII Replacement V3-MVP mode

[nina-xu]

Summary

This PR replaces the old NER / GLiNER / Jinja data_editor PII stack with a tabular, plan-driven replacer. Detection is heuristic (column names + value patterns); replacement is deterministic given a plan, seed, and persona backend. There is no GPU NER and no LLM in this release.

---

What users experience

When replace_pii is enabled (the default), the pipeline still runs PII replacement on the training dataframe before training. The mechanism is new:

1. Resolve a plan — auto-discover, load a YAML file, or use an inline plan.
2. Validate the plan against the dataframe (and, for user plans, also during --validate preflight).
3. Apply synthetic replacements, then propagate those values into free-text columns.
4. Write pii_replacement_plan.yaml into the run directory so the discovered (or used) plan can be edited and reused.

Turning it off is unchanged in spirit: replace_pii: null, --no-replace-pii, or .with_replace_pii(enable=False).

---

Config shape (replace_pii)

Old step/NER/classify YAML (PiiReplacerConfig) is gone. The new top-level object is ReplacePiiConfig:

Field	Behavior
replacement_plan	"auto_discovery" (default), a path to a plan YAML, or an inline plan (config file only)
replacement.locale / seed	Locale and RNG seed for generation
person.backend	managed (default, pre-generated personas), faker, or internal-only pgm
person.managed_assets_path / sdg_pgms_src	Asset / PGM source locations
discovery.replace_group_key	If true (default) and data.group_training_examples_by is set, that column is forced onto the plan as a standalone unique_identifier
llm_enhancement	Reserved; true raises NotImplementedError at apply time (preflight warns)

A plan has:

• scope: record | group | dataframe — how widely one original value keeps the same synthetic value
• persona_backed_columns: named personas; columns filled from one synthetic identity; optional match_persona_by (gender / ethnic_background) columns that are read only
• standalone_columns_to_replace: IDs, cards, free text, etc., replaced independently

Each column entry is column_name + entity_type + optional patterns (formats for writing values).

Important rule: the engine follows entity type, not YAML section. Putting an ID under a persona does not make it persona-sourced; putting a name only under standalone means it will not share a synthetic person with other name columns. Preflight/apply emit warnings for those mismatches; they do not rewrite the plan.

---

Discovery (auto plan)

With replacement_plan: auto_discovery, discovery:

• Uses column names (fuzzy), value regex coverage (≥ 85% of non-null values for a dominant pattern), cardinality (≤ 12 distinct → not treated as an ID from name alone), and optional group-constancy when a training group key is set.
• Groups person-like columns into personas; puts IDs/cards/IPs/etc. in standalone; marks long prose columns as free_text.
• Sets scope to group if group_training_examples_by is set, else dataframe.
• Marks generic temporal columns (date / datetime / time / duration) as identified, not replaced — logged, omitted from the plan.
• Skips free-text scanning when no structured replaceable columns exist (nothing to propagate).
• Infers patterns from observed value shapes (names/emails as persona placeholders; DOB as strftime; IDs/phones as character templates). IPv4/IPv6 get no templates.

No LLM or NER model is used.

---

Replacement (exact apply behavior)

Persona-sourced (from one synthetic identity when listed under a persona): names, email, street address, SSN, national ID. Phone is persona-sourced only under pgm; under managed/faker it is rebuilt from the column’s format like other identifiers.

Entity-driven (stable original→synthetic map, standalone path always): unique IDs, cards, API keys, IPs, DOB (age-preserving perturbation keeping format), and non-pgm phones.

Free text: does not replace the whole cell; substitutes values already replaced elsewhere in the row (plus name tokens for partial mentions like “Dr. Smith”). Word-boundary matching; original punctuation variants are normalized for matching.

Persona backends:

• managed — sample datasets/{locale}.parquet under NSS_MANAGED_ASSETS_PATH or ~/.data-designer/managed-assets; missing assets fall back to Faker with a warning.
• faker — Faker personas; locale must be supported.
• pgm — local sdg_pgms checkout; no fallback; en_US only; only backend that supplies persona phone numbers.

match_persona_by conditions which persona is drawn (e.g. sex→gender for name agreement). Condition columns are never replaced.

After apply, column stats for the evaluation report record transform methods (personas/Faker/PGM, pattern, perturbation, propagation).

---

Validation and preflight

User-supplied plans are checked in preflight (PiiPlanValidityCheck) so --validate catches bad plans without running replacement. Errors include unknown/duplicate columns, missing entity_type, entity_type: date, patterns that don’t match values or don’t apply to that entity, group scope without a group key, and persona/matcher conflicts. Auto-discovered plans are validated when replacement runs, not in that preflight check. Separate config checks cover Faker locale, managed assets presence (warning), and PGM source (error).

---

What was removed

• Entire NER stack, GLiNER/regex entity extractors, and NemoPII / data_editor (including Jinja/Faker template execution).
• Old replace_pii step/classify/NER config.
• Expression-based data actions (GenExpression, GenFaker, etc.) that depended on that engine.
• Dead fragment.py and NER metadata on field features.

Typed ActionExecutor actions (distributions, UUID, date constraints, etc.) remain but are not on the default PII path.

---

Net behavioral change vs main

Before	After
NER + optional LLM classification; transform steps / Jinja	Heuristic discovery + declarative plan + programmatic replace
Large entity vocabulary / NER models	Closed entity set above
GPU/model deps for detection	CPU heuristics only
Opaque transforms	Editable pii_replacement_plan.yaml round-trip

Default runs still replace PII automatically; the plan file in the run dir is the main new control surface for audit and override.

---

Testing Plan

Paths below assume datasets live in /root/datasets and configs in /root/configs. Artifacts land under safe-synthesizer-artifacts/<config>---<dataset>/<timestamp>/ (plan as pii_replacement_plan.yaml, transformed data as dataset/transformed_training.csv when PII changed anything).

For PII-only checks, prefer the SDK TabularPiiReplacer snippets (no GPU/train). Use CLI --validate for plan/preflight, and a tiny TinyLlama run only when you need end-to-end artifact wiring.

Dataset map

Dataset	Why use it	Behaviors exercised
/root/datasets/telco_churn.csv	Flat table: CustomerID, Name, Gender, city/state/zip	Auto-discovery; full_name + match_persona_by gender; ID patterns (CUST-…); dataframe scope
/root/datasets/patient_events.csv	Multi-row patients: patient_id, names, date_of_birth, sex, provider_name, notes/details	Group scope via group_training_examples_by=patient_id; DOB perturbation; dual personas (patient + provider); free-text propagation
/root/datasets/bike_sales.csv	Split name parts + address + phone + BirthDate + Gender	Persona-sourced first/middle/last; street address; phone as entity-driven under managed/faker; match_persona_by
/root/datasets/crm/Contact.csv	CRM person record: names, email, phone, mailing address	Email + phone patterns; street address; Salesforce-style IDs
/root/datasets/pii_dataset_no_label.csv or personal_bios.csv	Structured PII columns and a long text bio that mentions them	Free-text propagation of name/email/phone/address into prose
/root/datasets/call_transcripts.csv	Almost only free text (transcript)	No structured PII → free-text scan skipped; little/no replacement
/root/datasets/EHR.csv	Clinical IDs + demographics, little classic name PII	ID-heavy discovery; gender/ethnicity as matchers not replaced; temporal fields identified-not-replaced
/root/datasets/clinc_oos.csv	Intent utterances, no person columns	Empty / near-empty plan; confirms “nothing structured → skip free text”

1. Auto-discovery smoke (CLI)

Tiny config (/root/configs/quick-tinyllama-pii.yaml or equivalent) + Faker so managed assets are not required:

# Telco: flat tabular discovery
safe-synthesizer run \
  --config /root/configs/quick-tinyllama-pii.yaml \
  --data-source /root/datasets/telco_churn.csv \
  --replace_pii__person__backend faker \
  --replace_pii__replacement__seed 42 \
  --emit_telemetry false

# Patient events: group-scoped discovery
safe-synthesizer run \
  --config /root/configs/quick-tinyllama-pii.yaml \
  --data-source /root/datasets/patient_events.csv \
  --data__group_training_examples_by patient_id \
  --replace_pii__person__backend faker \
  --replace_pii__replacement__seed 42 \
  --emit_telemetry false

Check: run dir contains pii_replacement_plan.yaml. For patient events, expect scope: group, patient + provider personas, patient_id as standalone unique_identifier, notes/details as free_text, and DOB present. For telco, expect scope: dataframe, Name under a persona with Gender in match_persona_by, CustomerID standalone.

PII-only SDK equivalent:

import pandas as pd
from nemo_safe_synthesizer.config.data import DataParameters
from nemo_safe_synthesizer.config.pii_replacement import (
    AUTO_DISCOVERY, PiiPersonBackend, PiiPersonConfig, PiiReplacementSettings, ReplacePiiConfig,
)
from nemo_safe_synthesizer.pii_replacer import TabularPiiReplacer

df = pd.read_csv("/root/datasets/patient_events.csv")
cfg = ReplacePiiConfig(
    replacement_plan=AUTO_DISCOVERY,
    person=PiiPersonConfig(backend=PiiPersonBackend.faker),
    replacement=PiiReplacementSettings(seed=42),
)
replacer = TabularPiiReplacer(
    cfg,
    data_config=DataParameters(group_training_examples_by="patient_id"),
    workdir="/tmp/pii-patient-auto",
)
replacer.transform_df(df)
print(replacer.resolved_plan.model_dump())
print(replacer.result.transformed_df[["patient_id", "first_name", "last_name", "provider_name"]].head())

2. Plan round-trip (edit → reuse)

1. Run auto-discovery once (above) and copy pii_replacement_plan.yaml out of the run dir.
2. Edit: drop a column, move provider_name under its own persona, add/adjust patterns, or force scope.
3. Re-run pointing at the file:

safe-synthesizer run \
  --config /root/configs/quick-tinyllama.yaml \
  --data-source /root/datasets/patient_events.csv \
  --data__group_training_examples_by patient_id \
  --replace_pii__replacement_plan /path/to/edited_pii_replacement_plan.yaml \
  --replace_pii__person__backend faker \
  --emit_telemetry false

Or inline in a config (see /root/configs/new-pii-repl.yaml / quick-tinyllama-pii.yaml):

replace_pii:
  person:
    backend: faker
  replacement_plan:
    scope: group
    persona_backed_columns:
      - persona: patient
        columns_to_replace:
          - column_name: first_name
            entity_type: first_name
          - column_name: last_name
            entity_type: last_name
        match_persona_by:
          - persona_attribute: gender
            column_name: sex
      - persona: provider
        columns_to_replace:
          - column_name: provider_name
            entity_type: full_name
    standalone_columns_to_replace:
      - column_name: patient_id
        entity_type: unique_identifier
      - column_name: date_of_birth
        entity_type: date_of_birth
        patterns: ["%m/%d/%Y"]
      - column_name: notes
        entity_type: free_text

Check: dropped columns stay original; persona columns share one identity within scope; preflight/--validate rejects bad column names and bad patterns.

3. Preflight / --validate (no replacement apply)

# Good plan loads and validates against columns
safe-synthesizer run --validate \
  --config /root/configs/quick-tinyllama-pii.yaml \
  --data-source /root/datasets/telco_churn.csv

# Bad plan: expect pii_plan_column_not_found / pattern errors
# (point replacement_plan at a YAML that names MissingCol or sets patterns on ipv4)

Also exercise:

Intentional misconfig	Expected
scope: group without data.group_training_examples_by	pii_plan_group_scope_invalid
entity_type: date on a column	pii_plan_entity_type_invalid
llm_enhancement: true	today: preflight warning, apply NotImplementedError (candidate to harden to error)
person.backend: faker + bogus locale	pii_faker_locale_invalid
Missing managed parquet with backend: managed	pii_managed_assets_missing warning; apply falls back to Faker

4. Consistency scope

Use patient_events (many rows per patient_id):

Scope	How to set	What to assert
group (default when group key set)	data.group_training_examples_by: patient_id + auto plan	Same patient → same synthetic name across their rows; different patients independent
dataframe	Plan scope: dataframe (or no group key on telco/bike)	Same original value → same synthetic across the whole frame
record	Explicit plan scope: record on a small slice with duplicate names	Each row gets its own persona even if first_name repeats; also assert free-text notes agree with that row’s structured name (regression for H2)

# Record-scope regression sketch (duplicate Alice + notes)
df = pd.DataFrame({
    "first_name": ["Alice", "Alice"],
    "notes": ["Alice visited today", "Alice visited today"],
})
# plan: scope=record, persona first_name, standalone free_text notes
# assert transformed notes each contain that row's new first_name, not the other row's

For standalone IDs/phones under scope: group|record, today’s code keeps a dataframe-wide map (H1) — assert current behavior, then re-assert after the fix.

5. Free-text propagation

Best datasets: patient_events (notes mentioning clinical context; provider honorifics) and pii_dataset_no_label / personal_bios (text repeats name / email / phone / address).

safe-synthesizer run \
  --config /root/configs/quick-tinyllama.yaml \
  --data-source /root/datasets/pii_dataset_no_label.csv \
  --replace_pii__person__backend faker \
  --replace_pii__replacement__seed 42 \
  --training__num_input_records_to_sample 200 \
  --emit_telemetry false

Check: in transformed_training.csv, structured name/email/phone change and the same originals disappear from text. Partial mentions (Dr. <Last>) follow when only full_name is planned. Case variants (Smith vs SMITH) — document current case-sensitive behavior (M1), then retest after fix.

Negative control: call_transcripts.csv or clinc_oos.csv — plan should have little/no structured PII; free-text scan skipped or empty; transcripts largely unchanged.

6. Persona backends

# Faker (no assets)
safe-synthesizer run --config /root/configs/quick-tinyllama.yaml \
  --data-source /root/datasets/bike_sales.csv \
  --replace_pii__person__backend faker \
  --replace_pii__replacement__seed 7 \
  --emit_telemetry false

# Managed (needs ~/.data-designer/managed-assets/datasets/en_US.parquet or NSS_MANAGED_ASSETS_PATH)
safe-synthesizer run --config /root/configs/quick-tinyllama.yaml \
  --data-source /root/datasets/bike_sales.csv \
  --replace_pii__person__backend managed \
  --replace_pii__replacement__locale en_US \
  --emit_telemetry false

# Managed with assets deliberately missing → warning + Faker fallback; confirm logs/stats
NSS_MANAGED_ASSETS_PATH=/tmp/empty-assets \
safe-synthesizer run --config /root/configs/quick-tinyllama.yaml \
  --data-source /root/datasets/telco_churn.csv \
  --replace_pii__person__backend managed \
  --emit_telemetry false

Check on bike_sales: FirstName/LastName/AddressLine1/PhoneNumber/BirthDate replaced; Gender unchanged but conditions names; phone is not persona-tied under managed/faker (entity-driven / pattern). With pgm (if /root/sdg-pgms/src present), phone may come from the persona — only backend that does.

7. Gender / ethnicity conditioning

# telco: Name conditioned on Gender
# patient_events: names conditioned on sex
# bike_sales: names conditioned on Gender

After replacement, spot-check that synthetic given names are plausible for the row’s Gender/sex (Faker male/female generators). Ethnicity conditioning needs a column mapped via match_persona_by: ethnic_background (e.g. EHR ethnicity or patient race in an edited plan).

8. Disable PII / no-op paths

safe-synthesizer run --config /root/configs/quick-tinyllama.yaml \
  --data-source /root/datasets/telco_churn.csv \
  --no-replace-pii \
  --emit_telemetry false

SafeSynthesizer().with_data_source("/root/datasets/telco_churn.csv").with_replace_pii(enable=False).with_train()

Check: no pii_replacement_plan.yaml (or PII stage skipped); training CSV equals source holdout.

9. End-to-end SDK customization

from nemo_safe_synthesizer.sdk.library_builder import SafeSynthesizer

synth = (
    SafeSynthesizer()
    .with_data_source("/root/datasets/crm/Contact.csv")
    .with_replace_pii(
        person={"backend": "faker"},
        replacement={"locale": "en_US", "seed": 42},
        # after a first auto run, pass the edited plan path:
        # replacement_plan="/path/to/pii_replacement_plan.yaml",
    )
    .with_train(pretrained_model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", num_input_records_to_sample=200)
    .with_generate(num_records=50)
)
# synth.fit() / run per local SDK entrypoint

Check CRM Contact: FirstName/LastName/Email/Phone/MailingStreet replaced consistently as one persona; Id standalone.

10. Suggested matrix (quick vs deep)

Priority	Scenario	Dataset	Command style
P0	Auto-discovery + plan emit	telco_churn, patient_events	CLI tiny run or SDK replacer
P0	Group consistency	patient_events	SDK replacer + assert per patient_id
P0	Free-text propagation	pii_dataset_no_label, patient_events	SDK replacer; diff text/notes
P0	--validate bad plan	telco + hand-written YAML	CLI --validate
P1	Plan round-trip edit	patient_events	CLI with replacement_plan path
P1	Record-scope + notes	tiny synthetic DF	SDK unit-style
P1	Managed fallback	telco + empty assets path	CLI
P1	Phone routing managed vs pgm	bike_sales / Contact	SDK, compare backends
P2	No structured PII	call_transcripts, clinc_oos	CLI/SDK; empty plan
P2	Full TinyLlama E2E	telco or Contact	quick-tinyllama.yaml

What to inspect after each run

1. pii_replacement_plan.yaml — scope, personas, matchers, patterns, free_text columns
2. Logs — “Resolved PII replacement plan”, placement advisories, managed-fallback warnings
3. dataset/transformed_training.csv vs dataset/training.csv — columns changed, referential consistency within scope
4. Evaluation column stats / transform methods (personas, pattern, perturbation, propagation) when a full run completes

Pre-Review Checklist

Ensure that the following pass:

• mise run format && mise run check or via prek validation.
• mise run test passes locally
• mise run test:e2e passes locally
• mise run test:ci-container passes locally (recommended)
• GPU CI status check passes -- comment /sync on this PR to trigger a run (auto-triggers on ready-for-review)

Pre-Merge Checklist

• New or updated tests for any fix or new behavior
• Updated documentation for new features and behaviors, including docstrings for API docs.

Other Notes

• Closes #

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 96fd1d61-28ea-4e14-b236-6cbcab32e927

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch nina-xu/pii-repl-v3-alt-config

---

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

❌ Patch coverage is 88.63049% with 264 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/nemo_safe_synthesizer/pii_replacer/core.py	87.02%	144 Missing ⚠️
src/nemo_safe_synthesizer/pii_replacer/persona.py	61.22%	38 Missing ⚠️
.../nemo_safe_synthesizer/pii_replacer/replacement.py	88.62%	38 Missing ⚠️
src/nemo_safe_synthesizer/pii_replacer/plan.py	92.69%	16 Missing ⚠️
...rc/nemo_safe_synthesizer/pii_replacer/discovery.py	93.10%	8 Missing ⚠️
tests/smoke/test_pii_replacement_cpu.py	38.46%	8 Missing ⚠️
src/nemo_safe_synthesizer/preflight/checks/pii.py	93.10%	6 Missing ⚠️
src/nemo_safe_synthesizer/pii_replacer/replacer.py	97.22%	3 Missing ⚠️
tests/evaluation/test_render.py	0.00%	2 Missing ⚠️
...rc/nemo_safe_synthesizer/config/pii_replacement.py	99.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!